customizeable memristor model

started trying to implement a customizeable memristor model to pass to memTorch, porting a model I got from kaevin too fit the pyTorch framework

4 files changed, 403 insertions, 14 deletions

diff --git a/neurobench_testing/custom_memristor_model.py b/neurobench_testing/custom_memristor_model.py
new file mode 100644
index 0000000..8bdcde0
--- /dev/null
+++ b/neurobench_testing/custom_memristor_model.py
@@ -0,0 +1,42 @@
+import torch
+from memtorch.bh.memristor.Memristor import Memristor
+
+class MemtorchMemristor(Memristor):
+    def __init__(
+            self,
+            k_off       = 1.0,      # switching rate for off state
+            k_on        = -1.0,     # switching rate for on state
+            alpha_off   = 5,        # exponent controlling nonlinearity
+            alpha_on    = 5,        # exponent controlling nonlinearity
+            i_off       = 0.5e-3,   # threshhold current to trigger off state
+            i_on        = 0.5e-3,   # threshold current to trigger on state
+            r_on        = 1e3,      # maximum resistance
+            r_off       = 10e3,     # minimum resistance
+            p           = 2,        # window function exponent
+            **kwargs
+    ):
+        #initializing base memristor class
+        super(MemtorchMemristor, self).__init__(r_off=r_off, r_on=r_on, **kwargs)
+
+        # hyper parameters
+        self.k_off = k_off
+        self.k_on = k_on
+        self.alpha_off = alpha_off
+        self.alpha_on = alpha_on
+        self.i_on = i_on
+        self.i_off = i_off
+        self.p = p
+
+        # makes sure w starts in valid state
+        if not hasattr(self, 'w'):
+            self.w = torch.tensor(0.5)
+
+
+    """
+    Updates w and computes new resistance
+   
+    """ 
+    def step(self, v, dt):
+        i = v / self.r_curr
+
+
diff --git a/neurobench_testing/memTorch_testing.py b/neurobench_testing/memTorch_testing.py
index a44fd98..5c3d828 100644
--- a/neurobench_testing/memTorch_testing.py
+++ b/neurobench_testing/memTorch_testing.py
@@ -1,14 +1,19 @@
-from memtorch import memristor
+"""
+This is a program to test out using memTorch with neurobench.
+June 17th, 2026
+Author: Tanner Robison, 
+Teuscher Lab
+"""
 import torch
 import torch.nn as nn
 import snntorch as snn
+from snntorch import surrogate
 
 from torch.utils.data import DataLoader
 
 from neurobench.models import SNNTorchModel
 from neurobench.benchmarks import Benchmark, benchmark
 from neurobench.datasets import SpeechCommands
-
 from neurobench.metrics.workload import (
         ActivationSparsity,
         SynapticOperations,
@@ -29,8 +34,7 @@ from memtorch.map.Parameter import naive_map
 from memtorch.bh.memristor import VTEAM
 
 beta = 0.9
-
-class SimpleSNN(nn.Module):
+class SNN(nn.Module):
     def __init__(self):
         super().__init__()
     
@@ -54,36 +58,46 @@ class SimpleSNN(nn.Module):
         return spk2, mem2
 
 device = torch.device("cpu")
-net = SimpleSNN().to(device)
+spike_grad = surrogate.fast_sigmoid()
+net = nn.Sequential(
+        nn.Flatten(),
+        nn.Linear(20, 256),
+        snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True),
+        nn.Linear(256, 256),
+        snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True),
+        nn.Linear(256, 256),
+        snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True),
+        nn.Linear(256, 35),
+        snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True, output=True),
+)
 
 #memristor patch
 reference_memristor = VTEAM()
 
-print("Patching model to memristive crossbar")
+net.load_state_dict(torch.load("examples/gsc/model_data/s2s_gsc_snntorch", map_location=device))
+
 patched_net = patch_model(
         copy.deepcopy(net),
         memristor_model=reference_memristor,
-        memristor_model_params={},
+        memristor_model_params={'time_series_resolution': 1e-8},
         mapping_routine=naive_map,
         transistor=True,
+        tile_shape=(128, 128),
         ADC_resolution=8,
-        use_bindings=False
+        use_bindings=True
 )
 
-model = SNNTorchModel(patched_net)
-
 static_metrics = [Footprint, ConnectionSparsity]
 workload_metrics = [ActivationSparsity, SynapticOperations, ClassificationAccuracy]
 
 # data loader here maybe??
 test_set = SpeechCommands(path="data/SpeechCommands/", subset="testing")
