How to use the pyuvsim.mpi function in pyuvsim

return beam_model
        if beam_model.startswith('analytic'):
            bspl = beam_model.split('_')
            model = bspl[1]

            to_set = {}
            for extra in bspl[2:]:
                par, val = extra.split('=')
                full = self._float_params[par]
                to_set[full] = float(val)

            return AnalyticBeam(model, **to_set)

        path = beam_model  # beam_model = path to beamfits
        uvb = UVBeam()
        if use_shared_mem and (mpi.world_comm is not None):
            if mpi.rank == 0:
                uvb.read_beamfits(path)
                uvb.peak_normalize()
            for key, attr in uvb.__dict__.items():
                if not isinstance(attr, parameter.UVParameter):
                    continue
                if key == '_data_array':
                    uvb.__dict__[key].value = mpi.shared_mem_bcast(attr.value, root=0)
                else:
                    uvb.__dict__[key].value = mpi.world_comm.bcast(attr.value, root=0)
            mpi.world_comm.Barrier()
        else:
            uvb.read_beamfits(path)
        for key, val in self.uvb_params.items():
            setattr(uvb, key, val)
        uvb.extra_keywords['beam_path'] = path

return AnalyticBeam(model, **to_set)

        path = beam_model  # beam_model = path to beamfits
        uvb = UVBeam()
        if use_shared_mem and (mpi.world_comm is not None):
            if mpi.rank == 0:
                uvb.read_beamfits(path)
                uvb.peak_normalize()
            for key, attr in uvb.__dict__.items():
                if not isinstance(attr, parameter.UVParameter):
                    continue
                if key == '_data_array':
                    uvb.__dict__[key].value = mpi.shared_mem_bcast(attr.value, root=0)
                else:
                    uvb.__dict__[key].value = mpi.world_comm.bcast(attr.value, root=0)
            mpi.world_comm.Barrier()
        else:
            uvb.read_beamfits(path)
        for key, val in self.uvb_params.items():
            setattr(uvb, key, val)
        uvb.extra_keywords['beam_path'] = path
        return uvb

if args.paramsfile is None:
    raise ValueError("Parameter file required")

if args.profile is not None:
    pyuvsim.profiling.set_profiler(outfile_prefix=args.profile, dump_raw=args.raw_profile)

if not os.path.isdir(os.path.dirname(args.paramsfile)):
    args.paramsfile = os.path.join('.', args.paramsfile)

t0 = pytime.time()

pyuvsim.uvsim.run_uvsim(args.paramsfile, quiet=args.quiet)

if args.profile:
    dt = pytime.time() - t0
    maxrss = pyuvsim.mpi.get_max_node_rss()
    rtime = str(timedelta(seconds=dt))
    if isinstance(maxrss, float):
        print('\tRuntime: {} \n\tMaxRSS: {:.3f} GiB'.format(
            rtime, maxrss
        ))
    if hasattr(pyuvsim.profiling.prof, 'meta_file'):
        with open(pyuvsim.profiling.prof.meta_file, 'a') as afile:
            afile.write("Runtime \t {}\nMaxRSS \t {:.3f}\n".format(rtime, maxrss))
            afile.write("Date/Time \t {}".format(str(datetime.now())))

parser.add_argument('--beam', dest='beam', type=str, default='uniform')
parser.add_argument('--prof_out', dest='prof_out', type=str, default='time_profile.out')
parser.add_argument('--mem_out', dest='mem_out', type=str, default='memory_usage.out')
parser.add_argument('--time_out', dest='time_out', type=str, default='time_usage.out')

args = parser.parse_args()

with open(paramsfile, 'r') as pfile:
    params = yaml.safe_load(pfile)

params['config_path'] = os.path.dirname(paramsfile)

min_alt = 70  # Degrees

mpi.start_mpi()
rank = mpi.get_rank()

if rank == 0:
    t0 = pytime.time()

beam_list = None
beam_dict = None
input_uv = UVData()
mock_keywords = None
catalog = None
profiling.set_profiler(outfile_prefix=args.prof_out)

if rank == 0:
    print("{} freqs, {} times, {} bls, {} srcs, {} beam".format(
        args.Nfreqs, args.Ntimes, args.Nbls, args.Nsrcs, args.beam))
    params['freq']['Nfreqs'] = args.Nfreqs
    params['time']['Ntimes'] = args.Ntimes

parser.add_argument('--Nbls', dest='Nbls', type=int, default=1)
parser.add_argument('--beam', dest='beam', type=str, default='uniform')
parser.add_argument('--prof_out', dest='prof_out', type=str, default='time_profile.out')
parser.add_argument('--mem_out', dest='mem_out', type=str, default='memory_usage.out')
parser.add_argument('--time_out', dest='time_out', type=str, default='time_usage.out')

args = parser.parse_args()

with open(paramsfile, 'r') as pfile:
    params = yaml.safe_load(pfile)

params['config_path'] = os.path.dirname(paramsfile)

min_alt = 70  # Degrees

mpi.start_mpi()
rank = mpi.get_rank()

if rank == 0:
    t0 = pytime.time()

beam_list = None
beam_dict = None
input_uv = UVData()
mock_keywords = None
catalog = None
profiling.set_profiler(outfile_prefix=args.prof_out)

if rank == 0:
    print("{} freqs, {} times, {} bls, {} srcs, {} beam".format(
        args.Nfreqs, args.Ntimes, args.Nbls, args.Nsrcs, args.beam))
    params['freq']['Nfreqs'] = args.Nfreqs