Implement direction-independent calibration in peel subcommand

mdbook/src/user/peel/intro.md

@@ -114,9 +114,14 @@ hyperdrive also borrows from RTS, which used a gaussian weighting scheme to tape
 
 ## Source Counts
 
-Similar to [`vis-subtract`](../vis_subtract/intro.md), the total number of sources in the sky model to be subtracted is limited by `--num-sources`, but only the top `--iono-sub` brightest have their ionospheric constants measured and modeled during subtraction. By default, hyperdrive will include sources in the sky model after [vetoing](../vis_simulate/intro.md#vetoing) if `--num-sources` is not specified, and will ionospherically subtract all of these sources if `--iono-sub` is not specified. An error will occur if `--num-sources` is greater than the number of sources in the sky model.
+Similar to [`vis-subtract`](../vis_subtract/intro.md), the total number of sources in the sky model to be subtracted is limited by `--num-sources`, but only the top `--iono-sub` brightest have their ionospheric constants measured and modeled during subtraction. Additionally, the top `--peel` brightest sources will undergo full DI calibration peeling.
 
-Future versions of peel will include a `--peel` argument to specify the number of sources to peel.
+The relationship between these parameters is:
+- `--num-sources`: Total sources in sky model to subtract (default: all after vetoing)
+- `--iono-sub`: Number of sources for ionospheric subtraction (default: all sources, ≤ num-sources)  
+- `--peel`: Number of sources for full DI calibration peeling (default: 0, ≤ iono-sub)
+
+An error will occur if `--num-sources` is greater than the number of sources in the sky model, or if `--peel` > `--iono-sub` > `--num-sources`.
 
 ## High level overview
 
@@ -159,4 +164,39 @@ flowchart TD
 
 ## Peeling
 
-A full peel involves performing direction-independent calibration towards each source, this is currently a work in progress and not yet implemented.
+A full peel involves performing direction-independent calibration towards each source after ionospheric subtraction. This functionality has been implemented and can be controlled with the `--peel` argument.
+
+### Usage
+
+To perform full peeling on a subset of sources, use the `--peel` argument to specify how many of the brightest sources should be fully peeled:
+
+```shell
+hyperdrive peel --peel 3 --iono-sub 10 [other arguments...]
+```
+
+This example will:
+- Ionospherically subtract the 10 brightest sources
+- Perform full DI calibration peeling on the 3 brightest sources
+
+### DI Calibration Parameters
+
+The DI calibration step during peeling can be controlled with these additional parameters:
+
+- `--di-max-iterations`: Maximum number of DI calibration iterations (default: 50)
+- `--di-stop-threshold`: Stop threshold for convergence (default: 1e-8)  
+- `--di-min-threshold`: Minimum threshold for convergence (default: 1e-4)
+
+### Difference between Ionospheric Subtraction and Full Peeling
+
+**Ionospheric Subtraction** (`--iono-sub`):
+- Solves for ionospheric refraction parameters α and β (proportional to λ²)
+- Applies scalar gain correction
+- Subtracts the ionosphere-corrected model from visibilities
+
+**Full Peeling** (`--peel`):
+- Performs ionospheric subtraction (as above)
+- **Additionally** performs direction-independent calibration toward the source
+- Solves for complex gain corrections per antenna and channel
+- Corrects for scintillation and other phase/amplitude effects
+
+Full peeling provides more complete correction but is computationally more expensive and should typically be applied only to the brightest sources.

src/cli/error.rs

@@ -173,6 +173,7 @@ impl From<PeelArgsError> for HyperdriveError {
         match e {
             PeelArgsError::NoOutput
             | PeelArgsError::TooManyIonoSub { .. }
+            | PeelArgsError::TooManyPeel { .. }
             | PeelArgsError::ZeroPasses
             | PeelArgsError::ZeroLoops
             | PeelArgsError::ParseIonoTimeAverageFactor(_)

src/cli/peel/error.rs

@@ -12,6 +12,11 @@ pub(crate) enum PeelArgsError {
     )]
     TooManyIonoSub { total: usize, iono: usize },
 
+    #[error(
+        "The number of sources to subtract ({total}) is less than the number of sources to peel ({peel})"
+    )]
+    TooManyPeel { total: usize, peel: usize },
+
     #[error("The number of iono sub passes cannot be 0")]
     ZeroPasses,
 

src/cli/peel/mod.rs

