Skip to content

Commit

Permalink
Fix S1/S2 correction and gas gain when simulating S1/S2PE (#122)
Browse files Browse the repository at this point in the history
* Separate correction for S1PE/S2PE and cS1/cS2 using true/reconstructed position

* Add gas gain (x,y) dependence when simulating S2PE

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix parameter list of S2PE

* change gas gain map to relative map

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* add back gas gain to parameter

* modify the structure a bit

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* debug

* Update datastructure

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Zihao Xu <zihao.xu@columbia.edu>
Co-authored-by: Dacheng Xu <dx2227@columbia.edu>
  • Loading branch information
4 people authored Dec 20, 2023
1 parent 58d2843 commit e237adf
Show file tree
Hide file tree
Showing 4 changed files with 1,235 additions and 984 deletions.
1 change: 1 addition & 0 deletions appletree/maps/_gas_gain_relative.json
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
{"coordinate_lowers": [-70, -70], "coordinate_uppers": [70, 70], "coordinate_type": "regbin", "coordinate_name": ["x", "y"], "map": [[1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0]]}
63 changes: 34 additions & 29 deletions appletree/plugins/detector.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,52 +13,50 @@
@export
@takes_config(
Map(
name="s1_correction",
name="s1_lce",
default="_s1_correction.json",
help="S1 light collection efficiency correction",
help="S1 light collection efficiency",
),
)
class S1Correction(Plugin):
depends_on = ["rec_x", "rec_y", "rec_z"]
provides = ["s1_correction"]
class S1LCE(Plugin):
depends_on = ["x", "y", "z"]
provides = ["s1_lce"]

@partial(jit, static_argnums=(0,))
def simulate(self, key, parameters, rec_x, rec_y, rec_z):
pos = jnp.stack([rec_x, rec_y, rec_z]).T
s1_correction = self.s1_correction.apply(pos)
return key, s1_correction
def simulate(self, key, parameters, x, y, z):
pos_true = jnp.stack([x, y, z]).T
s1_lce = self.s1_lce.apply(pos_true)
return key, s1_lce


@export
@takes_config(
Map(
name="s2_correction",
name="s2_lce",
default="_s2_correction.json",
help="S2 light collection efficiency correction",
help="S2 light collection efficiency",
),
)
class S2Correction(Plugin):
depends_on = ["rec_x", "rec_y"]
provides = ["s2_correction"]
class S2LCE(Plugin):
depends_on = ["x", "y"]
provides = ["s2_lce"]

@partial(jit, static_argnums=(0,))
def simulate(self, key, parameters, rec_x, rec_y):
pos = jnp.stack([rec_x, rec_y]).T
s2_correction = self.s2_correction.apply(pos)
return key, s2_correction
def simulate(self, key, parameters, x, y):
pos_true = jnp.stack([x, y]).T
s2_lce = self.s2_lce.apply(pos_true)
return key, s2_lce


@export
class PhotonDetection(Plugin):
depends_on = ["num_photon", "s1_correction"]
depends_on = ["num_photon", "s1_lce"]
provides = ["num_s1_phd"]
parameters = ("g1", "p_dpe")

@partial(jit, static_argnums=(0,))
def simulate(self, key, parameters, num_photon, s1_correction):
g1_true_no_dpe = jnp.clip(
parameters["g1"] * s1_correction / (1.0 + parameters["p_dpe"]), 0, 1.0
)
def simulate(self, key, parameters, num_photon, s1_lce):
g1_true_no_dpe = jnp.clip(parameters["g1"] * s1_lce / (1.0 + parameters["p_dpe"]), 0, 1.0)
key, num_s1_phd = randgen.binomial(key, g1_true_no_dpe, num_photon)
return key, num_s1_phd

Expand Down Expand Up @@ -104,21 +102,28 @@ def simulate(self, key, parameters, num_electron, drift_survive_prob):
return key, num_electron_drifted


@takes_config(
Map(
name="gas_gain_relative",
default="_gas_gain_relative.json",
help="Gas gain as a function of (x, y) / mean gas gain",
),
)
@export
class S2PE(Plugin):
depends_on = ["num_electron_drifted", "s2_correction"]
depends_on = ["num_electron_drifted", "s2_lce", "x", "y"]
provides = ["num_s2_pe"]
parameters = ("g2", "gas_gain")

@partial(jit, static_argnums=(0,))
def simulate(self, key, parameters, num_electron_drifted, s2_correction):
extraction_eff = parameters["g2"] / parameters["gas_gain"]
g2_true = parameters["g2"] * s2_correction
gas_gain_true = g2_true / extraction_eff
def simulate(self, key, parameters, num_electron_drifted, s2_lce, x, y):
pos_true = jnp.stack([x, y]).T
gas_gain = parameters["gas_gain"] * self.gas_gain_relative.apply((pos_true))
extraction_eff = parameters["g2"] * s2_lce / gas_gain

key, num_electron_extracted = randgen.binomial(key, extraction_eff, num_electron_drifted)

mean_s2_pe = num_electron_extracted * gas_gain_true
mean_s2_pe = num_electron_extracted * gas_gain
key, num_s2_pe = randgen.truncate_normal(key, mean_s2_pe, jnp.sqrt(mean_s2_pe), vmin=0)

return key, num_s2_pe
38 changes: 38 additions & 0 deletions appletree/plugins/reconstruction.py
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,44 @@ def simulate(self, key, parameters, num_s2_pe):
return key, s2_area


@export
@takes_config(
Map(
name="s1_correction",
default="_s1_correction.json",
help="S1 xyz correction on reconstructed positions",
),
)
class S1Correction(Plugin):
depends_on = ["rec_x", "rec_y", "rec_z"]
provides = ["s1_correction"]

@partial(jit, static_argnums=(0,))
def simulate(self, key, parameters, rec_x, rec_y, rec_z):
pos_rec = jnp.stack([rec_x, rec_y, rec_z]).T
s1_correction = self.s1_correction.apply(pos_rec)
return key, s1_correction


@export
@takes_config(
Map(
name="s2_correction",
default="_s2_correction.json",
help="S2 xy correction on constructed positions",
),
)
class S2Correction(Plugin):
depends_on = ["rec_x", "rec_y"]
provides = ["s2_correction"]

@partial(jit, static_argnums=(0,))
def simulate(self, key, parameters, rec_x, rec_y):
pos_rec = jnp.stack([rec_x, rec_y]).T
s2_correction = self.s2_correction.apply(pos_rec)
return key, s2_correction


@export
class cS1(Plugin):
depends_on = ["s1_area", "s1_correction"]
Expand Down
Loading

0 comments on commit e237adf

Please sign in to comment.