Update script config parameters

.bumpversion.cfg

@@ -1,7 +1,7 @@
 [bumpversion]
 commit = True
 tag = True
-current_version = 3.0.47
+current_version = 3.0.48
 message = Bump version: {current_version} → {new_version} [skip ci]
 
 [bumpversion:file:pvnet/__init__.py]

README.md

@@ -101,7 +101,7 @@ conda create -n ocf_datapipes python=3.10
 Then go inside the ocf_datapipes repo to add packages
 
 ```bash
-pip install ".[dev]"
+pip install -r requirements.txt -r requirements-dev.txt
 ```
 
 Then exit this environment, and enter back into the pvnet conda environment and install ocf_datapies in editable mode (-e). This means the package is directly linked to the source code in the ocf_datapies repo.

pvnet/__init__.py

@@ -1,2 +1,2 @@
 """PVNet"""
-__version__ = "3.0.47"
+__version__ = "3.0.48"

pvnet/models/multimodal/linear_networks/networks.py

@@ -1,5 +1,5 @@
 """Linear networks used for the fusion model"""
-from torch import nn
+from torch import nn, rand
 
 from pvnet.models.multimodal.linear_networks.basic_blocks import (
     AbstractLinearNetwork,
@@ -316,6 +316,7 @@ def __init__(
             virtual_batch_size=virtual_batch_size,
             momentum=momentum,
             mask_type=mask_type,
+            group_attention_matrix=rand(4, in_features),
         )
 
         self.activation = nn.LeakyReLU(negative_slope=0.01)

pyproject.toml

@@ -15,7 +15,7 @@ dependencies = [
     "ipykernel",
     "h5netcdf",
     "torch>=2.0.0",
-    "lightning>=2.0.1",
+    "lightning",
     "torchvision",
     "pytest",
     "pytest-cov",

scripts/save_concurrent_batches.py

@@ -2,14 +2,16 @@
 Constructs batches where each batch includes all GSPs and only a single timestamp.
 
 Currently a slightly hacky implementation due to the way the configs are done. This script will use
-the same config file currently set to train the model.
+the same config file currently set to train the model. In the datamodule config it is possible
+to set the batch_output_dir and number of train/val batches, they can also be overriden in the
+command as shown in the example below.
 
 use:
 ```
 python save_concurrent_batches.py \
-    +batch_output_dir="/mnt/disks/nwp_rechunk/concurrent_batches_v3.9" \
-    +num_train_batches=20_000 \
-    +num_val_batches=4_000
+    datamodule.batch_output_dir="/mnt/disks/nwp_rechunk/concurrent_batches_v3.9" \
+    datamodule.num_train_batches=20_000 \
+    datamodule.num_val_batches=4_000
 ```
 
 """
@@ -157,12 +159,14 @@ def main(config: DictConfig):
     config_dm = config.datamodule
 
     # Set up directory
-    os.makedirs(config.batch_output_dir, exist_ok=False)
+    os.makedirs(config_dm.batch_output_dir, exist_ok=False)
 
-    with open(f"{config.batch_output_dir}/datamodule.yaml", "w") as f:
+    with open(f"{config_dm.batch_output_dir}/datamodule.yaml", "w") as f:
         f.write(OmegaConf.to_yaml(config.datamodule))
 
-    shutil.copyfile(config_dm.configuration, f"{config.batch_output_dir}/data_configuration.yaml")
+    shutil.copyfile(
+        config_dm.configuration, f"{config_dm.batch_output_dir}/data_configuration.yaml"
+    )
 
     dataloader_kwargs = dict(
         shuffle=False,
@@ -179,39 +183,39 @@ def main(config: DictConfig):
         persistent_workers=False,
     )
 
-    if config.num_val_batches > 0:
+    if config_dm.num_val_batches > 0:
         print("----- Saving val batches -----")
 
-        os.mkdir(f"{config.batch_output_dir}/val")
+        os.mkdir(f"{config_dm.batch_output_dir}/val")
 
         val_batch_pipe = _get_datapipe(
             config_dm.configuration,
             *config_dm.val_period,
-            config.num_val_batches,
+            config_dm.num_val_batches,
         )
 
         _save_batches_with_dataloader(
             batch_pipe=val_batch_pipe,
-            batch_dir=f"{config.batch_output_dir}/val",
-            num_batches=config.num_val_batches,
+            batch_dir=f"{config_dm.batch_output_dir}/val",
+            num_batches=config_dm.num_val_batches,
             dataloader_kwargs=dataloader_kwargs,
         )
 
-    if config.num_train_batches > 0:
+    if config_dm.num_train_batches > 0:
         print("----- Saving train batches -----")
 
-        os.mkdir(f"{config.batch_output_dir}/train")
+        os.mkdir(f"{config_dm.batch_output_dir}/train")
 
         train_batch_pipe = _get_datapipe(
             config_dm.configuration,
             *config_dm.train_period,
-            config.num_train_batches,
+            config_dm.num_train_batches,
         )
 
         _save_batches_with_dataloader(
             batch_pipe=train_batch_pipe,
-            batch_dir=f"{config.batch_output_dir}/train",
-            num_batches=config.num_train_batches,
+            batch_dir=f"{config_dm.batch_output_dir}/train",
+            num_batches=config_dm.num_train_batches,
             dataloader_kwargs=dataloader_kwargs,
         )