PolarArbitrary and doc updates

docs/user_guide/detector/index.md

@@ -13,7 +13,7 @@ includes effects pointed out by Adam Anthony in his [thesis](https://ezproxy.msu
 The detector system code simulates an event by going through a series of steps to 
 ultimately produce its point cloud. For each event, the detector system:
 
-1. Generate the trajectory of the each nucleus in the exit channel of the event by solving its equation of motion in the AT-TPC with a fine grained-timestep
+1. Generate the trajectory of the each nucleus in the exit channel (or the nuclei specified by the user) of the event by solving its equation of motion in the AT-TPC with a fine grained-timestep
 2. Determines how many electrons are created at each point of each nucleus' trajectory
 3. Converts the z coordinate of each point in each trajectory to a GET electronics sampling-frequency based Time Bucket using the electron drift velocity
 3. Transports each point of the nucleus' trajectory to the pad plane, applying diffusion if requested, to identify the sensing pad for each timestep
@@ -141,6 +141,17 @@ step reaction a(b,c)d a=0, b=1, c=2, d=3. In a two step, a(b,c)d->e+f, e=4, f=5,
 on for more complex scenarios. These labels are particularly useful for evaluating the
 performance of machine learning methods like clustering in downstream analyses.
 
+## Which nuclei get simulated in the detector simulation
+
+By default, the detector simulation considers all nuclei in the exit channel. That is,
+in a two step reaction a(b,c)d->e+f, the nuclei c, e, f are included in the detector
+simulation and used to generate a point cloud. However, this is not always desireable. 
+In some usecases, only on particle is interesting and all others simply generate extra
+data. As such, the `run_simulation` function contains an optional keyword argument 
+`indices` which can be used to specify the indices of the nuclei to include in the 
+detector simulation. Indices follow the same convention as the kinematics simulation,
+i.e. for our earlier example two step reaction a=0, b=1, c=2, d=3, e=4, and f=5.
+
 ## Why Point clouds
 
 It is worth elaborating why the detector system outputs point clouds instead of raw 

docs/user_guide/getting_started.md

@@ -1,6 +1,7 @@
 # Getting Started
 
-Once attpc_engine is installed it's time to setup our simulation project. Below is an outline of what we'd expect your project structure to look like
+Once attpc_engine is installed it's time to setup our simulation project. Below is an 
+outline of what we'd expect your project structure to look like
 
 ```txt
 |---my_sim
@@ -13,13 +14,18 @@ Once attpc_engine is installed it's time to setup our simulation project. Below
 |   |   |---detector
 ```
 
-You have a folder (in this case called `my_sim`) with two Python files (`generate_kinematics.py` and `apply_detector.py`) a virtual environment with attpc_engine installed (`.venv`) and an output directory with a folder for kinematics and detector effects. Note that you may want the output to be stored on a remote disk or other location as simulation files can be big at times (> 1 GB).
+You have a folder (in this case called `my_sim`) with two Python files 
+(`generate_kinematics.py` and `apply_detector.py`) a virtual environment with 
+attpc_engine installed (`.venv`) and an output directory with a folder for kinematics 
+and detector effects. Note that you may want the output to be stored on a remote disk 
+or other location as simulation files can be big at times (> 1 GB).
 
 Now lets talk about generating and sampling a kinematics phase space!
 
 ## Sampling Kinematics
 
-Below is an example script for running a kinematics sample, which in our example project we would write in `my_sim/generate_kinematics.py`
+Below is an example script for running a kinematics sample, which in our example 
+project we would write in `my_sim/generate_kinematics.py`
 
 ```python
 from attpc_engine.kinematics import (
@@ -70,10 +76,13 @@ if __name__ == "__main__":
     main()
 ```
 
-First we import all of our pieces from the attpc_engine library and its dependencies. We also import the Python standard library Path object to handle our file paths.
+First we import all of our pieces from the attpc_engine library and its dependencies. 
+We also import the Python standard library Path object to handle our file paths.
 
-We then start to define our kinematics configuration. First we define the output path to be an HDF5 file in our output directory. 
-We also load a spyral-utils [gas target](https://attpc.github.io/spyral-utils/api/nuclear/target) from a file path.
+We then start to define our kinematics configuration. First we define the output path 
+to be an HDF5 file in our output directory. 
+We also load a spyral-utils [gas target](https://attpc.github.io/spyral-utils/api/nuclear/target)
+from a file path.
 
 ```python
 target = load_target(target_path, nuclear_map)
@@ -90,7 +99,21 @@ nevents = 10000
 beam_energy = 184.131 # MeV
 ```
 
-Now were ready to define our kinematics Pipeline. The first argument of the Pipeline is a list of steps, where each step is either a Reaction or Deacy. The first element of the list must *always* be a Reaction, and all subsequent steps are *always* Decays. The residual of the previous step is *always* the parent of the next step. The Pipeline will attempt to validate this information for you. We also must define a list of ExcitationDistribution objects. These describe the state in the residual that is populated by each Reaction/Decay. There is exactly *one* distribution per step. There are two types of predefined ExcitationDistributions (ExcitationGaussian and ExcitationUniform), but others can be implemented by implementing the ExcitationDistribution protocol. Similarly, we use PolarUniform to define a uniform polar angle distribution from 0 degrees to 180 degrees in the center of mass frame (technically, this is from -1 to 1 in cos(polar)). We also define our KinematicsTargetMaterial, which includes our gas target, as well as the allowed space within the target for our reaction vertex (range in z in meters and standard deviation of cylindrical ρ in meters).
+Now were ready to define our kinematics Pipeline. The first argument of the Pipeline 
+is a list of steps, where each step is either a Reaction or Deacy. The first element 
+of the list must *always* be a Reaction, and all subsequent steps are *always* Decays. 
+The residual of the previous step is *always* the parent of the next step. The 
+Pipeline will attempt to validate this information for you. We also must define a list 
+of ExcitationDistribution objects. These describe the state in the residual that is 
+populated by each Reaction/Decay. There is exactly *one* distribution per step. There 
+are two types of predefined ExcitationDistributions (ExcitationGaussian and 
+ExcitationUniform), but others can be implemented by implementing the 
+ExcitationDistribution protocol. Similarly, we use PolarUniform to define a uniform 
+polar angle distribution from 0 degrees to 180 degrees in the center of mass frame 
+(technically, this is from -1 to 1 in cos(polar)). We also define our 
+KinematicsTargetMaterial, which includes our gas target, as well as the allowed space 
+within the target for our reaction vertex (range in z in meters and standard deviation 
+of cylindrical ρ in meters).
 
 ```python
 pipeline = KinematicsPipeline(
@@ -126,11 +149,13 @@ if __name__ == "__main__":
     main()
 ```
 
-That's it! This script will then sample 10000 events from the kinematic phase space of ${}^{16}C(d,d')$ and write the data out to an HDF5 file in our output directory.
+That's it! This script will then sample 10000 events from the kinematic phase space of 
+${}^{16}C(d,d')$ and write the data out to an HDF5 file in our output directory.
 
 ## Applying the Detector
 
-Below is an example script for running a kinematics sample, which in our example project we would write in `my_sim/apply_detector.py`
+Below is an example script for running a kinematics sample, which in our example 
+project we would write in `my_sim/apply_detector.py`
 
 ```python
 from attpc_engine.detector import (
@@ -193,7 +218,10 @@ if __name__ == "__main__":
     main()
 ```
 
-Just like in the kinematics script, we start off by importing a whole bunch of code. Next we define our kinematics input (which is the output of the kinematics script) and an output path. Note that for the output path we simply specify a directory; this is because our writer will handle breaking up the output data into reasonably sized files.
+Just like in the kinematics script, we start off by importing a whole bunch of code. 
+Next we define our kinematics input (which is the output of the kinematics script) and 
+an output path. Note that for the output path we simply specify a directory; this is 
+because our writer will handle breaking up the output data into reasonably sized files.
 
 ```python
 input_path = Path("output/kinematics/c16dd_d2_300Torr_184MeV.h5")
@@ -209,7 +237,8 @@ if not isinstance(gas, GasTarget):
     raise Exception(f"Could not load target data from {target_path}!")
 ```
 
-and finally we begin to define the detector specific configuration, which is ultimately stored in a Config object.
+and finally we begin to define the detector specific configuration, which is ultimately 
+stored in a Config object.
 
 ```python
 detector = DetectorParams(
@@ -237,13 +266,21 @@ pads = PadParams()
 config = Config(detector, electronics, pads)
 ```
 
-Note that by not passing any arguments to `PadParams` we are using the default pad description that is bundled with the package. See the [detector](./detector/index.md) guide for more details. For the output, we create a SpyralWriter object. This will take in the simulation data and convert it to match the format expected by the Spyral analysis. Note the final argument of the Writer; this is the maximum size of an individual file in events (here we've specified 5,000 events). The writer will then split our output up into many files, which will help later when trying to analyze the data with a framework like Spyral.
+Note that by not passing any arguments to `PadParams` we are using the default pad 
+description that is bundled with the package. See the [detector](./detector/index.md) 
+guide for more details. For the output, we create a SpyralWriter object. This will take
+in the simulation data and convert it to match the format expected by the Spyral 
+analysis. Note the final argument of the Writer; this is the maximum size of an 
+individual file in events (here we've specified 5,000 events). The writer will then 
+split our output up into many files, which will help later when trying to analyze the 
+data with a framework like Spyral.
 
 ```python
 writer = SpyralWriter(output_path, config, 5_000)
 ```
 
-Then, just like in the kinematics script we set up a main function and set it to be run when the script is processed
+Then, just like in the kinematics script we set up a main function and set it to be 
+run when the script is processed
 
 ```python
 def main():
@@ -257,7 +294,25 @@ if __name__ == "__main__":
     main()
 ```
 
-And just like that, we can now take our kinematic samples and apply detector effects!
+And just like that, we can now take our kinematic samples and apply detector effects! 
+Note that by default the simulation will apply detector effects and generate a point cloud
+containing *every single final product* from the kinematics simulation. If you only want
+to generate a point cloud containing specific particles you can utilize the `indices`
+argument of `run_simulation`:
+
+```python
+def main():
+    run_simulation(
+        config,
+        input_path,
+        writer,
+        indices=[2],
+    )
+```
+
+The `indices` argument is a list of particle indices to be included in the simulation.
+The indicies follow the same convention as the kinematics simulation: i.e. for a reaction
+a(b,c)d a=0, b=1, c=2, d=3. If you have a decay a(b,c)d->e+f e=4, f=5, and so on.
 
 ## More Details
 

src/attpc_engine/detector/simulator.py

@@ -56,7 +56,7 @@ def simulate(
     mass_numbers: np.ndarray,
     config: Config,
     rng: Generator,
-    indicies: list[int],
+    indices: list[int],
 ) -> tuple[np.ndarray, np.ndarray]:
     """Apply detector simulation to a kinematics event
 
@@ -78,7 +78,7 @@ def simulate(
         The detector simulation parameters
     rng: numpy.random.Generator
         A random number generator
-    indicies: list[int]
+    indices: list[int]
         The indicies in the list of nuclei which should be simulated.
         Typically this would be all final products of the reaction
 
@@ -93,7 +93,7 @@ def simulate(
     points = Dict.empty(
         key_type=types.int64, value_type=types.Tuple(types=[types.int64, types.int64])
     )
-    for idx in indicies:
+    for idx in indices:
         if proton_numbers[idx] == 0:
             continue
         nucleus = nuclear_map.get_data(proton_numbers[idx], mass_numbers[idx])
@@ -119,7 +119,7 @@ def run_simulation(
     config: Config,
     input_path: Path,
     writer: SimulationWriter,
-    indicies: list[int] | None = None,
+    indices: list[int] | None = None,
 ):
     """Run the detector simulation
 
@@ -134,7 +134,7 @@ def run_simulation(
         Path to HDF5 file containing kinematics
     writer: SimulationWriter
         An object which implements the SimulationWriter Protocol
-    indicies: list[int] | None
+    indices: list[int] | None
         List of which nuclei to include in the detector simulation. Nuclei are
         specified by index of which they occur in the kinematic arrays. For example,
         in a simple one step reaction, a(b,c)d 0=a, 1=b, 2=c, 3=d. For two step
@@ -150,8 +150,8 @@ def run_simulation(
 
     # Decide which nuclei to sim, either by user input or all reaction final products
     nuclei_to_sim = None
-    if indicies is not None:
-        nuclei_to_sim = indicies
+    if indices is not None:
+        nuclei_to_sim = indices
     else:
         # default nuclei to sim, all final outgoing particles
         nuclei_to_sim = [idx for idx in range(2, len(proton_numbers), 2)]

src/attpc_engine/kinematics/__init__.py

@@ -13,7 +13,7 @@
     ExcitationBreitWigner,
 )
 
-from .angle import PolarDistribution, PolarUniform
+from .angle import PolarDistribution, PolarUniform, PolarArbitrary
 
 from .reaction import Reaction, Decay
 
@@ -26,6 +26,7 @@
     "ExcitationUniform",
     "ExcitationBreitWigner",
     "PolarDistribution",
+    "PolarArbitrary",
     "PolarUniform",
     "Reaction",
     "Decay",

src/attpc_engine/kinematics/angle.py

@@ -26,7 +26,7 @@ def sample(self, rng: Generator) -> float:  # type: ignore
 
         Returns
         -------
-        float:
+        float
             The sampled angle in radians
         """
         pass
@@ -74,7 +74,79 @@ def sample(self, rng: Generator) -> float:
 
         Returns
         -------
-        float:
+        float
             The sampled angle in radians
         """
         return np.arccos(rng.uniform(self.cos_angle_min, self.cos_angle_max))
+
+
+class PolarArbitrary:
+    """An arbitrary, finite-precision polar angular distribution in the CM frame
+
+    Define an arbitrary probability distribution for the polar angle using an
+    array of uniformly-spaced angles and their probabilities, as well as the spacing
+    of the angle values. We then sample from the distribution and smear the output angle
+    within the spacing.
+
+    Note: If your distribution is defined using *any* of the pre-defined distributions,
+    you should favor using those over PolarArbitrary. Sampling of arbitrary functions
+    comes with a runtime performance penalty. That is: do not use this to sample from a
+    uniform distribution.
+
+    Parameters
+    ----------
+    angles: numpy.ndarray
+        The array of *lower* angle values. Each angle corresponds to a bin covering
+        angle + bin_width. Angles should be in units of radians.
+    probabilities: numpy.ndarray
+        The probability of a given angle bin. Should sum to 1.0
+    angle_bin_width: float
+        The width of the angle bins in radians
+
+    Attributes
+    ----------
+    angle_width: float
+        The angle bin width in radians
+    probs: numpy.ndarray
+        The probability of a given angle bin
+    angles: numpy.ndarray
+        The array of *lower* angle values. Each angle corresponds to a bin covering
+        angle + bin_width. Angles should be in units of radians.
+
+    Methods
+    -------
+    sample(rng)
+        Sample the distribution, returning a polar angle in radians
+    """
+
+    def __init__(
+        self,
+        angles: np.ndarray,
+        probabilities: np.ndarray,
+        angle_bin_width: float,
+    ):
+        if np.sum(probabilities) > 1.0:
+            raise ValueError(
+                f"The sum of the probabilities passed to PolarArbitrary should be 1.0. Yours sum to {np.sum(probabilities)}"
+            )
+        self.angle_width = angle_bin_width
+        self.probs = probabilities
+        self.angles = angles
+
+    def sample(self, rng: Generator) -> float:
+        """Sample the distribution, returning a polar angle in radians
+
+        Parameters
+        ----------
+        rng: numpy.random.Generator
+            The random number generator
+
+        Returns
+        -------
+        float
+            The sampled angle in radians
+        """
+        return (
+            rng.choice(self.angles, p=self.probs)
+            + rng.uniform(0.0, 1.0) * self.angle_width
+        )