[BUG] WaveDiff `TFPhysicalPolychromaticField` fails for arbitrary FOV positions due to exact-match lookup in `find_position_indices`

[jeipollack]

Describe the Bug
Inference fails when passing valid FOV positions that are not exactly present in obs_pos.
The failure occurs in:

find_position_indices(obs_pos, batch_positions)

Steps to Reproduce

inference_conf_dataset.yaml
inference:
  # Inference batch size
  batch_size: 16
  
  # Cycle to use for inference. Can be: 1, 2, ...
  cycle: 2

  # Dataset schema used to interpret input fields and control downstream conversion.
  # This does NOT modify dataset content; it only controls how fields are interpreted during conversion.
  # This affects required fields and available conversion handlers.
  #
  # Choices:
  # - INFERENCE: standard inference mode (requires positions and seds)
  # - EVALUATION: evaluation mode (may include additional fields such as sources or masks)
  schema_mode: INFERENCE
  
  # Paths to the configuration files and trained model directory
  configs:
    # Path to the directory containing the trained model
    trained_model_path: /sps/euclid/Users/jpollack/DR1/ppo_a0c6db68/workdir/

    # Subdirectory name of the trained model, e.g. psf_model
    model_subdir: checkpoint
  
    # Relative Path to the training configuration file used to train the model
    trained_model_config_path: config/training_config.yaml

    # Path to the data config file (this could contain prior information)
    data_config_path: config/data_config.yaml

  # The following parameters will overwrite the `model_params` in the training config file.
  model_params:
    # Num of wavelength bins to reconstruct polychromatic objects.
    n_bins_lambda: 20

    # Downsampling rate to match the oversampled model to the specified telescope's sampling.
    output_Q: 3

    # Dimension of the pixel PSF postage stamp
    output_dim: 32

    # Flag to perform centroid error correction
    correct_centroids: True

    # Flag to perform CCD misalignment error correction
    add_ccd_misalignments: True

    # Path to ccd_misalignments file
    ccd_misalignments_aux_path: /path/to/ccd_misalignments/tiles.npy

Code:

# Set path to inference configuration file
inference_config_path = "/path/to/yaml/inference_conf_dataset.yaml"

# Create PSFInference instance

psf_inferred = PSFInference(inference_config_path = inference_config_path,x_field = [-152.5],
                             y_field = [-160.596],
                             seds = seds,
                             sources = data["images"],
                             masks = masks)
                            
psf_inferred.prepare_configs()

Result: TensorFlow assertion fails with error:

assertion failed: [Some positions not found in obs_pos] [Condition x == y did not hold element-wise:] [x (All_1:0) = ] [0] [y (assert_equal_1[/y:0](https://notebook.cc.in2p3.fr/y#line=-1)) = ] [1]
	 [[{{node assert_equal_1[/Assert/AssertGuard/Assert](https://notebook.cc.in2p3.fr/Assert/AssertGuard/Assert)}}]] [Op:__inference_find_position_indices_402192]

Call arguments received by layer 'tf_physical_polychromatic_field_2' (type TFPhysicalPolychromaticField):
  • inputs=['tf.Tensor(shape=(1, 2), dtype=float32)', 'tf.Tensor(shape=(1, 20, 3), dtype=float32)']
  • training=False

Nearest-neighbour check confirms positions are within valid FOV range but do not exactly match any obs_pos entry

Expected Behaviour
Inference should support:

• arbitrary FOV positions
• continuous coordinate space

Screenshots

Your Setup
Tested in a Jupyter notebook with a conda environment with TF 2.15.0

[jeipollack]

The code generating the error is the following:

    def call(self, positions):
        """Calculate the prior Zernike coefficients for a batch of positions.

        This method calculates the Zernike coefficients for a batch of input positions
        based on the pre-computed Zernike coefficients for observed positions.

        Parameters
        ----------
        positions : tf.Tensor
            Tensor of shape (batch_size, 2) representing the positions.
            The first element represents the x-axis, and the second element represents the y-axis.

        Returns
        -------
        zernike_coeffs : tf.Tensor
            Tensor of shape (batch_size, n_zernikes, 1, 1) containing the prior Zernike coefficients
            corresponding to the input positions.

        Notes
        -----
        The method retrieves the Zernike coefficients for each input position from the pre-computed
        Zernike coefficients stored for observed positions. It matches each input position with
        the closest observed position and retrieves the corresponding Zernike coefficients.

        Before calling this method, ensure that the position polynomial matrix and the
        corresponding Zernike coefficients have been precomputed and set for the layer.


        Raises
        ------
        ValueError
            If the shape of the input `positions` tensor is not compatible.

        """
        # Find indices for all positions in one batch operation
        idx = find_position_indices(self.obs_pos, positions)

        # Gather the corresponding Zernike coefficients
        batch_zks = tf.gather(self.zks_prior, idx, axis=0)

        return batch_zks[:, :, tf.newaxis, tf.newaxis]

which is automatically called in the TFPhysical.__init__ constructor.

class TFPhysicalLayer(tf.keras.layers.Layer):
    def __init__(
        self,
        obs_pos,
        zks_prior,
        interpolation_type=None,
        interpolation_args=None,
        name="TF_physical_layer",
    ):
        super().__init__(name=name)
        self.obs_pos = obs_pos
        self.zks_prior = zks_prior

        if interpolation_args is None:
            self.interpolation_args = {"order": 2, "K": 50}

        # Define the prediction routine by default
        self.predict = self.call

        # Define the prediction routine

        if interpolation_type == "all":
            self.predict = self.interpolate_all
        elif interpolation_type == "top_K":
            self.predict = self.interpolate_top_K
        elif interpolation_type == "independent_Zk":
            self.predict = self.interpolate_independent_Zk

I tested the interpolation schemes below and they generated for poor results.

I consulted @tobias-liaudat , who informed me that in his investigations he found the interpolation schemes to be inadequate compared to directly using the data (prior data). For CCD misalignment contribution, the inference data should be directly used when re-constructing the trained PSF model, previously it was the training data.

So, I implemented some small changes:

@property
 def trained_psf_model(self):
         """
@@ -469,7 +470,7 @@ class PSFInference:
         # Load the trained PSF model
         return load_trained_psf_model(
             self.training_config,
-            self.model_data_adapter.complete_data,
+            self.inference_data_adapter.complete_data,
             weights_path_pattern,

and added helper for converting inference data to TensorFlow type

+    def _convert_inference_data_to_tensorflow(self):
+        # Convert to TensorFlow according to dataset schema mode
+        if self._inference_data_adapter.representation_state == RepresentationState.NUMPY:
+            self._inference_data_adapter.convert_to_tensorflow(
+                self.simPSF,
+                self.n_bins_lambda,
+                mode=self.config_handler.schema_mode,
+            )

I tested in a notebook generating PSFs with INFERENCE data schema (positions and SEDS with CCD misalignment corrections) and EVALUATION data schema (positions, SEDs, images, masks with Centroid and CCD misalignment corrections.

Below are some results evaluated on simulated data with masks.

Data Schema	Pixel RMSE +/- STD	Relative Pixel RMSE (%) +/- STD
INFERENCE.	3.76432062e-03 ± 1.62560930e-03	4.74398823e+01 ± 1.90543369e+01
EVALUATION	1.16268248e-04 ± 2.51358237e-05	1.50518900e+00 ± 4.42096629e-01

INFERENCE reconstruction

EVALUATION reconstruction

[jeipollack]

Evaluation results are consistent with previous iterations and the validation table

[jeipollack]

forgot to add inference result for random position