draft logic for backend isotope_fractionation routine; placeholder file for timescale attributes

Author: slayoo

PySDM/attributes/isotopes/timescales.py

@@ -0,0 +1,4 @@
+# TODO: register one class per isotope,
+#  and in each class call the relevant physics formula for a given isotope
+
+# TODO: code analogous to that in the delta.py file here

PySDM/backends/impl_numba/methods/isotope_methods.py

@@ -24,5 +24,54 @@ def body(output, ratio, reference_ratio):
     def isotopic_delta(self, output, ratio, reference_ratio):
         self._isotopic_delta_body(output.data, ratio.data, reference_ratio)
 
-    def isotopic_fractionation(self):
-        pass
+    @cached_property
+    def _isotopic_fractionation_body(self):
+        @numba.njit(**self.default_jit_flags)
+        def body(
+            *,
+            multiplicity,
+            signed_timescale,
+            moles,
+            dt,
+            ambient_isotope_mixing_ratio,
+            cell_id,
+            cell_volume,
+            dry_air_density,
+            molar_mass,
+        ):
+            for sd_id in numba.prange(multiplicity.shape[0]):
+                dn = dt / signed_timescale[sd_id] * moles[sd_id]
+                moles[sd_id] += dn
+                mass_of_dry_air = (
+                    dry_air_density[cell_id[sd_id]] * cell_volume[cell_id[sd_id]]
+                )
+                ambient_isotope_mixing_ratio[cell_id[sd_id]] -= (
+                    dn * multiplicity[sd_id] * molar_mass / mass_of_dry_air
+                )
+
+        return body
+
+    def isotopic_fractionation(
+        self,
+        *,
+        isotope,
+        multiplicity,
+        signed_timescale,
+        moles,
+        dt,
+        ambient_isotope_mixing_ratio,
+        cell_id,
+        cell_volume,
+        dry_air_density,
+    ):
+        self._isotopic_fractionation_body(
+            multiplicity=multiplicity.data,
+            signed_timescale=signed_timescale.data,
+            moles=moles.data,
+            dt=dt,
+            ambient_isotope_mixing_ratio=ambient_isotope_mixing_ratio.data,
+            cell_id=cell_id.data,
+            cell_volume=cell_volume.data,
+            dry_air_density=dry_air_density.data,
+            molar_mass=getattr(self.formulae.constants, f"M_{isotope}"),
+        )

PySDM/particulator.py

@@ -439,9 +439,23 @@ def calculate_displacement(
                 n_substeps=n_substeps,
             )
 
-    def isotopic_fractionation(self, heavy_isotopes: tuple):
-        self.backend.isotopic_fractionation()
+    def isotopic_fractionation(
+        self, heavy_isotopes: tuple, ambient_isotope_mixing_ratios
+    ):
         for isotope in heavy_isotopes:
+            self.backend.isotopic_fractionation(
+                isotope=isotope,
+                multiplicity=self.attributes["multiplicity"],
+                signed_timescale=self.attributes[f"isotope_timescale_{isotope}"],
+                moles=self.attributes[f"moles_{isotope}"],
+                dt=self.dt,
+                dry_air_density=self.environment[f"dry_air_density"],
+                cell_volume=self.environment.mesh.dv,
+                cell_id=self.attributes["cell id"],
+                ambient_isotope_mixing_ratio=self.environment[
+                    f"mixing_ratio_{isotope}"
+                ],
+            )
             self.attributes.mark_updated(f"moles_{isotope}")
 
     def seeding(