Field distortion and transverse diffusion map (#235)

* Included COMSOL-driven field distortion

* Transverse diffusion from COMSOL map included

* Remove useless argument make_map function

* Extended truth distorted x and y to fdc_3d

* Fix backwards compatibility

* Made xy truth a function

wfsim/core/rawdata.py

@@ -331,6 +331,8 @@ def get_truth(self, instruction, truth_buffer):
                 tb[f't_first_{quantum}'] = np.nan
                 tb[f't_last_{quantum}'] = np.nan
                 tb[f't_sigma_{quantum}'] = np.nan
+
+        self.get_mean_xy_electron(peak_type, instruction, tb)
 
         # Endtime is the end of the last pulse
         if np.isnan(tb['t_last_photon']):
@@ -368,6 +370,19 @@ def get_truth(self, instruction, truth_buffer):
 
         # Signal this row is now filled, so it won't be overwritten
         tb['fill'] = True
+
+    def get_mean_xy_electron(self, peak_type, instruction, tb):
+
+        if peak_type == 's2' and self.config.get('field_distortion_model', "none") in ['comsol', 'inverse_fdc']:
+            if self.config.get('field_distortion_model', "none") == 'comsol':
+                _, xy_tmp = self.pulses['s2'].field_distortion_comsol(instruction['x'], instruction['y'], instruction['z'], self.resource)
+            elif self.config.get('field_distortion_model', "none") == 'inverse_fdc':
+                _, xy_tmp = self.pulses['s2'].inverse_field_distortion_correction(instruction['x'], instruction['y'], instruction['z'], self.resource)
+            tb['x_mean_electron'] = np.mean(xy_tmp.T[0])
+            tb['y_mean_electron'] = np.mean(xy_tmp.T[1])
+        else:
+            tb['x_mean_electron'] = np.nan
+            tb['y_mean_electron'] = np.nan
 
     @staticmethod
     @njit

wfsim/core/s2.py

@@ -16,7 +16,7 @@
 @export
 class S2(Pulse):
     """
-    Given temperal inputs as well as number of electrons
+    Given temporal inputs as well as number of electrons
     Random generate photon timing and channel distribution.
     """
 
@@ -35,8 +35,11 @@ def __call__(self, instruction):
             np.array(v).reshape(-1) for v in zip(*instruction)]
 
         # Reverse engineerring FDC
-        if self.config['field_distortion_on']:
-            z_obs, positions = self.inverse_field_distortion(x, y, z, resource=self.resource)
+        if self.config['field_distortion_model'] == 'inverse_fdc':
+            z_obs, positions = self.inverse_field_distortion_correction(x, y, z, resource=self.resource)
+        # Reverse engineerring FDC
+        elif self.config['field_distortion_model'] == 'comsol':
+            z_obs, positions = self.field_distortion_comsol(x, y, z, resource=self.resource)
         else:
             z_obs, positions = z, np.array([x, y]).T
 
@@ -83,7 +86,7 @@ def get_s2_drift_time_params(z_obs, positions, config, resource):
             drift_velocity_liquid *= 1e-4  # cm/ns
         else:
             drift_velocity_liquid = config['drift_velocity_liquid']
-
+            
         if config['enable_field_dependencies']['diffusion_longitudinal_map']:
             diffusion_constant_longitudinal = resource.field_dependencies_map(z_obs, positions, map_name='diffusion_longitudinal_map')  # cm²/s
             diffusion_constant_longitudinal *= 1e-9  # cm²/ns
@@ -135,7 +138,7 @@ def get_electron_yield(n_electron, positions, z_obs, config, resource):
                                               config['electron_lifetime_liquid'])
         cy = config['electron_extraction_yield'] * electron_lifetime_correction
 
-        # Remove electrons in insensitive volumne
+        # Remove electrons in insensitive volume
         if config['enable_field_dependencies']['survival_probability_map']:
             survival_probability = resource.field_dependencies_map(z_obs, positions, map_name='survival_probability_map')
             cy *= survival_probability
@@ -146,8 +149,8 @@ def get_electron_yield(n_electron, positions, z_obs, config, resource):
         return n_electron
 
     @staticmethod
-    def inverse_field_distortion(x, y, z, resource):
-        """For 1T the pattern map is a data driven one so we need to reverse engineer field distortion
+    def inverse_field_distortion_correction(x, y, z, resource):
+        """For 1T the pattern map is a data driven one so we need to reverse engineer field distortion correction
         into the simulated positions
         :param x: 1d array of float
         :param y: 1d array of float
