LDCast implementation in MLCast (in progress)#8
Open
martinbo-meteo wants to merge 3 commits intomlcast-community:mainfrom
Open
LDCast implementation in MLCast (in progress)#8martinbo-meteo wants to merge 3 commits intomlcast-community:mainfrom
martinbo-meteo wants to merge 3 commits intomlcast-community:mainfrom
Conversation
added 3 commits
February 11, 2026 10:31
…sed, and to make the three main components (the autoencoder, the conditioner and the diffuser) distinct entities. Thanks to this, the way the data makes its way in these three components is clearer (see LDCast.ipynb). The files which were taken from the original code without change are: ldcast/utils.py (from ldcast/models/utils.py) ldcast/distributions.py (from ldcast/models/distributions.py) autoenc/auteonc.py autoenc/encoder.py blocks/afno.py blocks/attention.py blocks/resnet.py diffusion/ema.py diffusion/utils.py diffusion/unet.py (which was in the genforecast folder, even though unet is used only in the denoiser) diffusion/utils.py The changes I made are essentially: diffusion/diffusion.py: the LatentDiffusion was given the three parts: model (=denoiser), autoencoder and the context_encoder (= conditioner) and the interactions of these three was not clear. The main class is now DiffusionModel, which needs only the denoiser to be instantiated (the forward call still needs the context given by the conditioner). The interaction between DiffusionModel and the PLMSSampler could be improved (merge the two ?). I removed the ema scope for now, but it should taken care of. diffusion/plms.py: changed only the way the model (denoiser) is called I removed the genforecast folder of the original code: unet.py is now in diffusion, and analysis.py is now context/context.py. The context folder contains nowcast.py (which comes from nowcast/nowcast.py), so that the context folder contains everything to build the conditioner. I reworked the conext/nowcast.py file: I removed the Nowcaster, AFNONowcastNetBasic and the AFNONowcastNet classes (which were not used), and simplified a little the code of two remaining classes (some parts were not used either). The AFNONowcastNetBase class was also taking the autoencoder as input to build the conditioner, which I find very weird (this is why the data seemed to be decoded but not encoded in forecast.Forecast.__call__...). Now, the conditioner is built without the autoencoder. ldcast.py is a new file which will contain the classes subclassing the base classes of mlcast.
- some changes in /src/mlcast/models/base.py (attribute problem if the loss is a dict) - in /src/mlcast/models/ldcast/autoenc/autoenc.py: added the loss for the autoencoder; renamed the 'AutoencoderKL' class in 'AutoencoderKLNet' and set the encoder and decoder to default configurations (which were the ones used in the original code). Removed all the training logic from that class (it will be handled by the trainer) - in src/mlcast/models/ldcast/context/context.py: the timesteps passed to AFNONetCascade.forward were always [-3, -2, -1, 0], so I included them in this function - in /src/mlcast/models/ldcast/diffusion/diffusion.py: I tried to separate as much as possible what concerns the samplers from the rest. I replaced the LatentDiffusion class by the LatentNowcaster one. I could not manage to make this a subclass of NowcastingLightningModule but it would be nice to do so. I removed all the training logic which was contained in the LatentDiffusion class (it will be handled by the trainer) - /src/mlcast/models/ldcast/diffusion/plms.py: removed the score_corrector, corrector_kwargs and mask keywords, which were not used - /src/mlcast/models/ldcast/ldcast.py now contains the main LDCast class subclassing the NowcastingModelBase (only the predict method is implemented, partially) I did not manage to make LatentNowcaster a sublcass of NowcastingLightningModule because LatentNowcaster needs two nets (denoiser and conditioner) and because the training logic is not as straightforward as it is for the moment in NowcastingLightningModule. One should also take into account the fact that two different samplers are used for training and inference, so that the forward method can not just be self.net(x) It would be nice to have cleaner and consistent APIs for the samplers. For the moment, the PLMSSampler and the SimpleSampler are not totally consistent in their APIs, because the SimpleSampler (better/more common name for this one?) was only used during training, while the PLMSSampler was used during inference. The handling of the schedule of each sampler with respect to the schedule saved in the denoiser could also be clearer During training, an EMA scope was used for the weights of the denoiser, I removed this for the moment, but it should reincluded in some way. The variable sometimes refers to the timesteps of the diffusion process (= 1000 during training) and sometimes refers to the nowcasting timesteps (where each time step = 5 minutes). Better to have different names. An AutoencoderKLNet instance can now be passed to the NowcastingLightningModule with the autoenc_loss to handle the training In /src/mlcast/models/ldcast/diffusion/diffusion.py, one has to choose which sampler to use for testing
Author
|
I also included current questions/problems at the end of the README |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
I worked on the code of LDCast, trying to reimplement the code to make it clearer, and trying to make it compliant with the MLCast base classes.
Up to now, I tried essentially to make to make clearer the logic between the autoencoder, the conditioner and the denoiser. I think that the logic between the denoiser and the samplers could still be improved.
I removed also some parts of the code which were not used in the original code, again to make things clearer.
The README contains the current status of the code, and how it is organized. I also try to track all changes in messages of the commits.