Keep gradients finite when electrodes are driven to zero amplitude

dynaphos/simulator.py

@@ -12,6 +12,12 @@
                             get_deg2pix_coeff, set_deterministic,
                             print_stats, sigmoid, to_numpy, Map)
 
+# Lower bound on the argument of the sqrt size equation. Far below any
+# physical stimulation current, so the forward pass is unaffected; its only
+# purpose is to keep the gradient finite at zero amplitude.
+MIN_SIZE_ARG = 1e-30
+
+
 class State:
     def __init__(self, params: dict, shape: Tuple[int, ...],
                  verbose: Optional[bool] = False):
@@ -145,7 +151,14 @@ def __init__(self, params: dict, shape: Tuple[int, ...],
         p = self.params['size']
         if p['size_equation'] == 'sqrt':  # Tehovnik 2007
             def f(x):
-                return torch.sqrt(torch.div(x, p['current_spread']))
+                # The derivative of sqrt is infinite at zero, so an electrode
+                # driven to exactly 0 A would emit a NaN gradient that
+                # propagates silently through the whole model while the
+                # forward pass still looks correct. Clamping the argument
+                # keeps the derivative finite and leaves the forward output
+                # unchanged for any physically meaningful current.
+                return torch.sqrt(
+                    torch.div(x, p['current_spread']).clamp_min(MIN_SIZE_ARG))
         elif p['size_equation'] == 'sigmoid':  # Bosking et al., 2017
             def f(x):
                 return 0.5 * p['MD'] * sigmoid(p['slope_size'] *

test/test_gradients.py

@@ -0,0 +1,62 @@
+"""Gradients must stay finite when electrodes are switched off.
+
+The simulator is used for end-to-end optimisation, where an encoder routinely
+drives some electrodes to zero: sparse stimulation is the goal for a
+charge-limited device, sigmoid outputs underflow, and any charge penalty pushes
+amplitudes down. The sqrt size equation has an infinite derivative at zero, so
+without a lower bound those electrodes emit NaN gradients that propagate into
+every parameter while the forward pass still looks correct.
+"""
+
+import os
+
+import numpy as np
+import pytest
+import torch
+
+from dynaphos.cortex_models import \
+    get_visual_field_coordinates_probabilistically
+from dynaphos.simulator import GaussianSimulator
+from dynaphos.utils import load_params
+
+CONFIG = os.path.join(os.path.dirname(__file__), '..', 'config', 'params.yaml')
+
+
+def build(batch_size=2, num_phosphenes=64, resolution=32, seed=0):
+    params = load_params(CONFIG)
+    params['run']['resolution'] = [resolution, resolution]
+    params['run']['batch_size'] = batch_size
+    params['run']['seed'] = seed
+    coordinates = get_visual_field_coordinates_probabilistically(
+        params, num_phosphenes, np.random.default_rng(seed))
+    return GaussianSimulator(params, coordinates, rng=np.random.default_rng(seed))
+
+
+@pytest.mark.parametrize('fraction_off', [1.0, 0.5, 0.25])
+def test_gradient_finite_with_electrodes_off(fraction_off):
+    simulator = build()
+    amplitude = torch.full((2, simulator.num_phosphenes), 80e-6)
+    n_off = int(simulator.num_phosphenes * fraction_off)
+    amplitude[:, :n_off] = 0.0
+    amplitude.requires_grad_(True)
+
+    simulator.reset()
+    percept = simulator(amplitude)
+    percept.sum().backward()
+
+    assert torch.isfinite(percept).all(), 'forward pass produced non-finite values'
+    assert torch.isfinite(amplitude.grad).all(), (
+        f'{torch.isnan(amplitude.grad).float().mean():.0%} of amplitude '
+        f'gradients are NaN with {fraction_off:.0%} of electrodes off')
+
+
+def test_forward_unchanged_by_the_clamp():
+    """The bound sits far below any physical current, so percepts must not move."""
+    simulator = build()
+    generator = torch.Generator().manual_seed(0)
+    amplitude = torch.rand(
+        (2, simulator.num_phosphenes), generator=generator) * 100e-6
+    simulator.reset()
+    percept = simulator(amplitude)
+    assert torch.isfinite(percept).all()
+    assert percept.max() > 0, 'expected a visible percept at up to 100 uA'