Add an estimator setting for fixed-tau fits

flimkit_bridge/fitting.py

@@ -42,6 +42,9 @@
      'applies_to': ('roi', 'per_pixel'), 'advanced': True, 'default': True},
     {'key': 'fit_sigma', 'label': 'Fit IRF broadening', 'type': 'bool',
      'applies_to': ('roi', 'per_pixel'), 'advanced': True, 'default': False},
+    {'key': 'estimator', 'label': 'Fixed-tau amplitude estimator', 'type': 'choice',
+     'choices': ('poisson', 'nnls'), 'applies_to': ('per_pixel',),
+     'advanced': True, 'default': 'poisson'},
     {'key': 'free_tau', 'label': 'Free tau per pixel', 'type': 'bool',
      'applies_to': ('per_pixel',), 'advanced': True, 'default': False},
     {'key': 'lm_restarts', 'label': 'LM restarts', 'type': 'int', 'min': 1,
@@ -362,6 +365,7 @@ def fit_pixels(stack, tcspc_res, n_bins, params, irf_prompt=None, bands=8,
             use_gpu='auto' if params.get('use_gpu', True) else False,
             fit_model=fit_model,
             free_tau=bool(params.get('free_tau', False)),
+            estimator=params.get('estimator', 'poisson'),
             **pileup,
         )
         collected.append(maps)

flimkit_bridge/pipeline.py

@@ -24,6 +24,8 @@
      'type': 'int', 'min': 0, 'max': 2000, 'default': 120},
     {'key': 'correct_pileup', 'label': 'Correct pile-up', 'type': 'bool',
      'default': False},
+    {'key': 'estimator', 'label': 'Fixed-tau amplitude estimator', 'type': 'choice',
+     'choices': ('poisson', 'nnls'), 'default': 'poisson', 'advanced': True},
     {'key': 'channel', 'label': 'Channel', 'type': 'int', 'min': 0, 'max': 16,
      'default': None, 'advanced': True},
     {'key': 'pipeline', 'label': 'Pipeline', 'type': 'choice',
@@ -171,6 +173,7 @@ def build_args(container, tile_dir, output_dir, basename, params):
     args.estimate_irf = params['irf_strategy']
     args.machine_irf = params['irf_path'] or str(MACHINE_IRF_DEFAULT_PATH)
     args.correct_pileup = bool(params['correct_pileup'])
+    args.estimator = params['estimator']
     args.channel = params['channel']
     args.dist_type = 'discrete'
     args.mode = params['mode']

flimkit_bridge/zstack.py

@@ -21,6 +21,8 @@
      'default': False},
     {'key': 'correct_pileup', 'label': 'Correct pile-up', 'type': 'bool',
      'default': False},
+    {'key': 'estimator', 'label': 'Fixed-tau amplitude estimator', 'type': 'choice',
+     'choices': ('poisson', 'nnls'), 'default': 'poisson', 'advanced': True},
     {'key': 'volume_format', 'label': 'Volume format', 'type': 'choice',
      'choices': ('ome-zarr', 'ome-tiff'), 'default': 'ome-zarr'},
     {'key': 'channel', 'label': 'Channel', 'type': 'int', 'min': 0, 'max': 16,
@@ -152,6 +154,7 @@ def build_args(ptu_dir, output_dir, params):
     args.estimate_irf = params['irf_strategy']
     args.machine_irf = params['irf_path'] or str(MACHINE_IRF_DEFAULT_PATH)
     args.correct_pileup = bool(params['correct_pileup'])
+    args.estimator = params['estimator']
     args.channel = params['channel']
     args.bound_fraction = bool(params['bound_fraction'])
     args.ref_tau1 = params['ref_tau1_ns']

pyproject.toml

@@ -4,14 +4,14 @@ build-backend = 'setuptools.build_meta'
 
 [project]
 name = 'flimkit-bridge'
-version = '0.8.0'
+version = '0.9.0'
 description = 'The local HTTP server FLIMKit image analysis clients talk to'
 readme = 'README.md'
 requires-python = '>=3.12'
 license = 'MIT'
 authors = [{ name = 'Alex Hunt', email = 'alexander.hunt@ed.ac.uk' }]
 dependencies = [
-    'flimkit>=0.13.2',
+    'flimkit>=0.14.2',
     'numpy',
     'tifffile',
     'zarr>=2.16',

tests/test_fitting.py

@@ -311,6 +311,31 @@ def test_per_pixel_uses_the_gpu_unless_it_is_turned_off():
     assert fitting.merge_params({'use_gpu': False})['use_gpu'] is False
 
 
+def test_estimator_defaults_to_poisson_and_takes_nnls():
+    assert fitting.defaults()['values']['estimator'] == 'poisson'
+    assert fitting.merge_params({'estimator': 'nnls'})['estimator'] == 'nnls'
+    with pytest.raises(ValueError):
+        fitting.merge_params({'estimator': 'bogus'})
+
+
+@pytest.mark.parametrize('estimator', ['poisson', 'nnls'])
+def test_the_estimator_reaches_fit_per_pixel(monkeypatch, estimator):
+    import flimkit.FLIM.fitters as fitters
+    seen = {}
+    real = fitters.fit_per_pixel
+
+    def spy(*args, **kwargs):
+        seen['estimator'] = kwargs.get('estimator')
+        return real(*args, **kwargs)
+
+    monkeypatch.setattr(fitters, 'fit_per_pixel', spy)
+    stack, _ = _synthetic_stack(shape=(4, 4), photons=40000)
+    params = fitting.merge_params({'n_exp': 2, 'irf_strategy': 'gaussian',
+                                   'use_gpu': False, 'estimator': estimator})
+    fitting.fit_pixels(stack, 5e-11, 256, params, bands=2)
+    assert seen['estimator'] == estimator
+
+
 def test_accepts_reports_a_missing_keyword():
     from flimkit_bridge.fitting import accepts
 

tests/test_pipeline.py

@@ -163,3 +163,10 @@ def test_products_stay_float_when_counts_are_not_whole(tmp_path):
     written = tifffile.imread(products[0]['file'])
     assert written.dtype == np.float32
     assert np.array_equal(written, counts)
+
+
+def test_the_pipeline_passes_the_estimator_on():
+    assert pipeline.defaults()['values']['estimator'] == 'poisson'
+    params = pipeline.merge_params({'estimator': 'nnls'})
+    args = pipeline.build_args('c.xlif', 'tiles', 'out', 'base', params)
+    assert args.estimator == 'nnls'

tests/test_zstack.py

@@ -349,3 +349,10 @@ def test_a_name_with_xml_in_it_cannot_break_the_metadata(tmp_path):
     xml = (Path(written) / 'OME' / 'METADATA.ome.xml').read_text(encoding='utf-8')
     ElementTree.fromstring(xml)
     assert 'a<b&c"d' in xml
+
+
+def test_the_zstack_passes_the_estimator_on():
+    assert zstack.defaults()['values']['estimator'] == 'poisson'
+    params = zstack.merge_params({'estimator': 'nnls'})
+    args = zstack.build_args('tiles', 'out', params)
+    assert args.estimator == 'nnls'