Multiprocessing and minor changes in parameters

radiosim/jet.py

@@ -1,5 +1,5 @@
 import numpy as np
-from radiosim.utils import relativistic_boosting, pol2cart, zoom_on_source, zoom_out
+from radiosim.utils import relativistic_boosting, pol2cart
 from radiosim.gauss import twodgaussian
 from radiosim.flux_scaling import get_start_amp
 
@@ -39,6 +39,7 @@ def create_jet(grid, conf):
     targets = []
     for _ in grid:
         comps = np.random.randint(num_comps[0], num_comps[1] + 1)
+        comps4sim = np.random.uniform(comps, num_comps[1])
 
         amp = np.zeros(num_comps[1])
         x = np.zeros(num_comps[1])
@@ -51,12 +52,13 @@ def create_jet(grid, conf):
         # velocity in units of c_0, initialise velocity of first component
         beta = np.zeros(num_comps[1])
         beta[1] = np.random.uniform(0, 1)
-        y_rotation = np.random.uniform(0, np.pi)
+        y_rotation = np.random.uniform(0, 2 * np.pi)
         z_rotation = np.random.uniform(0, np.pi / 2)
+        expansion = np.random.uniform(0, 0.5)
 
         for i in range(comps):
             # amplitude decreases for more distant components, empirical
-            amp[i] = np.exp(-np.sqrt(i) * np.random.normal(1.3, 0.2))
+            amp[i] = np.exp(-np.sqrt(i) * np.random.normal(1.1, 0.2))
             if i >= 1 and np.random.rand() < 0.1:  # drop some components
                 amp[i] = 0
 
@@ -65,7 +67,7 @@ def create_jet(grid, conf):
                 beta[i] = beta[1] * np.exp(-np.sqrt(i - 1) * np.random.normal(0.5, 0.1))
 
             # curving the jet, empirical
-            y_rotation += np.random.normal(0, np.pi / 24)
+            y_rotation += np.random.normal(0, np.pi / 36)
 
             # distance between components, r_factor to fill the corners
             jet_angle_cos = np.abs(np.cos(y_rotation))
@@ -78,7 +80,7 @@ def create_jet(grid, conf):
                 r_factor = np.sqrt(2)
 
             # *0.7 so the last component is not on the edge
-            r = i / (comps - 1) * img_size / 2 * r_factor * np.sin(z_rotation) * 0.7
+            r = i / (comps4sim - 1) * img_size / 2 * r_factor * np.sin(z_rotation) * 0.7
 
             # get the cartesian coordinates
             x[i], y[i] = np.array(pol2cart(r, y_rotation)) + center
@@ -87,29 +89,33 @@ def create_jet(grid, conf):
             sx[i], sy[i] = np.sort(
                 (
                     img_size
-                    / comps
+                    / comps4sim
                     * r_factor
-                    * np.sqrt(i + 1)
-                    / np.random.uniform(3, 8, size=2)
+                    * (i + 1) ** expansion
+                    / np.random.uniform(3, 9, size=2)
                 )
             )[::-1]
 
-            # rotation aligned with the jet angle, empirical
+            # rotation of the gauss component, empirical
             if i >= 1:
                 rotation[i] = rotation[i - 1] + np.random.normal(0, np.pi / 18)
 
         # print('Velocity of the jet:', beta)
         boost_app, boost_rec = relativistic_boosting(z_rotation, beta)
 
-        center_shift_x = np.random.uniform(-img_size / 20, img_size / 20)
-        center_shift_y = np.random.uniform(-img_size / 20, img_size / 20)
+        center_shift_x = np.random.uniform(
+            -img_size / 20, img_size / 20
+        )
+        center_shift_y = np.random.uniform(
+            -img_size / 20, img_size / 20
+        )
 
         if conf["scaling"] == "mojave":
             amp *= get_start_amp("mojave")
 
         # mirror the data for the counter jet
         # random drop of counter jet, because the relativistic boosting only does not create clear one-sided jets
