Diffusion model #64

henrikjacobsenfys · 2026-01-06T14:02:18Z

Basic diffusion model.

I've kept it a bit simpler than some of the other PRs in terms of allowed inputs, to get things flowing faster.

Once the overall structure is approved, I have ~6 more diffusion models to implement following the same style.

codecov · 2026-01-06T14:05:00Z

Codecov Report

❌ Patch coverage is 99.18699% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 97.83%. Comparing base (5edc575) to head (5df823d).
⚠️ Report is 1 commits behind head on develop.

Files with missing lines	Patch %	Lines
...mple_model/diffusion_model/diffusion_model_base.py	97.22%	1 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff             @@
##           develop      #64      +/-   ##
===========================================
+ Coverage    97.58%   97.83%   +0.25%     
===========================================
  Files           13       16       +3     
  Lines          663      786     +123     
===========================================
+ Hits           647      769     +122     
- Misses          16       17       +1

Flag	Coverage Δ
unittests	`97.83% <99.18%> (+0.25%)`	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

henrikjacobsenfys · 2026-01-06T14:12:14Z

I realise I should probably not allow the inputs to be Parameters. This is a legacy from the old base object serialisation.

rozyczko

A few comments, but nothing serious enough to not approve.

rozyczko · 2026-01-08T10:50:24Z

src/easydynamics/sample_model/diffusion_model/brownian_translational_diffusion.py

+        if isinstance(Q, Numeric):
+            Q = np.array([Q])
+
+        if isinstance(Q, list):
+            Q = np.array(Q)
+
+        if not isinstance(Q, np.ndarray):
+            raise TypeError("Q must be a number, list, or numpy array.")
+


This is a copy of the check in calculate_width - consider refactoring out.
Something like this maybe?

def _validate_and_convert_Q(self, Q) -> np.ndarray: if isinstance(Q, Number): Q = np.array([Q]) if isinstance(Q, list): Q = np.array(Q) if not isinstance(Q, np.ndarray): raise TypeError("Q must be a number, list, or numpy array.") return Q

And consider adding the Q.ndim check below to that as well.

rozyczko · 2026-01-08T10:54:09Z

src/easydynamics/sample_model/diffusion_model/brownian_translational_diffusion.py

+            raise TypeError("diffusion_coefficient must be a number.")
+        self._diffusion_coefficient.value = diffusion_coefficient
+
+    def calculate_width(self, Q: np.ndarray) -> np.ndarray:


Q is type-hinted as np.ndarray, but the code manually handles Numeric and list.
Maybe accept ArrayLike to cover both lists and arrays? Then simply convert to np.array at the start.

I think I'll eventually also want to accept sc.Variable, but I didn't want to deal with all the checks and conversions

rozyczko · 2026-01-08T12:05:51Z

src/easydynamics/sample_model/diffusion_model/brownian_translational_diffusion.py

+            "angstrom": self._angstrom,
+        }
+
+    def __repr__(self):


This can easily be in the base class.

The dependency map expression is different for diferent models, and same for the __repr__ - other models have more Parameters than just D and scale

but in case the child doesn't override __repr__ we probably should have some default way of showing what the model is.

rozyczko · 2026-01-08T12:08:33Z

src/easydynamics/sample_model/diffusion_model/diffusion_model_base.py

+        -------
+        str or sc.Unit or None
+        """
+        return self._unit


This differs from def unit in ModelComponent. Consider standardizing.
Are we returning sc.Unit/str or just str?

Just the str, thanks for catching it

rozyczko · 2026-01-08T12:10:43Z

src/easydynamics/sample_model/diffusion_model/brownian_translational_diffusion.py

+        diffusion_coefficient : float or Parameter, optional
+            Diffusion coefficient D. If a number is provided, it is assumed to be in the unit given by diffusion_unit. Defaults to 1.0.
+        diffusion_unit : str, optional
+            Unit for the diffusion coefficient D. Default is "meV*Å**2". Options are 'meV*Å**2' or 'm**2/s'


Nope. the current default is m**2/s

rozyczko · 2026-01-08T12:15:16Z

src/easydynamics/sample_model/diffusion_model/brownian_translational_diffusion.py

+    def _write_width_dependency_map_expression(self) -> Dict[str, str]:
+        """


The return value is not Dict[str, str], since self.diffusion_coefficient is a Parameter and self._hbar and self._angstrom are Descriptors

Whoops, I blame copilot and myself for trusting copilot

rozyczko · 2026-01-08T12:18:13Z

