Merge pull request #1 from slimgroup/pack

Pack

README.md

@@ -3,3 +3,21 @@
 Implements a distributed version of [Fourier Neural Operators](https://arxiv.org/pdf/2010.08895.pdf)
 based off of the [code](https://github.com/zongyi-li/fourier_neural_operator) written by Zongyi Li, Shuhao Cao, and Jack Griffiths
 using [DistDL](https://github.com/distdl/distdl).
+
+## Installation
+
+
+To install this package you can either clone it and install it in dvelopper mode:
+
+
+```bash
+git clone https://github.com/slimgroup/dfno.git
+cd dfno
+pip install -e .
+```
+
+or directly install it with pip as a standard pythn package `pip install git+https://github.com/slimgroup/dfno.git`
+
+## Author
+
+This package is writeen by Thomas Grady <[email protected]> at Georgia Institute of Technology

dfno/__init__.py

@@ -0,0 +1,2 @@
+from .utils import *
+from .dfno import *

dfno.py → dfno/dfno.py

@@ -7,7 +7,7 @@
 
 from distdl.utilities.tensor_decomposition import *
 from distdl.utilities.torch import *
-from utils import create_root_partition
+from dfno.utils import create_root_partition
 
 class DXFFTN(nn.Module):
 

setup.py

@@ -0,0 +1,13 @@
+from setuptools import setup, find_packages
+
+requirements = ['distdl @ git+https://github.com/mloubout/distdl.git@master', 'mat73', 'matplotlib', 'mpi4py', 'numpy', 'scipy', 'torch==1.9']
+
+setup(
+    name = 'dfno',
+    version = '0.1',
+    author = 'Thomas Grady',
+    author_email = '[email protected]',
+    license = 'MIT',
+    install_requires=requirements,
+    packages = find_packages(),
+    description = 'Distributed Fourier Neural Operator based on DistDL')

experiment.py → training/experiment.py

@@ -9,19 +9,18 @@
 
 from argparse import ArgumentParser
 from distdl.utilities.torch import *
-from dfno import DistributedFNONd
+from dfno import *
 from mat73 import loadmat
 from matplotlib.animation import FuncAnimation
 from mpi4py import MPI
 from pathlib import Path
 from scipy import io
-from utils import *
 
 Partition = distdl.backend.backend.Partition
 
 parser = ArgumentParser()
 parser.add_argument('--input',                  '-i',  type=Path)
-parser.add_argument('--partition-shape',        '-ps', type=int,   nargs='+')
+parser.add_argument('--partition-shape',        '-ps', type=int,   default=(1,1,2,2,1), nargs=5)
 parser.add_argument('--num-data',               '-nd', type=int,   default=1000)
 parser.add_argument('--sampling-rate',          '-sr', type=int,   default=1)
 parser.add_argument('--in-timesteps',           '-it', type=int,   default=10)
@@ -30,7 +29,7 @@
 parser.add_argument('--num-gpus',               '-ng', type=int,   default=1)
 parser.add_argument('--train-split',            '-ts', type=float, default=0.8)
 parser.add_argument('--width',                  '-w',  type=int,   default=20)
-parser.add_argument('--modes',                  '-m',  type=int,   nargs='+')
+parser.add_argument('--modes',                  '-m',  type=int,   default=(4, 4, 4), nargs=3)
 parser.add_argument('--decomposition-order',    '-do', type=int,   default=1)
 parser.add_argument('--num-blocks',             '-nb', type=int,   default=4)
 parser.add_argument('--num-epochs',             '-ne', type=int,   default=500)
@@ -40,6 +39,10 @@
 
 args = parser.parse_args()
 
+if np.prod(args.partition_shape) != MPI.COMM_WORLD.size:
+    raise ValueError(f'The number of processes {MPI.COMM_WORLD.size} does not match the partition shape {args.partition_shape}.')
+
+
 P_world, P_x, P_0 = create_standard_partitions(args.partition_shape)
 dim = P_x.dim
 if args.device == 'cuda':