Default to data-driven diffusion map

wfsim/core/s2.py

@@ -166,8 +166,7 @@ def get_s2_drift_time_params(z_int, xy_int, config, resource):
         """
         drift_velocity_liquid = S2.get_avg_drift_velocity(z_int, xy_int, config, resource)
         if config['enable_field_dependencies']['diffusion_longitudinal_map']:
-            diffusion_constant_longitudinal = resource.field_dependencies_map(z_int, xy_int, map_name='diffusion_longitudinal_map')  # cm²/s
-            diffusion_constant_longitudinal *= 1e-9  # cm²/ns
+            diffusion_constant_longitudinal = resource.diffusion_longitudinal_map(z_int, xy_int)  # cm²/ns
         else:
             diffusion_constant_longitudinal = config['diffusion_constant_longitudinal']
 

wfsim/load_resource.py

@@ -100,6 +100,7 @@ def config_to_file(config):
                 'noise_file': 'x1t_noise_170203_0850_00_small.npz',
                 'fdc_3d': 'XnT_3D_FDC_xyt_dummy_all_zeros_v0.1.json.gz',
                 'field_dependencies_map': '',
+                "diffusion_longitudinal_map": '',
                 's1_time_spline': 'XENONnT_s1_proponly_va43fa9b_wires_20200625.json.gz',
                 's2_time_spline': '',
              })
@@ -337,6 +338,16 @@ def rz_map(z, xy, **kwargs):
                     return field_dependencies_map(np.array([r, z]).T, **kwargs)
                 self.field_dependencies_map = rz_map
 
+            # Data-driven longitudinal diffusion map
+            # TODO: Change to the best way to accommodate simulation/data-driven map
+            if config['enable_field_dependencies']["diffusion_longitudinal_map"]:
+                diffusion_longitudinal_map = make_map(files['diffusion_longitudinal_map'], fmt='json.gz',
+                                                  method='WeightedNearestNeighbors')
+                def _rz_map(z, xy, **kwargs):
+                    r = np.sqrt(xy[:, 0]**2 + xy[:, 1]**2)
+                    return diffusion_longitudinal_map(np.array([r, z]).T, **kwargs)
+                self.diffusion_longitudinal_map = _rz_map
+
             # Photon After Pulses
             if config.get('enable_pmt_afterpulses', False):
                 self.uniform_to_pmt_ap = straxen.get_resource(files['photon_ap_cdfs'], fmt='json.gz')