@@ -170,7 +173,25 @@ def inverse_field_distortion(x, y, z, resource):
 
         positions = np.array([x_obs, y_obs]).T 
         return z_obs, positions
-
+
+    @staticmethod
+    def field_distortion_comsol(x, y, z, resource):
+        """Field distortion from the COMSOL simulation for the given electrode configuration:
+        :param x: 1d array of float
+        :param y: 1d array of float
+        :param z: 1d array of float
+        :param resource: instance of resource class
+        returns z: 1d array, postions 2d array 
+        """
+        positions = np.array([np.sqrt(x**2 + y**2), z]).T
+        theta = np.arctan2(y, x)
+        r_obs = resource.fd_comsol(positions, map_name='r_distortion_map')
+        x_obs = r_obs * np.cos(theta)
+        y_obs = r_obs * np.sin(theta)
+
+        positions = np.array([x_obs, y_obs]).T 
+        return z, positions
+
     @staticmethod
     @njit
     def _luminescence_timings_simple(n, dG, E0, r, dr, rr, alpha, uE, p, n_photons):
@@ -272,7 +293,7 @@ def electron_timings(t, n_electron, drift_time_mean, drift_time_spread, sc_gain,
         :param n_electron:1 d array of ints
         :param drift_time_mean: 1d array of floats
         :param drift_time_spread: 1d array of floats
-        :param sc_gain: secondairy scintallation gain       
+        :param sc_gain: secondary scintillation gain       
         :param timings: empty array with length sum(n_electron)
         :param gains: empty array with length sum(n_electron)
         :param electron_trapping_time: configuration values
@@ -355,17 +376,34 @@ def s2_pattern_map_diffuse(n_electron, z, xy, config, resource):
         :param config: dict of the wfsim config
         :param resource: instance of the resource class
         """
-        drift_time_gate = config['drift_time_gate']
-        drift_velocity_liquid = config['drift_velocity_liquid']
-        diffusion_constant_transverse = getattr(config, 'diffusion_constant_transverse', 0)
-
         assert all(z < 0), 'All S2 in liquid should have z < 0'
+
+        if config['enable_field_dependencies']['drift_speed_map']:
+            drift_velocity_liquid = resource.field_dependencies_map(z, xy, map_name='drift_speed_map')  # mm/µs
+            drift_velocity_liquid *= 1e-4  # cm/ns
+        else:
+            drift_velocity_liquid = config['drift_velocity_liquid']
+
+        if config['enable_field_dependencies']['diffusion_transverse_map']:
+            diffusion_constant_radial = resource.field_dependencies_map(z, xy, map_name='diffusion_radial_map')  # cm²/s
+            diffusion_constant_azimuthal = resource.field_dependencies_map(z, xy, map_name='diffusion_azimuthal_map') # cm²/s
+            diffusion_constant_radial *= 1e-9  # cm²/ns
+            diffusion_constant_azimuthal *= 1e-9  # cm²/ns
+        else:
+            diffusion_constant_transverse = getattr(config, 'diffusion_constant_transverse', 0)
+            diffusion_constant_radial = diffusion_constant_transverse
+            diffusion_constant_azimuthal = diffusion_constant_transverse
 
-        drift_time_mean = - z / drift_velocity_liquid + drift_time_gate  # Add gate time for consistancy?
-        hdiff_stdev = np.sqrt(2 * diffusion_constant_transverse * drift_time_mean)
-
-        hdiff = np.random.normal(0, 1, (np.sum(n_electron), 2)) * \
-            np.repeat(hdiff_stdev, n_electron, axis=0).reshape((-1, 1))
+        drift_time_mean = - z / drift_velocity_liquid
+        hdiff_stdev_radial = np.sqrt(2 * diffusion_constant_radial * drift_time_mean)
+        hdiff_stdev_azimuthal = np.sqrt(2 * diffusion_constant_azimuthal * drift_time_mean)
+
+        hdiff_radial = np.random.normal(0, 1, np.sum(n_electron)) * np.repeat(hdiff_stdev_radial, n_electron, axis=0)
+        hdiff_azimuthal = np.random.normal(0, 1, np.sum(n_electron)) * np.repeat(hdiff_stdev_azimuthal, n_electron, axis=0)
+        hdiff = np.column_stack([hdiff_radial, hdiff_azimuthal])
+        theta = np.arctan2(xy[:,1], xy[:,0])
+        matrix = np.array([[np.cos(theta), np.sin(theta)], [-np.sin(theta), np.cos(theta)]]).T
+        hdiff = np.vstack([(matrix[i] @ np.split(hdiff, np.cumsum(n_electron))[:-1][i].T).T for i in range(len(matrix))])
         # Should we also output this xy position in truth?
         xy_multi = np.repeat(xy, n_electron, axis=0) + hdiff  # One entry xy per electron
         # Remove points outside tpc, and the pattern will be the average inside tpc

wfsim/load_resource.py

@@ -251,9 +251,12 @@ def __init__(self, config=None):
                 liquid_level = min(liquid_level_available, key=lambda x: abs(x - liquid_level))
                 self.s2_luminescence = s2_luminescence_map[s2_luminescence_map['ll'] == liquid_level]
 
-            if config.get('field_distortion_on', False):
+            if config.get('field_distortion_model', "none") == "inverse_fdc":
                 self.fdc_3d = make_map(files['fdc_3d'], fmt='json.gz')
 
+            if config.get('field_distortion_model', "none") == "comsol":
+                self.fd_comsol = make_map(config['field_distortion_comsol_map'], fmt='json.gz', method='RectBivariateSpline')
+
             # Gas gap warping map
             if config.get('enable_gas_gap_warping', False):
                 gas_gap_map = straxen.get_resource(files['gas_gap_map'], fmt='pkl')
@@ -262,7 +265,7 @@ def __init__(self, config=None):
             # Field dependencies 
             # This config entry a dictionary of 5 items
             if any(config['enable_field_dependencies'].values()):
-                field_dependencies_map = make_map(files['field_dependencies_map'], fmt='json.gz')
+                field_dependencies_map = make_map(files['field_dependencies_map'], fmt='json.gz', method='RectBivariateSpline')
 
                 def rz_map(z, xy, **kwargs):
                     r = np.sqrt(xy[:, 0]**2 + xy[:, 1]**2)

wfsim/strax_interface.py

@@ -24,9 +24,9 @@
 instruction_dtype = [(('Waveform simulator event number.', 'event_number'), np.int32),
                      (('Quanta type (S1 photons or S2 electrons)', 'type'), np.int8),
                      (('Time of the interaction [ns]', 'time'), np.int64),
-                     (('X position of the cluster[cm]', 'x'), np.float32),
-                     (('Y position of the cluster[cm]', 'y'), np.float32),
-                     (('Z position of the cluster[cm]', 'z'), np.float32),
+                     (('X position of the cluster [cm]', 'x'), np.float32),
+                     (('Y position of the cluster [cm]', 'y'), np.float32),
+                     (('Z position of the cluster [cm]', 'z'), np.float32),
                      (('Number of quanta', 'amp'), np.int32),
                      (('Recoil type of interaction.', 'recoil'), np.int8),
                      (('Energy deposit of interaction', 'e_dep'), np.float32),
@@ -50,6 +50,8 @@
     (('Arrival time of the last photon [ns]', 't_last_photon'), np.float64),
     (('Mean time of the photons [ns]', 't_mean_photon'), np.float64),
     (('Standard deviation of photon arrival times [ns]', 't_sigma_photon'), np.float64),
+    (('X field-distorted mean position of the electrons [cm]', 'x_mean_electron'), np.float32),
+    (('Y field-distorted mean position of the electrons [cm]', 'y_mean_electron'), np.float32),
     (('Arrival time of the first electron [ns]', 't_first_electron'), np.float64),
     (('Arrival time of the last electron [ns]', 't_last_electron'), np.float64),
     (('Mean time of the electrons [ns]', 't_mean_electron'), np.float64),
@@ -489,6 +491,10 @@ def set_config(self,):
         overrides = self.config['fax_config_override']
         if overrides is not None:
             self.config.update(overrides)
+
+        # backwards compatibility
+        if 'field_distortion_on' in self.config and not 'field_distortion_model' in self.config:
+            self.config.update({'field_distortion_model': "inverse_fdc" if self.config['field_distortion_on'] else "none"})
 
         # Update gains to the nT defaults
         self.to_pe = straxen.get_correction_from_cmt(self.run_id,