def setup_tracmass(self, seed_locations): tm.mod_seed.nqua = 4 tm.mod_param.partquant = 1 tm.mod_grid.subgrid = 0 tm.mod_param.kriva = 1 tm.mod_domain.timax = 3650.0 tm.loop.writedata = self.writedata tm.loop.readfields = self.readfields #Constants in this run tm.mod_param.lbt = 1 tm.mod_param.nend = tm.mod_param.lbt + 1 tm.mod_param.ntracmax = seed_locations.shape[0] + 1 tm.init_params2() tm.coordinat() tm.mod_time.intrun = self.ntimesteps - 1 if self.start > self.end: #Backward self.mod_seed.nff = 2 self.mod_time.intstep = -1 tm.mod_time.intstart = self.ntimesteps tm.mod_time.intend = 0 else: #Forward tm.mod_time.intstep = 1 tm.mod_seed.nff = 1 tm.mod_time.intstart = 0 tm.mod_time.intend = self.ntimesteps tm.mod_grid.kmt[:] = tm.mod_param.km if tm.mod_seed.nqua == 1: # number of trajectories (per time resolution) # num=NTRACMAX tm.mod_seed.num = tm.mod_param.partquant elif tm.mod_seed.nqua == 2: tm.mod_param.voltr = tm.mod_param.partquant elif tm.mod_seed.nqua == 3: tm.mod_param.voltr = tm.mod_param.partquant tm.mod_grid.kmt[:] = tm.mod_param.km tm.mod_time.intspin = 1 tm.mod_seed.idir = 0 tm.allocate_seed(seed_locations.shape[0]) tm.mod_seed.isec = 5 tm.mod_seed.seed_ijk[:, :] = 2 tm.mod_seed.seed_xyz[:, :] = seed_locations tm.mod_seed.seed_set[:] = [tm.mod_seed.isec, tm.mod_seed.idir] self.trajectory = [list() for _ in range(seed_locations.shape[0])]
def main(): parser = ArgumentParser() parser.add_argument('-r', '--run-params', type=FileType('r')) parser.add_argument('-g', '--grid-params', type=FileType('r')) parser.add_argument('-s', '--set', nargs=2, action='append') parser.add_argument('--seed') parser.add_argument('--exit') parser.add_argument('--path') parser.add_argument('--start') parser.add_argument('--end') parser.add_argument('--delta-hours') parser.add_argument('--delta-minutes') parser.add_argument('--delta-seconds') parser.add_argument('--print-state', action='store_true') args = parser.parse_args() read_params(args.run_params) read_params(args.grid_params) tm.init_params2() tm.coordinat() setup() tm.mod_grid.kmt[:] = tm.mod_param.km tm.allocate_seed(1) tm.mod_seed.isec = 5 tm.mod_seed.seed_ijk[:,:] = 2 tm.mod_seed.seed_xyz[:,:] = [[9.125, 49.125, 4.5]] tm.mod_seed.seed_set[:] = [tm.mod_seed.isec, tm.mod_seed.idir] if args.print_state: print_state() return print tm.mod_traj.trj.shape print tm.mod_traj.trj[:1,:] def writedata(*args): print tm.mod_traj.trj[:1,:] print tm.writedata2(*args) tm.loop.writedata = tm.writedata2 tm.loop.readfields = tm.tes_readfields tm.loop() print "Done!"