#!/usr/env/bin python
import os
# os.environ['OMP_NUM_THREADS'] = '1'
from newpoisson import poisson
import numpy as np
from fenics import set_log_level, File, RectangleMesh, Point

mesh = RectangleMesh(Point(0,0), Point(1,1), 36, 36)
# comm = mesh.mpi_comm()

set_log_level(40) # ERROR=40
# from mpi4py import MPI
# comm = MPI.COMM_WORLD
# rank = comm.Get_rank()

if __name__=='__main__':

    import argparse

    parser = argparse.ArgumentParser(description="Poisson Problem")
    parser.add_argument('-n', '--num', default = 10, type=int,
                       help="Number of samples")
    parser.add_argument('-o', '--outfile', default='results',
                       help="Output filename (no extension)")
    parser.add_argument('-i', '--input-dim', default=1, type=int)
    parser.add_argument('-d', '--dist', default='u', help='Distribution. `n` (normal), `u` (uniform, default)')
    args = parser.parse_args()

    num_samples = args.num
    dist = args.dist
    outfile = args.outfile.replace('.pkl','')
    inputdim = args.input_dim
    if inputdim == 1: # U[1,5]
        randsamples = 1 + 4*np.random.rand(num_samples)
    else: # N(0,1)
        if dist == 'n':
            randsamples = np.random.randn(num_samples, inputdim)
        elif dist == 'u':
            randsamples = -4*np.random.rand(num_samples, inputdim)
        else:
            raise ValueError("Improper distribution choice, use `n` (normal), `u` (uniform)")

    sample_seed_list = list(zip(range(num_samples), randsamples))

    def wrapper(sample, outfile):
        g=sample[1]
        u = poisson(gamma=g, mesh=mesh)
        # Save solution 
        fname = f"{outfile}-data/poisson-{int(sample[0]):06d}.xml"
        File(fname, 'w') << u
        return {int(sample[0]): {'u': fname, 'gamma': sample[1]}}


    results = []
    for sample in sample_seed_list:
        r = wrapper(sample, outfile)
        results.append(r)
#         print(results)
    import pickle
    pickle.dump(results, open(f'{outfile}.pkl','wb'))

[((261, 278), 'fenics.set_log_level', 'set_log_level', (['(40)'], {}), '(40)\n', (274, 278), False, 'from fenics import set_log_level, File, RectangleMesh, Point\n'), ((203, 214), 'fenics.Point', 'Point', (['(0)', '(0)'], {}), '(0, 0)\n', (208, 214), False, 'from fenics import set_log_level, File, RectangleMesh, Point\n'), ((215, 226), 'fenics.Point', 'Point', (['(1)', '(1)'], {}), '(1, 1)\n', (220, 226), False, 'from fenics import set_log_level, File, RectangleMesh, Point\n'), ((425, 479), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ([], {'description': '"""Poisson Problem"""'}), "(description='Poisson Problem')\n", (448, 479), False, 'import argparse\n'), ((1580, 1607), 'newpoisson.poisson', 'poisson', ([], {'gamma': 'g', 'mesh': 'mesh'}), '(gamma=g, mesh=mesh)\n', (1587, 1607), False, 'from newpoisson import poisson\n'), ((1206, 1244), 'numpy.random.randn', 'np.random.randn', (['num_samples', 'inputdim'], {}), '(num_samples, inputdim)\n', (1221, 1244), True, 'import numpy as np\n'), ((1708, 1724), 'fenics.File', 'File', (['fname', '"""w"""'], {}), "(fname, 'w')\n", (1712, 1724), False, 'from fenics import set_log_level, File, RectangleMesh, Point\n'), ((1109, 1136), 'numpy.random.rand', 'np.random.rand', (['num_samples'], {}), '(num_samples)\n', (1123, 1136), True, 'import numpy as np\n'), ((1300, 1337), 'numpy.random.rand', 'np.random.rand', (['num_samples', 'inputdim'], {}), '(num_samples, inputdim)\n', (1314, 1337), True, 'import numpy as np\n')]