source/workstation-adaptations.patch

diff --git a/main.py b/main.py
index 4c8a8f9..f45d630 100644
--- a/main.py
+++ b/main.py
@@ -19,13 +19,19 @@ from dataset import (CUB_PRESEGM_PARTS, CUB_PRESEGM_NUM_CLASSES, CUB_PRESEGM_DIR
 import torch
 from model import Model
 import lightning.pytorch as pl
+from lightning.fabric.plugins.environments.mpi import MPIEnvironment
 import os
 import argparse
 from utils import set_seed, compute_balanced_accuracy_loader, compute_acc_unordered, get_concept_balance
 from datetime import datetime
 from competitors import DNN, CBMDeep, CRM, CMR, _plot_intervenability
 from torchvision.utils import save_image
-from lightning.pytorch.loggers import WandbLogger
+from lightning.pytorch.loggers import CSVLogger
+
+# This workstation has an unrelated mpi4py package but no MPI runtime.  The
+# released experiment is single-device, so disable Lightning's auto-detection
+# probe without altering any scientific code path.
+MPIEnvironment.detect = staticmethod(lambda: False)
 
 
 def inspect_val_predictions_with_pause(model, val_loader, concept_names, output_folder, device):
@@ -542,10 +548,10 @@ def main(args):
     if test_only:
         log_name += "_TESTING"
     
-    # Save wandb logs to new_outputs/<timestamp>_<extra>/wandb
-    # You can change the 'entity' argument below to your team name (e.g. entity="my-team-name")
+    # Local logger: environment-only adaptation, with no change to the model,
+    # data, loss, optimizer, or evaluation paths.
     run_name = f"{timestamp}_{dataset}_{args.extra}" if args.extra != "" else f"{timestamp}_{dataset}"
-    logger = WandbLogger(project="HigherOrderCBMs", entity="conceptlords", save_dir=base_output_folder, name=run_name)
+    logger = CSVLogger(save_dir=base_output_folder, name=run_name)
     log_folder = os.path.join(base_output_folder, "checkpoints")
 
     # collect all hyperparameters and dataset data for logging
@@ -628,6 +634,8 @@ def main(args):
 
     if DEVICE == "cpu":
         trainer = pl.Trainer(accelerator='cpu', max_epochs=epochs, logger=logger, callbacks=selected_callbacks)
+    elif DEVICE == "mps":
+        trainer = pl.Trainer(accelerator='mps', devices=1, max_epochs=epochs, logger=logger, callbacks=selected_callbacks)
     else:
         trainer = pl.Trainer(accelerator='gpu', devices=[int(DEVICE[-1])], max_epochs=epochs, logger=logger, callbacks=selected_callbacks)
     if not test_only:
@@ -641,6 +649,8 @@ def main(args):
         print(f"Restoring best model from: {best_model_path}")  # pass val_loader and train_loader
         model = type(model).load_from_checkpoint(best_model_path, val_loader=val_loader, train_loader=train_loader)
         model.to(DEVICE)
+        model.ALWAYS_USE_TRUE_MASKS = True if always_use_true_masks else False
+        model.USE_INITIAL_PROTO_EMBS = True if use_initial_proto_embs else False
 
         # print epoch number of best model
         print(f"Best model epoch: {trainer.current_epoch}")
@@ -682,10 +692,8 @@ def main(args):
         "final_test_c_acc": test_c_acc
     })
 
