Merge pull request #96 from LemurPwned/feat/custom-drivers

driver updates

Author: LemurPwned

core/drivers.hpp

@@ -10,9 +10,10 @@
 #define _USE_MATH_DEFINES
 #include "cvector.hpp" // for CVector
 #include <cmath>       // for M_PI
-#include <stdexcept>   // for runtime_error
-#include <utility>     // for move
-#include <vector>      // for vector
+#include <pybind11/pybind11.h>
+#include <stdexcept> // for runtime_error
+#include <utility>   // for move
+#include <vector>    // for vector
 
 enum UpdateType {
   constant,
@@ -23,7 +24,8 @@ enum UpdateType {
   halfsine,
   trapezoid,
   gaussimpulse,
-  gaussstep
+  gaussstep,
+  custom = 100
 };
 
 template <typename T> class Driver {
@@ -65,24 +67,23 @@ template <typename T> class ScalarDriver : public Driver<T> {
 private:
   T edgeTime = 0;
   T steadyTime = 0;
+  pybind11::function m_callback;
 
 protected:
   T stepUpdate(T amplitude, T time, T timeStart, T timeStop) {
     if (time >= timeStart && time <= timeStop) {
       return amplitude;
-    } else {
-      return 0.0;
     }
+    return 0.0;
   }
   T pulseTrain(T amplitude, T time, T period, T cycle) {
     const int n = static_cast<int>(time / period);
     const T dT = cycle * period;
     const T nT = n * period;
     if (nT <= time && time <= (nT + dT)) {
       return amplitude;
-    } else {
-      return 0;
     }
+    return 0.0;
   }
 
   T trapezoidalUpdate(T amplitude, T time, T timeStart, T edgeTime,
@@ -110,7 +111,8 @@ template <typename T> class ScalarDriver : public Driver<T> {
   explicit ScalarDriver(UpdateType update = constant, T constantValue = 0,
                         T amplitude = 0, T frequency = -1, T phase = 0,
                         T period = -1, T cycle = -1, T timeStart = -1,
-                        T timeStop = -1, T edgeTime = -1, T steadyTime = -1)
+                        T timeStop = -1, T edgeTime = -1, T steadyTime = -1,
+                        pybind11::function m_callback = pybind11::function())
       : Driver<T>(update, constantValue, amplitude, frequency, phase, period,
                   cycle, timeStart, timeStop) {
     this->edgeTime = edgeTime;
@@ -122,6 +124,7 @@ template <typename T> class ScalarDriver : public Driver<T> {
       throw std::runtime_error(
           "Selected sine driver type but frequency was not set");
     }
+    this->m_callback = m_callback;
   }
 
   /**
@@ -241,6 +244,21 @@ template <typename T> class ScalarDriver : public Driver<T> {
                         t0, -1, sigma);
   }
 
+  static ScalarDriver getCustomDriver(pybind11::function callback) {
+    if (!callback) {
+      throw std::runtime_error("Callback function is not set");
+    }
+    // check if the callback is callable has one argument
+    // Check if the callback function has exactly one argument
+    // Using cast to int to avoid type mismatch error
+    if (pybind11::cast<int>(callback.attr("__code__").attr("co_argcount")) !=
+        1) {
+      throw std::runtime_error("Callback function must have one argument");
+    }
+
+    return ScalarDriver(custom, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, callback);
+  }
+
   T getCurrentScalarValue(T &time) override {
     T returnValue = this->constantValue;
     if (this->update == pulse) {
@@ -273,6 +291,15 @@ template <typename T> class ScalarDriver : public Driver<T> {
           0.5 * this->amplitude *
           (1 + std::erf((time - this->timeStart) / (sqrt(2) * this->edgeTime)));
       returnValue += gaussStep;
+    } else if (this->update == custom) {
+      // If it is, call the Python function
+      pybind11::gil_scoped_acquire gil;
+      try {
+        return pybind11::cast<double>(m_callback(time));
+      } catch (pybind11::error_already_set &e) {
+        std::cerr << "Error in Python callback: " << e.what() << std::endl;
+        throw std::runtime_error("Error in Python callback");
+      }
     }
     return returnValue;
   }

docs/api/drivers.md

@@ -45,14 +45,11 @@ from cmtj import (
 import numpy as np
 
 
-class MyDriver(ScalarDriver):
-    def getCurrentScalarValue(self, time: float) -> float:
-        return time * np.random.choice([-1, 1])
+def my_custom_function(time: float) -> float:
+    return time * np.random.choice([-1, 1])
 
-
-driver = MyDriver()
-for i in range(10):
-    print(driver.getCurrentScalarValue(i * 1e-9))
+# Create a driver with this function
+driver = ScalarDriver.getCustomDriver(my_custom_function)
 
 demag = [CVector(0, 0, 0), CVector(0, 0, 0), CVector(0, 0, 1)]
 layer = Layer(

python/cmtj.cpp

@@ -147,6 +147,18 @@ PYBIND11_MODULE(cmtj, m) {
           .value("DormandPrice", DORMAND_PRICE)
           .export_values();
 
+     py::enum_<UpdateType>(m, "UpdateType")
+          .value("constant", constant)
+          .value("pulse", pulse)
+          .value("sine", sine)
+          .value("step", step)
+          .value("posine", posine)
+          .value("halfsine", halfsine)
+          .value("trapezoid", trapezoid)
+          .value("gaussimpulse", gaussimpulse)
+          .value("gaussstep", gaussstep)
+          .value("custom", custom)
+          .export_values();
      // Driver Class
      py::class_<DScalarDriver>(m, "ScalarDriver")
           .def(py::init<>())
@@ -174,7 +186,9 @@ PYBIND11_MODULE(cmtj, m) {
                "amplitude"_a, "t0"_a, "sigma"_a)
           .def_static("getGaussianStepDriver",
                &DScalarDriver::getGaussianStepDriver, "constantValue"_a,
-               "amplitude"_a, "t0"_a, "sigma"_a);
+               "amplitude"_a, "t0"_a, "sigma"_a)
+          .def_static("getCustomDriver", &DScalarDriver::getCustomDriver,
+               "callback"_a);
 
      py::class_<DNullDriver, DScalarDriver>(m, "NullDriver")
           .def(py::init<>())

tests/test_drivers.py

@@ -1,6 +1,8 @@
-from cmtj import AxialDriver, CVector, Junction, Layer
+from cmtj import AxialDriver, CVector, Junction, Layer, ScalarDriver
 from cmtj import constantDriver, sineDriver
 import pytest
+import math
+
 
 def test_cvector_operators():
     vec1 = (1.0, 2.0, 3.0)
@@ -84,3 +86,51 @@ def test_junction_with_driver():
         AxialDriver(constantDriver(0), constantDriver(0), sineDriver(0, 1e3, 7e9, 0)),
     )
     junction.runSimulation(10e-9, 1e-13, 1e-13)
+
+
+def test_custom_driver():
+    def my_custom_function(time: float) -> float:
+        return math.sqrt(time)
+
+    driver = ScalarDriver.getCustomDriver(my_custom_function)
+    assert driver.getCurrentScalarValue(1e-9) == math.sqrt(1e-9)
+    assert driver.getCurrentScalarValue(3e-9) == math.sqrt(3e-9)
+
+    # Test with zero time value
+    assert my_custom_function(0) == 0, "Zero time value should return 0"
+
+    # Test error condition: Passing None should raise a TypeError (since None cannot be multiplied by an int)
+    with pytest.raises(TypeError):
+        my_custom_function(None)
+
+
+def test_custom_driver_exception():
+    # Test error condition: A callback that raises an exception
+    def faulty_callback(time: float) -> float:
+        raise ValueError("Intentional exception")
+
+    driver = ScalarDriver.getCustomDriver(faulty_callback)
+    with pytest.raises(RuntimeError):
+        driver.getCurrentScalarValue(1e-9)
+
+
+def test_no_argument_callback():
+    # Test error condition: A callback that raises an exception
+    def faulty_callback():
+        raise ValueError("Intentional exception")
+
+    with pytest.raises(RuntimeError):
+        driver = ScalarDriver.getCustomDriver(faulty_callback)
+
+
+def test_custom_on_junction(single_layer_mtj):
+    junction, _ = single_layer_mtj
+
+    def my_custom_function(time: float) -> float:
+        return 12345
+
+    driver = ScalarDriver.getCustomDriver(my_custom_function)
+    junction.setLayerAnisotropyDriver("all", driver)
+    junction.runSimulation(10e-9, 1e-13, 1e-13, calculateEnergies=True)
+    log = junction.getLog()
+    assert all(x == 12345 for x in log["free_K"])