diff --git a/math_explorer/src/ai/activations.rs b/math_explorer/src/ai/activations.rs
index 12dca42..7c53fc2 100644
--- a/math_explorer/src/ai/activations.rs
+++ b/math_explorer/src/ai/activations.rs
@@ -1,5 +1,72 @@
 // Implementation of various activation functions.
-use nalgebra::DMatrix;
+use nalgebra::{DMatrix, RealField};
+
+/// Trait for activation functions to allow interchangeability (OCP).
+pub trait ActivationFunction<T: RealField + Copy> {
+    /// Applies the activation function in-place.
+    fn apply(&self, x: &mut DMatrix<T>);
+}
+
+/// Rectified Linear Unit (ReLU).
+#[derive(Debug, Clone, Copy, Default)]
+pub struct ReLU;
+
+impl<T: RealField + Copy> ActivationFunction<T> for ReLU {
+    fn apply(&self, x: &mut DMatrix<T>) {
+        x.apply(|val| {
+            if *val < T::zero() {
+                *val = T::zero();
+            }
+        });
+    }
+}
+
+/// Leaky ReLU.
+#[derive(Debug, Clone, Copy)]
+pub struct LeakyReLU<T> {
+    pub alpha: T,
+}
+
+impl<T: RealField + Copy> LeakyReLU<T> {
+    pub fn new(alpha: T) -> Self {
+        Self { alpha }
+    }
+}
+
+impl<T: RealField + Copy> ActivationFunction<T> for LeakyReLU<T> {
+    fn apply(&self, x: &mut DMatrix<T>) {
+        x.apply(|val| {
+            if *val < T::zero() {
+                *val *= self.alpha;
+            }
+        });
+    }
+}
+
+/// Sigmoid function.
+#[derive(Debug, Clone, Copy, Default)]
+pub struct Sigmoid;
+
+impl<T: RealField + Copy> ActivationFunction<T> for Sigmoid {
+    fn apply(&self, x: &mut DMatrix<T>) {
+        x.apply(|val| {
+            // 1 / (1 + exp(-x))
+            *val = T::one() / (T::one() + (-*val).exp());
+        });
+    }
+}
+
+/// Hyperbolic Tangent (Tanh).
+#[derive(Debug, Clone, Copy, Default)]
+pub struct Tanh;
+
+impl<T: RealField + Copy> ActivationFunction<T> for Tanh {
+    fn apply(&self, x: &mut DMatrix<T>) {
+        x.apply(|val| {
+            *val = val.tanh();
+        });
+    }
+}
 
 /// Applies the Rectified Linear Unit (ReLU) activation function element-wise, in-place.
 /// ReLU(x) = max(0, x)
