Add mass stopping power calculation and bindings for protons in liqui…

src/particles/ions/ion.cpp

@@ -1,6 +1,8 @@
 #include "ion.h"
-#include "../particles.h"
 
+#include <stdexcept>
+
+#include "../particles.h"
 
 Ion::Ion(long long pdg_code) {
   nb::object p;

src/particles/ions/ion.h

@@ -1,4 +1,7 @@
 #pragma once
+#include <stdexcept>
+#include <string>
+
 #include "../particles.h"
 
 /**
@@ -18,6 +21,18 @@ class Ion : public Particle {
    */
   Ion();
 
-  long Z;  /**< Atomic number of the ion. */
-  long A;  /**< Mass number of the ion. */
+  long Z; /**< Atomic number of the ion. */
+  long A; /**< Mass number of the ion. */
+
+  /**
+   * @brief libamtrack particle number (1000*Z + A).
+   *
+   * Used when passing ions into calculation functions such as stopping power.
+   */
+  long get_particle_no() const {
+    if (Z < 1 || A < 1) {
+      throw std::invalid_argument("Ion is missing a valid Z and A");
+    }
+    return AT_particle_no_from_Z_and_A_single(Z, A);
+  }
 };

src/particles/particle_arguments.h

@@ -0,0 +1,118 @@
+#ifndef PARTICLE_ARGUMENTS_H
+#define PARTICLE_ARGUMENTS_H
+
+#include <nanobind/nanobind.h>
+#include <nanobind/ndarray.h>
+
+#include <stdexcept>
+#include <string>
+
+#include "ions/ion.h"
+#include "particles.h"
+
+namespace nb = nanobind;
+
+/**
+ * @brief Validates a particle argument, recursing into lists and NumPy arrays.
+ *
+ * Stopping-power calculations require an ion (Z and A). Elementary particles
+ * such as neutron and electron are rejected. Integer values are treated as
+ * libamtrack particle numbers (1000*Z + A). Boolean values are not accepted.
+ *
+ * @throws std::invalid_argument (ValueError) if a particle number does not
+ *         correspond to a known element, or an Ion is missing Z/A.
+ * @throws nb::type_error if the argument is not an int, Ion, list, or NumPy array.
+ */
+inline void validate_particle_argument(const nb::object& argument) {
+  if (nb::isinstance<nb::list>(argument)) {
+    nb::list values = nb::cast<nb::list>(argument);
+    for (size_t i = 0; i < values.size(); ++i) {
+      validate_particle_argument(values[i]);
+    }
+    return;
+  }
+
+  if (nb::isinstance<nb::ndarray<>>(argument)) {
+    validate_particle_argument(argument.attr("tolist")());
+    return;
+  }
+
+  if (nb::isinstance<Ion>(argument)) {
+    nb::cast<Ion>(argument).get_particle_no();
+    return;
+  }
+
+  if (nb::isinstance<Particle>(argument)) {
+    throw nb::type_error(
+        "particle must be an Ion (stopping power requires Z and A); "
+        "elementary particles are not supported");
+  }
+
+  long particle_no;
+  if (PyBool_Check(argument.ptr())) {
+    throw nb::type_error("particle must be an int (not bool), an Ion, or a list / NumPy array of either");
+  } else if (nb::isinstance<nb::int_>(argument)) {
+    particle_no = nb::cast<long>(argument);
+  } else {
+    throw nb::type_error("particle must be an int, an Ion, or a list / NumPy array of either");
+  }
+
+  const long Z = AT_Z_from_particle_no_single(particle_no);
+  const long A = AT_A_from_particle_no_single(particle_no);
+  if (A < 1) {
+    throw std::invalid_argument("invalid particle number: " + std::to_string(particle_no));
+  }
+  bool known_Z = false;
+  for (int i = 0; i < AT_Particle_Data.n; ++i) {
+    if (AT_Particle_Data.Z[i] == Z) {
+      known_Z = true;
+      break;
+    }
+  }
+  if (!known_Z) {
+    throw std::invalid_argument("invalid particle number: " + std::to_string(particle_no));
+  }
+}
+
+/**
+ * @brief Convert a particle argument to libamtrack particle number(s).
+ *
+ * Ion objects become 1000*Z + A. Integers are passed through. Lists and
+ * NumPy arrays are converted elementwise.
+ */
+inline nb::object parse_particle_argument(const nb::object& argument) {
+  if (nb::isinstance<Ion>(argument)) {
+    return nb::cast(nb::cast<Ion>(argument).get_particle_no());
+  }
+
+  if (nb::isinstance<Particle>(argument)) {
+    throw nb::type_error(
+        "particle must be an Ion (stopping power requires Z and A); "
+        "elementary particles are not supported");
+  }
+
+  if (PyBool_Check(argument.ptr())) {
+    throw nb::type_error("particle must be an integer (not bool), Ion, list, or NumPy array");
+  }
+
+  if (nb::isinstance<nb::int_>(argument)) {
+    return argument;
+  }
+
+  if (nb::isinstance<nb::list>(argument)) {
+    nb::list values = nb::cast<nb::list>(argument);
+    nb::list parsed_values;
+    for (size_t i = 0; i < values.size(); ++i) {
+      parsed_values.append(parse_particle_argument(values[i]));
+    }
+    return parsed_values;
+  }
+
+  if (nb::isinstance<nb::ndarray<>>(argument)) {
+    return parse_particle_argument(argument.attr("tolist")());
+  }
+
+  throw nb::type_error("particle must be an integer, Ion, list, or NumPy array");
+}
+
+#endif  // PARTICLE_ARGUMENTS_H

