subgroups implemented

Author: grzegorzbor

agnpy/time_evolution/_time_evolution_utils.py

@@ -1,8 +1,8 @@
 import numpy as np
-from astropy.units import Quantity
+from astropy.units import Quantity, Unit
 from pygments.styles import vs
-
-from agnpy import InterpolatedDistribution, ParticleDistribution
+from agnpy import InterpolatedDistribution
+from agnpy.time_evolution.types import BinsWithDensities, FnParams
 from agnpy.utils.conversion import mec2
 from numpy._typing import NDArray
 from scipy.interpolate import PchipInterpolator
@@ -16,36 +16,46 @@ def update_bin_ub(gamma_bins_from, gamma_bins_to, low_change_rates_mask = None):
     gamma_bins_to[~low_change_rates_mask] = gamma_bins_from[~low_change_rates_mask] * (1 + bin_size_factor)
 
 
-def recalc_new_rates(gm_bins, functions, unit, previous_rates=None, recalc_mask=None):
+def recalc_new_rates(gamma, functions, unit, dens, dens_groups, previous_rates=None, recalc_mask=None):
     """
     Calculates (or recalculates) the energy change rates or injection rates.
     If the mask is provided, only unmasked elements will be recalculated, and previous_rates will be used for masked elements.
-    gamma_bins - array of shape (N,) consisting of gamma values (lower bounds of gamma bins), or (2N,) consisting of inlined pairs of gamma values (lower and upper bounds of gamma bins)
+    gamma - array of shape (N,) consisting of gamma values (lower bounds of gamma bins), or (2N,) consisting of inlined pairs of gamma values (lower and upper bounds of gamma bins)
     previous_rates - optional dict of arrays of shape (N,) or (2N,), consisting of most recently calculated rates corresponding to gamma_bins values
-    low_change_rates_mask - optional array of shape (N,) a mask of bins that do not need recalculation
+    recalc_mask - optional array of shape (N,) a mask of bins that need recalculation
     """
     if previous_rates is None:
         previous_rates = {}
-    mask = np.ones_like(gm_bins, dtype=bool) if recalc_mask is None else\
-        np.repeat(recalc_mask, 2) if len(recalc_mask) * 2 == len(gm_bins) else recalc_mask
-    new_injection_rates = {}
+    mask_dens = np.ones_like(dens, dtype=bool) if recalc_mask is None else recalc_mask
+    dens_masked = dens[mask_dens]
+    dens_groups_masked = remap_subgroups_density(mask_to_mapping(mask_dens), dens_groups)
+    mask = np.ones_like(gamma, dtype=bool) if recalc_mask is None else \
+        np.repeat(recalc_mask, 2) if len(recalc_mask) * 2 == len(gamma) else recalc_mask
+    new_rates = {}
     for label, fn in functions.items():
-        new_injection_rates[label] = np.zeros_like(gm_bins) * unit \
+        new_rates[label] = np.zeros_like(gamma) * unit \
             if previous_rates.get(label) is None \
             else previous_rates[label].copy()
-        new_injection_rates[label][mask] = fn(gm_bins[mask]).to(unit)
-    return new_injection_rates
+        new_rates[label][mask] = fn(FnParams(gamma[mask], dens_masked, dens_groups_masked)).to(unit)
+    return new_rates
 
 
 def sum_change_rates_at_index(change_rates, index):
     return np.sum([arr[index].value for arr in change_rates.values()])
 
 
-def sum_change_rates(change_rates, shape, unit):
-    result = np.zeros(shape) * unit
-    for ch_rates in change_rates.values():
-        result += ch_rates.to(unit)
-    return result
+def sum_change_rates(change_rates, shape, unit, subgroups=None):
+    if subgroups is None:
+        result = np.zeros(shape) * unit
+        for ch_rates in change_rates.values():
+            result += ch_rates.to(unit)
+        return Quantity([result])
+    else:
+        return Quantity([
+            np.sum(Quantity([change_rates[x] for x in row if x in change_rates]), axis=0)
+            if any(x in change_rates for x in row) else np.zeros(shape) * unit
+            for row in subgroups
+        ])
 
 
 def remap(mapping, values, default_value=None):