diff --git a/math_explorer/src/applied/clinical_trials/design.rs b/math_explorer/src/applied/clinical_trials/design.rs
index 3d6d43b..d148f5c 100644
--- a/math_explorer/src/applied/clinical_trials/design.rs
+++ b/math_explorer/src/applied/clinical_trials/design.rs
@@ -175,7 +175,9 @@ impl<S: AllocationStrategy> StratifiedRandomizer<S> {
 #[deprecated(since = "0.2.0", note = "Use SimpleRandomizer struct instead")]
 pub fn simple_randomization(n_patients: usize) -> Vec<Group> {
     let mut rng = thread_rng();
-    SimpleRandomizer.assign(&mut rng, n_patients).unwrap()
+    SimpleRandomizer
+        .assign(&mut rng, n_patients)
+        .expect("SimpleRandomizer should never fail")
 }
 
 #[deprecated(since = "0.2.0", note = "Use BlockRandomizer struct instead")]
diff --git a/math_explorer/src/climate/autoencoder.rs b/math_explorer/src/climate/autoencoder.rs
index f5ff2cf..86d6b0b 100644
--- a/math_explorer/src/climate/autoencoder.rs
+++ b/math_explorer/src/climate/autoencoder.rs
@@ -1,22 +1,9 @@
 //! This module defines the autoencoder architecture for the CERA framework.
 
+use crate::ai::activations::{ActivationFunction, LeakyReLU};
 use crate::climate::tensor_ops::conv1d;
 use nalgebra::{DMatrix, DVector};
 
-/// A simple leaky ReLU activation function.
-///
-/// # Arguments
-///
-/// * `x` - The matrix to apply the activation to (in-place).
-/// * `alpha` - The negative slope coefficient.
-pub fn leaky_relu(x: &mut DMatrix<f32>, alpha: f32) {
-    x.iter_mut().for_each(|val| {
-        if *val < 0.0 {
-            *val *= alpha;
-        }
-    });
-}
-
 /// A single layer for the Encoder or Decoder, consisting of a convolution and activation.
 pub struct ConvLayer {
     /// The convolution kernel matrix.
@@ -57,48 +44,41 @@ impl ConvLayer {
 }
 
 /// The encoder component of the autoencoder.
-pub struct Encoder {
+pub struct EncoderGeneric<A: ActivationFunction<f32>> {
     /// The stack of convolutional layers.
     pub layers: Vec<ConvLayer>,
+    /// The activation function strategy.
+    pub activation: A,
 }
 
+/// Default Encoder type alias for backward compatibility.
+pub type Encoder = EncoderGeneric<LeakyReLU<f32>>;
+
 impl Encoder {
     /// Creates a new encoder with a hardcoded architecture.
-    /// Input (2 channels) -> 64 -> 64 -> Latent (3 channels)
-    ///
-    /// # Arguments
-    ///
-    /// * `in_channels` - Number of input channels.
-    /// * `latent_channels` - Dimension of the latent space.
-    ///
-    /// # Returns
-    ///
-    /// A new `Encoder`.
     pub fn new(in_channels: usize, latent_channels: usize) -> Self {
+        Self::new_with_activation(in_channels, latent_channels, LeakyReLU::new(0.01))
+    }
+}
+
+impl<A: ActivationFunction<f32>> EncoderGeneric<A> {
+    pub fn new_with_activation(in_channels: usize, latent_channels: usize, activation: A) -> Self {
         let layers = vec![
             ConvLayer::new(in_channels, 64),
             ConvLayer::new(64, 64),
             ConvLayer::new(64, latent_channels), // No activation on the latent layer
         ];
-        Self { layers }
+        Self { layers, activation }
     }
 
     /// Encodes the input data into a latent representation.
-    ///
-    /// # Arguments
-    ///
-    /// * `input` - The input data matrix.
-    ///
-    /// # Returns
-    ///
-    /// The latent representation matrix.
     pub fn forward(&self, input: &DMatrix<f32>) -> DMatrix<f32> {
         let mut x = input.clone();
         for (i, layer) in self.layers.iter().enumerate() {
             x = conv1d(&x, &layer.kernel, &layer.bias);
             // No activation on the final layer
             if i < self.layers.len() - 1 {
-                leaky_relu(&mut x, 0.01);
+                self.activation.apply(&mut x);
             }
         }
         x
@@ -106,47 +86,40 @@ impl Encoder {
 }
 
 /// The decoder component of the autoencoder.
-pub struct Decoder {
+pub struct DecoderGeneric<A: ActivationFunction<f32>> {
     /// The stack of convolutional layers.
     pub layers: Vec<ConvLayer>,
+    /// The activation function strategy.
+    pub activation: A,
 }
 
+/// Default Decoder type alias.
+pub type Decoder = DecoderGeneric<LeakyReLU<f32>>;
+
 impl Decoder {
     /// Creates a new decoder with a hardcoded architecture.
-    /// Latent (3 channels) -> 64 -> 64 -> Output (2 channels)
-    ///
-    /// # Arguments
-    ///
-    /// * `latent_channels` - Dimension of the latent space.
-    /// * `out_channels` - Number of output channels.
-    ///
-    /// # Returns
-    ///
-    /// A new `Decoder`.
     pub fn new(latent_channels: usize, out_channels: usize) -> Self {
+        Self::new_with_activation(latent_channels, out_channels, LeakyReLU::new(0.01))
+    }
+}
+
+impl<A: ActivationFunction<f32>> DecoderGeneric<A> {
+    pub fn new_with_activation(latent_channels: usize, out_channels: usize, activation: A) -> Self {
         let layers = vec![
             ConvLayer::new(latent_channels, 64),
             ConvLayer::new(64, 64),
             ConvLayer::new(64, out_channels), // No activation on the output layer
         ];
-        Self { layers }
+        Self { layers, activation }
     }
 
     /// Reconstructs the input data from the latent representation.
-    ///
-    /// # Arguments
-    ///
-    /// * `latent_representation` - The latent representation matrix.
-    ///
-    /// # Returns
-    ///
-    /// The reconstructed data matrix.
     pub fn forward(&self, latent_representation: &DMatrix<f32>) -> DMatrix<f32> {
         let mut x = latent_representation.clone();
         for (i, layer) in self.layers.iter().enumerate() {
             x = conv1d(&x, &layer.kernel, &layer.bias);
             if i < self.layers.len() - 1 {
-                leaky_relu(&mut x, 0.01);
+                self.activation.apply(&mut x);
             }
         }
         x
@@ -154,41 +127,32 @@ impl Decoder {
 }
 
 /// The autoencoder model for the CERA framework.
-pub struct Autoencoder {
+pub struct AutoencoderGeneric<A: ActivationFunction<f32>> {
     /// The encoder component.
-    pub encoder: Encoder,
+    pub encoder: EncoderGeneric<A>,
     /// The decoder component.
-    pub decoder: Decoder,
+    pub decoder: DecoderGeneric<A>,
 }
 
+/// Default Autoencoder type alias.
+pub type Autoencoder = AutoencoderGeneric<LeakyReLU<f32>>;
+
 impl Autoencoder {
     /// Creates a new autoencoder.
-    /// The paper specifies 2 input channels and 3 latent channels.
-    ///
-    /// # Arguments
-    ///
-    /// * `in_channels` - Number of input channels.
-    /// * `latent_channels` - Dimension of the latent space.
-    ///
-    /// # Returns
-    ///
-    /// A new `Autoencoder`.
     pub fn new(in_channels: usize, latent_channels: usize) -> Self {
-        // The decoder's input is the encoder's output, and vice versa.
-        let encoder = Encoder::new(in_channels, latent_channels);
-        let decoder = Decoder::new(latent_channels, in_channels);
+        Self::new_with_activation(in_channels, latent_channels, LeakyReLU::new(0.01))
+    }
+}
+
+impl<A: ActivationFunction<f32> + Clone> AutoencoderGeneric<A> {
+    pub fn new_with_activation(in_channels: usize, latent_channels: usize, activation: A) -> Self {
+        let encoder =
+            EncoderGeneric::new_with_activation(in_channels, latent_channels, activation.clone());
+        let decoder = DecoderGeneric::new_with_activation(latent_channels, in_channels, activation);
         Self { encoder, decoder }
     }
 
     /// Performs a forward pass through the autoencoder.
-    ///
-    /// # Arguments
-    ///
-    /// * `input` - The input data matrix.
-    ///
-    /// # Returns
-    ///
-    /// A tuple containing `(latent_representation, reconstruction)`.
     pub fn forward(&self, input: &DMatrix<f32>) -> (DMatrix<f32>, DMatrix<f32>) {
         let latent = self.encoder.forward(input);
         let reconstruction = self.decoder.forward(&latent);
@@ -220,12 +184,4 @@ mod tests {
         assert_eq!(reconstruction.nrows(), n_samples);
         assert_eq!(reconstruction.ncols(), in_channels);
     }
-
-    #[test]
-    fn test_leaky_relu() {
-        let mut matrix = DMatrix::from_row_slice(2, 2, &[-1.0, 2.0, -3.0, 0.0]);
-        leaky_relu(&mut matrix, 0.1);
-        let expected = DMatrix::from_row_slice(2, 2, &[-0.1, 2.0, -0.3, 0.0]);
-        assert!((matrix - expected).abs().max() < 1e-6);
-    }
 }
diff --git a/math_explorer/src/climate/predictor.rs b/math_explorer/src/climate/predictor.rs
index b501242..55edf1c 100644
--- a/math_explorer/src/climate/predictor.rs
+++ b/math_explorer/src/climate/predictor.rs
@@ -1,6 +1,7 @@
 //! This module defines the predictor model for the CERA framework.
 
-use crate::climate::autoencoder::{ConvLayer, leaky_relu};
+use crate::ai::activations::{ActivationFunction, LeakyReLU};
+use crate::climate::autoencoder::ConvLayer;
 use nalgebra::{DMatrix, DVector};
 
 /// A trait representing the predictor model interface.
@@ -18,9 +19,11 @@ pub trait PredictorModel {
 ///
 /// The predictor takes the flattened, aligned latent representation from the
 /// autoencoder's encoder and maps it to the target output variables.
-pub struct Predictor {
+pub struct Predictor<A: ActivationFunction<f32>> {
     /// The stack of layers (using `ConvLayer` for simplicity as dense layers).
     pub layers: Vec<ConvLayer>,
+    /// The activation function strategy.
+    pub activation: A,
     // Store dimensions for clarity
     #[allow(dead_code)]
     input_size: usize,
@@ -29,7 +32,7 @@ pub struct Predictor {
     output_size: usize,
 }
 
-impl Predictor {
+impl Predictor<LeakyReLU<f32>> {
     /// Creates a new predictor model with a hardcoded architecture.
     /// Input (60) -> 128 -> 128 -> 128 -> 128 -> Output (148)
     ///
@@ -42,6 +45,13 @@ impl Predictor {
     ///
     /// A new `Predictor` instance.
     pub fn new(input_size: usize, output_size: usize) -> Self {
+        Self::new_with_activation(input_size, output_size, LeakyReLU::new(0.01))
+    }
+}
+
+impl<A: ActivationFunction<f32>> Predictor<A> {
+    /// Creates a new predictor with a custom activation function.
+    pub fn new_with_activation(input_size: usize, output_size: usize, activation: A) -> Self {
         let layers = vec![
             ConvLayer::new(input_size, 128),
             ConvLayer::new(128, 128),
@@ -51,13 +61,14 @@ impl Predictor {
         ];
         Self {
             layers,
+            activation,
             input_size,
             output_size,
         }
     }
 }
 
-impl PredictorModel for Predictor {
+impl<A: ActivationFunction<f32>> PredictorModel for Predictor<A> {
     fn forward(&self, input: &DMatrix<f32>) -> DMatrix<f32> {
         let mut x = input.clone();
         for (i, layer) in self.layers.iter().enumerate() {
@@ -67,7 +78,7 @@ impl PredictorModel for Predictor {
             x = crate::climate::tensor_ops::conv1d(&x, &layer.kernel, &layer.bias);
             // No activation on the final layer
             if i < self.layers.len() - 1 {
-                leaky_relu(&mut x, 0.01);
+                self.activation.apply(&mut x);
             }
         }
         x
diff --git a/math_explorer/src/epidemiology/compartmental.rs b/math_explorer/src/epidemiology/compartmental.rs
index e777da8..70bdb09 100644
--- a/math_explorer/src/epidemiology/compartmental.rs
+++ b/math_explorer/src/epidemiology/compartmental.rs
@@ -1,5 +1,5 @@
-use crate::pure_math::analysis::ode::{OdeSystem, Solver, TimeStepper, VectorOperations};
-use std::ops::{Add, AddAssign, Mul, MulAssign};
+use crate::impl_vector_ops;
+use crate::pure_math::analysis::ode::{OdeSystem, Solver, TimeStepper};
 
 /// State for the SIR Model.
 #[derive(Debug, Clone, Copy, PartialEq)]
@@ -9,55 +9,7 @@ pub struct SIRState {
     pub r: f64,
 }
 
-impl Add for SIRState {
-    type Output = Self;
-    fn add(self, rhs: Self) -> Self {
-        Self {
-            s: self.s + rhs.s,
-            i: self.i + rhs.i,
-            r: self.r + rhs.r,
-        }
-    }
-}
-
-impl AddAssign for SIRState {
-    fn add_assign(&mut self, rhs: Self) {
-        self.s += rhs.s;
-        self.i += rhs.i;
-        self.r += rhs.r;
-    }
-}
-
-impl Mul<f64> for SIRState {
-    type Output = Self;
-    fn mul(self, scalar: f64) -> Self {
-        Self {
-            s: self.s * scalar,
-            i: self.i * scalar,
-            r: self.r * scalar,
-        }
-    }
-}
-
-impl MulAssign<f64> for SIRState {
-    fn mul_assign(&mut self, scalar: f64) {
-        self.s *= scalar;
-        self.i *= scalar;
-        self.r *= scalar;
-    }
-}
-
-impl VectorOperations for SIRState {
-    fn scale_add(&mut self, other: &Self, scale: f64) {
-        self.s += other.s * scale;
-        self.i += other.i * scale;
-        self.r += other.r * scale;
-    }
-
-    fn copy_from(&mut self, other: &Self) {
-        *self = *other;
-    }
-}
+impl_vector_ops!(SIRState, s, i, r);
 
 /// SIR Model: Susceptible, Infectious, Recovered.
 ///
@@ -134,60 +86,7 @@ pub struct SEIRState {
     pub r: f64,
 }
 
-impl Add for SEIRState {
-    type Output = Self;
-    fn add(self, rhs: Self) -> Self {
-        Self {
-            s: self.s + rhs.s,
-            e: self.e + rhs.e,
-            i: self.i + rhs.i,
-            r: self.r + rhs.r,
-        }
-    }
-}
-
-impl AddAssign for SEIRState {
-    fn add_assign(&mut self, rhs: Self) {
-        self.s += rhs.s;
-        self.e += rhs.e;
-        self.i += rhs.i;
-        self.r += rhs.r;
-    }
-}
-
-impl Mul<f64> for SEIRState {
-    type Output = Self;
-    fn mul(self, scalar: f64) -> Self {
-        Self {
-            s: self.s * scalar,
-            e: self.e * scalar,
-            i: self.i * scalar,
-            r: self.r * scalar,
-        }
-    }
-}
-
-impl MulAssign<f64> for SEIRState {
-    fn mul_assign(&mut self, scalar: f64) {
-        self.s *= scalar;
-        self.e *= scalar;
-        self.i *= scalar;
-        self.r *= scalar;
-    }
-}
-
-impl VectorOperations for SEIRState {
-    fn scale_add(&mut self, other: &Self, scale: f64) {
-        self.s += other.s * scale;
-        self.e += other.e * scale;
-        self.i += other.i * scale;
-        self.r += other.r * scale;
-    }
-
-    fn copy_from(&mut self, other: &Self) {
-        *self = *other;
-    }
-}
+impl_vector_ops!(SEIRState, s, e, i, r);
 
 /// SEIR Model: Susceptible, Exposed, Infectious, Recovered.
 ///
diff --git a/math_explorer/src/pure_math/analysis/ode.rs b/math_explorer/src/pure_math/analysis/ode.rs
index d08be36..5702249 100644
--- a/math_explorer/src/pure_math/analysis/ode.rs
+++ b/math_explorer/src/pure_math/analysis/ode.rs
@@ -291,3 +291,59 @@ pub trait TimeStepper<State: VectorOperations>: OdeSystem<State> {
         *self.get_state_mut() = new_state;
     }
 }
+
+/// Macro to implement arithmetic traits and `VectorOperations` for a struct with named f64 fields.
+#[macro_export]
+macro_rules! impl_vector_ops {
+    ($struct_name:ident, $($field:ident),+) => {
+        impl std::ops::Add for $struct_name {
+            type Output = Self;
+            fn add(self, rhs: Self) -> Self {
+                Self {
+                    $(
+                        $field: self.$field + rhs.$field,
+                    )+
+                }
+            }
+        }
+
+        impl std::ops::AddAssign for $struct_name {
+            fn add_assign(&mut self, rhs: Self) {
+                $(
+                    self.$field += rhs.$field;
+                )+
+            }
+        }
+
+        impl std::ops::Mul<f64> for $struct_name {
+            type Output = Self;
+            fn mul(self, scalar: f64) -> Self {
+                Self {
+                    $(
+                        $field: self.$field * scalar,
+                    )+
+                }
+            }
+        }
+
+        impl std::ops::MulAssign<f64> for $struct_name {
+            fn mul_assign(&mut self, scalar: f64) {
+                $(
+                    self.$field *= scalar;
+                )+
+            }
+        }
+
+        impl $crate::pure_math::analysis::ode::VectorOperations for $struct_name {
+            fn scale_add(&mut self, other: &Self, scale: f64) {
+                $(
+                    self.$field += other.$field * scale;
+                )+
+            }
+
+            fn copy_from(&mut self, other: &Self) {
+                *self = *other;
+            }
+        }
+    };
+}