@@ -82,13 +82,12 @@ pub(crate) struct PeelCliArgs {
     #[clap(long = "iono-sub", help_heading = "PEELING")]
     pub(super) num_sources_to_iono_subtract: Option<usize>,
 
-    // TODO: peel
-    // The number of sources to peel. Peel sources are treated the same as
-    // "ionospherically subtracted" sources, except before subtracting, a "DI
-    // calibration" is done between the iono-rotated model and the data. This
-    // allows for scintillation and any other phase shift to be corrected.
-    // #[clap(long = "peel", help_heading = "PEELING")]
-    // pub(super) num_sources_to_peel: Option<usize>,
+    /// The number of sources to peel. Peel sources are treated the same as
+    /// "ionospherically subtracted" sources, except before subtracting, a "DI
+    /// calibration" is done between the iono-rotated model and the data. This
+    /// allows for scintillation and any other phase shift to be corrected.
+    #[clap(long = "peel", help_heading = "PEELING")]
+    pub(super) num_sources_to_peel: Option<usize>,
     #[clap(long, help = NUM_PASSES_HELP.as_str(), help_heading = "PEELING")]
     pub(super) num_passes: Option<usize>,
 
@@ -113,6 +112,21 @@ pub(crate) struct PeelCliArgs {
     #[clap(long, help = CONVERGENCE_HELP.as_str(), help_heading = "PEELING")]
     pub(super) convergence: Option<f64>,
 
+    /// Maximum number of iterations for DI calibration during peeling. 
+    /// Default: 50
+    #[clap(long = "di-max-iterations", help_heading = "PEELING")]
+    pub(super) di_max_iterations: Option<u32>,
+
+    /// Stop threshold for DI calibration convergence during peeling.
+    /// Default: 1e-8
+    #[clap(long = "di-stop-threshold", help_heading = "PEELING")]
+    pub(super) di_stop_threshold: Option<f64>,
+
+    /// Minimum threshold for DI calibration convergence during peeling.
+    /// Default: 1e-4
+    #[clap(long = "di-min-threshold", help_heading = "PEELING")]
+    pub(super) di_min_threshold: Option<f64>,
+
     #[clap(short, long, multiple_values(true), help = VIS_OUTPUTS_HELP.as_str(), help_heading = "OUTPUT FILES")]
     pub(super) outputs: Option<Vec<PathBuf>>,
 
@@ -208,6 +222,7 @@ impl PeelArgs {
             peel_args:
                 PeelCliArgs {
                     num_sources_to_iono_subtract,
+                    num_sources_to_peel,
                     num_passes,
                     num_loops,
                     iono_time_average,
@@ -216,6 +231,9 @@ impl PeelArgs {
                     uvw_max,
                     short_baseline_sigma,
                     convergence,
+                    di_max_iterations,
+                    di_stop_threshold,
+                    di_min_threshold,
                     outputs,
                     output_vis_time_average,
                     output_vis_freq_average,
@@ -287,6 +305,16 @@ impl PeelArgs {
             }
         }
 
+        if let Some(ps) = num_sources_to_peel {
+            if ps > sky_model_source_count {
+                return Err(PeelArgsError::TooManyPeel {
+                    total: sky_model_source_count,
+                    peel: ps,
+                }
+                .into());
+            }
+        }
+
         let num_passes = NonZeroUsize::try_from(num_passes.unwrap_or(DEFAULT_NUM_PASSES))
             .map_err(|_| PeelArgsError::ZeroPasses)?;
 
@@ -537,16 +565,36 @@ impl PeelArgs {
 
         let num_sources_to_iono_subtract =
             num_sources_to_iono_subtract.unwrap_or(source_list.len());
+        let num_sources_to_peel = num_sources_to_peel.unwrap_or(0);
+
+        // Ensure num_sources_to_peel doesn't exceed num_sources_to_iono_subtract
+        if num_sources_to_peel > num_sources_to_iono_subtract {
+            return Err(PeelArgsError::TooManyPeel {
+                total: num_sources_to_iono_subtract,
+                peel: num_sources_to_peel,
+            }
+            .into());
+        }
+
+        let di_max_iterations = di_max_iterations.unwrap_or(50);
+        let di_stop_threshold = di_stop_threshold.unwrap_or(1e-8);
+        let di_min_threshold = di_min_threshold.unwrap_or(1e-4);
 
         let mut peel_printer = InfoPrinter::new("Peeling set up".into());
-        peel_printer.push_block(vec![
+        let mut print_blocks = vec![
             format!("Subtracting {} sources", source_list.len()).into(),
             format!(
                 "Ionospheric subtracting {} sources",
                 num_sources_to_iono_subtract
             )
             .into(),
-        ]);
+        ];
+        if num_sources_to_peel > 0 {
+            print_blocks.push(
+                format!("DI calibrating {} sources", num_sources_to_peel).into(),
+            );
+        }
+        peel_printer.push_block(print_blocks);
         if num_sources_to_iono_subtract > 0 {
             let mut block = vec![];
             block.push("Finding ionospheric offsets with data at:".into());
@@ -612,6 +660,10 @@ impl PeelArgs {
             peel_weight_params,
             peel_loop_params,
             num_sources_to_iono_subtract,
+            num_sources_to_peel,
+            di_max_iterations,
+            di_stop_threshold,
+            di_min_threshold,
         })
     }
 
@@ -636,6 +688,7 @@ impl PeelCliArgs {
             num_sources_to_iono_subtract: self
                 .num_sources_to_iono_subtract
                 .or(other.num_sources_to_iono_subtract),
+            num_sources_to_peel: self.num_sources_to_peel.or(other.num_sources_to_peel),
             num_passes: self.num_passes.or(other.num_passes),
             num_loops: self.num_loops.or(other.num_loops),
             iono_time_average: self.iono_time_average.or(other.iono_time_average),
@@ -644,6 +697,9 @@ impl PeelCliArgs {
             uvw_max: self.uvw_max.or(other.uvw_max),
             short_baseline_sigma: self.short_baseline_sigma.or(other.short_baseline_sigma),
             convergence: self.convergence.or(other.convergence),
+            di_max_iterations: self.di_max_iterations.or(other.di_max_iterations),
+            di_stop_threshold: self.di_stop_threshold.or(other.di_stop_threshold),
+            di_min_threshold: self.di_min_threshold.or(other.di_min_threshold),
             outputs: self.outputs.or(other.outputs),
             output_vis_time_average: self
                 .output_vis_time_average

src/params/peel/gpu.rs

@@ -29,6 +29,10 @@ pub(crate) fn peel_gpu(
     source_list: &SourceList,
     iono_consts: &mut [IonoConsts],
     source_weighted_positions: &[RADec],
+    num_sources_to_peel: usize,
+    di_max_iterations: u32,
+    di_stop_threshold: f64,
+    di_min_threshold: f64,
     peel_loop_params: &PeelLoopParams,
     chanblocks: &[Chanblock],
     low_res_lambdas_m: &[f64],
@@ -650,8 +654,11 @@ pub(crate) fn peel_gpu(
                 pb_trace!("{:?}: subtract_iono", start.elapsed());
 
                 // Peel?
-                let num_sources_to_peel = 0;
                 if pass == num_passes - 1 && i_source < num_sources_to_peel {
+                    multi_progress_bar.suspend(|| {
+                        debug!("Performing DI calibration for source {i_source}: {source_name}")
+                    });
+
                     // We currently can only do DI calibration on the CPU. Copy the visibilities back to the host.
                     let vis = d_high_res_vis_tfb.copy_from_device_new()?;
 
@@ -682,15 +689,37 @@ pub(crate) fn peel_gpu(
                         di_jones.view_mut(),
                         timeblock,
                         chanblocks,
-                        50,
-                        1e-8,
-                        1e-4,
+                        di_max_iterations,
+                        di_stop_threshold,
+                        di_min_threshold,
                         obs_context.polarisations,
                         pb,
                         true,
                     );
-                    if di_cal_results.into_iter().all(|r| r.converged) {
-                        // Apply.
+
+                    // Check if calibration converged for all channels
+                    let converged_count = di_cal_results.iter().filter(|r| r.converged).count();
+                    let total_channels = di_cal_results.len();
+
+                    multi_progress_bar.suspend(|| {
+                        info!(
+                            "DI calibration for {}: {}/{} channels converged", 
+                            source_name, converged_count, total_channels
+                        )
+                    });
+
+                    if converged_count > total_channels / 2 {
+                        multi_progress_bar.suspend(|| {
+                            debug!("DI calibration solutions computed for {} - application not yet implemented", source_name)
+                        });
+                        // TODO: Apply the DI Jones corrections
+                    } else {
+                        multi_progress_bar.suspend(|| {
+                            warn!(
+                                "DI calibration for {} failed to converge on enough channels ({}/{}) - not applying",
+                                source_name, converged_count, total_channels
+                            )
+                        });
                     }
                 }