Added source argument to dx to avoid repeated calculations.

src/tirl/costs/msd.py

@@ -195,7 +195,7 @@ class CostMSD(Cost):
         jac /= self.target.tsize
         return jac.reshape((source.domain.numel, -1))
 
-    def dx(self, dim=None):
+    def dx(self, dim=None, source=None):
         """
         Returns the voxelwise total derivative of the source with respect to the
         local displacement vector. The result has the following shape:
@@ -208,7 +208,8 @@ class CostMSD(Cost):
         image.
 
         """
-        source = self.source  # only evaluate once
+        if source is None:
+            source = self.source  # only evaluate once
         mask = self.combine_masks(self.target.mask, source.mask)
         # diff = source - self.target
         # diff.order = VOXEL_MAJOR