def copy_matching_fields(g_src,g_dest,eps=1.0, cell_fields=['depth'],edge_fields=['edge_depth']): # match cells by centroid in case centers have been # adjusted src_cc=g_src.cells_centroid() src_cc_kdt=KDTree(data=src_cc) dest_cc=g_dest.cells_centroid() for c in utils.progress(range(g_dest.Ncells())): dist,src_cell=src_cc_kdt.query(dest_cc[c],distance_upper_bound=eps) if not np.isfinite(dist): continue for fld in cell_fields: g_dest.cells[fld][c]=g_src.cells[fld][src_cell] src_ec=g_src.edges_center() src_ec_kdt=KDTree(data=src_ec) dest_ec=g_dest.edges_center() for j in utils.progress(range(g_dest.Nedges())): dist,src_edge=src_ec_kdt.query(dest_ec[j],distance_upper_bound=eps) if not np.isfinite(dist): continue for fld in edge_fields: g_dest.edges[fld][j]=g_src.edges[fld][src_edge]
def dem_to_cell_node_bathy(dem,g): cell_z=dem_to_cell_bathy(dem,g) node_z_to_cell_z=sparse.dok_matrix( (g.Ncells(),g.Nnodes()), np.float64 ) for c in utils.progress(range(g.Ncells())): nodes=g.cell_to_nodes(c) node_z_to_cell_z[c,nodes]=1./len(nodes) # A x = b # A: node_z_to_cell_z # x: node_z # b: cell_z # to better allow regularization, change this to a node elevation update. # A ( node_z0 + node_delta ) = cell_z # A*node_delta = cell_z - A*node_z0 node_z0=dem(g.nodes['x']) bad_nodes=np.isnan(node_z0) node_z0[bad_nodes]=0.0 # could come up with something better.. if np.any(bad_nodes): print("%d bad node elevations"%bad_nodes.sum()) b=cell_z - node_z_to_cell_z.dot(node_z0) # damp tries to keep the adjustments to O(2m) res=sparse.linalg.lsqr(node_z_to_cell_z.tocsr(),b,damp=0.05) node_delta, istop, itn, r1norm = res[:4] print("Adjustments to node elevations are %.2f to %.2f"%(node_delta.min(), node_delta.max())) final=node_z0+node_delta if np.any(np.isnan(final)): print("Bad news") import pdb pdb.set_trace() return final
def extract_particle_snapshot(t): # extract particle field at a given time snapshot=[] for traj_i,traj in enumerate(utils.progress(traj_data)): traj_t=traj['traj']['t'] if traj_t[0]>t: continue if traj_t[-1]<t: continue # time index along this trajectory matching up with t idx=np.searchsorted(traj_t,t) # x,y,Tfill,dTdt,fill_dist rec=[traj['traj']['x'][idx], traj['traj']['y'][idx], traj['Tfill'][idx], traj['dTdt'][idx], traj['fill_dist'][idx], traj_i] snapshot.append(rec) snap_struct=np.array(snapshot) snap_data=dict( x=snap_struct[:,0], y=snap_struct[:,1], Tfill=snap_struct[:,2], dTdt=snap_struct[:,3], fill_dist=snap_struct[:,4], traj_i=snap_struct[:,5].astype(np.int32) ) return snap_data
def group_trajectories(): trajectories = dict() # (grp,idx in group) => [ (t,xy), ..] times = range(0, ntimes) start_time = time.time() # by avoiding keeping references to particle arrays, and # sorting stuck particles on the fly, the memory usage is # kept in check for part_grp in range(len(ptm_groups)): print("Particle Group %d" % part_grp) # read the last time step to get the max particle # count, and final resting place. step_t, parts = ptm_data[part_grp].read_timestep(ntimes - 1) end_xy = parts['x'] # we output particles once when stuck, but then record # them here to avoid any more dead = np.zeros(len(end_xy), np.bool8) max_particles = len(parts) # use list to speed up references grp_trajectories = [[] for i in range(max_particles)] last_xy = np.nan * np.ones((max_particles, 2)) lastlast_xy = np.nan * np.ones((max_particles, 2)) for ti, t in enumerate(utils.progress(times)): step_t, parts = ptm_data[part_grp].read_timestep(t) step_t = utils.to_dt64(step_t) Nstep = len(parts) # move as much referencing outside the loop part_xy = parts['x'][:, :] # particles which do not move between this step and # the end of the simulation are presumed dead and # not procesed anymore. # it might be worth keeping the first dead location. # this is probably stripping out some otherwise useful # points. for part_idx in np.nonzero(~dead[:Nstep])[0]: # range(Nstep): traj = grp_trajectories[part_idx] # avoid any references back to parts # assumes that traj[-1] is same location as traj[-2] # probably safe. rec = [step_t, part_xy[part_idx, 0], part_xy[part_idx, 1]] if len(traj) >= 2 and (traj[-2][1] == rec[1]) and (traj[-2][2] == rec[2]): # stuck particles are just updated by the latest time/position traj[-1][0] = step_t else: traj.append(rec) # if a particle is stuck from here on out, remove it from play dead[:Nstep] = np.all(part_xy == end_xy[:Nstep, :], axis=1) for part_idx, traj in enumerate(grp_trajectories): trajectories[(part_grp, part_idx)] = traj return trajectories
def fig_dist(C,num=2,log=False,title="",local_max=False, direction=False): fig=plt.figure(num) fig.clf() fig.set_size_inches([6,9],forward=True) ax=fig.add_subplot(1,1,1) cax=fig.add_axes([0.05,0.6,0.03,0.35]) fig.subplots_adjust(left=0,right=1,top=1,bottom=0) if log: C=np.log10(C.clip(1e-10,np.inf)) label='log$_{10}$' else: label='linear' ccoll=g.plot_cells(values=C,ax=ax,cmap=cmap) ccoll.set_lw(0.05) ccoll.set_edgecolor('face') # plt.colorbar(ccoll,cax=cax,label=label) plot_utils.cbar(ccoll,cax=cax,label=label) if local_max: is_local_max=np.ones(g.Ncells(),np.bool8) e2c=g.edge_to_cells() internal=e2c.min(axis=1)>=0 c1=e2c[internal,0] c2=e2c[internal,1] c1_less=C[c1]<C[c2] is_local_max[ c1[c1_less] ]=False c2_less=C[c2]<C[c1] is_local_max[ c2[c2_less] ]=False cc=g.cells_center() ax.plot(cc[is_local_max,0],cc[is_local_max,1],'ko') if direction: nbrhood=4 idxs=np.arange(g.Ncells()) np.random.shuffle(idxs) samp_cells=idxs[:1000] cc=g.cells_center() XY=cc[samp_cells] UV=np.nan*XY for i,c in utils.progress(enumerate(samp_cells)): cells=[c] for _ in range(nbrhood): cells=np.unique(list(cells) + [c for c0 in cells for c in g.cell_to_cells(c0) if c>=0] ) #y=C[cells] #X=np.c_[cc[cells]-cc[cells].mean(axis=0), # np.ones(len(cells))] # location and bias #beta_hat=np.linalg.lstsq(X,y,rcond=None)[0] #UV[i,:]=beta_hat[:2] # [gradx,grady] low=cells[ np.argmin(C[cells]) ] high=cells[ np.argmax(C[cells]) ] UV[i,:]=(cc[high]-cc[low])/(C[high]-C[low]) ax.quiver( XY[:,0], XY[:,1], UV[:,0], UV[:,1],pivot='tip',scale=60,width=0.005) ax.xaxis.set_visible(0) ax.yaxis.set_visible(0) ax.axis('equal') ax.text(0.5,0.98,title,transform=ax.transAxes,va='top',ha='center') return fig
def find_all_pairs(g, max_segs=7): """ Automatically pull channel cross sections up to 7 segments wide. This does not use any hydro data, so relies on grid boundaries alone. For the CSC grid, it's right about 95% of the time, so don't assume that every channel section will be found, or that every node pair returns is a channel section. """ e2c = g.edge_to_cells() e_boundary = np.any(e2c < 0, axis=1) boundary_nodes = np.unique(g.edges['nodes'][e_boundary]) boundary_mask = np.zeros(g.Nnodes(), np.bool) boundary_mask[boundary_nodes] = True ## node_marks = np.zeros(g.Nnodes(), np.bool8) all_pairs = [] def boundary_weight(j): if e_boundary[j]: return 1.0 else: return np.nan def internal_weight(j): if e_boundary[j]: return np.nan else: return 1.0 for n in utils.progress(boundary_nodes, msg="find_all_pairs %s boundary nodes"): if node_marks[n]: continue node_marks[n] = True # search only boundary edges to rule out along-boundary neighbors my_nbrs = g.shortest_path(n, n2=boundary_nodes, edge_weight=boundary_weight, max_return=2 * max_segs) my_nbrs = [mn[0] for mn in my_nbrs] other_n2 = np.setdiff1d(boundary_nodes, my_nbrs) tran_path = g.shortest_path(n, n2=other_n2, edge_weight=internal_weight, max_return=1) if len(tran_path) == 0: continue n2 = tran_path[0][1][0] all_pairs.append([n, n2]) # this isn't a strictly commutative property, but close enough # for our purposes node_marks[n2] = True return all_pairs
def never_stuck(ptm_out, ntimes): t_a, state = ptm_out.read_timestep(ts=0) last_x = state['x'][:, :2] stuck_count = np.zeros(len(last_x), np.int32) for ti in utils.progress(range(1, ntimes)): t_a, state = ptm_out.read_timestep(ts=ti) this_x = state['x'][:, :2] stuck = np.all(last_x == this_x, axis=1) stuck_count[stuck] += 1 last_x = this_x return stuck_count == 0
def phase_map(val,fit_cos=fit_cos,fit_sin=fit_sin,small=1e-6): phases=np.zeros(Ncells,np.float64) for c in utils.progress(range(Ncells)): c_val = val[:,c] C=np.cov( np.c_[fit_cos,fit_sin,c_val].T ) if C[2,2] > small: # some signal at this cell phase=np.arctan2(C[0,2],C[1,2]) * 180/np.pi else: phase=np.nan phases[c]=phase return phases
def track_cloud(cloud_particles): # track cloud. this ti is an index into the time steps already extracted track_vars = [] for ti in utils.progress(time_steps, 1.0, "processing %s timesteps"): cloud_xy = particles[ti, cloud_particles, :] cloud_cells = np.array([g.select_cells_nearest(xy) for xy in cloud_xy]) pair_dists = pairwise_grid_distance(cloud_cells) variance_per_cell = (pair_dists**2).mean(axis=1) best_variance = variance_per_cell.min() track_vars.append(best_variance) track_vars = np.array(track_vars) mb = np.polyfit(track_time_s, track_vars, 1) return mb[0]
def get_particle_mass(pb,inflow,rel_stride=2): """ pb: a PtmBin object. inflow: dataset with Q, dnum rel_stride: how many ptm outputs go by for each release period. returns a {particle id: particle mass} """ ptm_out_dt_s=pb.dt_seconds() nsteps=pb.count_timesteps() # Associate weights with particle ids # At this stage, we just weight based on unit concentration # in the flow. part_mass={} # particle id => 'mass' for step in utils.progress(range(0,nsteps-1,rel_stride)): t,parts=pb.read_timestep(step) dnum=utils.to_dnum(t) new_ids=[ p['id'] for p in parts if p['id'] not in part_mass] # kludge - the quarter hour offset here gives more constant # particle mass. It's close enough, so I'm not going to worry # about replicating the integration further. Qnow=np.interp(dnum+ptm_out_dt_s*0.5/86400, inflow.dnum.values,inflow.Q) if len(new_ids)==0: if Qnow>0: print(f"WARNING: {t.strftime('%Y-%m-%d %H:%M')}: {len(new_ids):6d} new particles, Q={Qnow:9.2f} m3/s") continue mass_per_particle=max(0,Qnow*ptm_out_dt_s / len(new_ids)) if step%20==0: print(f"{t.strftime('%Y-%m-%d %H:%M')}: {len(new_ids):6d} new particles, Q={Qnow:9.2f} m3/s, mass/part {mass_per_particle:9.2f}") for i in new_ids: part_mass[i]=mass_per_particle if len(part_mass)==0: # like petaluma, has no flow in this period. return np.nan*np.ones(1) # And convert to array for faster indexing. max_id=np.max(list(part_mass.keys())) print("max_id: ",max_id) # leave unset values nan to detect issues later. part_mass_a=np.nan*np.zeros(max_id+1,np.float64) for k in part_mass: part_mass_a[k]=part_mass[k] return part_mass_a
def dem_to_cell_bathy(dem,g,fill_iters=20): cell_means=np.zeros(g.Ncells(),np.float64) for c in utils.progress(range(g.Ncells()),msg="dem_to_cell_bathy: %s"): #msk=dem.polygon_mask(g.cell_polygon(c)) #cell_means[c]=np.nanmean(dem.F[msk]) cell_means[c]=np.nanmean(dem.polygon_mask(g.cell_polygon(c),return_values=True)) for _ in range(fill_iters): missing=np.nonzero(np.isnan(cell_means))[0] if len(missing)==0: break new_depths=[] print("filling %d missing cell depths"%len(missing)) for c in missing: new_depths.append( np.nanmean(cell_means[g.cell_to_cells(c)]) ) cell_means[missing]=new_depths else: print("Filling still left %d nan cell elevations"%len(missing)) return cell_means
def particle_to_density(particles, grid, normalize='area'): """ particles: struct array with 'x' and 'mass' normalize: 'area' normalize by cell areas to get a mass/area 'volume': Not implemented yet. """ mass = np.zeros(grid.Ncells(), np.float64) for i in utils.progress(range(len(particles))): cell = grid.select_cells_nearest(particles['x'][i, :2]) if cell is None: continue mass[cell] += particles['mass'][i] if normalize == 'area': mass /= grid.cells_area() elif normalize == 'mass': pass else: raise Exception(f"Not ready for normalize={normalize}") return mass
def fill_by_divergence(gtri, Qtri): """ Update nan entries in Qtri based on divergence in adjacent cells. """ bad_cells = [] e2c = gtri.edge_to_cells() valid = np.isfinite(Qtri) for j in utils.progress(np.nonzero(~valid)[0]): c1, c2 = e2c[j] # flow out of each cell. # positive flow on an edge means out of c1, into c2. Qcs = [] Qabs = [] # to get a sense of roundoff scale # csgn relates the net flow out of that cell to what should # be put on j. so a net flow out of c1, Qc>0, means that # j should put flow into c1, which is a negative Q on edge # j. for csgn, c in [(-1, c1), (1, c2)]: if c < 0: continue Qc = cell_divergence(gtri, c, Qtri) Qcs.append(Qc * csgn) if len(Qcs) > 1: eps = 1e-3 abs_scale = 1e-2 abs_err = np.abs(Qcs[0] - Qcs[1]) rel_err = abs_err / (np.mean(np.abs(Qcs)) + eps) if (rel_err > eps) and (abs_err > abs_scale): print("Flows Q[c=%d]=%f Q[c=%d]=%f don't match" % (c1, Qcs[0], c2, Qcs[1])) bad_cells.append(c1) bad_cells.append(c2) Qtri[j] = np.mean(Qcs) assert np.isfinite(Qtri[j])
def track_cloud(cloud_particles): # track cloud. this ti is an index into the time steps already extracted track_vars=[] for ti in utils.progress(time_steps,1.0,"processing %s timesteps"): cloud_xy=particles[ti,cloud_particles,:] cloud_cells=np.array( [g.select_cells_nearest(xy) for xy in cloud_xy] ) if np.any( dead_cells[cloud_cells] ): # print("Particles hit boundary - truncate time") break pair_dists=pairwise_grid_distance(cloud_cells) variance_per_cell=(pair_dists**2).mean(axis=1) best_variance=variance_per_cell.min() track_vars.append(best_variance) track_vars=np.array(track_vars) t_s=track_time_s[:len(track_vars)] # give up if there is less than 12h or 5 data points. if len(t_s>5) and (t_s[-1]-t_s[0])>43200: mb=np.polyfit(t_s,track_vars,1) return mb[0] else: return np.nan
fig, ax = plt.subplots(num=1) ax.plot(parts['x'][:, 0], parts['x'][:, 1], 'g.') ax.plot(parts['x'][cloud_idxs, 0], parts['x'][cloud_idxs, 1], 'm.', ms=10) plot_wkb.plot_polygon(grid_poly, ax=ax, fc='0.8', ec='k', lw=0.5) ax.axis('equal') ax.axis(zoom) ## # build up just the 2D coordinates for all of the initial release particles part_locs = [] part_ts = [] for ti in utils.progress(range(0, 1000, 4)): t, parts = ptm_group.read_timestep(ti) part_locs.append(parts['x'][:, :2].copy()) part_ts.append(t) # This will fail is there are additional releases in this group part_locs = np.array(part_locs) part_ts = np.array([utils.to_dt64(t) for t in part_ts]) ## cloud_t_xy = part_locs[:, cloud_idxs, :] cloud_cc = cloud_t_xy.mean(axis=1) # In this case it remains more or less in the channel.
if 0: # Simplest option: # Put bathy on nodes, just direct sampling. z_node = dem(g.nodes['x']) if 1: # Bias deep name += "_deep" # Maybe a good match with bedlevtype=5. # BLT=5: edges get shallower node, cells get deepest edge. # So extract edge depths (min,max,mean), and nodes get deepest # edge. alpha = np.linspace(0, 1, 5) edge_data = np.zeros((g.Nedges(), 3), np.float64) # Find min/max/mean depth of each edge: for j in utils.progress(range(g.Nedges())): pnts = (alpha[:, None] * g.nodes['x'][g.edges['nodes'][j, 0]] + (1 - alpha[:, None]) * g.nodes['x'][g.edges['nodes'][j, 1]]) z = dem(pnts) edge_data[j, 0] = z.min() edge_data[j, 1] = z.max() edge_data[j, 2] = z.mean() z_node = np.zeros(g.Nnodes()) for n in utils.progress(range(g.Nnodes())): # This is the most extreme bias: nodes get the deepest # of the deepest points along adjacent edgse z_node[n] = edge_data[g.node_to_edges(n), 0].min() g.add_node_field('node_z_bed', z_node, on_exists='overwrite')
ntimes=ptm_data[0].count_timesteps() run_start=ptm_data[0].read_timestep(0)[0] run_stop =ptm_data[0].read_timestep(ntimes-1)[0] ## ptm_group=ptm_data[0] # initial release ## # build up just the 2D coordinates for all of the initial release particles part_locs=[] part_ts=[] for ti in utils.progress(range(0,1000,4)): t,parts=ptm_group.read_timestep(ti) part_locs.append( parts['x'][:,:2].copy() ) part_ts.append(t) # This will fail is there are additional releases in this group part_locs=np.array(part_locs) part_ts=np.array([utils.to_dt64(t) for t in part_ts]) ## from scipy import sparse from scipy.sparse import csgraph def grid_to_graph(g): # use cell-to-cell connections to make it easier to avoid
('AM2',AM2), ('AM3',AM3), ('AM4',AM4), ('AM6',AM6), ('AM7',AM7), ('AM8',AM8), ('AM9',AM9), ('SM1',SM1), ('SM2',SM2), ('SM3',SM3), ('SM4',SM4), ('SM7',SM7), ('SM8',SM8) ] all_detects=[pt.parse_tek(fn,name=name) for name,fn in utils.progress(all_receiver_fns)] ## all_detects_nomp=[pt.remove_multipath(d) for d in all_detects] ## # First cut: # Limit to 1 hour of detections. 3388 detections over this 1 hour. # one good tag - C56B (just 4 fixes) #t_clip=[np.datetime64("2019-03-25 02:00"), # np.datetime64("2019-03-25 03:00")] # one good tag - C535 (but just 4 fixes)
def parse_txts(txt_fns, pressure_range=[110e3, 225e3], name=None, beacon='auto', split_on_clock_change=True): """ Parse a collection of txt files, grab detections and optionally clock resync events. beacon: 'auto' set beacon tag id from most commonly received tag id. None: don't set a beacon tag id else: use the provided value as the beacon id. """ txt_fns = list(txt_fns) txt_fns.sort() dfs = [] for fn in utils.progress(txt_fns): df = parse_txt(fn) df['fn'] = fn dfs.append(df) df = pd.concat(dfs, sort=True) # sort columns to get alignment df = df.reset_index() # will use this for slicing later n_unknown = (df['type'] == 'unknown').sum() if n_unknown > 0: # So far, this is just corrupt lines. Might be able to # salvage the second part of corrupt lines, but it's # such a small number, and the second part is usually # just a NODE bootup msg anyway. log.warning("%d unknown line types in txt files" % n_unknown) # Do we need epoch? yes, it's used downstream epoch = utils.to_unix(df['time'].values) # It will get usec added, so be sure it's just the integer portion. assert np.all((epoch % 1.0)[np.isfinite(epoch)] == 0) df['epoch'] = epoch # Add microseconds to timestamps when t_usec is available sel = np.isfinite(df.t_usec.values) df.loc[sel, 'time'] = df.loc[ sel, 'time'] + (df.loc[sel, 't_usec'].values * 1e6).astype( np.int32) * np.timedelta64(1, 'us') df_det = df[df['type'] == 'DET'] # clean up time: bad_time = (df_det.time < np.datetime64('2018-01-01')) | ( df_det.time > np.datetime64('2022-01-01')) df2 = df_det[~bad_time].copy() # clean up temperature: df2.loc[df2.temp < -5, 'temp'] = np.nan df2.loc[df2.temp > 35, 'temp'] = np.nan # clean up pressure if pressure_range is not None: df2.loc[df2.pressure < pressure_range[0], 'pressure'] = np.nan df2.loc[df2.pressure > pressure_range[1], 'pressure'] = np.nan # trim to first/last valid pressure valid_idx = np.nonzero(np.isfinite(df2.pressure.values))[0] df3 = df2.iloc[valid_idx[0]:valid_idx[-1] + 1, :].copy() df3['tag'] = [s.strip() for s in df3.tag.values] # narrow that to the fields we actually care about: fields = [ 'rx_serial', 'tag', 'time', 't_usec', 'epoch', 'corrQ', 'nbwQ', 'pressure', 'temp', 'datetime_str', 'fn' ] ds = xr.Dataset.from_dataframe(df3.loc[:, fields]) ds['usec'] = ds['t_usec'] ds['cf2_filename'] = None if beacon == 'auto': beacon_id = df3.groupby('tag').size().sort_values().index[-1] ds['beacon_id'] = beacon_id ds['beacon_id'].attrs['source'] = 'received tags' print("auto_beacon: beacon_id inferred as ", beacon_id) # import pdb # pdb.set_trace() elif beacon is not None: ds['beacon_id'] = beacon ds['beacon_id'].attrs['source'] = 'Specified to parse_txts' ds.attrs['pressure_range'] = pressure_range if name is not None: ds['name'] = (), name if split_on_clock_change: # dice_by_clock_resets got crazy complicated. Rather than # re-living that experience, try something simple here, # but know that we may have to come back and refactor this # with dice_by_clock_resets. sync_sel = df['sync_status'].values == 1.0 sync_idxs = df.index.values[ sync_sel] # of course these won't actually be in df3! nonmono_sel = np.diff(df3.time.values) < np.timedelta64(0) # Should mean that each item is the first index of a new chunk. nonmono_idxs = df3.index.values[1:][nonmono_sel] all_breaks = np.unique(np.concatenate((sync_idxs, nonmono_idxs))) all_breaks = np.r_[0, all_breaks, len(df)] # as indices into ds.index ds_breaks = np.searchsorted(ds.index.values, all_breaks) diced = [] for start, stop in zip(ds_breaks[:-1], ds_breaks[1:]): if stop > start: diced.append(ds.isel(index=slice(start, stop))) else: # often will have a slice that's empty. pass return diced else: return ds
# g=unstructured_grid.UnstructuredGrid.read_ugrid(grid_fn) # g.renumber() ## g = unstructured_grid.UnstructuredGrid.read_ugrid( 'pesca_butano_existing_deep_bathy.nc') ## alpha = np.linspace(0, 1, 5) edge_data = np.zeros((g.Nedges(), 3), np.float64) # Find min/max/mean depth of each edge: for j in utils.progress(range(g.Nedges())): pnts = (alpha[:, None] * g.nodes['x'][g.edges['nodes'][j, 0]] + (1 - alpha[:, None]) * g.nodes['x'][g.edges['nodes'][j, 1]]) z = dem(pnts) edge_data[j, 0] = z.min() edge_data[j, 1] = z.max() edge_data[j, 2] = z.mean() ## # This is bedlevtype=5 z_edge = (g.nodes['node_z_bed'][g.edges['nodes']]).max(axis=1) plt.figure(1).clf() plt.plot(z_edge, edge_data[:, 0], 'g.') plt.plot(z_edge, edge_data[:, 1], 'b.')
def node_depths_edge_mean_opt(g, target_edge_depths, target_node_depths, section_weight=1.0, node_weight=0.1, nonsection_weight=0.0, max_segments_per_section=7): """ return per-node elevations such that the average of the endpoints per edge are close to the target edge depth. Edges which form channel cross-sections are given precedence. nodes per-edge to match the mean edge elevation from the DEM. """ # construct a linear system where we try to solve good elevations for # all of the nodes at once # First, do this but ignore boundary edges: e2c = g.edge_to_cells() boundary_edges = e2c.min(axis=1) < 0 boundary_nodes = np.unique(g.edges['nodes'][boundary_edges]) weight_by_pairs = True if section_weight == nonsection_weight: log.info( "Cross-section and non-cross-section weights the same. No tracing needed" ) else: all_pairs = find_all_pairs(g, max_segments_per_section) cross_edges = [] for n1, n2 in utils.progress(all_pairs, msg="Extracting section edges %s"): cross_edges.append(g.shortest_path(n1, n2, return_type='edges')) cross_edges = np.concatenate(cross_edges) edge_weights = nonsection_weight * np.ones(g.Nedges(), np.float64) edge_weights[cross_edges] = section_weight rows = [] cols = [] data = [] rhs = [] for j in utils.progress(range(g.Nedges())): row = len(rhs) if boundary_edges[j]: continue if edge_weights[j] == 0: continue n1, n2 = g.edges['nodes'][j] rows.append(row) cols.append(n1) data.append(0.5 * edge_weights[j]) rows.append(row) cols.append(n2) data.append(0.5 * edge_weights[j]) rhs.append(target_edge_depths[j] * edge_weights[j]) if node_weight > 0: node_weights = node_weight * np.ones(g.Nnodes(), np.float64) # boundary nodes we consider free node_weights[boundary_nodes] = 0 for n in range(g.Nnodes()): if node_weights[n] != 0: rows.append(len(rhs)) cols.append(n) data.append(node_weights[n]) rhs.append(node_weights[n] * target_node_depths[n]) M = sparse.coo_matrix((data, (rows, cols)), shape=(len(rhs), g.Nnodes())) node_depths_sparse, status, itn, r1norm, r2norm, anorm, acond, arnorm, xnorm, var_ = sparse.linalg.lsqr( M, rhs) return node_depths_sparse
# if a particle is stuck from here on out, remove it from play dead[:Nstep]=np.all(part_xy==end_xy[:Nstep,:],axis=1) for part_idx,traj in enumerate(grp_trajectories): trajectories[ (part_grp,part_idx) ] = traj return trajectories trajectories=group_trajectories() ## trimmed=dict() # Trim each trajectory, and convert to numpy array for k in utils.progress(trajectories.keys()): traj=trajectories[k] if len(traj)>2: t,x,y = zip(*traj) traj_np=np.zeros(len(traj),dtype=[('t','M8[us]'),('x',np.float64),('y',np.float64)]) traj_np['t']=t traj_np['x']=x traj_np['y']=y trimmed[k]=traj_np trimmed_l=list(trimmed.values()) ## # Detect when those particles pass through a footprint import matplotlib.path as mpltPath
def disp_array(self): self.hydro.infer_2d_elements() self.hydro.infer_2d_links() # first calculate all time steps, just in 2D. Q = np.zeros((len(self.hydro.t_secs), self.hydro.n_2d_links), np.float64) A = np.zeros((len(self.hydro.t_secs), self.hydro.n_2d_links), np.float64) for ti in utils.progress(range(len(self.hydro.t_secs))): t_sec = self.hydro.t_secs[ti] flows = [ hydro.flows(t_sec) for hydro in [self.hydro_tidal, self.hydro] ] flow_hp = flows[0] - flows[1] # depth-integrate flow_hor = flow_hp[:self.hydro_tidal.n_exch_x] link_flows = np.bincount( self.hydro.exch_to_2d_link['link'], self.hydro.exch_to_2d_link['sgn'] * flow_hor) Q[ti, :] = link_flows**2 A[ti, :] = np.bincount( self.hydro.exch_to_2d_link['link'], self.hydro.areas(t_sec)[:self.hydro.n_exch_x]) dt_s = np.median(np.diff(self.hydro.t_secs)) winsize = int(self.lowpass_days * 86400 / dt_s) # These are a little slow. 10s? # could streamline this some since we later only use a fraction of the values. # clip here is because in some cases the values are very low and # and some roundoff is creating negatives. Qlp = filters.lowpass_fir(Q, winsize=winsize, axis=0).clip(0) Alp = filters.lowpass_fir(A, winsize=winsize, axis=0).clip(0) rms_flows = np.sqrt(Qlp) mean_A = Alp Lexch = self.hydro.exchange_lengths.sum(axis=1)[:self.hydro.n_exch_x] L = [ Lexch[exchs[0]] for l, exchs in utils.enumerate_groups( self.hydro.exch_to_2d_link['link']) ] # This is just a placeholder. A proper scaling needs to account for # cell size. rms_flows has units of m3/s. probably that should be normalized # by dividing by average flux area, and possibly multiplying by the distance # between cell centers. that doesn't seem quite right. link_K = self.K_scale * rms_flows * L / mean_A # this is computed for every time step, but we can trim that down # it's lowpassed at winsize. Try stride of half winsize. # That was used for the first round of tests, but it looks a bit # sparse. K_stride = winsize // 4 K2D = link_K[::K_stride, :] K_t_secs = self.hydro.t_secs[::K_stride] if self.amp_factor != 1.0: Kbar = K2D.mean(axis=0) K2D = (Kbar[None, :] + self.amp_factor * (K2D - Kbar[None, :])).clip(0) K = np.zeros((len(K_t_secs), self.hydro.n_exch), np.float64) # and then project to 3D K[:, :self.hydro.n_exch_x] = K2D[:, self.hydro.exch_to_2d_link['link']] if 0: # DEBUGGING # verify that I can get back to the previous, constant in time # run. log.warning("Debugging K") Kconst = super(KautoUnsteady, self).disp_array() K[:, :] = Kconst[None, :] log.info("Median dispersion coefficient: %g" % (np.median(K))) return K_t_secs, K
def set_ic_from_hycom(model, hycom_ll_box, cache_dir, default_s=None, default_T=None): """ Update model.ic_ds with salinity and temperature from hycom. hycom_ll_box is like [-124.9, -121.7, 35.9, 39.0], and is specified here to make best use of cached data (so slightly mis-alignment between models doesn't require refetching all of the hycom files). where to save/find cached hycom files default_s, default_T: when a grid point does not intersect valid hycom data, what value to use. leave as None in order to leave the value in ic_ds alone. In the past used: default_s=33.4 default_T=10.0 """ fns = hycom.fetch_range( hycom_ll_box[:2], hycom_ll_box[2:], [model.run_start, model.run_start + np.timedelta64(1, 'D')], cache_dir=cache_dir) hycom_ic_fn = fns[0] hycom_ds = xr.open_dataset(hycom_ic_fn) if 'time' in hycom_ds.dims: hycom_ds = hycom_ds.isel(time=0) cc = model.grid.cells_center() cc_ll = model.native_to_ll(cc) # Careful - some experiments (such as 92.8) have lon in [0,360], # while later ones have lon in [-180,180] # this forces all to be [-180,180] hycom_ds.lon.values[:] = (hycom_ds.lon.values + 180) % 360.0 - 180.0 dlat = np.median(np.diff(hycom_ds.lat.values)) dlon = np.median(np.diff(hycom_ds.lon.values)) lat_i = utils.nearest(hycom_ds.lat.values, cc_ll[:, 1], max_dx=1.2 * dlat) lon_i = utils.nearest(hycom_ds.lon.values, cc_ll[:, 0], max_dx=1.2 * dlon) # make this positive:down to match hycom and make the interpolation sun_z = -model.ic_ds.z_r.values assert ('time', 'Nk', 'Nc') == model.ic_ds.salt.dims, "Workaround is fragile" for scal, hy_var, sun_var, default in [ ('s', 'salinity', 'salt', default_s), ('T', 'water_temp', 'temp', default_T) ]: if scal == 's' and float(model.config['beta']) == 0.0: continue if scal == 'T' and float(model.config['gamma']) == 0.0: continue for c in utils.progress(range(model.grid.Ncells()), msg="HYCOM initial condition %s %%s" % scal): sun_val = default if lat_i[c] < 0 or lon_i[c] < 0: print("Cell %d does not overlap HYCOM grid" % c) else: # top to bottom, depth positive:down val_profile = hycom_ds[hy_var].isel(lon=lon_i[c], lat=lat_i[c]).values valid = np.isfinite(val_profile) if not np.any(valid): # print("Cell %d is dry in HYCOM grid"%c) pass else: # could add bottom salinity if we really cared. sun_val = np.interp(sun_z, hycom_ds.depth.values[valid], val_profile[valid]) if sun_val is not None: model.ic_ds[sun_var].values[0, :, c] = sun_val
poly_rio = np.array([[615061., 4224762.], [616046., 4224084.], [615810., 4223645.], [614768., 4224230.]]) # fast lookups via matplotlib: dead_polys = [poly_geo, poly_rio] from matplotlib import path dead_paths = [path.Path(poly) for poly in dead_polys] # dead_cells |= .contains_points(ctrs) ## data_dir = "csvs" os.path.exists(data_dir) or os.mkdir(data_dir) for step in utils.progress(range(1000)): fn = os.path.join(data_dir, "combined_%04d.csv" % step) if os.path.exists(fn): continue dfs = [] for src_i, src in enumerate(ptm_data): t, parts = src.read_timestep(step) sel = np.ones(len(parts), np.bool8) for dp in dead_paths: sel &= ~dp.contains_points(parts['x'][:, :2]) df = pd.DataFrame() df['x'] = parts['x'][sel, 0] df['y'] = parts['x'][sel, 1] df['group'] = src_i dfs.append(df)
if np.sum(valid) == 0: return np.nan gam = LinearGAM( s(0, n_splines=4) + s(1, n_splines=5) + te(0, 1, n_splines=4)).gridsearch(sample_st, sample_z) z_pred = gam.predict(np.array([[0, 0]]))[0] return z_pred meth = 'gam' if meth == 'invdist': meth_pretty = "Inverse Distance\np=%g, $\\alpha$=%g" % (power, aniso) params = "p%g_aniso%g" % (-power, aniso) z_result = [interp_invdist(data) for data in utils.progress(all_data)] if meth == 'invdist_plane': meth_pretty = "IDW plane\np=%g, $\\alpha$=%g" % (power, aniso) params = "p%g_aniso%g" % (-power, aniso) z_result = [ interp_invdist_plane(data) for data in utils.progress(all_data) ] if meth == 'krige': meth_pretty = "Gaussian Process" params = "" z_result = [interp_krige(data) for data in utils.progress(all_data)] if meth == 'gam': meth_pretty = "GAM" params = "v2" z_result = [interp_gam(data) for data in utils.progress(all_data)]
def add_wind_coamps_sfei(model, cache_dir, pad=np.timedelta64(3 * 3600, 's'), coamps_buffer=30e3, air_temp=False): """ model: A HydroModel instance cache_dir: path for caching wind data pad: how far before/after the simulation the wind dataset should extend Combine SFEI interpolated winds and COAMPS winds. coamps_buffer: coamps samples within this distance of SFEI data are omitted. This method does not work so well with SUNTANS. The available interpolation methods (inverse distance and kriging) do not deal well with having two distinct, densely sampled datasets with a gap in between. air_temp: if 'coamps', fill in air temperature samples from coamps data. """ g = model.grid period_start = model.run_start - pad period_stop = model.run_stop + pad fields = ['wnd_utru', 'wnd_vtru', 'pres_msl'] if air_temp == 'coamps': # may add sol_rad at some point... fields += ['air_temp', 'rltv_hum'] coamps_ds = coamps.coamps_dataset(g.bounds(), period_start, period_stop, cache_dir=cache_dir, fields=fields) sfei_ds = xr.open_dataset('wind_natneighbor_WY2017.nc') # SFEI data is PST logging.info(sfei_ds.time.values[0]) sfei_ds.time.values[:] += np.timedelta64(8 * 3600, 's') logging.info(sfei_ds.time.values[0]) # just to be sure it took. sfei_ds.time.attrs['timezone'] = 'UTC' # Not ready for other years assert sfei_ds.time.values[ 0] <= period_start, "FIX: SFEI wind only setup for 2017" assert sfei_ds.time.values[ -1] >= period_stop, "FIX: SFEI wind only setup for 2017" # buffer out the model domain a bit to get a generous footprint g_poly = geometry.Polygon(g.boundary_polygon().exterior).buffer(3000) # For each of the sources, which points will be included? # Add a mask variable for each dataset exclude_poly = None for src_name, src_ds in [('SFEI', sfei_ds), ('COAMPS', coamps_ds)]: fld = field.SimpleGrid(extents=[ src_ds.x.values[0], src_ds.x.values[-1], src_ds.y.values[0], src_ds.y.values[-1] ], F=src_ds.wind_u.isel(time=0).values) logging.info("%s: gridded resolution: %.2f %.2f" % (src_name, fld.dx, fld.dy)) mask = fld.polygon_mask(g_poly) logging.info("%s: %d of %d samples fall within grid" % (src_name, mask.sum(), mask.size)) if exclude_poly is not None: omit = fld.polygon_mask(exclude_poly) mask = mask & (~omit) logging.info( "%s: %d of %d samples fall within exclusion poly, will use %d" % (src_name, omit.sum(), omit.size, mask.sum())) src_ds['mask'] = src_ds.wind_u.dims[1:], mask # Add these points to the exclusion polygon for successive sources X, Y = fld.XY() xys = np.c_[X[mask], Y[mask]] pnts = [geometry.Point(xy[0], xy[1]) for xy in xys] poly = cascaded_union([p.buffer(coamps_buffer) for p in pnts]) if exclude_poly is None: exclude_poly = poly else: exclude_poly = exclude_poly.union(poly) # Trim to the same period # SFEI time_slc = (sfei_ds.time.values >= period_start) & (sfei_ds.time.values <= period_stop) sfei_sub_ds = sfei_ds.isel(time=time_slc) # COAMPS time_slc = (coamps_ds.time.values >= period_start) & (coamps_ds.time.values <= period_stop) coamps_sub_ds = coamps_ds.isel(time=time_slc) # make sure that worked: assert np.all(sfei_sub_ds.time.values == coamps_sub_ds.time.values) times = sfei_sub_ds.time.values # Now we start to break down the interface with model, as wind is not really # ready to go. met_ds = model.zero_met(times=times) srcs = [sfei_sub_ds, coamps_sub_ds] src_counts = [src.mask.values.sum() for src in srcs] n_points = np.sum(src_counts) xcoords = [] ycoords = [] for src in srcs: X, Y = np.meshgrid(src.x.values, src.y.values) xcoords.append(X[src.mask.values]) ycoords.append(Y[src.mask.values]) xcoords = np.concatenate(xcoords) ycoords = np.concatenate(ycoords) # Replace placeholder coordinates for wind variables. for name in ['Uwind', 'Vwind']: del met_ds["x_" + name] del met_ds["y_" + name] del met_ds["z_" + name] del met_ds[name] met_ds["x_" + name] = ("N" + name), xcoords met_ds["y_" + name] = ("N" + name, ), ycoords met_ds["z_" + name] = ("N" + name, ), 10.0 * np.ones_like(xcoords) Uwind_t = [] Vwind_t = [] for ti in utils.progress(range(len(times)), msg="Compiling wind: %s"): Uwind = [] Vwind = [] for src in srcs: Uwind.append(src.wind_u.isel(time=ti).values[src.mask]) Vwind.append(src.wind_v.isel(time=ti).values[src.mask]) Uwind = np.concatenate(Uwind) Vwind = np.concatenate(Vwind) Uwind_t.append(Uwind) Vwind_t.append(Vwind) met_ds['Uwind'] = ('nt', "NUwind"), np.stack(Uwind_t) met_ds['Vwind'] = ('nt', "NVwind"), np.stack(Vwind_t) logging.info("New Met Dataset:") logging.info(str(met_ds)) model.met_ds = met_ds if int(model.config['metmodel']) not in [0, 4]: logging.warning("Adding wind, will override metmodel %s => %d" % (model.config['metmodel'], 4)) model.config['metmodel'] = 4 # wind only
def add_wind_preblended(model, cache_dir, pad=np.timedelta64(3 * 3600, 's')): """ model: A HydroModel instance cache_dir: path for caching wind data pad: how far before/after the simulation the wind dataset should extend Add wind data from pre-blended netcdf output """ g = model.grid period_start = model.run_start - pad period_stop = model.run_stop + pad # note that this is a bit different than the sfei data # - already in UTC # - larger footprint, coarser grid # - natural neighbors run on the observed + COAMPS (thinned). blended_ds = blended_dataset(period_start, period_stop) # Not ready for other years assert blended_ds.time.values[ 0] <= period_start, "FIX: pre-blended wind only set up for some of 2017" assert blended_ds.time.values[ -1] >= period_stop, "FIX: pre-blended wind only set up for some of 2017" # buffer out the model domain a bit to get a generous footprint g_poly = geometry.Polygon(g.boundary_polygon().exterior).buffer(3000) # For each of the sources, which points will be included? # Add a mask variable for each dataset exclude_poly = None for src_name, src_ds in [('BLENDED', blended_ds)]: fld = field.SimpleGrid(extents=[ src_ds.x.values[0], src_ds.x.values[-1], src_ds.y.values[0], src_ds.y.values[-1] ], F=src_ds.wind_u.isel(time=0).values) logging.info("%s: gridded resolution: %.2f %.2f" % (src_name, fld.dx, fld.dy)) mask = fld.polygon_mask(g_poly) logging.info("%s: %d of %d samples fall within grid" % (src_name, mask.sum(), mask.size)) if exclude_poly is not None: omit = fld.polygon_mask(exclude_poly) mask = mask & (~omit) logging.info( "%s: %d of %d samples fall within exclusion poly, will use %d" % (src_name, omit.sum(), omit.size, mask.sum())) src_ds['mask'] = src_ds.wind_u.dims[1:], mask # Trim to the same period time_slc = (blended_ds.time.values >= period_start) & (blended_ds.time.values <= period_stop) blended_sub_ds = blended_ds.isel(time=time_slc) times = blended_sub_ds.time.values # Now we start to break down the interface with model, as wind is not really # ready to go. met_ds = model.zero_met(times=times) srcs = [blended_sub_ds] src_counts = [src.mask.values.sum() for src in srcs] n_points = np.sum(src_counts) xcoords = [] ycoords = [] for src in srcs: X, Y = np.meshgrid(src.x.values, src.y.values) xcoords.append(X[src.mask.values]) ycoords.append(Y[src.mask.values]) xcoords = np.concatenate(xcoords) ycoords = np.concatenate(ycoords) # Replace placeholder coordinates for wind variables. for name in ['Uwind', 'Vwind']: del met_ds["x_" + name] del met_ds["y_" + name] del met_ds["z_" + name] del met_ds[name] met_ds["x_" + name] = ("N" + name), xcoords met_ds["y_" + name] = ("N" + name, ), ycoords met_ds["z_" + name] = ("N" + name, ), 10.0 * np.ones_like(xcoords) Uwind_t = [] Vwind_t = [] for ti in utils.progress(range(len(times)), msg="Compiling wind: %s"): Uwind = [] Vwind = [] for src in srcs: Uwind.append(src.wind_u.isel(time=ti).values[src.mask]) Vwind.append(src.wind_v.isel(time=ti).values[src.mask]) Uwind = np.concatenate(Uwind) Vwind = np.concatenate(Vwind) Uwind_t.append(Uwind) Vwind_t.append(Vwind) met_ds['Uwind'] = ('nt', "NUwind"), np.stack(Uwind_t) met_ds['Vwind'] = ('nt', "NVwind"), np.stack(Vwind_t) logging.info("New Met Dataset:") logging.info(str(met_ds)) model.met_ds = met_ds if int(model.config['metmodel']) not in [4, 5]: logging.warning("While adding wind, noticed metmodel %s" % (model.config['metmodel']))
if len(track)<2: results.append(None) continue track=track.copy() for fld in ['x','y','tnum']: fld_m=0.5*( track[fld].values[:-1] + track[fld].values[1:] ) track[fld+'_m']=np.r_[ fld_m, np.nan ] track['time_m']=utils.unix_to_dt64(track['tnum_m'].values) new_records=[] for i,seg in utils.progress( track.iloc[:-1,:].iterrows() ): rec={} # new fields to be added to segments. t=seg['time_m'].to_datetime64() run_i,t_i=quantize_time(t) rec['index']=i for slice_name,slice_def in z_slices: Uint=interpolator(run_i,t_i,**slice_def) # vorticity at centers x_samp=np.array( [ [seg['x_m'] , seg['y_m'] ], [seg['x_m']-eps, seg['y_m'] ], [seg['x_m']+eps, seg['y_m'] ], [seg['x_m'] , seg['y_m'] - eps], [seg['x_m'] , seg['y_m'] + eps]
def agg_decay_tracers(model, hydro_orig, agg_grid, force='auto'): def reckless_nanmean(*a, **kw): with warnings.catch_warnings(): warnings.simplefilter("ignore") return np.nanmean(*a, **kw) # pre-calculated the data on the original grid, # aggregate, write to netcdf. precalc_fn = os.path.join(model.base_path, 'decay_summary_v00.nc') age_orig_agg2ds = [] if os.path.exists(precalc_fn): if force == 'auto': # compare precalc_fn to map_fn force = utils.is_stale(precalc_fn, [model.map_ds_path()]) if force: os.unlink(precalc_fn) if not os.path.exists(precalc_fn): model_grid = model.load_hydro().grid() if ((agg_grid.Ncells() == model_grid.Ncells()) and (np.allclose( model_grid.cells_centroid(), agg_grid.cells_centroid()))): needs_aggregation = False else: needs_aggregation = True orig_map_ds = model.map_ds() layers = orig_map_ds.dims['layer'] # These could get large -- try to build it iteratively. ds = xr.Dataset() ds['time'] = orig_map_ds['time'] ds['t_sec'] = orig_map_ds['t_sec'] agg_grid.write_to_xarray(ds) ds.to_netcdf(precalc_fn) ds.close() nc = netCDF4.Dataset(precalc_fn, mode='r+') def condense(scal): """ take a 3D, potentially unaggregated scalar field, aggregate, nanmean, and condense to 2D. """ # then aggregate but avoid nan contamination if needs_aggregation: scal_agg = agger(hydro_orig, agg_grid).segment_aggregator( t_sec, scal.ravel(), nan_method='ignore').reshape( (layers, -1)) else: scal_agg = scal scal_agg2d = reckless_nanmean(scal_agg, axis=0) return scal_agg2d for tracer in range(5): # max number of decay tracers cTr = "cTr%d" % (tracer + 1) dTr = "dTr%d" % (tracer + 1) if cTr not in orig_map_ds: break log.info("%s: Processing tracers %s, %s" % (precalc_fn, cTr, dTr)) nc_var = nc.createVariable('age%d' % (tracer + 1), np.float32, ('time', 'face'), zlib=True, complevel=2) nc_var_cons = nc.createVariable(cTr, np.float32, ('time', 'face'), zlib=True, complevel=2) nc_var_decay = nc.createVariable(dTr, np.float32, ('time', 'face'), zlib=True, complevel=2) for t_idx in utils.progress(range(orig_map_ds.dims['time'])): t_sec = orig_map_ds['t_sec'].isel(time=t_idx).values cTr_orig = orig_map_ds[cTr].isel(time=t_idx).values dTr_orig = orig_map_ds[dTr].isel(time=t_idx).values age_orig = recalc_age(cTr_orig, dTr_orig) nc_var[t_idx, :] = condense(age_orig) nc_var_cons[t_idx, :] = condense(cTr_orig) nc_var_decay[t_idx, :] = condense(dTr_orig) nc.sync() nc.close() return xr.open_dataset(precalc_fn)