DM-41952: Add flux sampling of full covariance model to preKernel

python/lsst/cp/pipe/makeBrighterFatterKernel.py

@@ -109,6 +109,21 @@ class BrighterFatterKernelSolveConfig(pipeBase.PipelineTaskConfig,
         default=False,
     )
 
+    useCovModelSample = pexConfig.Field(
+        dtype=bool,
+        doc="Use the covariance matrix sampled from the full covariance model "
+        "(Astier et al. 2019 equation 20) instead of the average measured covariances?",
+        default=False,
+    )
+
+    covModelFluxSample = pexConfig.DictField(
+        keytype=str,
+        itemtype=float,
+        doc="Flux level in electrons at which to sample the full covariance"
+        "model if useCovModelSample=True. The same level is applied to all"
+        "amps if this parameter [`dict`] is passed as {'ALL_AMPS': value}",
+        default={'ALL_AMPS': 25000.0},
+    )
     maxIterSuccessiveOverRelaxation = pexConfig.Field(
         dtype=int,
         doc="The maximum number of iterations allowed for the successive over-relaxation method",
@@ -119,7 +134,6 @@ class BrighterFatterKernelSolveConfig(pipeBase.PipelineTaskConfig,
         doc="The target residual error for the successive over-relaxation method",
         default=5.0e-14
     )
-
     correlationQuadraticFit = pexConfig.Field(
         dtype=bool,
         doc="Use a quadratic fit to find the correlations instead of simple averaging?",
@@ -200,6 +214,17 @@ def run(self, inputPtc, dummy, camera, inputDims):
             detectorCorrList = list()
             detectorFluxes = list()
 
+        if not inputPtc.ptcFitType == "FULLCOVARIANCE" and self.config.useCovModelSample:
+            raise ValueError("ptcFitType must be FULLCOVARIANCE if useCovModelSample=True.")
+
+        # Get flux sample dictionary
+        fluxSampleDict = {ampName: 0.0 for ampName in inputPtc.ampNames}
+        for ampName in inputPtc.ampNames:
+            if 'ALL_AMPS' in self.config.covModelFluxSample:
+                fluxSampleDict[ampName] = self.config.covModelFluxSample['ALL_AMPS']
+            elif ampName in self.config.covModelFluxSample:
+                fluxSampleDict[ampName] = self.config.covModelFluxSample[ampName]
+
         bfk = BrighterFatterKernel(camera=camera, detectorId=detector.getId(), level=self.config.level)
         bfk.rawMeans = inputPtc.rawMeans  # ADU
         bfk.rawVariances = inputPtc.rawVars  # ADU^2
@@ -221,6 +246,7 @@ def run(self, inputPtc, dummy, camera, inputDims):
         for amp in detector:
             ampName = amp.getName()
             gain = bfk.gain[ampName]
+            noiseMatrix = inputPtc.noiseMatrix[ampName]
             mask = inputPtc.expIdMask[ampName]
             if gain <= 0:
                 # We've received very bad data.
@@ -237,6 +263,8 @@ def run(self, inputPtc, dummy, camera, inputDims):
             # covariances may now have the mask already applied.
             fluxes = np.array(bfk.rawMeans[ampName])[mask]
             variances = np.array(bfk.rawVariances[ampName])[mask]
+            covModelList = np.array(inputPtc.covariancesModel[ampName])
+
             xCorrList = np.array([np.array(xcorr) for xcorr in bfk.rawXcorrs[ampName]])
             if np.sum(mask) < len(xCorrList):
                 # Only apply the mask if needed.
@@ -305,6 +333,14 @@ def run(self, inputPtc, dummy, camera, inputDims):
                 # Use a quadratic fit to the correlations as a
                 # function of flux.
                 preKernel = self.quadraticCorrelations(corrList, fluxes, f"Amp: {ampName}")
+            elif self.config.useCovModelSample:
+                # Sample the full covariance model at a given flux.
+                # Use the non-truncated fluxes for this
+                mu = bfk.rawMeans[ampName]
+                covTilde = self.sampleCovModel(mu, noiseMatrix, gain,
+                                               covModelList, fluxSampleDict[ampName],
+                                               f"Amp: {ampName}")
+                preKernel = np.pad(self._tileArray(-1.0 * covTilde), ((1, 1)))
             else:
                 # Use a simple average of the measured correlations.
                 preKernel = self.averageCorrelations(scaledCorrList, f"Amp: {ampName}")
@@ -423,6 +459,50 @@ def quadraticCorrelations(self, xCorrList, fluxList, name):
 
         return meanXcorr
 
+    def sampleCovModel(self, fluxes, noiseMatrix, gain, covModelList, flux, name):
+        """Sample the correlation model and measure
+        widetile{C}_{ij} from Broughton et al. 2023 (eq. 4)
+
+        Parameters
+        ----------
+        fluxes : `list` [`float`]
+            List of fluxes (in ADU)
+        noiseMatrix : `numpy.array`, (N, N)
+            Noise matrix
+        gain : `float`
+            Amplifier gain
+        covModelList : `numpy.array`, (N, N)
+            List of covariance model matrices. These are
+            expected to be square arrays.
+        flux : `float`
+            Flux in electrons at which to sample the
+            covariance model.
+        name : `str`
+            Name for log messages.
+
+        Returns
+        -------
+        covTilde : `numpy.array`, (N, N)
+            The calculated C-tilde from Broughton et al. 2023 (eq. 4).
+        """
+
+        # Get the index of the flux sample
+        # (this must be done in electron units)
+        ix = np.argmin((fluxes*gain - flux)**2)
+        assert len(fluxes) == len(covModelList)
+
+        # Find the nearest measured flux level
+        # and the full covariance model at that point
+        nearestFlux = fluxes[ix]
+        covModelSample = covModelList[ix]
+
+        # Calculate flux sample
+        # covTilde returned in ADU units
+        covTilde = (covModelSample - noiseMatrix/gain**2)/(nearestFlux**2)
+        covTilde[0][0] -= (nearestFlux/gain)/(nearestFlux**2)
+
+        return covTilde
+
     @staticmethod
     def _tileArray(in_array):
         """Given an input quarter-image, tile/mirror it and return full image.

tests/test_brighterFatterKernel.py

@@ -63,8 +63,11 @@ def setUp(self):
                           [0.025 * residual, 0.015 * residual, 0.01 * residual]]
             self.ptc.covariances['amp 1'].append(covariance)
 
+        self.ptc.covariancesModel = self.ptc.covariances
         self.ptc.gain['amp 1'] = 1.0
         self.ptc.noise['amp 1'] = 5.0
+        self.ptc.noiseMatrix['amp 1'] = np.zeros((3, 3))
+        self.ptc.noiseMatrix['amp 1'][0][0] = self.ptc.noise['amp 1']
 
         # This is empirically determined from the above parameters.
         self.ptc.aMatrix['amp 1'] = np.array([[2.14329806e-06, -4.28659612e-07, -5.35824515e-08],
@@ -105,6 +108,40 @@ def test_aMatrix(self):
         results = task.run(self.ptc, ['this is a dummy exposure'], self.camera, {'detector': 1})
         self.assertFloatsAlmostEqual(results.outputBFK.ampKernels['amp 1'], self.expectation, atol=1e-5)
 
+    def test_covSample(self):
+        """Test solution from Broughton et al. 2024 eq. 4 preKernel
+        """
+        config = cpPipe.BrighterFatterKernelSolveConfig()
+        config.useCovModelSample = True
+        config.covModelFluxSample = {'ALL_AMPS': 30000.}
+        task = cpPipe.BrighterFatterKernelSolveTask(config=config)
+
+        results = task.run(self.ptc, ['this is a dummy exposure'], self.camera, {'detector': 1})
+
+        expectation = np.array([[2.19577206e-08, 4.22977941e-08, 5.54871324e-08,
+                                 5.85845588e-08, 5.54871324e-08, 4.22977941e-08,
+                                 2.19577206e-08],
+                                [4.55330882e-08, 9.17463235e-08, 1.21066176e-07,
+                                 1.23363971e-07, 1.21066176e-07, 9.17463235e-08,
+                                 4.55330882e-08],
+                                [6.84283088e-08, 1.48088235e-07, 1.98667279e-07,
+                                 1.67738971e-07, 1.98667279e-07, 1.48088235e-07,
+                                 6.84283088e-08],
+                                [8.00919118e-08, 1.83511029e-07, 2.57775735e-07,
+                                 -4.97426471e-08, 2.57775735e-07, 1.83511029e-07,
+                                 8.00919118e-08],
+                                [6.84283088e-08, 1.48088235e-07, 1.98667279e-07,
+                                 1.67738971e-07, 1.98667279e-07, 1.48088235e-07,
+                                 6.84283088e-08],
+                                [4.55330882e-08, 9.17463235e-08, 1.21066176e-07,
+                                 1.23363971e-07, 1.21066176e-07, 9.17463235e-08,
+                                 4.55330882e-08],
+                                [2.19577206e-08, 4.22977941e-08, 5.54871324e-08,
+                                 5.85845588e-08, 5.54871324e-08, 4.22977941e-08,
+                                 2.19577206e-08]])
+
+        self.assertFloatsAlmostEqual(results.outputBFK.ampKernels['amp 1'], expectation, atol=1e-5)
+
     def test_quadratic(self):
         """Test quadratic correlation solver.