From e3ee49e2503c1ad16d7f03919ba6b5124d8e152d Mon Sep 17 00:00:00 2001 From: mhliu Date: Sun, 29 Sep 2024 12:12:48 -0400 Subject: [PATCH 1/8] Add efield position dependence --- appletree/maps/_efield_map.json | 1 + appletree/plugins/detector.py | 30 +++++++++++++++++++++++++++++- appletree/plugins/er_microphys.py | 8 ++++---- appletree/plugins/nestv2.py | 11 +++-------- 4 files changed, 37 insertions(+), 13 deletions(-) create mode 100644 appletree/maps/_efield_map.json diff --git a/appletree/maps/_efield_map.json b/appletree/maps/_efield_map.json new file mode 100644 index 00000000..5eaa62e7 --- /dev/null +++ b/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]]} diff --git a/appletree/plugins/detector.py b/appletree/plugins/detector.py index e4292660..028c77f8 100644 --- a/appletree/plugins/detector.py +++ b/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) @@ -47,6 +47,34 @@ def simulate(self, key, parameters, x, y): s2_lce = self.s2_lce.apply(pos_true) 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_dependece.value: + return key, 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): diff --git a/appletree/plugins/er_microphys.py b/appletree/plugins/er_microphys.py index 2ee7ac5c..f6f74dea 100644 --- a/appletree/plugins/er_microphys.py +++ b/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"])) diff --git a/appletree/plugins/nestv2.py b/appletree/plugins/nestv2.py index 880bf929..8bb5cfa5 100644 --- a/appletree/plugins/nestv2.py +++ b/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 From 5403f3c74dc74e4437e925feb4f93154dad4cef4 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 29 Sep 2024 16:16:10 +0000 Subject: [PATCH 2/8] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- appletree/plugins/detector.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/appletree/plugins/detector.py b/appletree/plugins/detector.py index 028c77f8..ce596c04 100644 --- a/appletree/plugins/detector.py +++ b/appletree/plugins/detector.py @@ -47,6 +47,7 @@ def simulate(self, key, parameters, x, y): s2_lce = self.s2_lce.apply(pos_true) return key, s2_lce + @export @takes_config( Constant( @@ -65,13 +66,13 @@ 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_dependece.value: return key, parameters["field"] - r = jnp.sqrt(x**2+y**2) + r = jnp.sqrt(x**2 + y**2) pos_true = jnp.stack([r, z]).T return key, self.efield_map.apply(pos_true) From 9f9da114b82771485b683c08f1a533b0b0f5edbb Mon Sep 17 00:00:00 2001 From: mhliu Date: Sun, 29 Sep 2024 12:21:16 -0400 Subject: [PATCH 3/8] Fix typo --- appletree/instructs/literature_lyqy.json | 2 +- appletree/plugins/detector.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/appletree/instructs/literature_lyqy.json b/appletree/instructs/literature_lyqy.json index 30a7e7b2..d8f827b2 100644 --- a/appletree/instructs/literature_lyqy.json +++ b/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, diff --git a/appletree/plugins/detector.py b/appletree/plugins/detector.py index ce596c04..c562aaac 100644 --- a/appletree/plugins/detector.py +++ b/appletree/plugins/detector.py @@ -70,7 +70,7 @@ class EField(Plugin): @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_dependece.value: + if not self.efield_position_dependence.value: return key, parameters["field"] r = jnp.sqrt(x**2 + y**2) pos_true = jnp.stack([r, z]).T From 89676761e9561b020d5d9b259855310901c5b8e2 Mon Sep 17 00:00:00 2001 From: mhliu Date: Sun, 29 Sep 2024 12:27:53 -0400 Subject: [PATCH 4/8] Fix test --- appletree/components/yields.py | 2 ++ appletree/plugins/detector.py | 2 +- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/appletree/components/yields.py b/appletree/components/yields.py index d2f2ae28..da473a26 100644 --- a/appletree/components/yields.py +++ b/appletree/components/yields.py @@ -9,6 +9,7 @@ def __init__(self, *args, **kwargs): self.register(apt.plugins.MonoEnergiesSpectra) 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) @@ -21,6 +22,7 @@ def __init__(self, *args, **kwargs): self.register(apt.plugins.UniformEnergiesSpectra) 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) diff --git a/appletree/plugins/detector.py b/appletree/plugins/detector.py index c562aaac..2f04438c 100644 --- a/appletree/plugins/detector.py +++ b/appletree/plugins/detector.py @@ -65,7 +65,7 @@ def simulate(self, key, parameters, x, y): class EField(Plugin): depends_on = ["x", "y", "z"] provides = ["field"] - parameters = ["field"] + parameters = ("field",) @partial(jit, static_argnums=(0,)) def simulate(self, key, parameters, x, y, z): From 1ad7e3d068253356146c5e5f493e9f74c6354e89 Mon Sep 17 00:00:00 2001 From: mhliu Date: Sun, 29 Sep 2024 12:37:14 -0400 Subject: [PATCH 5/8] Add position plugins for test --- appletree/components/yields.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/appletree/components/yields.py b/appletree/components/yields.py index da473a26..927feef1 100644 --- a/appletree/components/yields.py +++ b/appletree/components/yields.py @@ -8,6 +8,7 @@ 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) @@ -21,6 +22,7 @@ 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) From 82f5718121076aa35aa77ad4dae05f4d99fa9052 Mon Sep 17 00:00:00 2001 From: mhliu Date: Sun, 29 Sep 2024 12:48:06 -0400 Subject: [PATCH 6/8] Fix test parameters; fix efield shape --- appletree/parameters/nestv2.json | 11 +++++++++++ appletree/parameters/nestv2_yield_only.json | 11 +++++++++++ appletree/plugins/detector.py | 2 +- 3 files changed, 23 insertions(+), 1 deletion(-) diff --git a/appletree/parameters/nestv2.json b/appletree/parameters/nestv2.json index 1fa52a86..77156565 100644 --- a/appletree/parameters/nestv2.json +++ b/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": { diff --git a/appletree/parameters/nestv2_yield_only.json b/appletree/parameters/nestv2_yield_only.json index c8b6d420..7acec608 100644 --- a/appletree/parameters/nestv2_yield_only.json +++ b/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": { diff --git a/appletree/plugins/detector.py b/appletree/plugins/detector.py index 2f04438c..3f3767ac 100644 --- a/appletree/plugins/detector.py +++ b/appletree/plugins/detector.py @@ -71,7 +71,7 @@ class EField(Plugin): 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, parameters["field"] + 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) From a9574393da132b5b770b943b10c99e36f16425d2 Mon Sep 17 00:00:00 2001 From: mhliu Date: Sun, 29 Sep 2024 15:13:44 -0400 Subject: [PATCH 7/8] Add test for position-dependent efield --- tests/test_component.py | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/tests/test_component.py b/tests/test_component.py index 70ed06b0..ed2e1e0a 100644 --- a/tests/test_component.py +++ b/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) From 28a203a3466c437adbe8cdc0df75ce18f36ee9ec Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 29 Sep 2024 19:13:56 +0000 Subject: [PATCH 8/8] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/test_component.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_component.py b/tests/test_component.py index ed2e1e0a..7ac3a7df 100644 --- a/tests/test_component.py +++ b/tests/test_component.py @@ -99,7 +99,7 @@ def benchmark(key): 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()