Merge pull request #20 from radionets-project/update_formatting

Fix formatting, update naming of output files

radiosim/flux_scaling.py

@@ -8,7 +8,9 @@ def get_start_amp(scale_type):
 
 
 def draw_from_mojave_dist():
-    """Values from a fit to the peak fluxes distribution in the MOJAVE data archive."""
+    """Values from a fit to the peak fluxes distribution
+    in the MOJAVE data archive.
+    """
     a = 0.8639672251677816
     b = 47.64189171625089
     loc = 0.09163404776954732

radiosim/gauss.py

@@ -38,6 +38,7 @@ def gauss(params: list, size: int):
         size=(size, size),
     ).data
     gaussian_scaled = gaussian_cutout / gaussian_cutout.max() * params[0]
+
     return gaussian_scaled
 
 
@@ -65,6 +66,8 @@ def twodgaussian(params: list, size: int):
     rotgauss: 2darray
         gaussian distribution in two dimensions
 
+    Notes
+    -----
     Short version of 'twodgaussian':
     https://github.com/keflavich/gaussfitter/blob/0891cd3605ab5ba000c2d5e1300dd21c15eee1fd/gaussfitter/gaussfitter.py#L75
     """

radiosim/jet.py

@@ -6,27 +6,33 @@
 
 def create_jet(grid, conf):
     """
-    Creates the clean jets with all its components written in a list. Dependend on the
-    'train_type' the components will be seperated or summed up.
+    Creates the clean jets with all its components written
+    in a list. Dependend on the 'train_type' the components
+    will be seperated or summed up.
 
     Parameters
     ----------
     grid: ndarray
-        input grid of shape [n, 1, img_size, img_size] or [1, img_size, img_size]
+        input grid of shape [n, 1, img_size, img_size] or
+        [1, img_size, img_size]
     conf:
         loaded config file
 
     Returns
     -------
     jets: ndarray
-        image of the full jet, sum over all components, shape: [n, 1, img_size, img_size]
+        image of the full jet, sum over all components,
+        shape: [n, 1, img_size, img_size]
     jet_comps: ndarray
-        images of each component and background, shape: [n, c*2, img_size, img_size]
-        with c being the max number of components. A jet without counter jet has c
-        components. A jet with counter jet has c*2-1 components, since the center
-        appears only once. Adding one channel for the backgound gives c*2 channels.
+        images of each component and background,
+        shape: [n, c*2, img_size, img_size] with c being the
+        max number of components. A jet without counter jet
+        has c components. A jet with counter jet has c*2-1
+        components, since the center appears only once. Adding
+        one channel for the backgound gives c*2 channels.
     source_lists: ndarray