-        if np.random.rand() < 0.3:
+        if np.random.rand() < 0.5:
             amp = np.concatenate((amp * boost_app, amp[1:] * boost_rec[1:]))
             x = np.concatenate((x + center_shift_x, img_size - x[1:] + center_shift_x))
             y = np.concatenate((y + center_shift_y, img_size - y[1:] + center_shift_y))
@@ -132,34 +138,6 @@ def create_jet(grid, conf):
         jet_comp = np.array(component_from_list(img_size, amp, x, y, sx, sy, rotation))
         jet_img = np.sum(jet_comp, axis=0)
 
-        # get values at the edge of the image
-        edge_list = [
-            jet_img[0, :-1],
-            jet_img[:-1, -1],
-            jet_img[-1, ::-1],
-            jet_img[-2:0:-1, 0],
-        ]
-        edges = np.concatenate(edge_list)
-        edge_threshold = 0.01
-        if edges.max() < edge_threshold:
-            # zoom on source to equalize size differences from z-rotation
-            jet_img, jet_comp, zoom_factor = zoom_on_source(
-                jet_img, jet_comp, max_amp=edge_threshold
-            )
-            x = img_size / 2 + (x - img_size / 2) * zoom_factor
-            y = img_size / 2 + (y - img_size / 2) * zoom_factor
-            sx *= zoom_factor
-            sy *= zoom_factor
-
-            # random zoom out for more variance
-            zoom_out_factor = np.random.uniform(1 / 2, 1)  # 1/8: pad eg. 128 -> 1024
-            pad_value = (1 / zoom_out_factor - 1) * img_size / 2
-            jet_img, jet_comp = zoom_out(jet_img, jet_comp, pad_value=pad_value)
-            x = img_size / 2 + (x - img_size / 2) * zoom_out_factor
-            y = img_size / 2 + (y - img_size / 2) * zoom_out_factor
-            sx *= zoom_out_factor
-            sy *= zoom_out_factor
-
         # normalisation
         if conf["scaling"] == "normalize":
             jet_max = jet_img.max()

radiosim/scripts/start_simulation.py

@@ -22,10 +22,10 @@ def main(configuration_path, mode):
     conf = read_config(config)
     print(conf, "\n")
 
-    outpath = config["paths"]["outpath"]
+    outpath = conf["outpath"]
     sim_sources = check_outpath(
         outpath,
-        quiet=config["general"]["quiet"],
+        quiet=conf["quiet"],
     )
 
     if mode == "simulate":

radiosim/simulations.py

@@ -1,39 +1,66 @@
 import click
+import multiprocessing
 from tqdm import tqdm
+
+from radiosim.jet import create_jet
+from radiosim.survey import create_survey
 from radiosim.utils import (
-    create_grid,
     add_noise,
-    adjust_outpath,
+    create_grid,
     save_sky_distribution_bundle,
 )
-from radiosim.jet import create_jet
-from radiosim.survey import create_survey
 
 
 def simulate_sky_distributions(conf):
     for opt in ["train", "valid", "test"]:
