Drew's attempt at reproducing Hurum's work

[drewoldag]

WIP - Attempting to reproduce Hurum's work in Hyrax

Notes

13 May

I'm using the two data files that are available on epyc here:

• /astro/store/epyc/ssd/users/htohfa/tp_combined_5_7sigma.npy
• /astro/store/epyc/ssd/users/htohfa/fp_combined_5_7sigma.npy
  Copied to local drive with scp.

I modified the KbmodStamps dataset class to not do normalization just to avoid the startup penalty.
Suggest that we move the normalization to the forward method to ensure that all data that goes in is properly normalized.

Note for Drew for later
Running h.train(trace=2) indicated that DataProvider.resolve_data took 0.2ms and collate took 0.1ms. These seem like very large numbers to retrieve only 2 samples. However, the runtime for get_stamps was 0.12ms and 0.28ms for the two samples.

I ran only 1 epoch (took about 14 minutes on my local machine) and noticed the telltail sign that we've been seeing in the loss curve of an immediate drop and then consistent noise, i.e. not learning anything after the first few batches.

Tomorrow:

• download the data again to the work laptop
• move normalization to the model
• try to understand what is going on in hurum_kbmod.ipynb cell 2 when rebalancing the input data again
• determine if there's a way to memory map the .npy files so we don't have to read them into memory when creating the dataset class instance.

[drewoldag]

The output from h.train(trace=2) with the setup from commit: 928942e