src/particles/particles.cpp

@@ -1,14 +1,13 @@
 #include "particles.h"
-#include "construct_utils.h"
-#include "ions/ion.h"
-#include "ions/ions.h"
 
 #include <algorithm>
 #include <cctype>
 #include <optional>
 #include <stdexcept>
 
-
+#include "construct_utils.h"
+#include "ions/ion.h"
+#include "ions/ions.h"
 
 /**
  * @brief Default constructor for Particle.
@@ -42,6 +41,9 @@ Particle::Particle(long long pdg_code) {
   }
 }
 
+long Particle::get_id() const {
+  return id;
+}
 
 /**
  * @brief Constructs a Particle from a string representation.
@@ -72,13 +74,14 @@ Particle::Particle(long long pdg_code) {
  * @throws std::invalid_argument or nb::value_error on invalid input.
  */
 nb::object from_string(const std::string& name) {
-
-  // logic for creating either Particle or Ion. Should be changed with further development of the class hierarchy. For now, we just check for special cases of proton, neutron and electron, which are not really ions, but we want to support them as special particles.
+  // logic for creating either Particle or Ion. Should be changed with further development of the class hierarchy. For
+  // now, we just check for special cases of proton, neutron and electron, which are not really ions, but we want to
+  // support them as special particles.
 
   if (name == "neutron" || name == "electron") {
     return nb::cast(create_particle(name));
   }
-  
+
   return nb::cast(create_ion(name));
 }
 
@@ -93,10 +96,10 @@ Particle create_particle(const std::string& name) {
   p.element_name = name;
   p.element_acronym = name[0];
   p.pdg = pdg_for_nonions[name];
-  
+
   return p;
 }
- 
+
 /**
  * @brief Create an Ion from element or isotope string.
  *
@@ -120,7 +123,6 @@ Ion create_ion(const std::string& name) {
     A_ = 4;
   }
 
-
   int it = -1;
   for (int i = 0; i < data.n; ++i) {
     if (to_lower_case(data.element_acronym[i]) == to_lower_case(symbol)) {
@@ -141,13 +143,14 @@ Ion create_ion(const std::string& name) {
   acronym = data.element_acronym[it];
 
   auto rng = isotope_A_range.find(acronym);
-  if (rng == isotope_A_range.end())
-      throw nb::value_error(("No isotope data for element " + acronym).c_str());
+  if (rng == isotope_A_range.end()) throw nb::value_error(("No isotope data for element " + acronym).c_str());
   auto [l, h] = rng->second;
 
   if (A_ != -1) {
     if (A_ < l || A_ > h) {
-      throw nb::value_error(("Invalid mass number A=" + std::to_string(A_) + " for element " + acronym + " (valid range: " + std::to_string(l) + "-" + std::to_string(h) + ")").c_str());
+      throw nb::value_error(("Invalid mass number A=" + std::to_string(A_) + " for element " + acronym +
+                             " (valid range: " + std::to_string(l) + "-" + std::to_string(h) + ")")
+                                .c_str());
     }
   } else if (A_ == -1) {
     A_ = most_popular_iso_A[acronym];
@@ -159,11 +162,10 @@ Ion create_ion(const std::string& name) {
   p.Z = data.Z[it];
   p.A = A_;
   p.pdg = calculatePDG(p.Z, p.A, 0, 0);
-  
+
   return p;
 }
 
-
 /**
  * @brief Create an Ion from atomic number Z and mass number A.
  *
@@ -182,11 +184,8 @@ Ion from_ZA(long long Z, long long A) {
         Ion ion = nb::cast<Ion>(result);
         return ion;
       } catch (const nb::cast_error&) {
-        throw std::invalid_argument(
-          "Expected Ion type for Z=" + std::to_string(Z) + 
-          ", A=" + std::to_string(A) + 
-          ", but got Particle instead"
-        );
+        throw std::invalid_argument("Expected Ion type for Z=" + std::to_string(Z) + ", A=" + std::to_string(A) +
+                                    ", but got Particle instead");
       }
     }
   }
@@ -205,7 +204,7 @@ nb::object from_pdg(long long pdg_code) {
   if (pdg_code == 2212) {
     return nb::cast(from_ZA(1, 1));
   }
-  
+
   if (pdg_code == 2112) {
     return nb::cast(from_string("neutron"));
   }
@@ -215,7 +214,6 @@ nb::object from_pdg(long long pdg_code) {
 
   const auto& data = AT_Particle_Data;
 
-
   if (pdg_code < 1000000000) {
     throw nb::value_error(("PDG code does not correspond to a nucleus: " + std::to_string(pdg_code)).c_str());
   }
@@ -257,18 +255,16 @@ std::vector<std::string> get_acronyms() {
   return acronyms;
 }
 
-
 /**
  * @brief Short string representation.
  */
 std::string Particle::str() const {
-    return element_acronym;  
+  return element_acronym;
 }
 
 /**
  * @brief Detailed string representation.
  */
 std::string Particle::repr() const {
-    return "Particle(name=\"" + element_name +
-           "\", acronym=\"" + element_acronym + "\")";
+  return "Particle(name=\"" + element_name + "\", acronym=\"" + element_acronym + "\")";
 }