@@ -70,6 +80,17 @@ def remap_interlaced(mapping, values, default_value = None):
         return interlace(remap(mapping, v1, default_value), remap(mapping, v2, default_value))
 
 
+def remap_subgroups_density(mapping, subgroups_density):
+    new_subgroups_density = subgroups_density[:,mapping]
+    new_subgroups_density[0, mapping == -1] = 1.0
+    new_subgroups_density[1:, mapping == -1] = 0.0
+    return new_subgroups_density
+
+
+def mask_to_mapping(mask):
+    return np.arange(len(mask))[mask]
+
+
 def combine_mapping(first_mapping, second_mapping):
     """ Combines two mappings, but with special treatment of "-1" value in the mapping:
         -1 does not mean the last element; it means a "new element", at must be retained across the mappings
@@ -83,11 +104,20 @@ def to_erg(gamma_bins):
     return (gamma_bins * mec2).to("erg")
 
 
-def recalc_gamma_bins_and_density(energy_bins, abs_energy_changes, density):
+def to_densities_grouped(density, subgroups_density):
+    density_grouped = density * subgroups_density if subgroups_density is not None else density.reshape(
+        (1, len(density)))
+    return density_grouped
+    
+def second_axis_size(a: np.ndarray) -> int:
+    return a.shape[1] if a.ndim > 1 else 1
+
+
+def recalc_gamma_bins_and_density(energy_bins, abs_energy_changes, density) -> BinsWithDensities:
     new_energy_bins, density_increase = recalc_energy_bins_and_density_increase(abs_energy_changes, energy_bins)
     new_density = density * density_increase
     new_gamma_bins = (new_energy_bins / mec2).to_value('')
-    return new_gamma_bins, new_density
+    return BinsWithDensities(new_gamma_bins, new_density)
 
 
 def recalc_energy_bins_and_density_increase(abs_energy_changes, energy_bins):
@@ -107,26 +137,35 @@ def recalc_energy_bins_and_density_increase(abs_energy_changes, energy_bins):
     return new_energy_bins, density_increase
 
 
-def calc_new_low_change_rates_mask(energy_bins, density, abs_energy_change, abs_particle_scaling, abs_particle_injection,
+def calc_new_low_change_rates_mask(energy_bins, density, subgroups_density, abs_energy_change_grouped, abs_particle_scaling_grouped, abs_particle_injection_grouped,
                                    max_energy_change_per_interval, max_density_change_per_interval, max_abs_injection_per_interval):
-    if np.any((density == 0) & (abs_particle_injection < 0)):
-        raise ValueError("Particles cannot escape when density is zero")
     # The obvious hard limit for escape is -1.0, which means "all particles escaped"; more negative value would clearly
     # mean we used too long time. The -0.5 is somewhat arbitrary, perhaps it should be configurable like other limits?
     # But it is tricky how to configure it, because e.g. -1.1 value makes no physical sense, on the other hand +1.1 is completely fine.
     max_escape = -0.5
-    relative_changes_bin_lb, relative_en_changes_bin_ub = deinterlace(abs_energy_change / energy_bins)
-    rel_density_increase_from_energy_change = recalc_energy_bins_and_density_increase(abs_energy_change, energy_bins)[1]
-    rel_density_change_from_injection = np.divide(density * abs_particle_scaling + abs_particle_injection, density,
-                                                  out=Quantity(np.zeros_like(density.value, dtype=float)),
-                                                  where=density!=0).value
-    rel_density_increase_combined = rel_density_change_from_injection * rel_density_increase_from_energy_change
-    new_low_change_rates_mask = ((abs(relative_en_changes_bin_ub) <= max_energy_change_per_interval) &
-                                 (~np.isnan(rel_density_increase_from_energy_change)) &
-                                 (abs(abs_particle_scaling) <= max_density_change_per_interval) &
-                                 (abs(abs_particle_injection) <= max_abs_injection_per_interval) &
-                                 (abs_particle_scaling >= max_escape) &
-                                 (abs(rel_density_increase_combined) <= max_density_change_per_interval))
+    new_low_change_rates_mask = np.ones(len(density), dtype=bool)
+    density_grouped = to_densities_grouped(density, subgroups_density)
+    for i in (range(subgroups_density.shape[0])):
+        if np.any((density_grouped[i] == 0) & (abs_particle_injection_grouped[i] < 0)):
+            raise ValueError("Particles cannot escape when density is zero")
+
+        abs_energy_change = abs_energy_change_grouped[i]
+        density = density_grouped[i]
+        abs_particle_scaling = abs_particle_scaling_grouped[i]
+        abs_particle_injection = abs_particle_injection_grouped[i]
+        relative_changes_bin_lb, relative_en_changes_bin_ub = deinterlace(abs_energy_change / energy_bins)
+        rel_density_increase_from_energy_change = recalc_energy_bins_and_density_increase(abs_energy_change, energy_bins)[1]
+        rel_density_change_from_injection = np.divide(density * abs_particle_scaling + abs_particle_injection, density,
+                                                      out=Quantity(np.zeros_like(density.value, dtype=float)),
+                                                      where=density!=0).value
+        rel_density_increase_combined = rel_density_change_from_injection * rel_density_increase_from_energy_change
+        group_mask = ((abs(relative_en_changes_bin_ub) <= max_energy_change_per_interval) &
+                                     (~np.isnan(rel_density_increase_from_energy_change)) &
+                                     (abs(abs_particle_scaling) <= max_density_change_per_interval) &
+                                     (abs(abs_particle_injection) <= max_abs_injection_per_interval) &
+                                     (abs_particle_scaling >= max_escape) &
+                                     (abs(rel_density_increase_combined) <= max_density_change_per_interval))
+        new_low_change_rates_mask &= group_mask
     return new_low_change_rates_mask
 
 
@@ -313,6 +352,35 @@ def update_distribution(gamma_array, n_array, blob):
     blob.n_e = InterpolatedDistribution(gamma_array_sorted, n_array_sorted, interpolator=PchipInterpolator)
 
 
+def apply_time_eval(en_bins, density, subgroups_density, abs_en_changes_grouped, total_injection_grouped):
+    new_bins_and_densities = [recalc_gamma_bins_and_density(en_bins, abs_en_changes_grouped[i], density*subgroups_density[i])
+                              for i in range(len(subgroups_density))]
+    for i in range(len(new_bins_and_densities)):
+        new_bins_and_densities[i] = BinsWithDensities(new_bins_and_densities[i].gamma_bins,
+                                                      new_bins_and_densities[i].densities + total_injection_grouped[i])
+
+    new_gm_bins = new_bins_and_densities[0].gamma_bins
+    new_gm_bins_lb, new_gm_bins_ub = deinterlace(new_gm_bins)
+
+    new_densities = np.zeros_like(subgroups_density) * Unit("cm-3")
+    new_densities[0] = new_bins_and_densities[0].densities
+
+    for i in range(1, len(new_bins_and_densities)):
+        new_gm_bins_subgroup, _ = deinterlace(new_bins_and_densities[i].gamma_bins)
+        new_dens_subgroup = new_bins_and_densities[i].densities
+        if np.all(new_dens_subgroup == 0) or np.all(new_bins_and_densities[i].gamma_bins == new_gm_bins):
+            new_densities[i] = new_dens_subgroup
+        else:
+            distribution = InterpolatedDistribution(new_gm_bins_subgroup, new_dens_subgroup)
+            densities_interpolated = distribution(new_gm_bins_lb)
+            new_densities[i] = densities_interpolated
+
+    new_dens = Quantity(new_densities.sum(axis=0))
+    subgroups_density_new = subgroups_density.copy()
+    mask = new_dens != 0
+    subgroups_density_new[:,mask] = (new_densities[:,mask] / new_dens[mask]).value
+    return new_gm_bins, new_gm_bins_lb, new_gm_bins_ub, new_dens, subgroups_density_new
+
 def interlace(gamma_bins_start: NDArray, gamma_bins_end: NDArray) -> NDArray:
     """ Combines two gamma arrays into one array, to make the calculations on them in one steps instead of two;
         the first array goes into odd indices, the second into even indices """