| 50 | |
| 51 | |
| 52 | def parse_args(): |
| 53 | parser = argparse.ArgumentParser() |
| 54 | |
| 55 | # Add your arguments here |
| 56 | parser.add_argument('--dataset', default="landoceandataset", type=str, choices=["checkerboard", |
| 57 | "inat2018", |
| 58 | "landoceandataset", |
| 59 | "era5dataset" |
| 60 | ]) |
| 61 | parser.add_argument('--pe', default=["sphericalharmonics"], type=str, nargs='+', help='positional encoder(s)', |
| 62 | choices=["sphericalharmonics", "theory", "grid", "spherec", "spherecplus", "direct", |
| 63 | "cartesian3d", "wrap", "spherem", "spheremplus"]) |
| 64 | parser.add_argument('--nn', default=["siren"], type=str, nargs='+', help='neural network(s)', |
| 65 | choices=["linear", "siren", "fcnet"]) |
| 66 | |
| 67 | # optional configs |
| 68 | parser.add_argument('--save-model', action="store_true", help='save model checkpoint to results-dir') |
| 69 | parser.add_argument('--log-wandb', action="store_true", help='log run to wandb') |
| 70 | parser.add_argument('--hparams', default="hparams.yaml", type=str, help='hypereparameter yaml') |
| 71 | parser.add_argument('--results-dir', default="results/train", type=str, help='results directory') |
| 72 | parser.add_argument('--expname', default=None, type=str, |
| 73 | help='experiment name. If specified, saves results in subfolder') |
| 74 | parser.add_argument('--seed', default=0, type=int, help='global random seed') |
| 75 | parser.add_argument('--max-epochs', default=None, type=int, |
| 76 | help='maximum number of epochs. If unset, uses value in hparams.yaml') |
| 77 | parser.add_argument('--gpus', default='-1', type=int, nargs='+', |
| 78 | help='which gpus to use; if unset uses -1 which we map to auto') |
| 79 | parser.add_argument('--accelerator', default='auto', type=str, |
| 80 | help='lightning accelerator') |
| 81 | |
| 82 | parser.add_argument('-r', '--resume-ckpt-from-results-dir', action="store_true", |
| 83 | help="searches through provided results dir and resumes from suitable checkpoint " |
| 84 | "that matches pe and nn") |
| 85 | parser.add_argument('--matplotlib', action="store_true", |
| 86 | help="plot maps with matplotlib") |
| 87 | parser.add_argument('--matplotlib-show', action="store_true", |
| 88 | help="shows matplotlib plots (can cause freezing when called remotely)") |
| 89 | |
| 90 | parser.add_argument('--use-expnamehps', default=False, type=bool, |
| 91 | help='whether expname is part of the hp file names') |
| 92 | |
| 93 | # checkerboard |
| 94 | parser.add_argument('--checkerboard-scale', default=1, type=float, |
| 95 | help="scales the number of support points for the checkerboard dataset (specificed in hparams.yaml) " |
| 96 | "by this factor. This is useful to vary the scale to test different resolutions of encoders") |
| 97 | |
| 98 | # overwrite certain hparams |
| 99 | parser.add_argument('--legendre-polys', default=None, type=int) |
| 100 | parser.add_argument('--min-radius', default=None, type=float) |
| 101 | parser.add_argument('--harmonics-calculation', default="analytic", type=str, |
| 102 | choices=["analytic", "closed-form", "discretized"], |
| 103 | help='calculation of spherical harmonics: ' + |
| 104 | 'analytic uses pre-computed equations. This is exact, but works only up to degree 50, ' + |
| 105 | 'closed-form uses one equation but is computationally slower (especially for high degrees)' + |
| 106 | 'discretized pre-computes harmonics on a grid and interpolates these later') |
| 107 | |
| 108 | args = parser.parse_args() |
| 109 | return args |