-        create_sky_distribution(
+        csd = create_sky_distribution(
             conf=conf,
             opt=opt,
         )
+        csd()
+
+
+class create_sky_distribution:
+    def __init__(self, conf, opt):
+        self.conf = conf
+        self.opt = opt
+
+    def __call__(self):
+        n_bundels = self.conf["bundles_" + self.opt]
+        n_cores = int(multiprocessing.cpu_count() * 0.5)  # use 50% of available cores
+        if n_cores == 1 or not self.conf["multiprocessing"]:
+            for i in tqdm(range(n_bundels)):
+                self.sky_distribution(i)
+        else:
+            with multiprocessing.Pool(n_cores) as p:
+                _ = list(
+                    tqdm(
+                        p.imap(self.sky_distribution, range(n_bundels)), total=n_bundels
+                    )
+                )
 
+    def sky_distribution(self, i: int):
+        """Create and save the sky distribution
 
-def create_sky_distribution(conf, opt: str):
-    for _ in tqdm(range(conf["bundles_" + opt])):
-        grid = create_grid(conf["img_size"], conf["bundle_size"])
-        if conf["mode"] == "jet":
-            sky, target = create_jet(grid, conf)
-        elif conf["mode"] == "survey":
-            sky, target = create_survey(grid, conf)
+        Parameters
+        ----------
+        i: int
+            n-th sky distribution to be saved
+        """
+        grid = create_grid(self.conf["img_size"], self.conf["bundle_size"])
+        if self.conf["mode"] == "jet":
+            sky, target = create_jet(grid, self.conf)
+        elif self.conf["mode"] == "survey":
+            sky, target = create_survey(grid, self.conf)
         else:
             click.echo("Given mode not found. Choose 'survey' or 'jet' in config file")
 
         sky_bundle = sky.copy()
         target_bundle = target.copy()
-        if conf["noise"] and conf["noise_level"] > 0:
-            sky_bundle = add_noise(sky_bundle, conf["noise_level"])
+        if self.conf["noise"] and self.conf["noise_level"] > 0:
+            sky_bundle = add_noise(sky_bundle, self.conf["noise_level"])
             for img in sky_bundle:
                 img -= img.min()
                 img /= img.max()
-        path = adjust_outpath(conf["outpath"], "/samp_" + opt)
+        path = self.conf["outpath"] + "/samp_" + self.opt + "_" + str(i) + ".h5"
         save_sky_distribution_bundle(path, sky_bundle, target_bundle)

radiosim/utils.py

@@ -57,94 +57,14 @@ def relativistic_boosting(theta, beta):
     boost_rec: float
         boosting factor for the receding jet
     """
-    gamma = 1 / np.sqrt(1 - beta ** 2)  # Lorentz factor
+    gamma = 1 / np.sqrt(1 - beta**2)  # Lorentz factor
     mu = np.cos(theta)
 
     boost_app = 1 / (gamma * (1 - beta * mu))
     boost_rec = 1 / (gamma * (1 + beta * mu))
     return boost_app, boost_rec
 
 
-def zoom_on_source(img, comp=None, max_amp=0.01):
-    """
-    Zoom on source to cut out irrelevant area. Shape will stay equal.
-
-    Parameters
-    ----------
-    img: 2D array
-        Image of the sky used to zoom on
-    comp: 3D array (n, (img))
-        Images of the components, same zooming as on img
-    max_amp: float
-        Maximal amplitude which will be at the edge of the image
-
-    Returns
-    -------
-    zoomed_img: ndarray
-        Image after zooming
-    zoom_factor: float
-        Zooming factor
-    """
-    # find farest outside column or row with amplitude > max_amp
-    mask = img > max_amp
-    mask_flip = np.flip(mask)
-
-    idx_left = np.argmax(np.sum(mask, axis=0) > 0)
-    idx_right = np.argmax(np.sum(mask_flip, axis=0) > 0)
-    idx_bottom = np.argmax(np.sum(mask, axis=1) > 0)
-    idx_top = np.argmax(np.sum(mask_flip, axis=1) > 0)
-    # print(idx_left, idx_right, idx_bottom, idx_top)
-    idx = np.min([idx_left, idx_right, idx_bottom, idx_top])
-    size = img.shape[0]
-    zoom_factor = size / (size - 2 * idx)
-
-    # crop the source
-    cropped_img = img[idx:size-idx, idx:size-idx]
-    zoomed_img = cv2.resize(cropped_img, dsize=(size, size), interpolation=cv2.INTER_LINEAR)
-
-    if comp is not None:
-        cropped_comp = comp[:, idx:size-idx, idx:size-idx]
-        zoomed_comp = np.empty_like(comp)
-        for i, component in enumerate(cropped_comp):
-            zoomed_comp[i] = cv2.resize(component, dsize=(size, size), interpolation=cv2.INTER_LINEAR)
-        return zoomed_img, zoomed_comp, zoom_factor
-
-    return zoomed_img, zoom_factor
-
-
-def zoom_out(img, comp=None, pad_value=0):
-    """
-    Zoom out of an image. Padding edges with zeros.
-
-    Parameters
-    ----------
-    img: 2D array
-        Image of the sky
-    comp: 3D array (n, (img))
-        Images of the components
-    pad_value: int
-        Number of pixels to pad around the source
-
-    Returns
-    -------
-    zoomed_img: ndarray
-        Image after zooming
-    zoomed_comp: ndarray
-        Componets after zooming
-    """
-    if not isinstance(pad_value, int):
-        pad_value = np.int64(pad_value)
-    size = img.shape[0]
-    img = cv2.resize(np.pad(img, pad_value), dsize=(size, size), interpolation=cv2.INTER_LINEAR)
-
-    if comp is not None:
-        for component in comp:
-            component = cv2.resize(np.pad(component, pad_value), dsize=(size, size), interpolation=cv2.INTER_LINEAR)
-        return img, comp
-
-    return img
-
-
 def check_outpath(outpath, quiet=False):
     """
     Check if outpath exists. Check for existing source_bundle files.
@@ -206,7 +126,9 @@ def read_config(config):
         unpacked configurations
     """
     sim_conf = {}
+    sim_conf["quiet"] = config["general"]["quiet"]
     sim_conf["mode"] = config["general"]["mode"]
+    sim_conf["multiprocessing"] = config["general"]["multiprocessing"]
     sim_conf["outpath"] = config["paths"]["outpath"]
     sim_conf["training_type"] = config["jet"]["training_type"]
     sim_conf["num_jet_components"] = config["jet"]["num_jet_components"]
@@ -277,32 +199,6 @@ def call_noise(kernels, strengths):
     return image_noised
 
 
-def adjust_outpath(path, option, form="h5"):
-    """
-    Add number to out path when filename already exists.
-
-    Parameters
-    ----------
-    path: str
-        path to save directory
-    option: str
-        additional keyword to add to path
-    form: str
-        file extension
-
-    Returns
-    -------
-    out: str
-        adjusted path
-    """
-    counter = 0
-    filename = str(path) + (option + "_{}." + form)
-    while os.path.isfile(filename.format(counter)):
-        counter += 1
-    out = filename.format(counter)
-    return out
-
-
 def save_sky_distribution_bundle(path, x, y, name_x="x", name_y="y"):
     """
     Write images created in analysis to h5 file.
@@ -344,7 +240,7 @@ def cart2pol(x: float, y: float):
     phi: float
         angle in radian
     """
-    r = np.sqrt(x ** 2 + y ** 2)
+    r = np.sqrt(x**2 + y**2)
     phi = np.arctan2(y, x)
     return (r, phi)
 

rc/default_simulation.toml

@@ -4,7 +4,8 @@ title = "Simulation configuration"
 
 [general]
 quiet = true
-mode = "jet"   # jet or survey
+mode = "jet"            # jet or survey
+multiprocessing = true
 
 [paths]
 outpath = "./build/example_data/"

setup.py

@@ -2,7 +2,7 @@
 
 setup(
     name="radiosim",
-    version="0.0.2",
+    version="0.0.3",
     description="Simulation of radio skies to create astrophysical data sets",
     url="https://github.com/Kevin2/radionets",
     author="Kevin Schmidt, Felix Geyer, Paul-Simon Blomenkamp, Stefan Fröse",

tests/simulate.toml

@@ -4,7 +4,8 @@ title = "Simulation configuration"
 
 [general]
 quiet = true
-mode = "survey"   # jet or survey
+mode = "survey"         # jet or survey
+multiprocessing = true
 
 [paths]
 outpath = "./tests/build/data/"