collect.py

import time
from multiprocessing import Pool

import boto3
import botocore
import meep as mp


def make_dft_data(flx_reg=None, n2f_reg=None, frc_reg=None, fldc=None, flds=None, fldw=None, fld_cmp=None):
    dft_data = {
        'flux_regions': flx_reg,
        'n2f_regions': n2f_reg,
        'force_regions': frc_reg,
        'fields_center': fldc,
        'fields_size': flds,
        'fields_where': fldw,
        'fields_components': fld_cmp
    }
    return dft_data


def run(b,
        eps_offd=mp.Vector3(),
        mu_offd=mp.Vector3(),
        echi2=mp.Vector3(),
        hchi3=mp.vector3(),
        esus=[],
        hsus=[],
        dcond=mp.Vector3(),
        bcond=mp.Vector3(),
        dfts={}):

    resolution = 10
    cell_size = mp.Vector3(10, 10)
    mat = mp.Medium(
        epsilon=12,
        epsilon_offdiag=eps_offd,
        mu_offdiag=mu_offd,
        E_chi2_diag=echi2,
        H_chi3_diag=hchi3,
        E_susceptibilities=esus,
        H_susceptibilities=hsus,
        D_conductivity_diag=dcond,
        B_conductivity_diag=bcond,
    )
    geom = [mp.Block(size=mp.Vector3(x=b, y=b), material=mat)]

    source = [mp.Source(src=mp.GaussianSource(0.5, 0.25), component=mp.Ex, center=mp.Vector3())]

    start = time.time()
    sim = mp.Simulation(cell_size=cell_size,
                        geometry=geom,
                        resolution=resolution,
                        sources=source)

    if dfts:
        if dfts['flux_regions']:
            sim.add_flux(1, 0.5, 5, *dfts['flux_regions'])
        if dfts['n2f_regions']:
            sim.add_near2far(1, 0.5, 7, *dfts['n2f_regions'])
        if dfts['force_regions']:
            sim.add_force(1, 0.5, 9, *dfts['force_regions'])
        if dfts['fields_components']:
            sim.add_dft_fields(dfts['fields_components'], 0, 1, 5, where=dfts['fields_where'],
                               center=dfts['fields_center'], size=dfts['fields_size'])

    sim.run(until=200)
    total_time = time.time() - start

    assert len(sim.fragment_stats) == 1
    stats = sim.fragment_stats[0]
    results = "{},{},{},{},{},{},{:.4f}\n"

    return results.format(stats.num_anisotropic_eps_pixels,
                          stats.num_anisotropic_mu_pixels,
                          stats.num_nonlinear_pixels,
                          stats.num_susceptibility_pixels,
                          stats.num_nonzero_conductivity_pixels,
                          stats.num_dft_pixels,
                          total_time)


def get_s3_file(fn, bucket):
    s3 = boto3.resource('s3')

    try:
        s3.Bucket(bucket).download_file(fn, fn)
    except botocore.exceptions.ClientError as e:
        if e.response['Error']['Code'] == 404:
            print("Object {} does not exist in bucket {}".format(fn, bucket))
        else:
            raise
    print("Downloaded {} from S3".format(fn))


def put_s3_file(fn, bucket):
    s3 = boto3.client('s3')
    s3.upload_file(fn, bucket, fn)
    print("Uploaded {} to S3".format(fn))


if __name__ == '__main__':
    fn = 'fragment_stats.csv'
    bucket = 'hogan-fragment-stats'
    get_s3_file(fn, bucket)

    sizes = [2, 6, 10]
    vecs = [mp.Vector3(), mp.Vector3(2), mp.Vector3(2, 2)]
    sus = [[],
           [mp.LorentzianSusceptibility(0.5, 0.02)],
           [mp.LorentzianSusceptibility(0.5, 0.02), mp.DrudeSusceptibility(0.5, 0.02)]]
    dfts = [
        {},
        make_dft_data([mp.FluxRegion(mp.Vector3(), size=mp.Vector3(10, 10), direction=mp.X)]),
        make_dft_data([mp.FluxRegion(mp.Vector3(), direction=mp.X, size=mp.Vector3(5, 5))],
                      n2f_reg=[mp.Near2FarRegion(mp.Vector3(0, 10), size=mp.Vector3(10, 10), direction=mp.X)]),
    ]

    pool = Pool()

    args = [(b, eps_offd, mu_offd, echi2, hchi3, esus, hsus, dcond, bcond, dft)
            for b in sizes
            for eps_offd in vecs
            for mu_offd in vecs
            for echi2 in vecs
            for hchi3 in vecs
            for esus in sus
            for hsus in sus
            for dcond in vecs
            for bcond in vecs
            for dft in dfts]

    res = pool.starmap(run, args)

    with open(fn, 'a') as f:
        for line in res:
            f.write(line)

    put_s3_file(fn, bucket)