{"id":20520285,"url":"https://github.com/plasmacontrol/d3d_loaders","last_synced_at":"2026-05-27T13:35:22.324Z","repository":{"id":109403398,"uuid":"500835186","full_name":"PlasmaControl/d3d_loaders","owner":"PlasmaControl","description":"Pytorch style dataloader for D3D data in HDF5 format","archived":false,"fork":false,"pushed_at":"2023-03-31T01:11:49.000Z","size":1292,"stargazers_count":1,"open_issues_count":2,"forks_count":0,"subscribers_count":4,"default_branch":"main","last_synced_at":"2025-03-02T18:12:57.679Z","etag":null,"topics":[],"latest_commit_sha":null,"homepage":null,"language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"mit","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/PlasmaControl.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2022-06-07T12:40:35.000Z","updated_at":"2023-03-08T20:47:50.000Z","dependencies_parsed_at":"2023-07-27T21:02:18.613Z","dependency_job_id":null,"html_url":"https://github.com/PlasmaControl/d3d_loaders","commit_stats":null,"previous_names":[],"tags_count":0,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/PlasmaControl%2Fd3d_loaders","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/PlasmaControl%2Fd3d_loaders/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/PlasmaControl%2Fd3d_loaders/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/PlasmaControl%2Fd3d_loaders/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/PlasmaControl","download_url":"https://codeload.github.com/PlasmaControl/d3d_loaders/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":242123561,"owners_count":20075385,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":[],"created_at":"2024-11-15T22:18:59.170Z","updated_at":"2026-05-27T13:35:17.286Z","avatar_url":"https://github.com/PlasmaControl.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"D3D Loaders\n===========\n\nThis package implements Pytorch-style [Iterable Datasets](https://pytorch.org/docs/stable/data.html#)\nfor D3D data. It contains It contains helper function to fetch data from MDS and store it locally\nin HDF5 files. It also contains custom [samplers](https://pytorch.org/docs/stable/data.html#torch.utils.data.Sampler) to easily load batch random sequences from the data. These can be used in conjunction with\n[data loaders](https://pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader).\n\n\n\nLoad Multiple Signals in Single Dataset\n=======================================\n\nImport and instantiate the loader like this:\n```python\n\u003e\u003e\u003e from d3d_loaders.d3d_loaders import D3D_dataset\n\u003e\u003e\u003e from d3d_loaders.samplers import  RandomBatchSequenceSampler, collate_fn_random_batch_seq\n\u003e\u003e\u003e shotnr = 172337\n\u003e\u003e\u003e t_params = {'tstart' : 200.0,\n                'tend'   : 1000.0,\n                'tsample': 1.0\n\u003e\u003e\u003e }\n\u003e\u003e\u003e shift_targets = {'ae_prob': 10.0}\n\u003e\u003e\u003e device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n\n\n\u003e\u003e\u003e my_ds = D3D_dataset(shotnr, t_params,\n\u003e\u003e\u003e                     predictors=[\"pinj\", \"tinj\", \"iptipp\", \"dstdenp\", \"doutu\", \"dssdenest\", \"ae_prob\"],\n\u003e\u003e\u003e                     targets=[\"ae_prob\"],\n\u003e\u003e\u003e                     shift_targets=shift_targets,\n\u003e\u003e\u003e                     datapath=\"/projects/EKOLEMEN/d3dloader/test\",\n\u003e\u003e\u003e                     device=device)\n```\n\nHere `shotnr` gives the shot from which the signals are used. Together `tstart`, `tend`, and\n`tsample` define a common time base for all signals. The signals are resampled in a manner that\nconserves causality, that is, no future information is used to generate a single sample.\n\nThe dataset defines two groups of signals. `Predictors` are to be used as input to a model and\n`targets` are to be used as ground-truth outputs of the model. The dictionary `shift_targets`\ndefines an offset by which targets are shifted into the future. The name of the signals\nrefer to either `MDS` nodes or `PTDATA` point names. The names of the signals and their\nlocation in `MDS` or `PTDATA` is defined in the file `downloading.py`. A custom signal,\n`ae_prob` refers to output of Aza's Alfven Eigenmode probability model.\n\nPredictor and target signals are stored in the dictionaries `my_ds.predictors` and `my_ds.target` respectively. \n\nThe `datapath` argument defines the directory where the D3D data will be cached\nin `HDF5` files. There is one `HDF5` file per shot, which stores the individual signals within\na group. If a signal is not found in that group, it will be fetched automatically and added to\nthe file. The layout of the files are \n\n```julia\njulia\u003e fid = h5open(\"/projects/EKOLEMEN/d3dloader/test/172337.h5\", \"r\")\n🗂️ HDF5.File: (read-only) /projects/EKOLEMEN/d3dloader/test/172337.h5\n├─ 📂 ali\n│  ├─ 🏷️ origin\n│  ├─ 🔢 xdata\n│  │  └─ 🏷️ xunits\n│  └─ 🔢 zdata\n│     └─ 🏷️ zunits\n├─ 📂 doutu\n│  ├─ 🏷️ origin\n│  ├─ 🔢 xdata\n│  │  └─ 🏷️ xunits\n│  └─ 🔢 zdata\n│     └─ 🏷️ zunits\n├─ 📂 dssdenest\n│  ├─ 🏷️ origin\n│  ├─ 🔢 xdata\n│  └─ 🔢 zdata\n├─ 📂 dstdenp\n│  ├─ 🏷️ origin\n│  ├─ 🔢 xdata\n│  └─ 🔢 zdata\n...\n```\n\nEach group is named after the shortname signal and contains `xdata`, `zdata` nodes returned from `gadata.py`.\nAdditionally, the origin of the data is stored in the `origin` attribute, and the units of the signals\nare stored in the `xunits` and `zunits` attributes, when available.\n\n\nAdditionally, a `device` can be passed. If set to be the GPU, the `predictor` and \n`target` tensors will be stored on the GPU. This avoids having to load batches into GPU memory in the\ntraining data loop.\n\n\n\nIterating over signal sequences in a single shot\n================================================\nTo fetch data sequences from a single shot, we can instantiate a `DataLoader`. This package\nimplements a `RandomBatchSequenceSampler` that fetches linear sequences from both, `predictor`\nand `target` tensors of the dataset. `RandomBatchSequenceSampler` fetches a batch of sequences,\nall of the same length and starting at the same random point. The sample below shows how \n`D3D_dataset` and `RandomBatchSequenceSampler` work together to allow iteration over the sequences:\n\n```python\n\u003e\u003e\u003e len(ds)\n    800\n\u003e\u003e\u003e batch_size = 32\n\u003e\u003e\u003e seq_length = 512\n\u003e\u003e\u003e sampler = BatchedSampler(len(ds), seq_length=seq_length, batch_size=batch_size, shuffle=True)\n\u003e\u003e\u003e loader_train = torch.utils.data.DataLoader(my_ds, num_workers=0, \n                                               batch_sampler=sampler,\n                                               collate_fn = collate_fn_batched)\n\n\u003e\u003e\u003e for x_b, y_b in loader_train:\n        print(f\"x_b.shape={x_b.shape}, y_b.shape={y_b.shape}\")\n\nx_b.shape=torch.Size([32, 513, 11]), y_b.shape=torch.Size([32, 513, 5])\nx_b.shape=torch.Size([32, 513, 11]), y_b.shape=torch.Size([32, 513, 5])\nx_b.shape=torch.Size([32, 513, 11]), y_b.shape=torch.Size([32, 513, 5])\nx_b.shape=torch.Size([32, 513, 11]), y_b.shape=torch.Size([32, 513, 5])\nx_b.shape=torch.Size([32, 513, 11]), y_b.shape=torch.Size([32, 513, 5])\nx_b.shape=torch.Size([32, 513, 11]), y_b.shape=torch.Size([32, 513, 5])\nx_b.shape=torch.Size([32, 513, 11]), y_b.shape=torch.Size([32, 513, 5])\nx_b.shape=torch.Size([32, 513, 11]), y_b.shape=torch.Size([32, 513, 5])\nx_b.shape=torch.Size([32, 513, 11]), y_b.shape=torch.Size([32, 513, 5])\n```\n\nIn each iteration, `loader_train` returns a tuple of tensors. The first tensor contains\n`32=batch_size` sequences of length `seq_length+1`. Each sequence has `11` features, which\nis the sum of the feature dimension of each `predictor` signal. The second tensor has the\nsame batch size and sequence length, but the feature dimension is `5`, corresponding to\nthe probability of each Alfven Eigenmode.\n\nThe length of the sequences is `seq_length+1=513`. Note also, that the `predictor` signals are\nshifted `10ms` into the future. That is the samples at `y_b[:, -1, :]`  are `11ms` ahead\nof the samples in `x_b[:, -2, :]`. \n\nIn the example above, `collate_fn_batched` reshapes the returned data\nfrom `d3d_dataset.__getidx__` into a tuple of 2 tensors. The call semantics are explained in the\ncode example [here](https://pytorch.org/docs/stable/data.html#disable-automatic-batching).\nNote that `RandomBatchSequenceSampler` takes care of batching. Automatic batching in the pytorch\n`DataLoader` needs to be disabled.\n\n\nThe `BatchedSampler` can either sample sequences linearly, starting at 0, or shuffle\nthe starting points of the sequences.\n\nFor `shuffle=False`, `BatchedSampler` generates sequences where the starting index shifts by 1 for\neach sequence in a batch. Using a `BatchedSampler` in the example above would have\nthe first sequence start at index 0 and extending to 512. The second sequence\nwould start at index 1 and extend to 513. And so on.\n\nThis is useful for inference, where we may want to predict a target sequentially over the \nentire shot. We input samples of all predictors at time index 0...511 into the model and get a \nprediction for the target at time index 512. Then we input the predictors at time index\n1..512 to predict the target at time index 513. And so in. Thus, by iterating over a\nDataLoader which uses a `BatchedSampler`, we can easily get a reconstruction of the\npredicted target for the entire shot.\n\nSetting `shuffle=True`, `BatchedSampler` generates sequences that start at a random point\nin the dataset. Looking at the code above, the first of the 32 sequences in the batch\nmay start at index 18 (and extend to index 18+513=631) of the entire shot.\nThe second sequence may start at index 788 and extend to index 788+513. \nSo every starting index is at random. This is useful for training.\n\n\nThere are also the versions of the sampler `BatchedSampler_dist`. These should\nbe used for distributed training. They effectively distribute all available samples\nin the dataset across ranks.\n\n\n\nMulti-shot datasets\n===================\n\nTHIS IS A BIT OUTDATED. THE INTERFACE TO THE MULTISHOT DATASET HAS CHANGED.\nSEE [https://github.com/PPPLDeepLearning/frnn_examples](https://github.com/PPPLDeepLearning/frnn_examples) FOR PROPER USAGE.\n\nThe class `Multishot_dataset` defines a dataset that spans multiple shots. It can be instantiated in\na very similar manner:\n\n```python\n\n\u003e\u003e\u003e from d3d_loaders.d3d_loaders import Multishot_dataset\n\u003e\u003e\u003e from d3d_loaders.samplers import BatchedSampler_multi, collate_fn_batched\n\u003e\u003e\u003e shot_list_train = [172337, 172339] \n\u003e\u003e\u003e tstart = 110.0 # Time of first sample for upper triangularity is 100.0\n\u003e\u003e\u003e tend = 2000.0\n\u003e\u003e\u003e t_params = {\"tstart\": tstart, \"tend\": tend, \"tsample\": 1.0}\n\u003e\u003e\u003e t_shift = 10.0\n\u003e\u003e\u003e device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n\u003e\u003e\u003e seq_length = 512\n\u003e\u003e\u003e batch_size = 4\n\u003e\u003e\u003e pred_list = [\"pinj\",  # Injected power\n\u003e\u003e\u003e              \"tinj\", # Injected torque\n\u003e\u003e\u003e              \"iptipp\",  # Target current\n\u003e\u003e\u003e              \"dstdenp\", # Target density\n\u003e\u003e\u003e              \"doutu\", # Top triangularity\n\u003e\u003e\u003e              \"dssdenest\", # Line-averaged density\n\u003e\u003e\u003e              \"ae_prob\"] # AE mode probability\n\u003e\u003e\u003e targ_list = [\"ae_prob\"]\n\n\u003e\u003e\u003e ds_train = Multishot_dataset(shot_list_train, t_params, pred_list, targ_list,\n\u003e\u003e\u003e                              {\"ae_prob\": t_shift}, \n\u003e\u003e\u003e                              datapath=\"/projects/EKOLEMEN/d3dloader/test\", device)\n```\n\nTo iterate over a multi-shot database, we can again instantiate a DataLoader using the custom\n`RandomBatchSequenceSampler_multids` sampler as well as a custom `collate_fn`:\n\n```python\n\u003e\u003e\u003e for xb, yb in loader_train_b:\n\u003e\u003e\u003e     print(xb.shape, yb.shape)\n\u003e\u003e\u003e    break\n\ntorch.Size([513, 4, 11]) torch.Size([513, 4, 5])\n```\n\nAgain, the tensor tuple `xb` and `yb` contain a batch of random sequences. But the sequences are picked\nat random from the list of shots used to construct the `Multishot_dataset`. All sequences are constructed\nfrom data of only a single shot.\n\n\nDownloading signals\n===================\nMDS and PTDATA can be downloaded with the script `downloading.py`. This script fetches data for\na given shot from D3D's MDSplus server and stores in in a HDF5 file.\n\nThe current predictors are:\n\n| Predictor             | Key             |\n|-----------------------|-----------------|\n| AE Mode probabilities | ae_prob         |\n| Pinj                  | pinj            |\n| Neutron Rate          | neutronsrate    |\n| Injected Power        | ip              |\n| ECH                   | ech             |\n| q95                   | q95             |\n| kappa                 | kappa           |\n| Density Profile       | dens            |\n| Pressure Profile      | pres            |\n| Temperature Profile   | temp            |\n| q Profile             | q               |\n| Lower Triangularity   | doutl           |\n| Upper Triangularity   | doutl           |\n| Raw ECE Channels      | raw_ece         |\n| Raw CO2 dp Channels   | raw_co2_dp      |\n| Raw CO2 pl Channels   | raw_co2_pl      |\n| Raw MPI Channels      | raw_mpi         |\n| Raw BES Channels      | raw_bes         |\n| Bill's AE Labels (+-250ms window) | uci_label |\n\nNOTE: the shape signal of kappa and shape profiles of upper and lower triangularity don't have data even though\nthe keys exist in the hdf5 files. You should get a helpful error telling you this, however you may also get\nan error telling you of an array splicing error. Also when these signals do exist, the first signal is very \nlate, so make sure `tstart` is sufficiently large. \n\nAt a given index, the output is a list of 2 tensors.\n```\n\nIn [119]: my_ds[0]\nOut[119]: \n(tensor([-0.5374, -0.4721, -0.5583, -0.5301, -0.4714, -1.1228, -1.7159]),\n tensor([-0.0207, -0.0328, -0.0478, -0.0331, -0.0119]))\n```\n\nThe first tensor has a number of elements fixed by the number of signals being used. \nThe order will match the order of the signals in `predictors`. \nEach of these signals is normalized to their separate mean and standard deviation.\n\nCurrently, the only target is `ae_prob_delta`, calculated using AE probabilities from Aza's RCN model. \nThe output tensor (2nd tensor) is the change in the Alfven Eigenmode probabilities over the interval given in\n`shift_targets` with the `ae_prob_delta` key,\ni.e. AE mode probability at t0 + 10ms as calculated using ECE data. This output is also\nnormalized to zero mean and unit standard deviation.\n\n\n\nIteration over the dataset is done in pytorch-style:\n```\nfor i in torch.utils.data.DataLoader(my_ds):\n    # use i\n```\n\nAn example of how to use this dataloader for predictive modelling is in 'runme.py'\n\nLoad Full Resolution Single Signal\n==================================\nBy setting `tsample=-1`, the full resolution will return with the `tstart` and `tend` \nfound closest to the true measurement values (forwards or backwards in time). This will\nnot load as a dataset since you cannot temporally align the full resolution signals. \nAn example for loading  neutron rate into a numpy array is below\n\n```python\nfrom d3d_loaders.signal0d import signal_neut\nshotnr = 169113\nt_params = {'tstart' : tstart,\n            'tend'   : tend,\n            'tsample': -1\n}\n\nsig_neut = signal_neut(shotnr, t_params)\n\ndata = sig_neut.data.numpy()\n```\n\n\n\nFull training examples\n======================\n\nTraining examples can be found in [this repository](https://github.com/PPPLDeepLearning/frnn_examples).\nThe dataset these examples operate on can be found [here](https://github.com/PPPLDeepLearning/dataset_D3D_100).\n\n\n\n\nThe folder `examples` contains several notebooks that illustrate how to train predictive\nmodels using the infrastructure provided in this package. In particular, the notebook\n`AE_pred_LSTMtest_1723xx.ipynb` shows how to train a LSTM-based model for predictive\ntasks. And the nodebook `AE_pred_transformer_154xxx.ipynb` shows how to train a transformer-based\nmodel for predictive tasks. Both notebooks rely on the dataset and dataloader infrastructure\nimplemented in this package. \n\nThe datasets used by the notebooks are located on Princeton's gpfs, accessible on\n`stellar` and `traverse`: `stellar:/projects/EKOLEMEN/AE_datasets`. They have been compiled\nusing the [d3d_signals](https://github.com/PlasmaControl/d3d_signals) package. Definitions of\nthe datasets in `yaml` format is [here](https://github.com/PlasmaControl/AE_datasets). \nThere is an additional `README` with notes on issues that arose when compiling the dataset. \n\n\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fplasmacontrol%2Fd3d_loaders","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fplasmacontrol%2Fd3d_loaders","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fplasmacontrol%2Fd3d_loaders/lists"}