-    if isinstance(model, Model) and dataset != "cub":
-        model.generate_plots(logger=logger, output_folder=output_folder)
-
-    _plot_intervenability(model, logger=logger, nb_runs=1, specific_loader=test_loader, specific_epoch_nb="final_test", is_pgcm=args.model=="PGCM")
+    # WandB-only plots are omitted; the audit computes and stores the numeric
+    # prototype/intervention metrics in a separate deterministic driver.
 
     if dataset == "celebamask" and args.model == "CBM":
         inspect_val_predictions_with_pause(
diff --git a/model.py b/model.py
index 029fb85..5ca81eb 100644
--- a/model.py
+++ b/model.py
@@ -244,7 +244,8 @@ class Model(pl.LightningModule):
         if self.reconstruction and not self.swapped and not self.use_pretrained_autoencoder:
             prototypes_as_features = self.map_proto_to_image(prototype_embeddings)
             prototype_emb_rec = self.map_proto_image_to_proto(prototypes_as_features) if not self.USE_INITIAL_PROTO_EMBS else prototype_embeddings
-            prototype_emb_loss = F.mse_loss(prototype_emb_rec, prototype_embeddings)
+            prototype_embeddings_batch = prototype_embeddings.unsqueeze(0).expand_as(prototype_emb_rec)
+            prototype_emb_loss = F.mse_loss(prototype_emb_rec, prototype_embeddings_batch)
             concept_logits_per_proto = self.map_proto_to_concepts(prototype_emb_rec)
             concept_probs_per_proto = torch.sigmoid(concept_logits_per_proto)
             prototypes_as_images = prototypes_as_features.unsqueeze(0) if prototypes_as_features.dim() == 2 else prototypes_as_features
@@ -325,7 +326,16 @@ class Model(pl.LightningModule):
         if self.is_cubEMB:
             return self._forward_cubEMB(x, interventions_mask=interventions_mask, standard_interventions_mask=standard_interventions_mask)
 
-        if self.use_pretrained_segmenter:
+        if self.ALWAYS_USE_TRUE_MASKS and not self.use_pretrained_segmenter:
+            # Native ColorMNIST supplies exact object masks.  This ablation
+            # isolates the released PGCM/prototype path from the separately
+            # pretrained segmenter artifact used in the paper environment.
+            X, M_test, C, Y = x
+            M = M_test
+            generated_M = M_test
+            generated_M_logits = None
+            segmentation_loss = torch.tensor(0.0, device=self.device)
+        elif self.use_pretrained_segmenter:
             segmentation, generated_M, C, Y, X = x
             generated_M_logits = None
             M = generated_M
@@ -382,7 +392,8 @@ class Model(pl.LightningModule):
             # prototype_emb_rec = prototype_emb_rec.view(b, nb_proto, -1)  # (b, nb_proto, proto_size)
             # prototype_embeddings_batch = prototype_embeddings.unsqueeze(0).expand(b, -1, -1)  # (b, nb_proto, proto_size)
             # prototype_emb_loss = F.mse_loss(prototype_emb_rec, prototype_embeddings_batch)
-            prototype_emb_loss = F.mse_loss(prototype_emb_rec, prototype_embeddings)
+            prototype_embeddings_batch = prototype_embeddings.unsqueeze(0).expand_as(prototype_emb_rec)
+            prototype_emb_loss = F.mse_loss(prototype_emb_rec, prototype_embeddings_batch)
 
             concept_logits_per_proto = self.map_proto_to_concepts(prototype_emb_rec.view(-1, prototype_emb_rec.shape[-1])) # (nb_proto, nb_concepts)
             concept_probs_per_proto = torch.sigmoid(concept_logits_per_proto)  # (nb_proto, nb_concepts)
@@ -744,6 +755,12 @@ class Model(pl.LightningModule):
 
                 prototypes_as_images = prototypes_as_images.unsqueeze(1)  # (b, 1, nb_proto, 3, H, W)
 
+                target_shape = (
+                    masked_images.shape[0], masked_images.shape[1],
+                    prototypes_as_images.shape[2], *masked_images.shape[3:]
+                )
+                prototypes_as_images = prototypes_as_images.expand(target_shape)
+                masked_images = masked_images.expand(target_shape)
                 recons_error = F.mse_loss(prototypes_as_images, masked_images, reduction='none')  # (b, nb_objects, nb_proto, 3, H, W)
                 recons_error = recons_error.mean(dim=[3, 4, 5])  # (b, nb_objects, nb_proto)
 
@@ -909,6 +926,12 @@ class Model(pl.LightningModule):
 
                 prototypes_as_images = prototypes_as_images.unsqueeze(1)  # (b, 1, nb_proto, 3, H, W)
 
+                target_shape = (
+                    masked_images.shape[0], masked_images.shape[1],
+                    prototypes_as_images.shape[2], *masked_images.shape[3:]
+                )
+                prototypes_as_images = prototypes_as_images.expand(target_shape)
+                masked_images = masked_images.expand(target_shape)
                 recons_error = F.mse_loss(prototypes_as_images, masked_images, reduction='none')  # (b, nb_objects, nb_proto, 3, H, W)
                 recons_error = recons_error.mean(dim=[3, 4, 5])  # (b, nb_objects, nb_proto)
 
@@ -1056,8 +1079,9 @@ class Model(pl.LightningModule):
             used_indices = get_used_prototypes_indices(self, self.val_loader)
             print(f"Used {len(used_indices)} prototypes out of {self.prototypes.weight.shape[0]}")
             print(f"Used prototypes: {used_indices}")
-        if not self.is_cubEMB and self.current_epoch % self.plot_frequency == 0:
-            self.generate_plots(logger=None)
+        # Visualization is omitted in this local audit because the released
+        # helper is hard-wired to a WandB logger.  Numeric training,
+        # validation, prototype swapping, and checkpoints are unchanged.
 
         if self.current_epoch == self.trainer.max_epochs // 2:
 
diff --git a/neural_networks.py b/neural_networks.py
index f74b02c..18350c2 100644
--- a/neural_networks.py
+++ b/neural_networks.py
@@ -237,7 +237,7 @@ class ResNetUNetSegmenter(nn.Module):
         
         # 1. Encoder (Pre-trained ResNet18 or ResNet34 is usually sufficient)
         # We grab the layers to access intermediate features for skip connections
-        base_model = models.resnet18(pretrained=True)
+        base_model = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1)
         self.base_layers = list(base_model.children())
         
         self.layer0 = nn.Sequential(*self.base_layers[:3]) # size=(N, 64, x.H/2, x.W/2)