-test_set_loader = DataLoader(test_set, batch_size=50, shuffle=True)
+test_set_loader = DataLoader(test_set, batch_size=500, shuffle=True)
 
 pre_processor = [S2SPreProcessor(device=device)]
 post_processor = [ChooseMaxCount()]
 
-# print("Layer 1 Beta:", net[2].beta)
-
+model = SNNTorchModel(net)
 benchmark = Benchmark(
         model,
         test_set_loader,
@@ -93,7 +107,46 @@ benchmark = Benchmark(
 )
 
 results = benchmark.run()
-print(results)
+print("\n\n----- IDEAL BENCHMARK -----")
+print(f"Footprint: {results['Footprint']}")
+print(f"Connection Sparsity: {results['ConnectionSparsity']}")
+print(f"Activation Sparsity: {results['ActivationSparsity']}")
+print(f"Synaptic Operations: {results['SynapticOperations']}")
+print(f"Classification Accuracy: {results['ClassificationAccuracy']}")
+
+#energy calculations
+ENERGY_PER_MAC = 0.9e-12
+ENERGY_PER_AC = 0.1e-12
+
+macs = results['SynapticOperations']['Effective_MACs']
+acs = results['SynapticOperations']['Effective_ACs']
+
+total_energy = (macs * ENERGY_PER_MAC) + (acs * ENERGY_PER_AC)
+
+print("Energy Report:")
+print(f"Total Operations: {macs} MACS, {acs} acs ")
+print(f"Calculated Energy cost: {total_energy} joules per batch")
+
+
+model = SNNTorchModel(patched_net)
+benchmark = Benchmark(
+        model,
+        test_set_loader,
+        pre_processor,
+        post_processor,
+        [static_metrics, workload_metrics]
+)
+
+with torch.no_grad(): #makes sure its in inference mode
+                      #otherwise you get memory leaks
+    results = benchmark.run()
+    
+print("\n\n----- MEMRISTOR BENCHMARK -----")
+print(f"Footprint: {results['Footprint']}")
+print(f"Connection Sparsity: {results['ConnectionSparsity']}")
+print(f"Activation Sparsity: {results['ActivationSparsity']}")
+print(f"Synaptic Operations: {results['SynapticOperations']}")
+print(f"Classification Accuracy: {results['ClassificationAccuracy']}")
 
 #energy calculations
 ENERGY_PER_MAC = 0.9e-12
diff --git a/neurobench_testing/memTorch_testing_custom_memristor.py b/neurobench_testing/memTorch_testing_custom_memristor.py
new file mode 100644
index 0000000..64f338e
--- /dev/null
+++ b/neurobench_testing/memTorch_testing_custom_memristor.py
@@ -0,0 +1,160 @@
+from memtorch import memristor
+import torch
+import torch.nn as nn
+import snntorch as snn
+from snntorch import surrogate
+
+from torch.utils.data import DataLoader
+
+from neurobench.models import SNNTorchModel
+from neurobench.benchmarks import Benchmark, benchmark
+from neurobench.datasets import SpeechCommands
+
+from neurobench.metrics.workload import (
+        ActivationSparsity,
+        SynapticOperations,
+        ClassificationAccuracy,
+)
+
+from neurobench.metrics.static import (
+        Footprint,
+        ConnectionSparsity,
+)
+
+from neurobench.processors.preprocessors import S2SPreProcessor
+from neurobench.processors.postprocessors import ChooseMaxCount
+
+import copy
+from memtorch.mn.Module import patch_model
+from memtorch.map.Parameter import naive_map 
+from memtorch.bh.memristor import VTEAM
+
+from kaevin_memristor import TEAMMemristor
+
+beta = 0.9
+class SNN(nn.Module):
+    def __init__(self):
+        super().__init__()
+    
+        #standard layers
+        self.fc1 = nn.Linear(20, 128)
+        self.fc2 = nn.Linear(128, 35)
+
+        #Spiking neurons
+        self.lif1 = snn.Leaky(beta=beta, init_hidden=True)
+        self.lif2 = snn.Leaky(beta=beta, init_hidden=True, output=True)
+
+    def forward(self, x):
+        x = x.view(x.size(0), -1)
+
+        cur1 = self.fc1(x)
+        spk1 = self.lif1(cur1)
+
+        cur2 = self.fc2(spk1)
+        spk2, mem2 = self.lif2(cur2)
+
+        return spk2, mem2
+
+device = torch.device("cpu")
+spike_grad = surrogate.fast_sigmoid()
+net = nn.Sequential(
+        nn.Flatten(),
+        nn.Linear(20, 256),
+        snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True),
+        nn.Linear(256, 256),
+        snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True),
+        nn.Linear(256, 256),
+        snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True),
+        nn.Linear(256, 35),
+        snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True, output=True),
+)
+
+#memristor patch
+reference_memristor = TEAMMemristor
+
+net.load_state_dict(torch.load("examples/gsc/model_data/s2s_gsc_snntorch", map_location=device))
+
+patched_net = patch_model(
+        copy.deepcopy(net),
+        memristor_model=reference_memristor,
+        memristor_model_params={'time_series_resolution': 1e-8},
+        mapping_routine=naive_map,
+        transistor=True,
+        tile_shape=(128, 128),
+        ADC_resolution=8,
+        use_bindings=True
+)
+
+static_metrics = [Footprint, ConnectionSparsity]
+workload_metrics = [ActivationSparsity, SynapticOperations, ClassificationAccuracy]
+
+# data loader here maybe??
+test_set = SpeechCommands(path="data/SpeechCommands/", subset="testing")
+test_set_loader = DataLoader(test_set, batch_size=500, shuffle=True)
+
+pre_processor = [S2SPreProcessor(device=device)]
+post_processor = [ChooseMaxCount()]
+
+model = SNNTorchModel(net)
+benchmark = Benchmark(
+        model,
+        test_set_loader,
+        pre_processor,
+        post_processor,
+        [static_metrics, workload_metrics]
+)
+
+results = benchmark.run()
+print("\n\n----- IDEAL BENCHMARK -----")
+print(f"Footprint: {results['Footprint']}")
+print(f"Connection Sparsity: {results['ConnectionSparsity']}")
+print(f"Activation Sparsity: {results['ActivationSparsity']}")
+print(f"Synaptic Operations: {results['SynapticOperations']}")
+print(f"Classification Accuracy: {results['ClassificationAccuracy']}")
+
+#energy calculations
+ENERGY_PER_MAC = 0.9e-12
+ENERGY_PER_AC = 0.1e-12
+
+macs = results['SynapticOperations']['Effective_MACs']
+acs = results['SynapticOperations']['Effective_ACs']
+
+total_energy = (macs * ENERGY_PER_MAC) + (acs * ENERGY_PER_AC)
+
+print("Energy Report:")
+print(f"Total Operations: {macs} MACS, {acs} acs ")
+print(f"Calculated Energy cost: {total_energy} joules per batch")
+
+
+model = SNNTorchModel(patched_net)
+benchmark = Benchmark(
+        model,
+        test_set_loader,
+        pre_processor,
+        post_processor,
+        [static_metrics, workload_metrics]
+)
+
+with torch.no_grad(): #makes sure its in inference mode
+                      #otherwise you get memory leaks
+    results = benchmark.run()
+    
+print("\n\n----- MEMRISTOR BENCHMARK -----")
+print(f"Footprint: {results['Footprint']}")
+print(f"Connection Sparsity: {results['ConnectionSparsity']}")
+print(f"Activation Sparsity: {results['ActivationSparsity']}")
+print(f"Synaptic Operations: {results['SynapticOperations']}")
+print(f"Classification Accuracy: {results['ClassificationAccuracy']}")
+
+#energy calculations
+ENERGY_PER_MAC = 0.9e-12
+ENERGY_PER_AC = 0.1e-12
+
+macs = results['SynapticOperations']['Effective_MACs']
+acs = results['SynapticOperations']['Effective_ACs']
+
+total_energy = (macs * ENERGY_PER_MAC) + (acs * ENERGY_PER_AC)
+
+print("Energy Report:")
+print(f"Total Operations: {macs} MACS, {acs} acs ")
+print(f"Calculated Energy cost: {total_energy} joules per batch")
diff --git a/references/kaevin_memristor.py b/references/kaevin_memristor.py
new file mode 100644
index 0000000..b4a4364
--- /dev/null
+++ b/references/kaevin_memristor.py
@@ -0,0 +1,134 @@
+import numpy as np
+from scipy.integrate import solve_ivp
+
+# TEAM (Threshold Adaptive Memristor) model as a reusable class
+class TEAMMemristor:
+
+    def __init__(
+        self,
+        k_off=1,          # switching rate for off-state
+        k_on=-1,          # switching rate for on-state
+        alpha_off=5,      # exponent controlling nonlinearity when switching off
+        alpha_on=5,       # exponent controlling nonlinearity when switching on
+        i_off=0.5e-3,     # threshold current to trigger off state switching
+        i_on=-0.5e-3,     # threshold current to trigger on state switching
+        g_on=1/1e3,       # maximum conductance (1 kohm = 1 ms)
+        g_off=1/10e3,     # minimum conductance (10 kohm = 0.1 ms)
+        w_init=0.5,       # initial state variable (0=off, 1=on)
+        p=2               # window function exponent
+    ):
+        self.k_off = k_off
+        self.k_on = k_on
+        self.alpha_off = alpha_off
+        self.alpha_on = alpha_on
+        self.i_off = i_off
+        self.i_on = i_on
+        self.G_on = G_on
+        self.G_off = G_off
+        self.w_init = w_init
+        self.p = p
+
+    def set_state(self, w):
+        # Update initial condition for next simulation
+        self.w_init = np.clip(w, 0, 1)
+
+    def window(self, w, i):
+        # Nonlinear window function: reduces switching rate near the boundaries
+        w = np.clip(w, 0.0, 1.0)
+        if i >= 0:
+            return 1 - w**(2*self.p)  # switching off: slower near w=1
+        else:
+            return 1 - (1-w)**(2*self.p)  # switching on: slower near w=0
+
+    def conductance(self, w):
+        # Linear interpolation between off and on conductance based on state w
+        w = np.clip(w, 0.0, 1.0)
+        return self.G_off + w*(self.G_on - self.G_off)
+
+    def dw_dt(self, w, i):
+        # TEAM state dynamics: dw/dt depends on current magnitude and direction
+        w = np.clip(w, 0.0, 1.0)
+        
+        if i >= self.i_off:  # positive current above threshold = switch off
+            dw = (
+                self.k_off
+                * ((i/self.i_off)-1)**self.alpha_off
+                * self.window(w, i)
+            )
+        elif i <= self.i_on:  # negative current below threshold = switch on
+            dw = (
+                self.k_on
+                * (((-i)/abs(self.i_on))-1)**self.alpha_on
+                * self.window(w, i)
+            )
+        else:  # between thresholds = no switching
+            dw = 0.0
+        
+        # Enforce physical bounds: prevent state from leaving [0,1]
+        if w <= 0 and dw < 0:
+            dw = 0
+        if w >= 1 and dw > 0:
+            dw = 0
+        
+        return dw
+
+    def simulate(self,
+                 freq=1,          # excitation frequency (Hz)
+                 V_amp=1.5,       # sinusoid amplitude (V)
+                 cycles=3):       # number of periods to simulate
+        # Solve the memristor ODE for given frequency and voltage amplitude
+        
+        def voltage(t):  # sinusoidal excitation signal
+            return V_amp*np.sin(2*np.pi*freq*t)
+        
+        def ode(t, y):  # dy/dt: current through memristor
+            w = y[0]
+            v = voltage(t)
+            G = self.conductance(w)
+            i = G*v  # Ohm's law: i = G*v
+            return [self.dw_dt(w, i)]
+        
+        T = 1/freq  # period
+        t_end = cycles*T
+        t_eval = np.linspace(0, t_end, 10000)  # dense time grid for smooth curves
+        
+        # Solve with RK45, small max step
+        sol = solve_ivp(
+            ode,
+            [0, t_end],
+            [self.w_init],
+            t_eval=t_eval,
+            method='RK45',
+            max_step=T/1000,  # max step keeps resolution within one period
+            rtol=1e-8,
+            atol=1e-10
+        )
+        
+        raw_w = sol.y[0]
+        
+        # Check if numerical solver violated physical bounds
+        eps = 1e-6
+        if np.any(raw_w < -eps) or np.any(raw_w > 1+eps):
+            print(
+                "WARNING: solver left bounds "
+                f"min={raw_w.min():.12f}, "
+                f"max={raw_w.max():.12f}"
+            )
+        
+        w = np.clip(raw_w, 0, 1)  # enforce bounds just in case
+        t = sol.t
+        v = voltage(t)
+        G = self.conductance(w)
+        i = G*v  # current throughout simulation
+        
+        return t, w, v, i
+
+    def resistance(self, w):
+        # Compute resistance as reciprocal of conductance
+        return 1/self.conductance(w)
+    
+    def reset(self):
+        # Reset state to default initial condition
+        self.w_init = 0.5
+
+