Trace Stages {
	dataset_getter: [
		data__get_classification(index) -> classification duration=0.00585 ms
		inputs:
			index = 1
		outputs:
			classification = 1
		
		data__get_object_id(index) -> object_id duration=0.00413 ms
		inputs:
			index = 1
		outputs:
			object_id = '0000001'
		
		data__get_stamps(index) -> stamps duration=0.119 ms
		inputs:
			index = 1
		outputs:
			stamps = <numpy.ndarray shape=(3, 21, 21) hash=13806450671840395264 device=cpu>
		
		data__get_classification(index) -> classification duration=0.00135 ms
		inputs:
			index = 2
		outputs:
			classification = 1
		
		data__get_object_id(index) -> object_id duration=0.00199 ms
		inputs:
			index = 2
		outputs:
			object_id = '0000002'
		
		data__get_stamps(index) -> stamps duration=0.284 ms
		inputs:
			index = 2
		outputs:
			stamps = <numpy.ndarray shape=(3, 21, 21) hash=4579189852410478592 device=cpu>
		
	]
	resolve_data: [
		DataProvider__resolve_data(index) -> data_dict duration=0.215 ms
		inputs:
			index = 1
		outputs:
			data_dict = {
				data: {
					classification: 1 
					object_id: '0000001' 
					stamps: <numpy.ndarray shape=(3, 21, 21) hash=13806450671840395264 device=cpu> 
				} 
				object_id: '0000001' 
			}
		
		DataProvider__resolve_data(index) -> data_dict duration=0.349 ms
		inputs:
			index = 2
		outputs:
			data_dict = {
				data: {
					classification: 1 
					object_id: '0000002' 
					stamps: <numpy.ndarray shape=(3, 21, 21) hash=4579189852410478592 device=cpu> 
				} 
				object_id: '0000002' 
			}
		
	]
	collate: [
		DataProvider__collate(batch_dicts) -> batch_dict duration=0.127 ms
		inputs:
			batch_dicts = <list len=2> [
				{
					data: {
						classification: 1 
						object_id: '0000001' 
						stamps: <numpy.ndarray shape=(3, 21, 21) hash=13806450671840395264 device=cpu> 
					} 
					object_id: '0000001' 
				}
				{
					data: {
						classification: 1 
						object_id: '0000002' 
						stamps: <numpy.ndarray shape=(3, 21, 21) hash=4579189852410478592 device=cpu> 
					} 
					object_id: '0000002' 
				}
			]
		outputs:
			batch_dict = {
				object_id: <numpy.ndarray shape=(2,) hash=4392798328429168925 device=cpu> 
				data: {
					classification: <numpy.ndarray shape=(2,) hash=0 device=cpu> 
					object_id: <numpy.ndarray shape=(2,) hash=4392798328429168925 device=cpu> 
					stamps: <numpy.ndarray shape=(2, 3, 21, 21) hash=9229794646123610112 device=cpu> 
				} 
			}
		
		DataProvider__handle_nans(batch_dict) -> batch_dict_no_nan duration=0.0629 ms
		inputs:
			batch_dict = {
				object_id: <numpy.ndarray shape=(2,) hash=4392798328429168925 device=cpu> 
				data: {
					classification: <numpy.ndarray shape=(2,) hash=0 device=cpu> 
					object_id: <numpy.ndarray shape=(2,) hash=4392798328429168925 device=cpu> 
					stamps: <numpy.ndarray shape=(2, 3, 21, 21) hash=9229794646123610112 device=cpu> 
				} 
			}
		outputs:
			batch_dict_no_nan = {
				object_id: <numpy.ndarray shape=(2,) hash=4392798328429168925 device=cpu> 
				data: {
					classification: <numpy.ndarray shape=(2,) hash=0 device=cpu> 
					object_id: <numpy.ndarray shape=(2,) hash=4392798328429168925 device=cpu> 
					stamps: <numpy.ndarray shape=(2, 3, 21, 21) hash=9229794646123610112 device=cpu> 
				} 
			}
		
	]
	prepare_inputs: [
		KBModNet__prepare_inputs(batch_dict) -> batch_ndarray duration=0.011 ms
		inputs:
			batch_dict = {
				object_id: <numpy.ndarray shape=(2,) hash=4392798328429168925 device=cpu> 
				data: {
					classification: <numpy.ndarray shape=(2,) hash=0 device=cpu> 
					object_id: <numpy.ndarray shape=(2,) hash=4392798328429168925 device=cpu> 
					stamps: <numpy.ndarray shape=(2, 3, 21, 21) hash=9229794646123610112 device=cpu> 
				} 
			}
		outputs:
			batch_ndarray = (
				<numpy.ndarray shape=(2, 3, 21, 21) hash=9229794646123610112 device=cpu>,
				<numpy.ndarray shape=(2,) hash=0 device=cpu>,
			)
		
		KBModNet__prepare_inputs(batch_dict) -> batch_ndarray duration=0.00559 ms
		inputs:
			batch_dict = {
				object_id: <numpy.ndarray shape=(2,) hash=4392798328429168925 device=cpu> 
				data: {
					classification: <numpy.ndarray shape=(2,) hash=0 device=cpu> 
					object_id: <numpy.ndarray shape=(2,) hash=4392798328429168925 device=cpu> 
					stamps: <numpy.ndarray shape=(2, 3, 21, 21) hash=9229794646123610112 device=cpu> 
				} 
			}
		outputs:
			batch_ndarray = (
				<numpy.ndarray shape=(2, 3, 21, 21) hash=9229794646123610112 device=cpu>,
				<numpy.ndarray shape=(2,) hash=0 device=cpu>,
			)
		
		KBModNet__prepare_inputs(batch_dict) -> batch_ndarray duration=0.00177 ms
		inputs:
			batch_dict = {
				object_id: <numpy.ndarray shape=(2,) hash=4392798328429168925 device=cpu> 
				data: {
					classification: <numpy.ndarray shape=(2,) hash=0 device=cpu> 
					object_id: <numpy.ndarray shape=(2,) hash=4392798328429168925 device=cpu> 
					stamps: <numpy.ndarray shape=(2, 3, 21, 21) hash=9229794646123610112 device=cpu> 
				} 
			}
		outputs:
			batch_ndarray = (
				<numpy.ndarray shape=(2, 3, 21, 21) hash=9229794646123610112 device=cpu>,
				<numpy.ndarray shape=(2,) hash=0 device=cpu>,
			)
		
	]
	evaluation: [
		KBModNet__train_batch(batch) -> loss_dict duration=54 ms
		inputs:
			batch = (
				<torch.Tensor shape=(2, 3, 21, 21) hash=9341934248820473856 device=cuda:0>,
				<torch.Tensor shape=(2,) hash=0 device=cuda:0>,
			)
		outputs:
			loss_dict = {
				loss: 0.20155265927314758 
			}
		
		KBModNet__train_batch(batch) -> loss_dict duration=54 ms
		inputs:
			batch = (
				<torch.Tensor shape=(2, 3, 21, 21) hash=9341934248820473856 device=cuda:0>,
				<torch.Tensor shape=(2,) hash=0 device=cuda:0>,
			)
		outputs:
			loss_dict = {
				loss: 0.20155265927314758 
			}
		
		KBModNet__train_batch(batch) -> loss_dict duration=54 ms
		inputs:
			batch = (
				<torch.Tensor shape=(2, 3, 21, 21) hash=9341934248820473856 device=cuda:0>,
				<torch.Tensor shape=(2,) hash=0 device=cuda:0>,
			)
		outputs:
			loss_dict = {
				loss: 0.20155265927314758 
			}
		
		KBModNet__forward(batch) -> batch_results duration=19.7 ms
		inputs:
			batch = <torch.Tensor shape=(2, 3, 21, 21) hash=9341934248820473856 device=cuda:0>
		outputs:
			batch_results = <torch.Tensor shape=(2, 2) hash=23331930409795584 device=cuda:0>
		
		KBModNet__forward(batch) -> batch_results duration=19.7 ms
		inputs:
			batch = <torch.Tensor shape=(2, 3, 21, 21) hash=9341934248820473856 device=cuda:0>
		outputs:
			batch_results = <torch.Tensor shape=(2, 2) hash=23331930409795584 device=cuda:0>
		
		KBModNet__forward(batch) -> batch_results duration=19.7 ms
		inputs:
			batch = <torch.Tensor shape=(2, 3, 21, 21) hash=9341934248820473856 device=cuda:0>
		outputs:
			batch_results = <torch.Tensor shape=(2, 2) hash=23331930409795584 device=cuda:0>
		
	]
}

