Add efield position dependence

appletree/components/yields.py

@@ -8,7 +8,9 @@ def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
         self.register(apt.plugins.MonoEnergiesSpectra)
+        self.register(apt.plugins.PositionSpectra)
         self.register(apt.plugins.TotalQuanta)
+        self.register(apt.plugins.EField)
         self.register(apt.plugins.ThomasImelBox)
         self.register(apt.plugins.QyNR)
         self.register(apt.plugins.LyNR)
@@ -20,7 +22,9 @@ def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
         self.register(apt.plugins.UniformEnergiesSpectra)
+        self.register(apt.plugins.PositionSpectra)
         self.register(apt.plugins.TotalQuanta)
+        self.register(apt.plugins.EField)
         self.register(apt.plugins.ThomasImelBox)
         self.register(apt.plugins.QyNR)
         self.register(apt.plugins.LyNR)

appletree/instructs/literature_lyqy.json

@@ -30,7 +30,7 @@
             "ly_lux_2022": 2.5,
             "qy_lux_2022": 2.5
         },
-        "literature_field": {
+        "field": {
             "ly_lux_2016": 180.0,
             "qy_lenardo_2019": 220.0,
             "ly_lux_2022": 400.0,

appletree/maps/_efield_map.json

@@ -0,0 +1 @@
+{"coordinate_lowers": [-140, 0], "coordinate_uppers": [0, 70], "coordinate_type": "regbin", "coordinate_name": ["r", "z"], "map": [[23.0, 23.0, 23.0, 23.0, 23.0], [23.0, 23.0, 23.0, 23.0, 23.0], [23.0, 23.0, 23.0, 23.0, 23.0], [23.0, 23.0, 23.0, 23.0, 23.0], [23.0, 23.0, 23.0, 23.0, 23.0]]}

appletree/parameters/nestv2.json

@@ -10,6 +10,17 @@
         "unit": "g/cm^3",
         "doc": "Density of the lXe"
     },
+    "field": {
+        "prior_type": "fixed",
+        "prior_args": {
+            "val": 23.0
+        },
+        "allowed_range": null,
+        "init_mean": null,
+        "init_std": null,
+        "unit": "V/cm",
+        "doc": "Drift field"
+    },
     "alpha": {
         "prior_type": "twohalfnorm",
         "prior_args": {

appletree/parameters/nestv2_yield_only.json

@@ -10,6 +10,17 @@
         "unit": "g/cm^3",
         "doc": "Density of the lXe"
     },
+    "field": {
+        "prior_type": "fixed",
+        "prior_args": {
+            "val": 23.0
+        },
+        "allowed_range": null,
+        "init_mean": null,
+        "init_std": null,
+        "unit": "V/cm",
+        "doc": "Drift field"
+    },
     "alpha": {
         "prior_type": "twohalfnorm",
         "prior_args": {

appletree/plugins/detector.py

@@ -4,7 +4,7 @@
 
 from appletree import randgen
 from appletree.plugin import Plugin
-from appletree.config import takes_config, Map
+from appletree.config import takes_config, Map, Constant
 from appletree.utils import exporter
 
 export, __all__ = exporter(export_self=False)
@@ -48,6 +48,35 @@ def simulate(self, key, parameters, x, y):
         return key, s2_lce
 
 
+@export
+@takes_config(
+    Constant(
+        name="efield_position_dependence",
+        type=bool,
+        default=False,
+        help="Whether to enable position-dependent efield",
+    ),
+    Map(
+        name="efield_map",
+        default="_efield_map.json",
+        help="Electric field map",
+    ),
+)
+class EField(Plugin):
+    depends_on = ["x", "y", "z"]
+    provides = ["field"]
+    parameters = ("field",)
+
+    @partial(jit, static_argnums=(0,))
+    def simulate(self, key, parameters, x, y, z):
+        # Safe to use "if" because we expect config to be fixed
+        if not self.efield_position_dependence.value:
+            return key, jnp.ones(x.shape) * parameters["field"]
+        r = jnp.sqrt(x**2 + y**2)
+        pos_true = jnp.stack([r, z]).T
+        return key, self.efield_map.apply(pos_true)
+
+
 @export
 class PhotonDetection(Plugin):
     depends_on = ["num_photon", "s1_lce"]

appletree/plugins/er_microphys.py

@@ -41,18 +41,18 @@ def simulate(self, key, parameters, num_quanta):
 
 @export
 class mTI(Plugin):
-    depends_on = ["energy"]
+    depends_on = ["energy", "field"]
     provides = ["recomb_mean"]
-    parameters = ("w", "nex_ni_ratio", "py0", "py1", "py2", "py3", "py4", "field")
+    parameters = ("w", "nex_ni_ratio", "py0", "py1", "py2", "py3", "py4")
 
     @partial(jit, static_argnums=(0,))
-    def simulate(self, key, parameters, energy):
+    def simulate(self, key, parameters, energy, field):
         ni = energy / parameters["w"] / (1.0 + parameters["nex_ni_ratio"])
         ti = (
             ni
             * parameters["py0"]
             * jnp.exp(-energy / parameters["py1"])
-            * parameters["field"] ** parameters["py2"]
+            * field ** parameters["py2"]
             / 4.0
         )
         fd = 1.0 / (1.0 + jnp.exp(-(energy - parameters["py3"]) / parameters["py4"]))

appletree/plugins/nestv2.py

@@ -89,20 +89,15 @@ def simulate(self, key, parameters, energy):
 
 
 @export
-@takes_config(
-    Constant(
-        name="literature_field", type=float, default=23.0, help="Drift field in each literature"
-    ),
-)
 class ThomasImelBox(Plugin):
-    depends_on = ["energy"]
+    depends_on = ["energy", "field"]
     provides = ["ThomasImel"]
     parameters = ("gamma", "delta", "liquid_xe_density")
 
     @partial(jit, static_argnums=(0,))
-    def simulate(self, key, parameters, energy):
+    def simulate(self, key, parameters, energy, field):
         ThomasImel = jnp.ones(shape=jnp.shape(energy))
-        ThomasImel *= parameters["gamma"] * self.literature_field.value ** parameters["delta"]
+        ThomasImel *= parameters["gamma"] * field ** parameters["delta"]
         ThomasImel *= (parameters["liquid_xe_density"] / 2.9) ** 0.3
         return key, ThomasImel
 

tests/test_component.py

@@ -96,6 +96,18 @@ def benchmark(key):
     )
     er.compile()
     er.lineage_hash
-    er.simulate_hist(key, batch_size, parameters)
+    key, h = er.simulate_hist(key, batch_size, parameters)
     with pytest.raises(RuntimeError):
         key, r = er.multiple_simulations(key, batch_size, parameters, 5, apply_eff=True)
+
+    # test efield dependence
+    er.set_config({"efield_position_dependence": True})
+    _cached_functions.clear()
+    er.deduce(
+        data_names=("cs1", "cs2"),
+        func_name="er_sim",
+        force_no_eff=True,
+    )
+    er.compile()
+    er.lineage_hash
+    key, h = er.simulate_hist(key, batch_size, parameters)