adding find_voltages script (#9)

* adding testing notebook

* add test files

* update test ipynb

* Add a file containing the function to calculate the voltage corrections

* updating script

* changing script file

* add tests for find_votlages script

* remove temp test files

* add dependencies to pyproject.toml

---------

Co-authored-by: Luke Fiddy <[email protected]>
Co-authored-by: root <[email protected]>
Co-authored-by: Luke Fiddy <[email protected]>

pyproject.toml

@@ -12,7 +12,9 @@ classifiers = [
     "Programming Language :: Python :: 3.12",
 ]
 description = "package to calibrate bimorph mirrors to focus a beamline to a target position"
-dependencies = [] # Add project dependencies here, e.g. ["click", "numpy"]
+dependencies = [
+    "numpy",
+] # Add project dependencies here, e.g. ["click", "numpy"]
 dynamic = ["version"]
 license.file = "LICENSE"
 readme = "README.md"
@@ -47,7 +49,7 @@ version_file = "src/bimorph_mirror_analysis/_version.py"
 
 [tool.pyright]
 typeCheckingMode = "standard"
-reportMissingImports = false # Ignore missing stubs in imported modules
+reportMissingImports = false  # Ignore missing stubs in imported modules
 
 [tool.pytest.ini_options]
 # Run pytest with all our checkers, and don't spam us with massive tracebacks on error

src/bimorph_mirror_analysis/maths.py

@@ -0,0 +1,51 @@
+import numpy as np
+
+
+def find_voltages(
+    data: np.ndarray, v: float, baseline_voltage_scan: int = 0
+) -> np.ndarray:
+    """Calculate voltage corrections to apply to bimorph.
+
+    Given a matrix of beamline centroid data, with columns of beamline scans at
+    different actuator voltages and rows of slit positions, calculate the necessary
+    voltages corrections to achive the target centroid position.
+
+    Args:
+        data: A matrix of beamline centroid data, with rows of different slit positions
+            and columns of pencil beam scans at different actuator voltages
+        v: The voltage increment applied to the actuators between pencil beam scans
+        baseline_voltage_scan: The pencil beam scan to use as the baseline for the
+            centroid calculation. 0 is the first scan, 1 is the second scan, etc.
+            -1 can be used for the last scan and -2 for the second to last scan etc.
+
+    Returns:
+        An array of voltage corrections required to move the centroid of each pencil
+        beam scan to the target position.
+    """
+
+    if not isinstance(baseline_voltage_scan, int):
+        raise TypeError("baseline_voltage_scan must be an integer")
+    if baseline_voltage_scan < -data.shape[1] or baseline_voltage_scan >= data.shape[1]:
+        raise IndexError(
+            f"baseline_voltage_scan is out of range, it must be between\
+                  {-1*data.shape[1]} and {data.shape[1]-1}"
+        )
+
+    responses = np.diff(
+        data, axis=1
+    )  # calculate the response of each actuator by subtracting previous pencil beam
+
+    H = responses / v  # response per unit charge
+    H = np.hstack((np.ones((H.shape[0], 1)), H))  # add columns of 1's to the left of H
+    H_inv = np.linalg.pinv(H)  # calculate the Moore-Penrose pseudo inverse of H
+
+    baseline_voltage_beamline_positions = data[:, baseline_voltage_scan]
+
+    target = np.mean(baseline_voltage_beamline_positions)
+    Y = target - baseline_voltage_beamline_positions
+
+    voltage_corrections = np.matmul(
+        H_inv, Y
+    )  # calculate the voltage required to move the centroid to the target position
+
+    return voltage_corrections[1:]  # return the voltages

tests/data/16_actuator_output.txt

@@ -0,0 +1,16 @@
+3.754035246300606445e+02
+3.173552610735841313e+02
+-8.871216249116981700e+00
+9.823656308341850263e+01
+1.240974691415672169e+02
+-1.921633058651442241e+02
+-1.580320457126290767e-01
+-6.235823640873921647e+01
+2.134018588040029556e+02
+-9.588821999790660300e+01
+4.260153919587435922e+01
+-6.367412313316757633e+01
+1.210759471343801295e+02
+-4.776920420905791786e+01
+1.272677250651050542e+02
+7.358411715637718942e+01

tests/data/8_actuator_output.txt

@@ -0,0 +1,8 @@
+1.038570862242364683e+02
+-7.851295652142359138e+00
+3.810159077901251123e+01
+2.169648645139839260e+01
+-8.453195968470640764e+01
+5.327945430609133126e+01
+-6.701602106542451054e+01
+2.399563025205714073e+00

tests/script_test.py

@@ -0,0 +1,34 @@
+import numpy as np
+import pytest
+
+from bimorph_mirror_analysis.maths import find_voltages
+
+# from bimorph_mirror_analysis.maths import find_voltages
+
+
+@pytest.mark.parametrize(
+    ["input_path", "output_path"],
+    [
+        ["tests/data/8_actuator_data.txt", "tests/data/8_actuator_output.txt"],
+        ["tests/data/16_actuator_data.txt", "tests/data/16_actuator_output.txt"],
+    ],
+)
+def test_find_voltages_correct_output(input_path, output_path):
+    data = np.loadtxt(input_path, delimiter=",")
+    v = -100
+    expected_output = np.loadtxt(output_path, delimiter=",")
+    np.testing.assert_almost_equal(
+        find_voltages(data, v, baseline_voltage_scan=-1), expected_output
+    )
+
+
+def test_find_voltages_index_error_throw():
+    data = np.loadtxt("tests/data/8_actuator_data.txt", delimiter=",")
+    with pytest.raises(IndexError):
+        find_voltages(data, -100, baseline_voltage_scan=12)
+
+
+def test_find_voltages_type_error_throw():
+    data = np.loadtxt("tests/data/8_actuator_data.txt", delimiter=",")
+    with pytest.raises(TypeError):
+        find_voltages(data, -100, baseline_voltage_scan="12")  # type: ignore