[Copilot]

Pull request overview

This PR is a WIP attempt to reproduce Hurum’s KBMOD ML work using the Hyrax framework by migrating existing Fibad-based components, adding a Hurum-inspired network, and updating dataset/config/notebook wiring for Hyrax-driven training.

Changes:

• Migrates model/dataset registration from Fibad to Hyrax and updates README/packaging references accordingly.
• Adds a new KBModNet (Hurum-style) model and a new Hyrax-compatible KbmodStamps dataset implementation (including optional augmentation and metadata table creation).
• Adds an exploratory notebook (drews_attempt.ipynb) and updates the docs training notebook toward Hyrax usage.

Reviewed changes

Copilot reviewed 3 out of 4 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
user_config.toml	Removes the local user config file (notebook/docs still reference it).
src/kbmod_ml/models/resnet.py	Switches to Hyrax registry and changes model to ResNet18; updates training target handling.
src/kbmod_ml/models/hurum_net.py	Adds Hurum-inspired KBModNet model with focal loss and augmentation.
src/kbmod_ml/models/cnn.py	Switches CNN registration from Fibad to Hyrax.
src/kbmod_ml/default_config.toml	Simplifies kbmod-ml config and adds augment flag.
src/kbmod_ml/data_sets/kbmod_stamps.py	Reworks KbmodStamps into a HyraxDataset-compatible dataset with augmentation and metadata.
scripts/true_positive_stamp_generation.py	Minor formatting/import ordering adjustments.
scripts/random_selections.py	Minor import ordering/whitespace cleanup.
scripts/false_positive_stamp_generation.py	Minor import ordering cleanup.
README.md	Updates dependency instructions from Fibad to Hyrax.
pyproject.toml	Updates dependency comment from Fibad to Hyrax.
drews_attempt.ipynb	Adds an exploratory Hyrax training notebook (currently environment/path-specific with outputs).
docs/notebooks/train_model.ipynb	Updates notebook toward Hyrax but contains several now-broken references.

Comments suppressed due to low confidence (3)

src/kbmod_ml/data_sets/kbmod_stamps.py:52