@@ -299,6 +299,20 @@ class ResNetUNetSegmenter(nn.Module):
 class ResNetUNetSegmenterMNIST(nn.Module):
     def __init__(self, n_class=1):
         super().__init__()
+
+        # The released forward method references the same four ResNet stages
+        # and three decoder blocks as the generic segmenter, but the MNIST
+        # constructor omitted their initialization.  Restore those literal
+        # modules with channel dimensions implied by the released forward.
+        base_model = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1)
+        self.base_layers = list(base_model.children())
+        self.layer0 = nn.Sequential(*self.base_layers[:3])
+        self.layer1 = nn.Sequential(*self.base_layers[3:5])
+        self.layer2 = self.base_layers[5]
+        self.layer3 = self.base_layers[6]
+        self.up3 = self._up_block(256, 128)
+        self.up2 = self._up_block(128 + 128, 64)
+        self.up1 = self._up_block(64 + 64, 64)
         
         self.final_up = nn.Sequential(
             nn.Conv2d(64, 32, kernel_size=3, padding=1),
diff --git a/configs/config_mnist_slot8_cbm.yml b/configs/config_mnist_slot8_cbm.yml
new file mode 100644
index 0000000..a5a5e42
--- /dev/null
+++ b/configs/config_mnist_slot8_cbm.yml
@@ -0,0 +1,27 @@
+dataset: "mnist"
+model: "CBM"
+epochs: 100
+batch_size: 128
+num_workers: 0
+lr: 0.001
+warmup_epochs: 10
+embedding_size: 128
+rule_emb_size: 100
+concepts_to_task: "thresholding"
+use_weights: true
+plot_frequency: 1000
+intv_prob: 0.2
+use_balanced_accuracy: true
+use_linear_task_predictor: false
+segmentation_method: "mask"
+use_pretrained_segmenter: false
+presegmented_datasets_path: null
+use_pretrained_autoencoder: false
+fixed_lr: false
+always_use_true_masks: true
+use_initial_proto_embs: false
+noisy_prob: 0.0
+noisy_digit: null
+noisy_target_digit: null
+noisy_part: null
+noisy_target_part: null
diff --git a/configs/config_mnist_slot8_cmr.yml b/configs/config_mnist_slot8_cmr.yml
new file mode 100644
index 0000000..3b61bf4
--- /dev/null
+++ b/configs/config_mnist_slot8_cmr.yml
@@ -0,0 +1,15 @@
+dataset: "mnist"
+model: "CMR"
+epochs: 100
+batch_size: 128
+num_workers: 0
+lr: 0.001
+warmup_epochs: 10
+embedding_size: 128
+rule_emb_size: 100
+concepts_to_task: "thresholding"
+use_weights: true
+plot_frequency: 1000
+intv_prob: 0.20
+use_balanced_accuracy: true
+use_linear_task_predictor: false
diff --git a/configs/config_mnist_slot8_crm.yml b/configs/config_mnist_slot8_crm.yml
new file mode 100644
index 0000000..d8d1def
--- /dev/null
+++ b/configs/config_mnist_slot8_crm.yml
@@ -0,0 +1,15 @@
+dataset: "mnist"
+model: "CRM"
+epochs: 100
+batch_size: 128
+num_workers: 0
+lr: 0.001
+warmup_epochs: 10
+embedding_size: 128
+rule_emb_size: 100
+concepts_to_task: "thresholding"
+use_weights: true
+plot_frequency: 1000
+intv_prob: 0.20
+use_balanced_accuracy: true
+use_linear_task_predictor: false
diff --git a/configs/config_mnist_slot8_full.yml b/configs/config_mnist_slot8_full.yml
new file mode 100644
index 0000000..a29477a
--- /dev/null
+++ b/configs/config_mnist_slot8_full.yml
@@ -0,0 +1,37 @@
+dataset: "mnist"
+model: "PGCM"
+epochs: 100
+batch_size: 128
+num_workers: 0
+lr: 0.001
+warmup_epochs: 10
+fixed_lr: false
+protosize: 64
+embedding_size: 128
+nb_proto: 30
+lam_entropy: 0.01
+lam_batch_entropy: 0.01
+decay_lam_entropy: true
+lam_reconstruction: 5.0
+lam_kl: 0.0
+lam_segmentation: 1.0
+lam_orth: 0.0
+lam_proto_emb: 0.0
+segmentation_method: "mask"
+plot_frequency: 1000
+concepts_to_task: "thresholding"
+use_weights: true
+always_use_true_masks: true
+use_initial_proto_embs: false
+use_balanced_accuracy: true
+use_linear_task_predictor: false
+use_pretrained_autoencoder: false
+autoencoder_path: null
+use_pretrained_segmenter: false
+presegmented_datasets_path: null
+intv_prob: 0.2
+noisy_prob: 0.0
+noisy_digit: null
+noisy_target_digit: null
+noisy_part: null
+noisy_target_part: null
diff --git a/configs/config_mnist_slot8_joint_full.yml b/configs/config_mnist_slot8_joint_full.yml
new file mode 100644
index 0000000..27f67d2
--- /dev/null
+++ b/configs/config_mnist_slot8_joint_full.yml
@@ -0,0 +1,37 @@
+dataset: "mnist"
+model: "PGCM"
+epochs: 100
+batch_size: 128
+num_workers: 0
+lr: 0.001
+warmup_epochs: 10
+fixed_lr: false
+protosize: 64
+embedding_size: 128
+nb_proto: 30
+lam_entropy: 0.01
+lam_batch_entropy: 0.01
+decay_lam_entropy: true
+lam_reconstruction: 5.0
+lam_kl: 0.0
+lam_segmentation: 1.0
+lam_orth: 0.0
+lam_proto_emb: 0.0
+segmentation_method: "mask"
+plot_frequency: 1000
+concepts_to_task: "thresholding"
+use_weights: true
+always_use_true_masks: false
+use_initial_proto_embs: false
+use_balanced_accuracy: true
+use_linear_task_predictor: false
+use_pretrained_autoencoder: false
+autoencoder_path: null
+use_pretrained_segmenter: false
+presegmented_datasets_path: null
+intv_prob: 0.2
+noisy_prob: 0.0
+noisy_digit: null
+noisy_target_digit: null
+noisy_part: null
+noisy_target_part: null
diff --git a/configs/config_mnist_slot8_joint_noisy.yml b/configs/config_mnist_slot8_joint_noisy.yml
new file mode 100644
index 0000000..6d10c33
--- /dev/null
+++ b/configs/config_mnist_slot8_joint_noisy.yml
@@ -0,0 +1,37 @@
+dataset: "mnist"
+model: "PGCM"
+epochs: 100
+batch_size: 128
+num_workers: 0
+lr: 0.001
+warmup_epochs: 10
+fixed_lr: false
+protosize: 64
+embedding_size: 128
+nb_proto: 30
+lam_entropy: 0.01
+lam_batch_entropy: 0.01
+decay_lam_entropy: true
+lam_reconstruction: 5.0
+lam_kl: 0.0
+lam_segmentation: 1.0
+lam_orth: 0.0
+lam_proto_emb: 0.0
+segmentation_method: "mask"
+plot_frequency: 1000
+concepts_to_task: "thresholding"
+use_weights: true
+always_use_true_masks: false
+use_initial_proto_embs: false
+use_balanced_accuracy: true
+use_linear_task_predictor: false
+use_pretrained_autoencoder: false
+autoencoder_path: null
+use_pretrained_segmenter: false
+presegmented_datasets_path: null
+intv_prob: 0.2
+noisy_prob: 0.3
+noisy_digit: [3, 4]
+noisy_target_digit: [1, 8]
+noisy_part: null
+noisy_target_part: null
diff --git a/configs/config_mnist_slot8_joint_smoke.yml b/configs/config_mnist_slot8_joint_smoke.yml
new file mode 100644
index 0000000..f87c40f
--- /dev/null
+++ b/configs/config_mnist_slot8_joint_smoke.yml
@@ -0,0 +1,37 @@
+dataset: "mnist"
+model: "PGCM"
+epochs: 2
+batch_size: 128
+num_workers: 0
+lr: 0.001
+warmup_epochs: 1
+fixed_lr: false
+protosize: 64
+embedding_size: 128
+nb_proto: 30
+lam_entropy: 0.01
+lam_batch_entropy: 0.01
+decay_lam_entropy: true
+lam_reconstruction: 5.0
+lam_kl: 0.0
+lam_segmentation: 1.0
+lam_orth: 0.0
+lam_proto_emb: 0.0
+segmentation_method: "mask"
+plot_frequency: 1000
+concepts_to_task: "thresholding"
+use_weights: true
+always_use_true_masks: false
+use_initial_proto_embs: false
+use_balanced_accuracy: true
+use_linear_task_predictor: false
+use_pretrained_autoencoder: false
+autoencoder_path: null
+use_pretrained_segmenter: false
+presegmented_datasets_path: null
+intv_prob: 0.2
+noisy_prob: 0.0
+noisy_digit: null
+noisy_target_digit: null
+noisy_part: null
+noisy_target_part: null
diff --git a/configs/config_mnist_slot8_noisy.yml b/configs/config_mnist_slot8_noisy.yml
new file mode 100644
index 0000000..800ba31
--- /dev/null
+++ b/configs/config_mnist_slot8_noisy.yml
@@ -0,0 +1,37 @@
+dataset: "mnist"
+model: "PGCM"
+epochs: 100
+batch_size: 128
+num_workers: 0
+lr: 0.001
+warmup_epochs: 10
+fixed_lr: false
+protosize: 64
+embedding_size: 128
+nb_proto: 30
+lam_entropy: 0.01
+lam_batch_entropy: 0.01
+decay_lam_entropy: true
+lam_reconstruction: 5.0
+lam_kl: 0.0
+lam_segmentation: 1.0
+lam_orth: 0.0
+lam_proto_emb: 0.0
+segmentation_method: "mask"
+plot_frequency: 1000
+concepts_to_task: "thresholding"
+use_weights: true
+always_use_true_masks: true
+use_initial_proto_embs: false
+use_balanced_accuracy: true
+use_linear_task_predictor: false
+use_pretrained_autoencoder: false
+autoencoder_path: null
+use_pretrained_segmenter: false
+presegmented_datasets_path: null
+intv_prob: 0.2
+noisy_prob: 0.3
+noisy_digit: [3, 4]
+noisy_target_digit: [1, 8]
+noisy_part: null
+noisy_target_part: null
diff --git a/configs/config_mnist_slot8_smoke.yml b/configs/config_mnist_slot8_smoke.yml
new file mode 100644
index 0000000..7ce8c9c
--- /dev/null
+++ b/configs/config_mnist_slot8_smoke.yml
@@ -0,0 +1,37 @@
+dataset: "mnist"
+model: "PGCM"
+epochs: 2
+batch_size: 128
+num_workers: 0
+lr: 0.001
+warmup_epochs: 1
+fixed_lr: false
+protosize: 64
+embedding_size: 128
+nb_proto: 30
+lam_entropy: 0.01
+lam_batch_entropy: 0.01
+decay_lam_entropy: true
+lam_reconstruction: 5.0
+lam_kl: 0.0
+lam_segmentation: 1.0
+lam_orth: 0.0
+lam_proto_emb: 0.0
+segmentation_method: "mask"
+plot_frequency: 1000
+concepts_to_task: "thresholding"
+use_weights: true
+always_use_true_masks: true
+use_initial_proto_embs: false
+use_balanced_accuracy: true
+use_linear_task_predictor: false
+use_pretrained_autoencoder: false
+autoencoder_path: null
+use_pretrained_segmenter: false
+presegmented_datasets_path: null
+intv_prob: 0.2
+noisy_prob: 0.0
+noisy_digit: null
+noisy_target_digit: null
+noisy_part: null
+noisy_target_part: null