[sevmag]

Addresses #753, following the design @RasmusOrsoe proposed in his review of #759. Picks up where @mobra7's PR left off.

Problem

FourierEncoder bakes in two assumptions:

1. Column layout the input must be (x, y, z, time, charge, auxiliary), with the slicing hardcoded (x[:, :, :3], x[:, :, 3], …) and the presence of charge/auxiliary inferred from the feature count. This is the index error reported in FourierEncoder not flexible enough for different detectors #753 for FEATURES.ICECUBE86.
2. Numeric range — the multipliers 4096 (xyz, time) and 1024 (charge) are fixed to the Kaggle normalisation.

Change

As proposed in the review of #759:

• FourierEncoderEPJC — the existing class, behaviour unchanged, with a docstring noting it is the graphnet implementation of the EPJ-C solution (arXiv:2310.15674) and naming the assumptions that make it hard to use elsewhere. IceMix/DeepIce and ISeeCube now reference it, so nothing currently on main changes behaviour.
• FourierEncoder — new, taking schema: Dict[int, float] mapping column index to multiplier, e.g. {1: 1024, 3: 4096}. Columns absent from the schema are not embedded. Nothing is assumed about column order or meaning, and the boolean embedding and MLP projection are left to the model, both being problem-specific. Exposes output_dim so the model can size what follows.

The schema docstring states the wavelength relation, so the next person choosing a value has the arithmetic rather than a constant copied from another detector.

Not included

The review also proposed renaming DeepIce to DeepIceEPJC with a refactored DeepIce owning the MLP, dropping the auxiliary embedding, and updating examples/04_training/06_train_icemix_model.py. Left for a follow-up: dropping the auxiliary embedding is a behaviour change worth deciding on its own, whereas this half is strictly additive. Happy to carry it through in this PR instead if you would rather have it in one piece.

[sevmag]

Also I think we should move DeepIce out of the models/gnn folder and rename the icemix.py file to deepice.py . It is really confusing

[sevmag]

@pweigel I think your opinion on this is valuable!

[sevmag]

@Aske-Rosted and I had the concern that renaming the Kaggle version of the FourierEncoder to a different name will break backwards compatibility with older pretrained state dicts. I thought about this a bit more, and I actually do not think this is an issue, since the name of the nn.Modules is determined by the attribute name and not the class name.

Example:
The FourierEncoder is initialized like this:
self.fourier_ext = FourierEncoder(
The torch state_dicts will then have a key called fourier_ext holding its weights. Changing FourierEncoder into FourierEncoderEPJC is not an issue at all.

[Aske-Rosted]

I think the description is a bit verbose and historical. I think it's better to provide just a description of what the class does but exclude information about how it came to exist. If this was a special scenario of the more general class then it might make sense to have that description but not the other way around.

[Aske-Rosted]

Apart from the doc string looks fine to me.

[felixyu7]

We've built Neptune (#923) on this FourierEncoder, and the per-column (multiplier, n_freq) schema covers everything it needs. Thanks! One request: under bf16-true / 16-true, the input and freqs are cast to low precision, which is too coarse for the phases at large multipliers. Neptune works around this by calling the encoder in float32 with autocast off. Could SinusoidalPosEmb compute the phases in float32 itself, so each model doesn't need its own guard?

[sevmag]

We've built Neptune (#923) on this FourierEncoder, and the per-column (multiplier, n_freq) schema covers everything it needs. Thanks! One request: under bf16-true / 16-true, the input and freqs are cast to low precision, which is too coarse for the phases at large multipliers. Neptune works around this by calling the encoder in float32 with autocast off. Could SinusoidalPosEmb compute the phases in float32 itself, so each model doesn't need its own guard?

good point, checkout the way I handled it @felixyu7