• super().__init__ is called twice (once with only config, and again with config, metadata_table). This is redundant at best and can overwrite HyraxDataset state at worst. Initialize the base class once after you’ve built the metadata table (or only once with the correct arguments).

    def __init__(self, config, data_location=None):
        super().__init__(config)

        # data_dir = config["general"]["data_dir"]
        true_positive_file_name = config["kbmod_ml"]["true_positive_file_name"]
        false_positive_file_name = config["kbmod_ml"]["false_positive_file_name"]

        true_data_path = os.path.join(data_location, true_positive_file_name)
        false_data_path = os.path.join(data_location, false_positive_file_name)

        if not os.path.isfile(true_data_path):
            raise ValueError(f"Could not find {true_positive_file_name} in provided {data_location}")
        if not os.path.isfile(false_data_path):
            raise ValueError(f"could not find {false_positive_file_name} in provided {data_location}")

        true_positive_samples = np.load(true_data_path)
        logger.warning(f"Loaded {len(true_positive_samples)} true positive samples from {true_data_path}")
        false_positive_samples = np.load(false_data_path)
        logger.warning(f"Loaded {len(false_positive_samples)} false positive samples from {false_data_path}")

        n_tp = len(true_positive_samples)
        n_fp = len(false_positive_samples)

        raw_data = np.concatenate([true_positive_samples, false_positive_samples[:, 0:3]])
        raw_labels = np.concatenate([
            np.ones(n_tp, dtype=np.int64),
            np.zeros(n_fp, dtype=np.int64),
        ])

        self._data = raw_data.astype(np.float32)
        # self._data = self._normalize(raw_data)
        self._labels = raw_labels

        seed = config["data_set"]["seed"] if config["data_set"]["seed"] else 42
        # self._arrange_for_hyrax_splits(n_tp, n_fp, seed)

        self.augment = config["kbmod_ml"]["augment"]

        metadata_table = self._read_metadata()
        super().__init__(config, metadata_table)

src/kbmod_ml/data_sets/kbmod_stamps.py:204

• _read_metadata() stores object_id as integer indices (np.arange(...)), but get_object_id() returns a zero-padded string. If Hyrax uses the metadata table for primary_id_field='object_id', this type/format mismatch can break joins/lookups. Store object_id consistently (e.g., padded strings everywhere).

    def ids(self):
        return np.arange(len(self._data))

    def shape(self):
        width, height = self._data[0][0].shape
        return (3, width, height)

    def get_object_id(self, idx):
        # Return a 0-padded string of the index with the number of digits needed
        # for the largest index, e.g. "000123" for index 123 if there are less
        # than 1 million samples
        num_digits = len(str(len(self._data) - 1))
        return f"{idx:0{num_digits}d}"


    def get_classification(self, idx):
        return self._labels[idx]

    def get_stamps(self, idx):
        x = self._data[idx]
        if self.augment:
            # reproduce this code using numpy instead of torch
            x = np.rot90(x, k=np.random.randint(0, 4), axes=(1, 2))
            if np.random.rand() > 0.5: x = np.flip(x, axis=2)
            if np.random.rand() > 0.5: x = np.flip(x, axis=1)
            x = x + np.random.randn(*x.shape).astype(np.float32) * 0.05
        return x

    def _read_metadata(self):
        from astropy.table import Table
        return Table({
            "object_id": self.ids(),
            "classification": self._labels,
        })

src/kbmod_ml/data_sets/kbmod_stamps.py:20

• data_location defaults to None, but the constructor immediately calls os.path.join(data_location, ...), which will raise a TypeError with an unclear message. If data_location is required, validate it early and raise a clear ValueError when it’s missing.

    def __init__(self, config, data_location=None):
        super().__init__(config)

        # data_dir = config["general"]["data_dir"]
        true_positive_file_name = config["kbmod_ml"]["true_positive_file_name"]
        false_positive_file_name = config["kbmod_ml"]["false_positive_file_name"]

        true_data_path = os.path.join(data_location, true_positive_file_name)
        false_data_path = os.path.join(data_location, false_positive_file_name)

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.