src/easydynamics/sample_model/diffusion_model/brownian_translational_diffusion.py

+        for Q_value in Q:
+            # Q is given as a float, so we need to divide by angstrom**2 to get the right units
+            width = (
+                self._hbar
+                * self.diffusion_coefficient
+                * Q_value**2
+                / (self._angstrom**2)
+            )
+            width.convert_unit(self.unit)
+            width_list.append(width.value)
+        width = np.array(width_list)


We really should be able to vectorize this.

This would be awesome to have:

width_var = ( self._hbar * self.diffusion_coefficient * Q_arr**2 / (self._angstrom**2) )

It would be awesome. But to do that we need to extend our arithmetic dunder methods to allow arrays.

Then let's have it done in the core...

unit_conversion_factor = ( self._hbar * self.diffusion_coefficient / (self._angstrom**2) ) unit_conversion_factor.convert_unit(self.unit) width = Q**2 * unit_conversion_factor.value

This may be more complicated with other diffusion models, but at least here it's simple.

damskii9992

Mostly minor comments, was one of your easier PRs to review :)

damskii9992 · 2026-01-08T12:28:58Z

src/easydynamics/sample_model/__init__.py

    "DeltaFunction",
    "DampedHarmonicOscillator",
    "Polynomial",
+    "BrownianTranslationalDiffusion",


Just a note, maybe you don't want everything exposed at the top-level?
More obscure/rare functionality makes sense to require a full path to import, just like how it makes sense to import your components from the components group:

from easydynamic.sample_model.components import Gaussian, Lorentzian

Rather than:

from easydynamics import Gaussian

Or whatever your modules are called. Just a note for consideration :)

damskii9992 · 2026-01-08T12:33:43Z

src/easydynamics/sample_model/diffusion_model/brownian_translational_diffusion.py

+            raise TypeError("diffusion_coefficient must be a Parameter or a number.")
+
+        if not isinstance(diffusion_unit, str):
+            raise TypeError("diffusion_unit must be 'meV*Å**2' or 'm**2/s'.")


The error message here should be: "diffusion_unit must be a valid string, such as: 'meV*Å2' or 'm2/s'."

Well right now it's only allowed to be one of those two specific strings :)

damskii9992 · 2026-01-08T12:38:44Z

src/easydynamics/sample_model/diffusion_model/brownian_translational_diffusion.py

+            raise ValueError("diffusion_unit must be 'meV*Å**2' or 'm**2/s'.")
+
+        if not isinstance(scale, Parameter):
+            scale = Parameter(name="scale", value=float(scale), fixed=False, min=0.0)


If you allow users to pass Parameters as well, then you also need to check if the given parameter has the correct unit. And maybe also "min" in this case?
And the same for the diffusion_cofficient.

I'm filing this under the same "don't allow parameters" issue as the components. I might just remove it since I don't think I need it.

damskii9992 · 2026-01-08T12:42:56Z

src/easydynamics/sample_model/diffusion_model/diffusion_model_base.py

+        """
+
+        if not (unit is None or isinstance(unit, (str, sc.Unit))):
+            raise TypeError("unit must be None, a string, or a scipp Unit")


It's not enough to check its type, you also need to check if this is actually a valid unit. Otherwise passing stuff like: "test" would work, even though that is clearly not a valid unit. Consider adding a:

try: test = Descriptor(unit=unit) test.to('meV') except exception as e: raise UnitError('Not a valid unit etc.) from e

damskii9992 · 2026-01-08T12:46:07Z

src/easydynamics/sample_model/diffusion_model/brownian_translational_diffusion.py

+        unit: str | sc.Unit = "meV",
+        scale: Parameter | Numeric = 1.0,
+        diffusion_coefficient: Parameter | Numeric = 1.0,
+        diffusion_unit: str = "m**2/s",


It's a bit confusing that diffusion_unit cannot be a sc.Unit, I would standardize, either always allow both, or always only allow a string.
And this is of course not just for this PR, but throughout your codebase.

Yeah I guess it can be a sc.Unit. The issue is that units are a mess with diffusion models since people (myelf included) tend to absorb hbar here and there and not be explicit about it.

damskii9992 · 2026-01-08T13:04:54Z

src/easydynamics/sample_model/diffusion_model/brownian_translational_diffusion.py

+        if isinstance(Q, Numeric):
+            Q = np.array([Q])
+
+        if isinstance(Q, list):
+            Q = np.array(Q)
+
+        if not isinstance(Q, np.ndarray):
+            raise TypeError("Q must be a number, list, or numpy array.")
+


And consider adding the Q.ndim check below to that as well.

damskii9992 · 2026-01-08T13:07:43Z

src/easydynamics/sample_model/diffusion_model/brownian_translational_diffusion.py

+        QISF = self.calculate_QISF(Q)
+
+        # Create a Lorentzian component for each Q-value, with width D*Q^2 and area equal to scale. No delta function, as the EISF is 0.
+        for i in range(len(Q)):


use enumerate instead.

damskii9992 · 2026-01-08T13:11:34Z

src/easydynamics/sample_model/diffusion_model/brownian_translational_diffusion.py

+        Write the dependency expression for the width as a function of Q to make dependent Parameters.
+        """
+        if not isinstance(Q, (float)):
+            raise TypeError("Q must be a float.")


