c645e6222d
CRITICAL FINDINGS from 3-trial validation: - 85,120 gradient clipping warnings (81.6% of logs) - REGRESSION - Rainbow features DISABLED: use_dueling=false, use_distributional=false, use_noisy_nets=false - Negative Q-values confirmed: HOLD -1000 to -3250 - Performance: Sharpe 0.29 (target 0.77) Changes: - Fixed N-Step compilation (7/7 tests passing) - Fixed Distributional compilation (6/6 tests passing) - Fixed Dueling CUDA errors (10/10 tests passing) - Added TDD validation for state_dim=225 - Total: 23/23 Wave 11 tests passing (100%) Issues requiring investigation: 1. Why are Dueling/Distributional/Noisy disabled in hyperopt? 2. Why gradient explosion despite previous fixes? 3. Test coverage gaps - unit tests pass but integration fails 🤖 Generated with Claude Code Co-Authored-By: Claude <noreply@anthropic.com>
596 lines
19 KiB
Rust
596 lines
19 KiB
Rust
//! Rainbow DQN Network Architecture Validation Tests
|
||
//!
|
||
//! Validates that the implementation matches the paper specification:
|
||
//! "Rainbow: Combining Improvements in Deep Reinforcement Learning" (Hessel et al., 2017)
|
||
//!
|
||
//! Key components verified:
|
||
//! 1. Noisy Linear layers (not standard Linear)
|
||
//! 2. Dueling architecture (value + advantage streams)
|
||
//! 3. C51 distributional output (num_actions × num_atoms)
|
||
//! 4. Forward pass shapes
|
||
//! 5. Softmax over atoms dimension
|
||
|
||
use anyhow::Result;
|
||
use candle_core::{DType, Device, Tensor};
|
||
use candle_nn::{Module, VarBuilder, VarMap};
|
||
use ml::dqn::{
|
||
CategoricalDistribution, DistributionalConfig, RainbowNetwork, RainbowNetworkConfig,
|
||
};
|
||
use ml::MLError;
|
||
|
||
/// Test 1: Network initialization with correct parameter counts
|
||
#[test]
|
||
fn test_rainbow_network_initialization() -> Result<(), MLError> {
|
||
let device = Device::Cpu;
|
||
let varmap = VarMap::new();
|
||
let vs = VarBuilder::from_varmap(&varmap, DType::F32, &device);
|
||
|
||
let config = RainbowNetworkConfig {
|
||
input_size: 128,
|
||
hidden_sizes: vec![512, 512],
|
||
num_actions: 3,
|
||
activation: ml::dqn::rainbow_network::ActivationType::ReLU,
|
||
dropout_rate: 0.1,
|
||
distributional: DistributionalConfig {
|
||
num_atoms: 51,
|
||
v_min: -10.0,
|
||
v_max: 10.0,
|
||
},
|
||
use_noisy_layers: true,
|
||
dueling: true,
|
||
};
|
||
|
||
let network = RainbowNetwork::new(&vs, config)?;
|
||
|
||
// Verify network was created successfully
|
||
assert_eq!(network.config().input_size, 128);
|
||
assert_eq!(network.config().num_actions, 3);
|
||
assert_eq!(network.config().distributional.num_atoms, 51);
|
||
assert!(network.config().use_noisy_layers);
|
||
assert!(network.config().dueling);
|
||
|
||
Ok(())
|
||
}
|
||
|
||
/// Test 2: Forward pass shape validation for single sample
|
||
#[test]
|
||
fn test_forward_pass_single_sample() -> Result<(), MLError> {
|
||
let device = Device::Cpu;
|
||
let varmap = VarMap::new();
|
||
let vs = VarBuilder::from_varmap(&varmap, DType::F32, &device);
|
||
|
||
let config = RainbowNetworkConfig {
|
||
input_size: 128,
|
||
hidden_sizes: vec![256, 256],
|
||
num_actions: 3,
|
||
activation: ml::dqn::rainbow_network::ActivationType::ReLU,
|
||
dropout_rate: 0.0,
|
||
distributional: DistributionalConfig {
|
||
num_atoms: 51,
|
||
v_min: -10.0,
|
||
v_max: 10.0,
|
||
},
|
||
use_noisy_layers: true,
|
||
dueling: true,
|
||
};
|
||
|
||
let network = RainbowNetwork::new(&vs, config)?;
|
||
|
||
// Input: [batch=1, state_dim=128]
|
||
let input = Tensor::randn(0.0_f32, 1.0_f32, (1, 128), &device)?;
|
||
|
||
// Forward pass
|
||
let output = network
|
||
.forward(&input)
|
||
.map_err(|e| MLError::ModelError(format!("Forward pass failed: {}", e)))?;
|
||
|
||
// Output should be [batch=1, num_actions=3, num_atoms=51]
|
||
assert_eq!(output.shape().dims(), &[1, 3, 51]);
|
||
|
||
Ok(())
|
||
}
|
||
|
||
/// Test 3: Forward pass shape validation for batched input
|
||
#[test]
|
||
fn test_forward_pass_batch() -> Result<(), MLError> {
|
||
let device = Device::Cpu;
|
||
let varmap = VarMap::new();
|
||
let vs = VarBuilder::from_varmap(&varmap, DType::F32, &device);
|
||
|
||
let config = RainbowNetworkConfig {
|
||
input_size: 128,
|
||
hidden_sizes: vec![512, 512],
|
||
num_actions: 3,
|
||
activation: ml::dqn::rainbow_network::ActivationType::ReLU,
|
||
dropout_rate: 0.0,
|
||
distributional: DistributionalConfig {
|
||
num_atoms: 51,
|
||
v_min: -10.0,
|
||
v_max: 10.0,
|
||
},
|
||
use_noisy_layers: true,
|
||
dueling: true,
|
||
};
|
||
|
||
let network = RainbowNetwork::new(&vs, config)?;
|
||
|
||
// Input: [batch=32, state_dim=128]
|
||
let batch_size = 32;
|
||
let input = Tensor::randn(0.0_f32, 1.0_f32, (batch_size, 128), &device)?;
|
||
|
||
// Forward pass
|
||
let output = network
|
||
.forward(&input)
|
||
.map_err(|e| MLError::ModelError(format!("Forward pass failed: {}", e)))?;
|
||
|
||
// Output should be [batch=32, num_actions=3, num_atoms=51]
|
||
assert_eq!(output.shape().dims(), &[batch_size, 3, 51]);
|
||
|
||
Ok(())
|
||
}
|
||
|
||
/// Test 4: C51 distributional output - verify softmax over atoms dimension
|
||
#[test]
|
||
fn test_c51_output_is_probability_distribution() -> Result<(), MLError> {
|
||
let device = Device::Cpu;
|
||
let varmap = VarMap::new();
|
||
let vs = VarBuilder::from_varmap(&varmap, DType::F32, &device);
|
||
|
||
let config = RainbowNetworkConfig {
|
||
input_size: 128,
|
||
hidden_sizes: vec![256],
|
||
num_actions: 3,
|
||
activation: ml::dqn::rainbow_network::ActivationType::ReLU,
|
||
dropout_rate: 0.0,
|
||
distributional: DistributionalConfig {
|
||
num_atoms: 51,
|
||
v_min: -10.0,
|
||
v_max: 10.0,
|
||
},
|
||
use_noisy_layers: true,
|
||
dueling: true,
|
||
};
|
||
|
||
let network = RainbowNetwork::new(&vs, config)?;
|
||
|
||
let input = Tensor::randn(0.0_f32, 1.0_f32, (4, 128), &device)?;
|
||
|
||
// Forward pass
|
||
let output = network
|
||
.forward(&input)
|
||
.map_err(|e| MLError::ModelError(format!("Forward pass failed: {}", e)))?;
|
||
|
||
// Output shape: [batch=4, actions=3, atoms=51]
|
||
assert_eq!(output.shape().dims(), &[4, 3, 51]);
|
||
|
||
// For each (batch, action) pair, the distribution over atoms should sum to 1.0
|
||
for batch_idx in 0..4 {
|
||
for action_idx in 0..3 {
|
||
let dist = output
|
||
.get(batch_idx)
|
||
.map_err(|e| MLError::ModelError(format!("Failed to get batch: {}", e)))?
|
||
.get(action_idx)
|
||
.map_err(|e| MLError::ModelError(format!("Failed to get action: {}", e)))?;
|
||
|
||
let sum_tensor = dist
|
||
.sum_all()
|
||
.map_err(|e| MLError::ModelError(format!("Failed to sum: {}", e)))?;
|
||
|
||
// Handle both [] and [1] shapes
|
||
let sum: f32 = if sum_tensor.rank() == 0 {
|
||
sum_tensor.to_scalar().map_err(|e| {
|
||
MLError::ModelError(format!("Failed to convert to scalar: {}", e))
|
||
})?
|
||
} else {
|
||
sum_tensor
|
||
.squeeze(0)
|
||
.map_err(|e| MLError::ModelError(format!("Failed to squeeze: {}", e)))?
|
||
.to_scalar()
|
||
.map_err(|e| {
|
||
MLError::ModelError(format!("Failed to convert to scalar: {}", e))
|
||
})?
|
||
};
|
||
|
||
// Check sum is approximately 1.0 (allow small numerical error)
|
||
assert!(
|
||
(sum - 1.0).abs() < 1e-4,
|
||
"Distribution sum should be 1.0, got {}",
|
||
sum
|
||
);
|
||
|
||
// Check all probabilities are non-negative
|
||
let min_tensor = dist
|
||
.min_keepdim(0)
|
||
.map_err(|e| MLError::ModelError(format!("Failed to get min: {}", e)))?;
|
||
|
||
// Handle both [] and [1] shapes
|
||
let min_val: f32 = if min_tensor.rank() == 0 {
|
||
min_tensor.to_scalar().map_err(|e| {
|
||
MLError::ModelError(format!("Failed to convert to scalar: {}", e))
|
||
})?
|
||
} else {
|
||
min_tensor
|
||
.squeeze(0)
|
||
.map_err(|e| MLError::ModelError(format!("Failed to squeeze: {}", e)))?
|
||
.to_scalar()
|
||
.map_err(|e| {
|
||
MLError::ModelError(format!("Failed to convert to scalar: {}", e))
|
||
})?
|
||
};
|
||
|
||
assert!(
|
||
min_val >= -1e-6,
|
||
"All probabilities should be non-negative, got min={}",
|
||
min_val
|
||
);
|
||
}
|
||
}
|
||
|
||
Ok(())
|
||
}
|
||
|
||
/// Test 5: Dueling architecture - verify value and advantage streams are used
|
||
#[test]
|
||
fn test_dueling_architecture() -> Result<(), MLError> {
|
||
let device = Device::Cpu;
|
||
let varmap = VarMap::new();
|
||
let vs = VarBuilder::from_varmap(&varmap, DType::F32, &device);
|
||
|
||
// Create two networks: one with dueling, one without
|
||
let config_dueling = RainbowNetworkConfig {
|
||
input_size: 64,
|
||
hidden_sizes: vec![128],
|
||
num_actions: 3,
|
||
activation: ml::dqn::rainbow_network::ActivationType::ReLU,
|
||
dropout_rate: 0.0,
|
||
distributional: DistributionalConfig {
|
||
num_atoms: 51,
|
||
v_min: -10.0,
|
||
v_max: 10.0,
|
||
},
|
||
use_noisy_layers: false, // Disable noise for deterministic test
|
||
dueling: true,
|
||
};
|
||
|
||
let config_no_dueling = RainbowNetworkConfig {
|
||
dueling: false,
|
||
..config_dueling.clone()
|
||
};
|
||
|
||
let network_dueling = RainbowNetwork::new(&vs.pp("dueling"), config_dueling)?;
|
||
let network_standard = RainbowNetwork::new(&vs.pp("standard"), config_no_dueling)?;
|
||
|
||
let input = Tensor::randn(0.0_f32, 1.0_f32, (2, 64), &device)?;
|
||
|
||
// Both should produce valid outputs
|
||
let output_dueling = network_dueling
|
||
.forward(&input)
|
||
.map_err(|e| MLError::ModelError(format!("Dueling forward failed: {}", e)))?;
|
||
|
||
let output_standard = network_standard
|
||
.forward(&input)
|
||
.map_err(|e| MLError::ModelError(format!("Standard forward failed: {}", e)))?;
|
||
|
||
// Both should have correct shape [batch=2, actions=3, atoms=51]
|
||
assert_eq!(output_dueling.shape().dims(), &[2, 3, 51]);
|
||
assert_eq!(output_standard.shape().dims(), &[2, 3, 51]);
|
||
|
||
// Outputs should be different (dueling combines value+advantage, standard doesn't)
|
||
let diff = output_dueling
|
||
.sub(&output_standard)
|
||
.map_err(|e| MLError::ModelError(format!("Failed to compute diff: {}", e)))?;
|
||
|
||
let diff_norm: f32 = diff
|
||
.sqr()
|
||
.map_err(|e| MLError::ModelError(format!("Failed to square: {}", e)))?
|
||
.sum_all()
|
||
.map_err(|e| MLError::ModelError(format!("Failed to sum: {}", e)))?
|
||
.to_scalar()
|
||
.map_err(|e| MLError::ModelError(format!("Failed to convert to scalar: {}", e)))?;
|
||
|
||
// Should be significantly different
|
||
assert!(
|
||
diff_norm > 1e-3,
|
||
"Dueling and standard networks should produce different outputs"
|
||
);
|
||
|
||
Ok(())
|
||
}
|
||
|
||
/// Test 6: Noisy layers - verify noise sampling changes outputs
|
||
#[test]
|
||
fn test_noisy_layers_exploration() -> Result<(), MLError> {
|
||
let device = Device::Cpu;
|
||
let varmap = VarMap::new();
|
||
let vs = VarBuilder::from_varmap(&varmap, DType::F32, &device);
|
||
|
||
let config = RainbowNetworkConfig {
|
||
input_size: 64,
|
||
hidden_sizes: vec![128],
|
||
num_actions: 3,
|
||
activation: ml::dqn::rainbow_network::ActivationType::ReLU,
|
||
dropout_rate: 0.0,
|
||
distributional: DistributionalConfig {
|
||
num_atoms: 51,
|
||
v_min: -10.0,
|
||
v_max: 10.0,
|
||
},
|
||
use_noisy_layers: true,
|
||
dueling: true,
|
||
};
|
||
|
||
let network = RainbowNetwork::new(&vs, config)?;
|
||
|
||
let input = Tensor::randn(0.0_f32, 1.0_f32, (1, 64), &device)?;
|
||
|
||
// First forward pass
|
||
let output1 = network
|
||
.forward(&input)
|
||
.map_err(|e| MLError::ModelError(format!("First forward failed: {}", e)))?;
|
||
|
||
// Second forward pass (noise should be different if reset)
|
||
let output2 = network
|
||
.forward(&input)
|
||
.map_err(|e| MLError::ModelError(format!("Second forward failed: {}", e)))?;
|
||
|
||
// Both should have correct shape
|
||
assert_eq!(output1.shape().dims(), &[1, 3, 51]);
|
||
assert_eq!(output2.shape().dims(), &[1, 3, 51]);
|
||
|
||
// Note: In this test, noise is NOT reset between forward passes,
|
||
// so outputs might be the same. This test verifies that the network
|
||
// CAN produce outputs (noise sampling doesn't crash).
|
||
// A proper test would require access to noise reset functionality.
|
||
|
||
Ok(())
|
||
}
|
||
|
||
/// Test 7: Q-value extraction from distributions
|
||
#[test]
|
||
fn test_q_value_extraction() -> Result<(), MLError> {
|
||
let device = Device::Cpu;
|
||
let varmap = VarMap::new();
|
||
let vs = VarBuilder::from_varmap(&varmap, DType::F32, &device);
|
||
|
||
let config = RainbowNetworkConfig {
|
||
input_size: 64,
|
||
hidden_sizes: vec![128],
|
||
num_actions: 3,
|
||
activation: ml::dqn::rainbow_network::ActivationType::ReLU,
|
||
dropout_rate: 0.0,
|
||
distributional: DistributionalConfig {
|
||
num_atoms: 51,
|
||
v_min: -10.0,
|
||
v_max: 10.0,
|
||
},
|
||
use_noisy_layers: false,
|
||
dueling: true,
|
||
};
|
||
|
||
let network = RainbowNetwork::new(&vs, config)?;
|
||
|
||
let input = Tensor::randn(0.0_f32, 1.0_f32, (2, 64), &device)?;
|
||
|
||
// Get distributions
|
||
let distributions = network
|
||
.forward(&input)
|
||
.map_err(|e| MLError::ModelError(format!("Forward failed: {}", e)))?;
|
||
|
||
// Extract Q-values
|
||
let q_values = network
|
||
.get_q_values(&distributions)
|
||
.map_err(|e| MLError::ModelError(format!("Q-value extraction failed: {}", e)))?;
|
||
|
||
// Q-values should have shape [batch=2, actions=3]
|
||
assert_eq!(q_values.shape().dims(), &[2, 3]);
|
||
|
||
// Q-values should be within reasonable range (-10.0 to 10.0 given v_min/v_max)
|
||
let q_min_tensor = q_values
|
||
.min_keepdim(1)
|
||
.map_err(|e| MLError::ModelError(format!("Failed to get min: {}", e)))?
|
||
.min_keepdim(0)
|
||
.map_err(|e| MLError::ModelError(format!("Failed to get min: {}", e)))?;
|
||
|
||
// Handle both [] and [1, 1] shapes
|
||
let q_min: f32 = if q_min_tensor.rank() == 0 {
|
||
q_min_tensor
|
||
.to_scalar()
|
||
.map_err(|e| MLError::ModelError(format!("Failed to convert to scalar: {}", e)))?
|
||
} else {
|
||
q_min_tensor
|
||
.flatten_all()
|
||
.map_err(|e| MLError::ModelError(format!("Failed to flatten: {}", e)))?
|
||
.get(0)
|
||
.map_err(|e| MLError::ModelError(format!("Failed to get index: {}", e)))?
|
||
.to_scalar()
|
||
.map_err(|e| MLError::ModelError(format!("Failed to convert to scalar: {}", e)))?
|
||
};
|
||
|
||
let q_max_tensor = q_values
|
||
.max_keepdim(1)
|
||
.map_err(|e| MLError::ModelError(format!("Failed to get max: {}", e)))?
|
||
.max_keepdim(0)
|
||
.map_err(|e| MLError::ModelError(format!("Failed to get max: {}", e)))?;
|
||
|
||
// Handle both [] and [1, 1] shapes
|
||
let q_max: f32 = if q_max_tensor.rank() == 0 {
|
||
q_max_tensor
|
||
.to_scalar()
|
||
.map_err(|e| MLError::ModelError(format!("Failed to convert to scalar: {}", e)))?
|
||
} else {
|
||
q_max_tensor
|
||
.flatten_all()
|
||
.map_err(|e| MLError::ModelError(format!("Failed to flatten: {}", e)))?
|
||
.get(0)
|
||
.map_err(|e| MLError::ModelError(format!("Failed to get index: {}", e)))?
|
||
.to_scalar()
|
||
.map_err(|e| MLError::ModelError(format!("Failed to convert to scalar: {}", e)))?
|
||
};
|
||
|
||
assert!(
|
||
q_min >= -10.0 && q_max <= 10.0,
|
||
"Q-values should be within v_min/v_max range, got [{}, {}]",
|
||
q_min,
|
||
q_max
|
||
);
|
||
|
||
Ok(())
|
||
}
|
||
|
||
/// Test 8: Different activation functions
|
||
#[test]
|
||
fn test_activation_functions() -> Result<(), MLError> {
|
||
let device = Device::Cpu;
|
||
let varmap = VarMap::new();
|
||
let vs = VarBuilder::from_varmap(&varmap, DType::F32, &device);
|
||
|
||
let activations = vec![
|
||
ml::dqn::rainbow_network::ActivationType::ReLU,
|
||
ml::dqn::rainbow_network::ActivationType::LeakyReLU,
|
||
ml::dqn::rainbow_network::ActivationType::Swish,
|
||
ml::dqn::rainbow_network::ActivationType::ELU,
|
||
];
|
||
|
||
for (idx, activation) in activations.iter().enumerate() {
|
||
let config = RainbowNetworkConfig {
|
||
input_size: 64,
|
||
hidden_sizes: vec![128],
|
||
num_actions: 3,
|
||
activation: *activation,
|
||
dropout_rate: 0.0,
|
||
distributional: DistributionalConfig {
|
||
num_atoms: 51,
|
||
v_min: -10.0,
|
||
v_max: 10.0,
|
||
},
|
||
use_noisy_layers: false,
|
||
dueling: true,
|
||
};
|
||
|
||
let network = RainbowNetwork::new(&vs.pp(&format!("net_{}", idx)), config)?;
|
||
|
||
let input = Tensor::randn(0.0_f32, 1.0_f32, (2, 64), &device)?;
|
||
|
||
let output = network.forward(&input).map_err(|e| {
|
||
MLError::ModelError(format!("Forward failed for {:?}: {}", activation, e))
|
||
})?;
|
||
|
||
// Should produce correct shape regardless of activation
|
||
assert_eq!(output.shape().dims(), &[2, 3, 51]);
|
||
}
|
||
|
||
Ok(())
|
||
}
|
||
|
||
/// Test 9: Network with different hidden layer configurations
|
||
#[test]
|
||
fn test_different_hidden_layer_configs() -> Result<(), MLError> {
|
||
let device = Device::Cpu;
|
||
let varmap = VarMap::new();
|
||
let vs = VarBuilder::from_varmap(&varmap, DType::F32, &device);
|
||
|
||
let hidden_configs = vec![
|
||
vec![128],
|
||
vec![256, 256],
|
||
vec![512, 512, 256],
|
||
vec![64, 128, 256, 128],
|
||
];
|
||
|
||
for (idx, hidden_sizes) in hidden_configs.iter().enumerate() {
|
||
let config = RainbowNetworkConfig {
|
||
input_size: 64,
|
||
hidden_sizes: hidden_sizes.clone(),
|
||
num_actions: 3,
|
||
activation: ml::dqn::rainbow_network::ActivationType::ReLU,
|
||
dropout_rate: 0.0,
|
||
distributional: DistributionalConfig {
|
||
num_atoms: 51,
|
||
v_min: -10.0,
|
||
v_max: 10.0,
|
||
},
|
||
use_noisy_layers: false,
|
||
dueling: true,
|
||
};
|
||
|
||
let network = RainbowNetwork::new(&vs.pp(&format!("config_{}", idx)), config)?;
|
||
|
||
let input = Tensor::randn(0.0_f32, 1.0_f32, (2, 64), &device)?;
|
||
|
||
let output = network.forward(&input).map_err(|e| {
|
||
MLError::ModelError(format!(
|
||
"Forward failed for hidden config {:?}: {}",
|
||
hidden_sizes, e
|
||
))
|
||
})?;
|
||
|
||
// Should produce correct shape regardless of hidden layer configuration
|
||
assert_eq!(output.shape().dims(), &[2, 3, 51]);
|
||
}
|
||
|
||
Ok(())
|
||
}
|
||
|
||
/// Test 10: Categorical distribution support values
|
||
#[test]
|
||
fn test_categorical_distribution_support() -> Result<(), MLError> {
|
||
let config = DistributionalConfig {
|
||
num_atoms: 51,
|
||
v_min: -10.0,
|
||
v_max: 10.0,
|
||
};
|
||
|
||
let device = Device::Cpu;
|
||
let dist = CategoricalDistribution::new(&config, &device)?;
|
||
|
||
// Verify support tensor has correct size
|
||
assert_eq!(dist.num_atoms(), 51);
|
||
|
||
let support = dist.support();
|
||
assert_eq!(support.shape().dims(), &[51]);
|
||
|
||
// Check first and last values
|
||
let first: f32 = support
|
||
.get(0)
|
||
.map_err(|e| MLError::ModelError(format!("Failed to get first: {}", e)))?
|
||
.to_scalar()
|
||
.map_err(|e| MLError::ModelError(format!("Failed to convert to scalar: {}", e)))?;
|
||
|
||
let last: f32 = support
|
||
.get(50)
|
||
.map_err(|e| MLError::ModelError(format!("Failed to get last: {}", e)))?
|
||
.to_scalar()
|
||
.map_err(|e| MLError::ModelError(format!("Failed to convert to scalar: {}", e)))?;
|
||
|
||
assert!(
|
||
(first - (-10.0_f32)).abs() < 1e-5,
|
||
"First support value should be v_min=-10.0, got {}",
|
||
first
|
||
);
|
||
assert!(
|
||
(last - 10.0_f32).abs() < 1e-5,
|
||
"Last support value should be v_max=10.0, got {}",
|
||
last
|
||
);
|
||
|
||
// Verify atoms are evenly spaced
|
||
let delta_z = (10.0 - (-10.0)) / 50.0; // (v_max - v_min) / (num_atoms - 1)
|
||
for i in 0..50 {
|
||
let expected = -10.0 + i as f32 * delta_z;
|
||
let actual: f32 = support
|
||
.get(i)
|
||
.map_err(|e| MLError::ModelError(format!("Failed to get atom {}: {}", i, e)))?
|
||
.to_scalar()
|
||
.map_err(|e| MLError::ModelError(format!("Failed to convert to scalar: {}", e)))?;
|
||
|
||
assert!(
|
||
(actual - expected).abs() < 1e-4,
|
||
"Atom {} should be {}, got {}",
|
||
i,
|
||
expected,
|
||
actual
|
||
);
|
||
}
|
||
|
||
Ok(())
|
||
}
|