-        array which stores all (seven) properties of each component, shape: [n, c*2-1, 7]
+        array which stores all (seven) properties of each component,
+        shape: [n, c*2-1, 7]
     """
     num_comps = conf["num_jet_components"]
 
@@ -83,7 +89,8 @@ def create_jet(grid, conf):
             # get the cartesian coordinates
             x[i], y[i] = np.array(pol2cart(r, y_rotation)) + center
 
-            # width of gaussian, empirical, sx > sy because rotation up to pi 'changes' this property - fixed to have consistency
+            # width of gaussian, empirical, sx > sy because rotation up
+            # to pi 'changes' this property - fixed to have consistency
             sx[i], sy[i] = np.sort(
                 (
                     img_size
@@ -108,7 +115,8 @@ def create_jet(grid, conf):
             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
+        # random drop of counter jet, because the relativistic
+        # boosting only does not create clear one-sided jets
         if np.random.rand() < 0.3:
             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))
@@ -212,6 +220,7 @@ def apply_train_type(train_type, jet_img, jet_comp, source_list):
         y = np.concatenate((jet_comp, list_to_add))
     if train_type == "clean":
         y = jet_img[None]
+
     return y
 
 
@@ -241,4 +250,5 @@ def component_from_list(size, amp, x, y, sx, sy, rotation):
                 size,
             )
             jet_comp += [g]
+
     return jet_comp

radiosim/mojave.py

@@ -26,12 +26,12 @@ def create_mojave(conf, rng):
     if threads == "none":
         threads = 1
     elif not isinstance(threads, int):
-        raise ValueError("threads has to be int >0 or \"none\"")
+        raise ValueError('threads has to be int >0 or "none"')
     elif threads <= 0:
-        raise ValueError("threads has to be int >0 or \"none\"")
-    
+        raise ValueError('threads has to be int >0 or "none"')
+
     size = conf["img_size"]
-        
+
     # calculate amount of each class to generate per bundle
     bundle_size = conf["bundle_size"]
     ratio = np.array(conf["class_ratio"])
@@ -42,11 +42,17 @@ def create_mojave(conf, rng):
         amount[amount.argmax()] += 1
 
     jets = []
-    compact = Parallel(n_jobs=threads)(delayed(gen_compact)(rng=child, size=size) for child in rng.spawn(amount[0]))
-    one_jet = Parallel(n_jobs=threads)(delayed(gen_one_jet)(rng=child, size=size) for child in rng.spawn(amount[1]))
-    two_jet = Parallel(n_jobs=threads)(delayed(gen_two_jet)(rng=child, size=size) for child in rng.spawn(amount[2]))
+    compact = Parallel(n_jobs=threads)(
+        delayed(gen_compact)(rng=child, size=size) for child in rng.spawn(amount[0])
+    )
+    one_jet = Parallel(n_jobs=threads)(
+        delayed(gen_one_jet)(rng=child, size=size) for child in rng.spawn(amount[1])
+    )
+    two_jet = Parallel(n_jobs=threads)(
+        delayed(gen_two_jet)(rng=child, size=size) for child in rng.spawn(amount[2])
+    )
     jets = np.array([*compact, *one_jet, *two_jet])
-    
+
     glx_class = []
     for glx_type, amt in zip([0, 1, 2], amount):
         glx_class.extend([glx_type] * amt)
@@ -63,9 +69,12 @@ def create_mojave(conf, rng):
 
     return data, data_name
 
-def gen_jet(size: int, amplitude: float, width: float, length: float, a: float) -> ArrayLike:
+
+def gen_jet(
+    size: int, amplitude: float, width: float, length: float, a: float
+) -> ArrayLike:
     """
-    Generate jet from skewed 2d normal distributiuon. 
+    Generate jet from skewed 2d normal distributiuon.
 
     Parameters
     ----------
@@ -89,14 +98,30 @@ def gen_jet(size: int, amplitude: float, width: float, length: float, a: float)
     jet : ArrayLike
         jet for source
     """
-
-    jet = skewed_gauss(size=size, x=size/2, y=size/2, amp=amplitude, width=width, length=length, a=a)
+
+    jet = skewed_gauss(
+        size=size,
+        x=size / 2,
+        y=size / 2,
+        amp=amplitude,
+        width=width,
+        length=length,
+        a=a,
+    )
 
     jet[np.isclose(jet, 0)] = 0
 
     return jet
 
-def gen_shocks(jet: ArrayLike, shock_comp: int, length: float, width: float, sx: float, rng: "np.Generator") -> tuple[ArrayLike, tuple, int]:
+
+def gen_shocks(
+    jet: ArrayLike,
+    shock_comp: int,
+    length: float,
+    width: float,
+    sx: float,
+    rng: "np.Generator",
+) -> tuple[ArrayLike, tuple, int]:
     """
     Generate equally spaced shocks in jet.
 