You don't REALLY need checks on internal functions inputs, as you yourself is the only user (or other developers).

damskii9992 · 2026-01-08T13:14:49Z

src/easydynamics/sample_model/diffusion_model/brownian_translational_diffusion.py

+            )
+
+            # Resolving the dependency can do weird things to the units, so we make sure it's correct.
+            lorentzian_component.width.convert_unit(self.unit)


Uhh, if resolving the dependency makes the unit weird, this line only solves it temporarily. As soon as the fit runs, this unit will be weird again.

Yikes. Halp?

Maybe you can add the convert_unit to the dependency expression? That way it will be evaluated everytime the dependency is updated as well.

You might want to make a unit test to check if this works though. So make a diffusion_model and its component collection, then update the value of the diffusion_constant, and then check the unit/value of the component_collections Lorentzian.

Okay, this is a serious problem. The dependency expression is
f"hbar * D* {Q} **2*1/(angstrom**2)"
hbar is in mev*s, D is in m**2/s and angstrom is in m**2, so the resulting unit should be meV. However, it becomes J. This is easy to fix with convert_unit, but of course problematic if it changes every fit. Does this mean I need to rethink how I do the dependency, or do we need to add a unit option to dependencies?

It doesn't seem to work. I tried
return f"(hbar * D* {Q} **2*1/(angstrom**2)).convert_unit('meV')"
But get
A Parameter with unique_name meV does not exist. Please check your dependency expression.

Everything in a string is assumed to be a unique name by the parser, which makes sense. It seems the best solution is to allow specifying a desired unit for dependencies. I'll bring it up at standup tomorrow :)

damskii9992 · 2026-01-08T13:20:36Z

tests/unit_tests/sample_model/diffusion_model/test_brownian_translational_diffusion.py

+            "angstrom": brownian_diffusion_model._angstrom,
+        }
+
+        assert expression_map == expected_map


You shouldn't test internal methods unless you're afraid that they will break. Internal methods are implementation details so such tests will test your implementation rather than your functionality.
Remember, a 100% tested codebase is over-tested.

You should very much test internal methods, but preferably through the external API.

I'm doing both :) We can always delete the test if they get annoying, but last time I added a test to go from 98 to 100% coverage I found a bug, so even though I agree in principle, I fear I need to over-test my code at least for the time being.

henrikjacobsenfys added 2 commits January 6, 2026 14:51

Brownian diffusion with example and test

774c0a8

polishing

fab4aee

henrikjacobsenfys added this to the First release milestone Jan 6, 2026

henrikjacobsenfys assigned rozyczko and damskii9992 Jan 6, 2026

henrikjacobsenfys added [scope] enhancement Adds/improves features (major.MINOR.patch) [priority] medium Normal/default priority labels Jan 6, 2026

Move tests to where they belong

a118a31

rozyczko approved these changes Jan 8, 2026

View reviewed changes

damskii9992 approved these changes Jan 8, 2026

View reviewed changes

henrikjacobsenfys added 3 commits January 8, 2026 14:44

Updates based on PR comments

c5fafaf

Remove parameters as input

ff2ddca

Updates based on PR comments

5df823d

henrikjacobsenfys merged commit 2ef5011 into develop Jan 9, 2026
16 checks passed

		def _write_width_dependency_map_expression(self) -> Dict[str, str]:
		"""

Diffusion model #64

Diffusion model #64

Uh oh!

Conversation

henrikjacobsenfys commented Jan 6, 2026

Uh oh!

codecov bot commented Jan 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

henrikjacobsenfys commented Jan 6, 2026

Uh oh!

rozyczko left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

damskii9992 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

codecov bot commented Jan 6, 2026 •

edited

Loading