From f17debf73d2a4089b2a6f414e2bbbb95dc9e0467 Mon Sep 17 00:00:00 2001
From: Lanqing Yuan <yuanlq@uchicago.edu>
Date: Mon, 26 Aug 2024 00:49:33 -0500
Subject: [PATCH] Explicitly assign local field in AmBe/YBe generator (#150)

---
 saltax/instructions/generator.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/saltax/instructions/generator.py b/saltax/instructions/generator.py
index b32aa77..ef1b16d 100644
--- a/saltax/instructions/generator.py
+++ b/saltax/instructions/generator.py
@@ -325,6 +325,7 @@ def generator_ambe(
     rate=1e9 / SALT_TIME_INTERVAL,
     time_mode="uniform",
     ambe_instructions_file=AMBE_INSTRUCTIONS_FILE,
+    fmap=FIELD_MAP,
     **kwargs
 ):
     """Generate instructions for a run with AmBe source.
@@ -340,6 +341,7 @@ def generator_ambe(
     :param time_mode: 'uniform' or 'realistic', default: 'uniform'
     :param ambe_instructions_file: file containing ambe instructions,
         default: AMBE_INSTRUCTIONS_FILE
+    :param fmap: field map, default: FIELD_MAP, defined above
     :return: instructions in numpy array
     """
     # determine time offsets to shift ambe instructions
@@ -372,6 +374,8 @@ def generator_ambe(
         instr_i["e_dep"] = selected_ambe["e_dep"]
         instr_i["amp"] = selected_ambe["amp"]
         instr_i["n_excitons"] = selected_ambe["n_excitons"]
+        instr_i["local_field"] = fmap(np.array([np.sqrt(selected_ambe["x"]**2 + selected_ambe["y"]**2), 
+                                                selected_ambe["z"]]).T).repeat(2)
 
         # concatenate instr
         instr = np.concatenate((instr, instr_i))
@@ -388,6 +392,7 @@ def generator_ybe(
     rate=1e9 / SALT_TIME_INTERVAL,
     time_mode="uniform",
     ybe_instructions_file=YBE_INSTRUCTIONS_FILE,
+    fmap=FIELD_MAP,
     **kwargs
 ):
     """Generate instructions for a run with YBe source.
@@ -403,6 +408,7 @@ def generator_ybe(
     :param time_mode: 'uniform' or 'realistic', default: 'uniform'
     :param ybe_instructions_file: file containing ybe instructions,
         default: YBE_INSTRUCTIONS_FILE
+    :param fmap: field map, default: FIELD_MAP, defined above
     :return: instructions in numpy array
     """
     # determine time offsets to shift ybe instructions
@@ -433,6 +439,8 @@ def generator_ybe(
         instr_i["e_dep"] = selected_ybe["e_dep"]
         instr_i["amp"] = selected_ybe["amp"]
         instr_i["n_excitons"] = selected_ybe["n_excitons"]
+        instr_i["local_field"] = fmap(np.array([np.sqrt(selected_ybe["x"]**2 + selected_ybe["y"]**2), 
+                                                selected_ybe["z"]]).T).repeat(2)
 
         # concatenate instr
         instr = np.concatenate((instr, instr_i))