@@ -114,7 +139,7 @@ def gen_shocks(jet: ArrayLike, shock_comp: int, length: float, width: float, sx:
         sx component of main gaussian
     rng : np.Generator
         numpy random generator
-    
+
     Returns
     -------
     jet : ArrayLike
@@ -175,13 +200,17 @@ def gen_shocks(jet: ArrayLike, shock_comp: int, length: float, width: float, sx:
         ps_orientation.append("r")
 
         s_dist += s_dist0
+
     return (
         jet,
         (ps_orientation, ps_amp, ps_x, ps_y, ps_sx, ps_sy, ps_rot, ps_dist, ps_a),
         skipped,
     )
 
-def add_swirl(jet: ArrayLike, rng: "np.Generator", first_jet_params: tuple | None = None) -> tuple[ArrayLike, tuple]:
+
+def add_swirl(
+    jet: ArrayLike, rng: "np.Generator", first_jet_params: tuple | None = None
+) -> tuple[ArrayLike, tuple]:
     """
     Add swirl distortion to the jet.
 
@@ -195,7 +224,7 @@ def add_swirl(jet: ArrayLike, rng: "np.Generator", first_jet_params: tuple | Non
         Use when generate two jet sources.
         Tuple of parameters to generate similar swirl for second jet.
         Contains the returned parameters from the first jet.
-    
+
     Returns
     -------
     swirled_jet : ArrayLike
@@ -218,7 +247,8 @@ def add_swirl(jet: ArrayLike, rng: "np.Generator", first_jet_params: tuple | Non
 
     return swirled_jet, (strength, radius, center)
 
-def gen_two_jet(rng : int, size : int = 1024, printout : bool = False) -> ArrayLike:
+
+def gen_two_jet(rng: int, size: int = 1024, printout: bool = False) -> ArrayLike:
     """
     Generate a two jet galaxy.
 
@@ -260,10 +290,10 @@ def gen_two_jet(rng : int, size : int = 1024, printout : bool = False) -> ArrayL
 
     sx = rng.uniform(dimensions.min(), dimensions.max()) * 1.5  # too large
     sy = rng.uniform(dimensions.min(), dimensions.max()) * 1.5
-    while np.any(
-            [sx / sy < 1 / 2, sx / sy > 2]
-        ):  # redraw if gaussian is too elliptical
-            sy = rng.uniform(dimensions.min(), dimensions.max()) * 1.5
+
+    # redraw if gaussian is too elliptical
+    while np.any([sx / sy < 1 / 2, sx / sy > 2]):
+        sy = rng.uniform(dimensions.min(), dimensions.max()) * 1.5
     rot = rng.uniform(0, 2 * np.pi)
 
     a = 4
@@ -275,10 +305,14 @@ def gen_two_jet(rng : int, size : int = 1024, printout : bool = False) -> ArrayL
     shock_comp_l = int(np.round(rng.power(0.5), decimals=1) * 10)
 
     if shock_comp_r > 0:
-        r_jet, r_printout, r_skipped = gen_shocks(r_jet, shock_comp_r, r_length, r_width, sx, rng)
+        r_jet, r_printout, r_skipped = gen_shocks(
+            r_jet, shock_comp_r, r_length, r_width, sx, rng
+        )
 
     if shock_comp_l > 0:
-        l_jet, l_printout, l_skipped = gen_shocks(l_jet, shock_comp_l, l_length, l_width, sx, rng)
+        l_jet, l_printout, l_skipped = gen_shocks(
+            l_jet, shock_comp_l, l_length, l_width, sx, rng
+        )
 
     r_jet, r_swirl_comp = add_swirl(r_jet, rng=rng)
     l_jet, l_swirl_comp = add_swirl(l_jet, rng=rng, first_jet_params=r_swirl_comp)
@@ -346,7 +380,8 @@ def gen_two_jet(rng : int, size : int = 1024, printout : bool = False) -> ArrayL
 
     return glx
 
-def gen_one_jet(rng : int, size : int = 1024, printout : bool = False) -> ArrayLike:
+
+def gen_one_jet(rng: int, size: int = 1024, printout: bool = False) -> ArrayLike:
     """
     Generate a one jet galaxy.
 
@@ -379,12 +414,12 @@ def gen_one_jet(rng : int, size : int = 1024, printout : bool = False) -> ArrayL
 
     dimensions = np.sort([width / 2, length / 5])
 
-    sx = rng.uniform(*dimensions) * 1.5 # too large
+    sx = rng.uniform(*dimensions) * 1.5  # too large
     sy = rng.uniform(*dimensions) * 1.5
-    while np.any(
-            [sx / sy < 1 / 2, sx / sy > 2]
-        ):  # redraw if gaussian is too elliptical
-            sy = rng.uniform(dimensions.min(), dimensions.max()) * 1.5
+
+    # redraw if gaussian is too elliptical
+    while np.any([sx / sy < 1 / 2, sx / sy > 2]):
+        sy = rng.uniform(dimensions.min(), dimensions.max()) * 1.5
     rot = rng.uniform(0, 2 * np.pi)
 
     a = 4
@@ -439,7 +474,8 @@ def gen_one_jet(rng : int, size : int = 1024, printout : bool = False) -> ArrayL
 
     return glx
 
-def gen_compact(rng : int, size : int = 1024, printout : bool = False) -> ArrayLike:
+
+def gen_compact(rng: int, size: int = 1024, printout: bool = False) -> ArrayLike:
     """Generate a compact galaxy.
 
     Parameters
@@ -468,7 +504,7 @@ def gen_compact(rng : int, size : int = 1024, printout : bool = False) -> ArrayL
 
     amp = expon.rvs(*amp_params, random_state=rng)
 
-    jet_amp = amp * rng.normal(2/3, 0.05)
+    jet_amp = amp * rng.normal(2 / 3, 0.05)
 
     y = np.arange(size)
 
@@ -503,7 +539,7 @@ def gen_compact(rng : int, size : int = 1024, printout : bool = False) -> ArrayL
             x = 0
             y = 0
         if i == 0:
-            sx = rng.uniform(*dimensions) * 1.7 # too large
+            sx = rng.uniform(*dimensions) * 1.7  # too large
             j = 0
             while sx > 15:
                 if j > 10:
@@ -523,7 +559,9 @@ def gen_compact(rng : int, size : int = 1024, printout : bool = False) -> ArrayL
                 j += 1
 
             j = 0
-            while np.any([sx/sy < 0.33, sx/sy>3]) or np.all([sx/sy > 0.75, sx/sy < 1.33]):
+            while np.any([sx / sy < 0.33, sx / sy > 3]) or np.all(
+                [sx / sy > 0.75, sx / sy < 1.33]
+            ):
                 if j > 10:
                     sy = sx * rng.choice([rng.uniform(0.4, 0.8), rng.uniform(1.2, 2.5)])
                     if sy > 15:
@@ -537,25 +575,25 @@ def gen_compact(rng : int, size : int = 1024, printout : bool = False) -> ArrayL
                         break
                     sy = rng.uniform(*dimensions) * 1.7
                     k += 1
-                j+=1
+                j += 1
         else:
             sx = comp_sx[-1] * deviation[0]
             sy = comp_sy[-1] * deviation[1]
-        
+
         if i == 0:
             if sx > sy:
                 vertical = True
             else:
                 vertical = False
-        
+
         try:
             rot = rng.uniform(comp_rot[-1] - 0.1, comp_rot[-1] + 0.1)
         except IndexError:
             if vertical:
-                rot = rng.uniform(-np.pi/4, np.pi/4)
+                rot = rng.uniform(-np.pi / 4, np.pi / 4)
             else:
-                rot = rng.uniform(-np.pi/4+np.pi/2, np.pi/4+np.pi/2)
-        
+                rot = rng.uniform(-np.pi / 4 + np.pi / 2, np.pi / 4 + np.pi / 2)
+
         c_amp = amp * rng.uniform(0.5, 1.2)
 
         gauss = twodgaussian([c_amp, size / 2 - x, size / 2 - y, sx, sy, rot], size)
@@ -578,7 +616,7 @@ def gen_compact(rng : int, size : int = 1024, printout : bool = False) -> ArrayL
     glx = np.roll(glx, shift=shift, axis=[0, 1])
 
     glx[np.isclose(glx, 0)] = 0
-    
+
     glx /= glx.max()
     glx *= amp
 

radiosim/simulations.py

@@ -29,10 +29,10 @@ def create_sky_distribution(conf, opt: str) -> None:
         elif isinstance(seed, int):
             rng = default_rng(seed)
         else:
-            raise TypeError("seed has to be int or \"none\"")
-        
+            raise TypeError('seed has to be int or "none"')
+
     for _ in tqdm(range(conf["bundles_" + opt])):
-        path = adjust_outpath(conf["outpath"], "/samp_" + opt)
+        path = adjust_outpath(conf["outpath"], f"/data_{conf['mode']}_" + opt)
         grid = create_grid(conf["img_size"], conf["bundle_size"])
         if conf["mode"] == "jet":
             sky, target = create_jet(grid, conf)
@@ -52,4 +52,5 @@ def create_sky_distribution(conf, opt: str) -> None:
             for img in sky_bundle:
                 img -= img.min()
                 img /= img.max()
+
         save_sky_distribution_bundle(path, sky_bundle, target_bundle)