def __init__(self, directory, domain, model_type="earth_model"): """ The init function. :param directory: The directory where the earth model or kernel is located. :param model_type: Determined the type of model loaded. Currently two are supported: * earth_model - The standard SES3D model files (default) * kernel - The kernels. Identifies by lots of grad_* files. * wavefield - The raw wavefields. """ self.directory = directory self.boxfile = os.path.join(self.directory, "boxfile") if not os.path.exists(self.boxfile): msg = "boxfile not found. Wrong directory?" raise ValueError(msg) # Read the boxfile. self.setup = self._read_boxfile() self.domain = domain self.model_type = model_type self.one_d_model = OneDimensionalModel("ak135-f") if model_type == "earth_model": # Now check what different models are available in the directory. # This information is also used to infer the degree of the used # lagrange polynomial. components = ["A", "B", "C", "lambda", "mu", "rhoinv", "Q"] self.available_derived_components = ["vp", "vsh", "vsv", "rho"] self.components = {} self.parsed_components = {} for component in components: files = glob.glob( os.path.join(directory, "%s[0-9]*" % component)) if len(files) != len(self.setup["subdomains"]): continue # Check that the naming is continuous. all_good = True for _i in range(len(self.setup["subdomains"])): if os.path.join(directory, "%s%i" % (component, _i)) in files: continue all_good = False break if all_good is False: msg = "Naming for component %s is off. It will be skipped." warnings.warn(msg) continue # They also all need to have the same size. if len(set([os.path.getsize(_i) for _i in files])) != 1: msg = ("Component %s has the right number of model files " "but they are not of equal size") % component warnings.warn(msg) continue # Sort the files by ascending number. files.sort(key=lambda x: int( re.findall(r"\d+$", (os.path.basename(x)))[0])) self.components[component] = {"filenames": files} elif model_type == "wavefield": components = ["vz", "vx", "vy", "vz"] self.available_derived_components = [] self.components = {} self.parsed_components = {} for component in components: files = glob.glob(os.path.join(directory, "%s_*_*" % component)) if not files: continue timesteps = collections.defaultdict(list) for filename in files: timestep = int(os.path.basename(filename).split("_")[-1]) timesteps[timestep].append(filename) for timestep, filenames in timesteps.items(): self.components["%s %s" % (component, timestep)] = \ {"filenames": sorted( filenames, key=lambda x: int( os.path.basename(x).split("_")[1]))} elif model_type == "kernel": # Now check what different models are available in the directory. # This information is also used to infer the degree of the used # lagrange polynomial. components = ["grad_cp", "grad_csh", "grad_csv", "grad_rho"] self.available_derived_components = [] self.components = {} self.parsed_components = {} for component in components: files = glob.glob( os.path.join(directory, "%s_[0-9]*" % component)) if len(files) != len(self.setup["subdomains"]): continue if len(set([os.path.getsize(_i) for _i in files])) != 1: msg = ("Component %s has the right number of model files " "but they are not of equal size") % component warnings.warn(msg) continue # Sort the files by ascending number. files.sort(key=lambda x: int( re.findall(r"\d+$", (os.path.basename(x)))[0])) self.components[component] = {"filenames": files} else: msg = "model_type '%s' not known." % model_type raise ValueError(msg) # All files for a single component have the same size. Now check that # all files have the same size. unique_filesizes = len( list( set([ os.path.getsize(_i["filenames"][0]) for _i in self.components.values() ]))) if unique_filesizes != 1: msg = ("The different components in the folder do not have the " "same number of samples") raise ValueError(msg) # Now calculate the lagrange polynomial degree. All necessary # information is present. size = os.path.getsize( list(self.components.values())[0]["filenames"][0]) sd = self.setup["subdomains"][0] x, y, z = sd["index_x_count"], sd["index_y_count"], sd["index_z_count"] self.lagrange_polynomial_degree = \ int(round(((size * 0.25) / (x * y * z)) ** (1.0 / 3.0) - 1)) self._calculate_final_dimensions() # Setup the boundaries. self.lat_bounds = [ rotations.colat2lat(_i) for _i in self.setup["physical_boundaries_x"][::-1] ] self.lng_bounds = self.setup["physical_boundaries_y"] self.depth_bounds = [ 6371 - _i / 1000.0 for _i in self.setup["physical_boundaries_z"] ] self.collocation_points_lngs = self._get_collocation_points_along_axis( self.lng_bounds[0], self.lng_bounds[1], self.setup["point_count_in_y"]) self.collocation_points_lats = self._get_collocation_points_along_axis( self.lat_bounds[0], self.lat_bounds[1], self.setup["point_count_in_x"]) self.collocation_points_depth = \ self._get_collocation_points_along_axis( self.depth_bounds[1], self.depth_bounds[0], self.setup["point_count_in_z"])[::-1]
def __init__(self, directory, domain, model_type="earth_model"): """ The init function. :param directory: The directory where the earth model or kernel is located. :param model_type: Determined the type of model loaded. Currently two are supported: * earth_model - The standard SES3D model files (default) * kernel - The kernels. Identifies by lots of grad_* files. * wavefield - The raw wavefields. """ self.directory = directory self.boxfile = os.path.join(self.directory, "boxfile") if not os.path.exists(self.boxfile): msg = "boxfile not found. Wrong directory?" raise ValueError(msg) # Read the boxfile. self.setup = self._read_boxfile() self.domain = domain self.model_type = model_type self.one_d_model = OneDimensionalModel("ak135-f") if model_type == "earth_model": # Now check what different models are available in the directory. # This information is also used to infer the degree of the used # lagrange polynomial. components = ["A", "B", "C", "lambda", "mu", "rhoinv", "Q"] self.available_derived_components = ["vp", "vsh", "vsv", "rho"] self.components = {} self.parsed_components = {} for component in components: files = glob.glob( os.path.join(directory, "%s[0-9]*" % component)) if len(files) != len(self.setup["subdomains"]): continue # Check that the naming is continuous. all_good = True for _i in xrange(len(self.setup["subdomains"])): if os.path.join(directory, "%s%i" % (component, _i)) in files: continue all_good = False break if all_good is False: msg = "Naming for component %s is off. It will be skipped." warnings.warn(msg) continue # They also all need to have the same size. if len(set([os.path.getsize(_i) for _i in files])) != 1: msg = ("Component %s has the right number of model files " "but they are not of equal size") % component warnings.warn(msg) continue # Sort the files by ascending number. files.sort(key=lambda x: int(re.findall(r"\d+$", (os.path.basename(x)))[0])) self.components[component] = {"filenames": files} elif model_type == "wavefield": components = ["vz", "vx", "vy", "vz"] self.available_derived_components = [] self.components = {} self.parsed_components = {} for component in components: files = glob.glob(os.path.join(directory, "%s_*_*" % component)) if not files: continue timesteps = collections.defaultdict(list) for filename in files: timestep = int(os.path.basename(filename).split("_")[-1]) timesteps[timestep].append(filename) for timestep, filenames in timesteps.iteritems(): self.components["%s %s" % (component, timestep)] = \ {"filenames": sorted( filenames, key=lambda x: int( os.path.basename(x).split("_")[1]))} elif model_type == "kernel": # Now check what different models are available in the directory. # This information is also used to infer the degree of the used # lagrange polynomial. components = ["grad_cp", "grad_csh", "grad_csv", "grad_rho"] self.available_derived_components = [] self.components = {} self.parsed_components = {} for component in components: files = glob.glob( os.path.join(directory, "%s_[0-9]*" % component)) if len(files) != len(self.setup["subdomains"]): continue if len(set([os.path.getsize(_i) for _i in files])) != 1: msg = ("Component %s has the right number of model files " "but they are not of equal size") % component warnings.warn(msg) continue # Sort the files by ascending number. files.sort(key=lambda x: int( re.findall(r"\d+$", (os.path.basename(x)))[0])) self.components[component] = {"filenames": files} else: msg = "model_type '%s' not known." % model_type raise ValueError(msg) # All files for a single component have the same size. Now check that # all files have the same size. unique_filesizes = len(list(set([ os.path.getsize(_i["filenames"][0]) for _i in self.components.itervalues()]))) if unique_filesizes != 1: msg = ("The different components in the folder do not have the " "same number of samples") raise ValueError(msg) # Now calculate the lagrange polynomial degree. All necessary # information is present. size = os.path.getsize(self.components.values()[0]["filenames"][0]) sd = self.setup["subdomains"][0] x, y, z = sd["index_x_count"], sd["index_y_count"], sd["index_z_count"] self.lagrange_polynomial_degree = \ int(round(((size * 0.25) / (x * y * z)) ** (1.0 / 3.0) - 1)) self._calculate_final_dimensions() # Setup the boundaries. self.lat_bounds = [ rotations.colat2lat(_i) for _i in self.setup["physical_boundaries_x"][::-1]] self.lng_bounds = self.setup["physical_boundaries_y"] self.depth_bounds = [ 6371 - _i / 1000.0 for _i in self.setup["physical_boundaries_z"]] self.collocation_points_lngs = self._get_collocation_points_along_axis( self.lng_bounds[0], self.lng_bounds[1], self.setup["point_count_in_y"]) self.collocation_points_lats = self._get_collocation_points_along_axis( self.lat_bounds[0], self.lat_bounds[1], self.setup["point_count_in_x"]) self.collocation_points_depth = \ self._get_collocation_points_along_axis( self.depth_bounds[1], self.depth_bounds[0], self.setup["point_count_in_z"])[::-1]
class RawSES3DModelHandler(object): """ Class able to deal with a directory of raw SES3D 4.0 models. An earth model directory is defined by containing: * boxfile -> File describing the discretization * A0-A[XX] -> Anisotropic parameters (one per CPU) * B0-B[XX] -> Anisotropic parameters (one per CPU) * C0-C[XX] -> Anisotropic parameters (one per CPU) * lambda0 - lambda[XX] -> Elastic parameters (one per CPU) * mu0 - mu[XX] -> Elastic parameters (one per CPU) * rhoinv0 - rhoinv[XX] -> Elastic parameters (one per CPU) * Q0 - Q[XX] -> Elastic parameters (one per CPU) A kernel directory is defined by containing: * boxfile -> File describing the discretization * grad_cp_[XX]_[XXXX] * grad_csh_[XX]_[XXXX] * grad_csv_[XX]_[XXXX] * grad_rho_[XX]_[XXXX] A wavefield directory contains the following files: * boxfile -> File describing the discretization * vx_[xx]_[timestep] * vy_[xx]_[timestep] * vz_[xx]_[timestep] """ def __init__(self, directory, domain, model_type="earth_model"): """ The init function. :param directory: The directory where the earth model or kernel is located. :param model_type: Determined the type of model loaded. Currently two are supported: * earth_model - The standard SES3D model files (default) * kernel - The kernels. Identifies by lots of grad_* files. * wavefield - The raw wavefields. """ self.directory = directory self.boxfile = os.path.join(self.directory, "boxfile") if not os.path.exists(self.boxfile): msg = "boxfile not found. Wrong directory?" raise ValueError(msg) # Read the boxfile. self.setup = self._read_boxfile() self.domain = domain self.model_type = model_type self.one_d_model = OneDimensionalModel("ak135-f") if model_type == "earth_model": # Now check what different models are available in the directory. # This information is also used to infer the degree of the used # lagrange polynomial. components = ["A", "B", "C", "lambda", "mu", "rhoinv", "Q"] self.available_derived_components = ["vp", "vsh", "vsv", "rho"] self.components = {} self.parsed_components = {} for component in components: files = glob.glob( os.path.join(directory, "%s[0-9]*" % component)) if len(files) != len(self.setup["subdomains"]): continue # Check that the naming is continuous. all_good = True for _i in range(len(self.setup["subdomains"])): if os.path.join(directory, "%s%i" % (component, _i)) in files: continue all_good = False break if all_good is False: msg = "Naming for component %s is off. It will be skipped." warnings.warn(msg) continue # They also all need to have the same size. if len(set([os.path.getsize(_i) for _i in files])) != 1: msg = ("Component %s has the right number of model files " "but they are not of equal size") % component warnings.warn(msg) continue # Sort the files by ascending number. files.sort(key=lambda x: int( re.findall(r"\d+$", (os.path.basename(x)))[0])) self.components[component] = {"filenames": files} elif model_type == "wavefield": components = ["vz", "vx", "vy", "vz"] self.available_derived_components = [] self.components = {} self.parsed_components = {} for component in components: files = glob.glob(os.path.join(directory, "%s_*_*" % component)) if not files: continue timesteps = collections.defaultdict(list) for filename in files: timestep = int(os.path.basename(filename).split("_")[-1]) timesteps[timestep].append(filename) for timestep, filenames in timesteps.items(): self.components["%s %s" % (component, timestep)] = \ {"filenames": sorted( filenames, key=lambda x: int( os.path.basename(x).split("_")[1]))} elif model_type == "kernel": # Now check what different models are available in the directory. # This information is also used to infer the degree of the used # lagrange polynomial. components = ["grad_cp", "grad_csh", "grad_csv", "grad_rho"] self.available_derived_components = [] self.components = {} self.parsed_components = {} for component in components: files = glob.glob( os.path.join(directory, "%s_[0-9]*" % component)) if len(files) != len(self.setup["subdomains"]): continue if len(set([os.path.getsize(_i) for _i in files])) != 1: msg = ("Component %s has the right number of model files " "but they are not of equal size") % component warnings.warn(msg) continue # Sort the files by ascending number. files.sort(key=lambda x: int( re.findall(r"\d+$", (os.path.basename(x)))[0])) self.components[component] = {"filenames": files} else: msg = "model_type '%s' not known." % model_type raise ValueError(msg) # All files for a single component have the same size. Now check that # all files have the same size. unique_filesizes = len( list( set([ os.path.getsize(_i["filenames"][0]) for _i in self.components.values() ]))) if unique_filesizes != 1: msg = ("The different components in the folder do not have the " "same number of samples") raise ValueError(msg) # Now calculate the lagrange polynomial degree. All necessary # information is present. size = os.path.getsize( list(self.components.values())[0]["filenames"][0]) sd = self.setup["subdomains"][0] x, y, z = sd["index_x_count"], sd["index_y_count"], sd["index_z_count"] self.lagrange_polynomial_degree = \ int(round(((size * 0.25) / (x * y * z)) ** (1.0 / 3.0) - 1)) self._calculate_final_dimensions() # Setup the boundaries. self.lat_bounds = [ rotations.colat2lat(_i) for _i in self.setup["physical_boundaries_x"][::-1] ] self.lng_bounds = self.setup["physical_boundaries_y"] self.depth_bounds = [ 6371 - _i / 1000.0 for _i in self.setup["physical_boundaries_z"] ] self.collocation_points_lngs = self._get_collocation_points_along_axis( self.lng_bounds[0], self.lng_bounds[1], self.setup["point_count_in_y"]) self.collocation_points_lats = self._get_collocation_points_along_axis( self.lat_bounds[0], self.lat_bounds[1], self.setup["point_count_in_x"]) self.collocation_points_depth = \ self._get_collocation_points_along_axis( self.depth_bounds[1], self.depth_bounds[0], self.setup["point_count_in_z"])[::-1] def _read_single_box(self, component, file_number): """ This function reads Ses3ds raw binary files, e.g. 3d velocity field snapshots, as well as model parameter files or sensitivity kernels. It returns the field as an array of rank 3 with shape (nx*lpd+1, ny*lpd+1, nz*lpd+1), discarding the duplicates by default. """ # Get the file and the corresponding domain. filename = self.components[component]["filenames"][file_number] domain = self.setup["subdomains"][file_number] lpd = self.lagrange_polynomial_degree shape = (domain["index_x_count"], domain["index_y_count"], domain["index_z_count"], lpd + 1, lpd + 1, lpd + 1) # Take care: The first and last four bytes in the arrays are invalid # due to them being written by Fortran. with open(filename, "rb") as open_file: field = np.ndarray(shape, buffer=open_file.read()[4:-4], dtype="float32", order="F") # Calculate the new shape by multiplying every dimension with lpd + 1 # value for every dimension. new_shape = [_i * _j for _i, _j in zip(shape[:3], shape[3:])] # Reorder the axes: # v[x, y, z, lpd + 1 ,lpd + 1, lpd + 1] -> v[x, lpd + 1, y, z, # lpd + 1, lpd + 1] -> v[x, lpd + 1, y, lpd + 1, z, lpd + 1] field = np.rollaxis(np.rollaxis(field, 3, 1), 3, lpd + 1) # Reshape the data: # v[nx,lpd+1,ny,lpd+1,nz,lpd+1] to v[nx*(lpd+1),ny*(lpd+1),nz*(lpd+1)] field = field.reshape(new_shape, order="C") # XXX: Attempt to do this in one step. for axis, count in enumerate(new_shape): # Mask the duplicate values. mask = np.ones(count, dtype="bool") mask[::lpd + 1][1:] = False # Remove them by compressing the array. field = field.compress(mask, axis=axis) return field[:, :, ::-1] def parse_component(self, component): """ Helper function parsing a whole component. :param component: The component name. :type component: basestring """ # If a real component, if component in list(self.components.keys()): self._parse_component(component) return elif component not in self.available_derived_components: msg = "Component %s is unknown" % component raise ValueError(msg) if component == "vp": self._parse_component("lambda") self._parse_component("mu") self._parse_component("rhoinv") lambda_ = self.parsed_components["lambda"] mu = self.parsed_components["mu"] rhoinv = self.parsed_components["rhoinv"] self.parsed_components["vp"] = \ np.sqrt(((lambda_ + 2.0 * mu) * rhoinv)) / 1000.0 elif component == "vsh": self._parse_component("mu") self._parse_component("rhoinv") mu = self.parsed_components["mu"] rhoinv = self.parsed_components["rhoinv"] self.parsed_components["vsh"] = np.sqrt((mu * rhoinv)) / 1000.0 elif component == "vsv": self._parse_component("mu") self._parse_component("rhoinv") self._parse_component("B") mu = self.parsed_components["mu"] rhoinv = self.parsed_components["rhoinv"] b = self.parsed_components["B"] self.parsed_components["vsv"] = \ np.sqrt((mu + b) * rhoinv) / 1000.0 elif component == "rho": self._parse_component("rhoinv") rhoinv = self.parsed_components["rhoinv"] self.parsed_components["rho"] = 1.0 / rhoinv def _parse_component(self, component): """ Parses the specified component. """ if component in self.parsed_components: return # Allocate empty array with the necessary dimensions. data = np.empty( (self.setup["point_count_in_x"], self.setup["point_count_in_y"], self.setup["point_count_in_z"]), dtype="float32") for _i, domain in enumerate(self.setup["subdomains"]): x_min, x_max = domain["boundaries_x"] y_min, y_max = domain["boundaries_y"] z_min, z_max = domain["boundaries_z"] # Minimum indices x_min, y_min, z_min = [ self.lagrange_polynomial_degree * _j for _j in (x_min, y_min, z_min) ] # Maximum indices x_max, y_max, z_max = [ self.lagrange_polynomial_degree * (_j + 1) for _j in (x_max, y_max, z_max) ] # Merge into data. data[x_min: x_max + 1, y_min: y_max + 1, z_min: z_max + 1] = \ self._read_single_box(component, _i) self.parsed_components[component] = data def _calculate_final_dimensions(self): """ Calculates the total number of elements and points and also the size of the final array per component. """ # Elements in each directions. x, y, z = (self.setup["boundaries_x"][1] - self.setup["boundaries_x"][0] + 1, self.setup["boundaries_y"][1] - self.setup["boundaries_y"][0] + 1, self.setup["boundaries_z"][1] - self.setup["boundaries_z"][0] + 1) self.setup["total_element_count"] = x * y * z self.setup["total_point_count"] = (x * y * z) * \ (self.lagrange_polynomial_degree + 1) ** 3 # This is the actual point count with removed duplicates. Much # smaller because each element shared a large number of elements # with its neighbours. self.setup["point_count_in_x"] = \ (x * self.lagrange_polynomial_degree + 1) self.setup["point_count_in_y"] = \ (y * self.lagrange_polynomial_degree + 1) self.setup["point_count_in_z"] = \ (z * self.lagrange_polynomial_degree + 1) self.setup["total_point_count_without_duplicates"] = ( self.setup["point_count_in_x"] * self.setup["point_count_in_y"] * self.setup["point_count_in_z"]) def _get_collocation_points_along_axis(self, min_value, max_value, count): """ Calculates count collocation points from min_value to max_value as they would be used by a simulation. The border elements are not repeated. :param min_value: The minimum value. :param max_value: The maximum value. :param count: The number of values. """ # Normalize from 0.0 to 1.0 coll_points = get_lpd_sampling_points(self.lagrange_polynomial_degree) coll_points += 1.0 coll_points /= 2.0 # Some view and reshaping tricks to ensure performance. Probably not # needed. points = np.empty(count) lpd = self.lagrange_polynomial_degree p_view = points[:count - 1].view().reshape(((count - 1) / lpd, lpd)) for _i in range(p_view.shape[0]): p_view[_i] = _i points[-1] = points[-2] + 1 for _i, value in enumerate(coll_points[1:-1]): points[_i + 1::lpd] += value # Normalize from min_value to max_value. points /= points.max() points *= abs(max_value - min_value) points += min_value return points def get_closest_gll_index(self, coord, value): if coord == "depth": return np.argmin(np.abs(self.collocation_points_depth - value)) elif coord == "longitude": return np.argmin(np.abs(self.collocation_points_lngs - value)) elif coord == "latitude": return np.argmin(np.abs(self.collocation_points_lats[::-1] - value)) def plot_depth_slice(self, component, depth_in_km, m, absolute_values=True): """ Plots a depth slice. :param component: The component to plot. :type component: basestring :param depth_in_km: The depth in km to plot. If the exact depth does not exists, the nearest neighbour will be plotted. :type depth_in_km: integer or float :param m: Basemap instance. """ depth_index = self.get_closest_gll_index("depth", depth_in_km) # No need to do anything if the currently plotted slice is already # plotted. This is useful for interactive use when the desired depth # is changed but the closest GLL collocation point is still the same. if hasattr(m, "_plotted_depth_slice"): # Use a tuple of relevant parameters. if m._plotted_depth_slice == (self.directory, depth_index, component, absolute_values): return None data = self.parsed_components[component] depth = self.collocation_points_depth[depth_index] lngs = self.collocation_points_lngs lats = self.collocation_points_lats # Rotate data if needed. lon, lat = np.meshgrid(lngs, lats) if hasattr(self.domain, "rotation_axis") and \ self.domain.rotation_axis and \ self.domain.rotation_angle_in_degree: lon_shape = lon.shape lat_shape = lat.shape lon.shape = lon.size lat.shape = lat.size lat, lon = rotations.rotate_lat_lon( lat, lon, self.domain.rotation_axis, self.domain.rotation_angle_in_degree) lon.shape = lon_shape lat.shape = lat_shape x, y = m(lon, lat) depth_data = data[::-1, :, depth_index] # Plot values relative to AK135. if not absolute_values: cmp_map = {"rho": "density", "vp": "vp", "vsh": "vs", "vsv": "vs"} factor = { "rho": 1000.0, "vp": 1.0, "vsh": 1.0, "vsv": 1.0, } if component not in cmp_map: vmin, vmax = depth_data.min(), depth_data.max() vmedian = np.median(depth_data) offset = max(abs(vmax - vmedian), abs(vmedian - vmin)) if vmax - vmin == 0: offset = 0.01 vmin = vmedian - offset vmax = vmedian + offset else: reference_value = self.one_d_model.get_value( cmp_map[component], depth) * factor[component] depth_data = (depth_data - reference_value) / reference_value depth_data *= 100.0 offset = np.abs(depth_data) try: offset = offset[offset < 50].max() except BaseException: offset = offset.max() vmin = -offset vmax = offset else: vmin, vmax = depth_data.min(), depth_data.max() vmedian = np.median(depth_data) offset = max(abs(vmax - vmedian), abs(vmedian - vmin)) min_delta = abs(vmax * 0.005) if (vmax - vmin) < min_delta: offset = min_delta vmin = vmedian - offset vmax = vmedian + offset # Remove an existing pcolormesh if it exists. This does not hurt in # any case but is useful for interactive use. if hasattr(m, "_depth_slice"): m._depth_slice.remove() del m._depth_slice im = m.pcolormesh(x, y, depth_data, cmap=tomo_colormap, vmin=vmin, vmax=vmax) m._depth_slice = im # Store what is currently plotted. m._plotted_depth_slice = (self.directory, depth_index, component, absolute_values) return {"depth": depth, "mesh": im, "data": depth_data} def get_depth_profile(self, component, latitude, longitude): """ Returns a depth profile of the model at the requested at the GLL points closest do latitude and longitude. :param component: The component of the model. :param latitude: The latitude. :param longitude: The longitude. """ # Need to rotate latitude and longitude. if hasattr(self.domain, "rotation_axis") and \ self.domain.rotation_axis and \ self.domain.rotation_angle_in_degree: latitude, longitude = rotations.rotate_lat_lon( latitude, longitude, self.domain.rotation_axis, -1.0 * self.domain.rotation_angle_in_degree) x_index = self.get_closest_gll_index("latitude", latitude) y_index = self.get_closest_gll_index("longitude", longitude) data = self.parsed_components[component] depths = self.collocation_points_depth values = data[x_index, y_index, :] lat = self.collocation_points_lats[::-1][x_index] lng = self.collocation_points_lngs[y_index] # Rotate back. if hasattr(self.domain, "rotation_axis") and \ self.domain.rotation_axis and \ self.domain.rotation_angle_in_degree: lat, lng = rotations.rotate_lat_lon( lat, lng, self.domain.rotation_axis, self.domain.rotation_angle_in_degree) return { "depths": depths, "values": values, "latitude": lat, "longitude": lng } def __str__(self): """ Eases interactive working. """ ret_str = "Raw SES3D Model (split in %i parts)\n" % \ (self.setup["total_cpu_count"]) ret_str += "\tSetup:\n" ret_str += "\t\tLatitude: {:.2f} - {:.2f}\n".format(*[ rotations.colat2lat(_i) for _i in self.setup["physical_boundaries_x"][::-1] ]) ret_str += "\t\tLongitude: %.2f - %.2f\n" % \ self.setup["physical_boundaries_y"] ret_str += "\t\tDepth in km: {:.2f} - {:.2f}\n".format(*[ 6371 - _i / 1000 for _i in self.setup["physical_boundaries_z"][::-1] ]) ret_str += "\t\tTotal element count: %i\n" % \ self.setup["total_element_count"] ret_str += "\t\tTotal collocation point count: %i (without " \ "duplicates: %i)\n" % ( self.setup["total_point_count"], self.setup["total_point_count_without_duplicates"]) ret_str += "\tMemory requirement per component: %.1f MB\n" % \ ((self.setup["total_point_count_without_duplicates"] * 4) / (1024.0 ** 2)) ret_str += "\tAvailable components: %s\n" % (", ".join( sorted(self.components.keys()))) ret_str += "\tAvailable derived components: %s\n" % (", ".join( sorted(self.available_derived_components))) ret_str += "\tParsed components: %s" % (", ".join( sorted(self.parsed_components.keys()))) return ret_str def _read_boxfile(self): setup = {"subdomains": []} with open(self.boxfile, "rU") as fh: # The first 14 lines denote the header lines = fh.readlines()[14:] # Strip lines and remove empty lines. lines = [_i.strip() for _i in lines if _i.strip()] # The next 4 are the global CPU distribution. setup["total_cpu_count"] = int(lines.pop(0)) setup["cpu_count_in_x_direction"] = int(lines.pop(0)) setup["cpu_count_in_y_direction"] = int(lines.pop(0)) setup["cpu_count_in_z_direction"] = int(lines.pop(0)) if set(lines[0]) == set("-"): lines.pop(0) # Small sanity check. if setup["total_cpu_count"] != setup["cpu_count_in_x_direction"] *\ setup["cpu_count_in_y_direction"] * \ setup["cpu_count_in_z_direction"]: msg = ("Invalid boxfile. Total and individual processor " "counts do not match.") raise ValueError(msg) # Now parse the rest of file which contains the subdomains. def subdomain_generator(data): """ Simple generator looping over each defined box and yielding a dictionary for each. :param data: The text. """ while data: subdom = {} # Convert both indices to 0-based indices subdom["single_index"] = int(data.pop(0)) - 1 subdom["multi_index"] = [ int(x) - 1 for x in data.pop(0).split() ] subdom["boundaries_x"] = list(map(int, data.pop(0).split())) subdom["boundaries_y"] = list(map(int, data.pop(0).split())) subdom["boundaries_z"] = list(map(int, data.pop(0).split())) # Convert radians to degree. subdom["physical_boundaries_x"] = [ math.degrees(float(x)) for x in data.pop(0).split() ] subdom["physical_boundaries_y"] = [ math.degrees(float(x)) for x in data.pop(0).split() ] # z is in meter. subdom["physical_boundaries_z"] = \ list(map(float, data.pop(0).split())) for component in ("x", "y", "z"): idx = "boundaries_%s" % component index_count = subdom[idx][1] - subdom[idx][0] + 1 subdom["index_%s_count" % component] = index_count # The boxfiles are slightly awkward in that the indices # are not really continuous. For example if one box # has 22 as the last index, the first index of the next # box will also be 22, even though it should be 23. The # next snippet attempts to fix this deficiency. offset = int( round(subdom[idx][0] / float(index_count - 1))) subdom[idx][0] += offset subdom[idx][1] += offset # Remove separator_line if existent. if set(lines[0]) == set("-"): lines.pop(0) yield subdom # Sort them after with the single index. setup["subdomains"] = sorted(list(subdomain_generator(lines)), key=lambda x: x["single_index"]) # Do some more sanity checks. if len(setup["subdomains"]) != setup["total_cpu_count"]: msg = ("Invalid boxfile. Number of processors and subdomains " "to not match.") raise ValueError(msg) for component in ("x", "y", "z"): idx = "index_%s_count" % component if len(set([_i[idx] for _i in setup["subdomains"]])) != 1: msg = ("Invalid boxfile. Unequal %s index count across " "subdomains.") % component raise ValueError(msg) # Now generate the absolute indices for the whole domains. for component in ("x", "y", "z"): setup["boundaries_%s" % component] = (min([ _i["boundaries_%s" % component][0] for _i in setup["subdomains"] ]), max([ _i["boundaries_%s" % component][1] for _i in setup["subdomains"] ])) setup["physical_boundaries_%s" % component] = (min([ _i["physical_boundaries_%s" % component][0] for _i in setup["subdomains"] ]), max([ _i["physical_boundaries_%s" % component][1] for _i in setup["subdomains"] ])) return setup
class RawSES3DModelHandler(object): """ Class able to deal with a directory of raw SES3D 4.0 models. An earth model directory is defined by containing: * boxfile -> File describing the discretization * A0-A[XX] -> Anisotropic parameters (one per CPU) * B0-B[XX] -> Anisotropic parameters (one per CPU) * C0-C[XX] -> Anisotropic parameters (one per CPU) * lambda0 - lambda[XX] -> Elastic parameters (one per CPU) * mu0 - mu[XX] -> Elastic parameters (one per CPU) * rhoinv0 - rhoinv[XX] -> Elastic parameters (one per CPU) * Q0 - Q[XX] -> Elastic parameters (one per CPU) A kernel directory is defined by containing: * boxfile -> File describing the discretization * grad_cp_[XX]_[XXXX] * grad_csh_[XX]_[XXXX] * grad_csv_[XX]_[XXXX] * grad_rho_[XX]_[XXXX] A wavefield directory contains the following files: * boxfile -> File describing the discretization * vx_[xx]_[timestep] * vy_[xx]_[timestep] * vz_[xx]_[timestep] """ def __init__(self, directory, domain, model_type="earth_model"): """ The init function. :param directory: The directory where the earth model or kernel is located. :param model_type: Determined the type of model loaded. Currently two are supported: * earth_model - The standard SES3D model files (default) * kernel - The kernels. Identifies by lots of grad_* files. * wavefield - The raw wavefields. """ self.directory = directory self.boxfile = os.path.join(self.directory, "boxfile") if not os.path.exists(self.boxfile): msg = "boxfile not found. Wrong directory?" raise ValueError(msg) # Read the boxfile. self.setup = self._read_boxfile() self.domain = domain self.model_type = model_type self.one_d_model = OneDimensionalModel("ak135-f") if model_type == "earth_model": # Now check what different models are available in the directory. # This information is also used to infer the degree of the used # lagrange polynomial. components = ["A", "B", "C", "lambda", "mu", "rhoinv", "Q"] self.available_derived_components = ["vp", "vsh", "vsv", "rho"] self.components = {} self.parsed_components = {} for component in components: files = glob.glob( os.path.join(directory, "%s[0-9]*" % component)) if len(files) != len(self.setup["subdomains"]): continue # Check that the naming is continuous. all_good = True for _i in xrange(len(self.setup["subdomains"])): if os.path.join(directory, "%s%i" % (component, _i)) in files: continue all_good = False break if all_good is False: msg = "Naming for component %s is off. It will be skipped." warnings.warn(msg) continue # They also all need to have the same size. if len(set([os.path.getsize(_i) for _i in files])) != 1: msg = ("Component %s has the right number of model files " "but they are not of equal size") % component warnings.warn(msg) continue # Sort the files by ascending number. files.sort(key=lambda x: int(re.findall(r"\d+$", (os.path.basename(x)))[0])) self.components[component] = {"filenames": files} elif model_type == "wavefield": components = ["vz", "vx", "vy", "vz"] self.available_derived_components = [] self.components = {} self.parsed_components = {} for component in components: files = glob.glob(os.path.join(directory, "%s_*_*" % component)) if not files: continue timesteps = collections.defaultdict(list) for filename in files: timestep = int(os.path.basename(filename).split("_")[-1]) timesteps[timestep].append(filename) for timestep, filenames in timesteps.iteritems(): self.components["%s %s" % (component, timestep)] = \ {"filenames": sorted( filenames, key=lambda x: int( os.path.basename(x).split("_")[1]))} elif model_type == "kernel": # Now check what different models are available in the directory. # This information is also used to infer the degree of the used # lagrange polynomial. components = ["grad_cp", "grad_csh", "grad_csv", "grad_rho"] self.available_derived_components = [] self.components = {} self.parsed_components = {} for component in components: files = glob.glob( os.path.join(directory, "%s_[0-9]*" % component)) if len(files) != len(self.setup["subdomains"]): continue if len(set([os.path.getsize(_i) for _i in files])) != 1: msg = ("Component %s has the right number of model files " "but they are not of equal size") % component warnings.warn(msg) continue # Sort the files by ascending number. files.sort(key=lambda x: int( re.findall(r"\d+$", (os.path.basename(x)))[0])) self.components[component] = {"filenames": files} else: msg = "model_type '%s' not known." % model_type raise ValueError(msg) # All files for a single component have the same size. Now check that # all files have the same size. unique_filesizes = len(list(set([ os.path.getsize(_i["filenames"][0]) for _i in self.components.itervalues()]))) if unique_filesizes != 1: msg = ("The different components in the folder do not have the " "same number of samples") raise ValueError(msg) # Now calculate the lagrange polynomial degree. All necessary # information is present. size = os.path.getsize(self.components.values()[0]["filenames"][0]) sd = self.setup["subdomains"][0] x, y, z = sd["index_x_count"], sd["index_y_count"], sd["index_z_count"] self.lagrange_polynomial_degree = \ int(round(((size * 0.25) / (x * y * z)) ** (1.0 / 3.0) - 1)) self._calculate_final_dimensions() # Setup the boundaries. self.lat_bounds = [ rotations.colat2lat(_i) for _i in self.setup["physical_boundaries_x"][::-1]] self.lng_bounds = self.setup["physical_boundaries_y"] self.depth_bounds = [ 6371 - _i / 1000.0 for _i in self.setup["physical_boundaries_z"]] self.collocation_points_lngs = self._get_collocation_points_along_axis( self.lng_bounds[0], self.lng_bounds[1], self.setup["point_count_in_y"]) self.collocation_points_lats = self._get_collocation_points_along_axis( self.lat_bounds[0], self.lat_bounds[1], self.setup["point_count_in_x"]) self.collocation_points_depth = \ self._get_collocation_points_along_axis( self.depth_bounds[1], self.depth_bounds[0], self.setup["point_count_in_z"])[::-1] def _read_single_box(self, component, file_number): """ This function reads Ses3ds raw binary files, e.g. 3d velocity field snapshots, as well as model parameter files or sensitivity kernels. It returns the field as an array of rank 3 with shape (nx*lpd+1, ny*lpd+1, nz*lpd+1), discarding the duplicates by default. """ # Get the file and the corresponding domain. filename = self.components[component]["filenames"][file_number] domain = self.setup["subdomains"][file_number] lpd = self.lagrange_polynomial_degree shape = (domain["index_x_count"], domain["index_y_count"], domain["index_z_count"], lpd + 1, lpd + 1, lpd + 1) # Take care: The first and last four bytes in the arrays are invalid # due to them being written by Fortran. with open(filename, "rb") as open_file: field = np.ndarray(shape, buffer=open_file.read()[4:-4], dtype="float32", order="F") # Calculate the new shape by multiplying every dimension with lpd + 1 # value for every dimension. new_shape = [_i * _j for _i, _j in zip(shape[:3], shape[3:])] # Reorder the axes: # v[x, y, z, lpd + 1 ,lpd + 1, lpd + 1] -> v[x, lpd + 1, y, z, # lpd + 1, lpd + 1] -> v[x, lpd + 1, y, lpd + 1, z, lpd + 1] field = np.rollaxis(np.rollaxis(field, 3, 1), 3, lpd + 1) # Reshape the data: # v[nx,lpd+1,ny,lpd+1,nz,lpd+1] to v[nx*(lpd+1),ny*(lpd+1),nz*(lpd+1)] field = field.reshape(new_shape, order="C") # XXX: Attempt to do this in one step. for axis, count in enumerate(new_shape): # Mask the duplicate values. mask = np.ones(count, dtype="bool") mask[::lpd + 1][1:] = False # Remove them by compressing the array. field = field.compress(mask, axis=axis) return field[:, :, ::-1] def parse_component(self, component): """ Helper function parsing a whole component. :param component: The component name. :type component: basestring """ # If a real component, if component in self.components.keys(): self._parse_component(component) return elif component not in self.available_derived_components: msg = "Component %s is unknown" % component raise ValueError(msg) if component == "vp": self._parse_component("lambda") self._parse_component("mu") self._parse_component("rhoinv") lambda_ = self.parsed_components["lambda"] mu = self.parsed_components["mu"] rhoinv = self.parsed_components["rhoinv"] self.parsed_components["vp"] = \ np.sqrt(((lambda_ + 2.0 * mu) * rhoinv)) / 1000.0 elif component == "vsh": self._parse_component("mu") self._parse_component("rhoinv") mu = self.parsed_components["mu"] rhoinv = self.parsed_components["rhoinv"] self.parsed_components["vsh"] = np.sqrt((mu * rhoinv)) / 1000.0 elif component == "vsv": self._parse_component("mu") self._parse_component("rhoinv") self._parse_component("B") mu = self.parsed_components["mu"] rhoinv = self.parsed_components["rhoinv"] b = self.parsed_components["B"] self.parsed_components["vsv"] = \ np.sqrt((mu + b) * rhoinv) / 1000.0 elif component == "rho": self._parse_component("rhoinv") rhoinv = self.parsed_components["rhoinv"] self.parsed_components["rho"] = 1.0 / rhoinv def _parse_component(self, component): """ Parses the specified component. """ if component in self.parsed_components: return # Allocate empty array with the necessary dimensions. data = np.empty(( self.setup["point_count_in_x"], self.setup["point_count_in_y"], self.setup["point_count_in_z"]), dtype="float32") for _i, domain in enumerate(self.setup["subdomains"]): x_min, x_max = domain["boundaries_x"] y_min, y_max = domain["boundaries_y"] z_min, z_max = domain["boundaries_z"] # Minimum indices x_min, y_min, z_min = [self.lagrange_polynomial_degree * _j for _j in (x_min, y_min, z_min)] # Maximum indices x_max, y_max, z_max = [self.lagrange_polynomial_degree * (_j + 1) for _j in (x_max, y_max, z_max)] # Merge into data. data[x_min: x_max + 1, y_min: y_max + 1, z_min: z_max + 1] = \ self._read_single_box(component, _i) self.parsed_components[component] = data def _calculate_final_dimensions(self): """ Calculates the total number of elements and points and also the size of the final array per component. """ # Elements in each directions. x, y, z = ( self.setup["boundaries_x"][1] - self.setup["boundaries_x"][0] + 1, self.setup["boundaries_y"][1] - self.setup["boundaries_y"][0] + 1, self.setup["boundaries_z"][1] - self.setup["boundaries_z"][0] + 1) self.setup["total_element_count"] = x * y * z self.setup["total_point_count"] = (x * y * z) * \ (self.lagrange_polynomial_degree + 1) ** 3 # This is the actual point count with removed duplicates. Much # smaller because each element shared a large number of elements # with its neighbours. self.setup["point_count_in_x"] = \ (x * self.lagrange_polynomial_degree + 1) self.setup["point_count_in_y"] = \ (y * self.lagrange_polynomial_degree + 1) self.setup["point_count_in_z"] = \ (z * self.lagrange_polynomial_degree + 1) self.setup["total_point_count_without_duplicates"] = ( self.setup["point_count_in_x"] * self.setup["point_count_in_y"] * self.setup["point_count_in_z"]) def _get_collocation_points_along_axis(self, min_value, max_value, count): """ Calculates count collocation points from min_value to max_value as they would be used by a simulation. The border elements are not repeated. :param min_value: The minimum value. :param max_value: The maximum value. :param count: The number of values. """ # Normalize from 0.0 to 1.0 coll_points = get_lpd_sampling_points(self.lagrange_polynomial_degree) coll_points += 1.0 coll_points /= 2.0 # Some view and reshaping tricks to ensure performance. Probably not # needed. points = np.empty(count) lpd = self.lagrange_polynomial_degree p_view = points[:count - 1].view().reshape(((count - 1) / lpd, lpd)) for _i in xrange(p_view.shape[0]): p_view[_i] = _i points[-1] = points[-2] + 1 for _i, value in enumerate(coll_points[1:-1]): points[_i + 1::lpd] += value # Normalize from min_value to max_value. points /= points.max() points *= abs(max_value - min_value) points += min_value return points def get_closest_gll_index(self, coord, value): if coord == "depth": return np.argmin(np.abs(self.collocation_points_depth - value)) elif coord == "longitude": return np.argmin(np.abs(self.collocation_points_lngs - value)) elif coord == "latitude": return np.argmin(np.abs(self.collocation_points_lats[::-1] - value)) def plot_depth_slice(self, component, depth_in_km, m, absolute_values=True): """ Plots a depth slice. :param component: The component to plot. :type component: basestring :param depth_in_km: The depth in km to plot. If the exact depth does not exists, the nearest neighbour will be plotted. :type depth_in_km: integer or float :param m: Basemap instance. """ depth_index = self.get_closest_gll_index("depth", depth_in_km) # No need to do anything if the currently plotted slice is already # plotted. This is useful for interactive use when the desired depth # is changed but the closest GLL collocation point is still the same. if hasattr(m, "_plotted_depth_slice"): # Use a tuple of relevant parameters. if m._plotted_depth_slice == (self.directory, depth_index, component, absolute_values): return None data = self.parsed_components[component] depth = self.collocation_points_depth[depth_index] lngs = self.collocation_points_lngs lats = self.collocation_points_lats # Rotate data if needed. lon, lat = np.meshgrid(lngs, lats) if hasattr(self.domain, "rotation_axis") and \ self.domain.rotation_axis and \ self.domain.rotation_angle_in_degree: lon_shape = lon.shape lat_shape = lat.shape lon.shape = lon.size lat.shape = lat.size lat, lon = rotations.rotate_lat_lon( lat, lon, self.domain.rotation_axis, self.domain.rotation_angle_in_degree) lon.shape = lon_shape lat.shape = lat_shape x, y = m(lon, lat) depth_data = data[::-1, :, depth_index] # Plot values relative to AK135. if not absolute_values: cmp_map = { "rho": "density", "vp": "vp", "vsh": "vs", "vsv": "vs" } factor = { "rho": 1000.0, "vp": 1.0, "vsh": 1.0, "vsv": 1.0, } if component not in cmp_map: vmin, vmax = depth_data.min(), depth_data.max() vmedian = np.median(depth_data) offset = max(abs(vmax - vmedian), abs(vmedian - vmin)) if vmax - vmin == 0: offset = 0.01 vmin = vmedian - offset vmax = vmedian + offset else: reference_value = self.one_d_model.get_value( cmp_map[component], depth) * factor[component] depth_data = (depth_data - reference_value) / reference_value depth_data *= 100.0 offset = np.abs(depth_data) try: offset = offset[offset < 50].max() except: offset = offset.max() vmin = -offset vmax = offset else: vmin, vmax = depth_data.min(), depth_data.max() vmedian = np.median(depth_data) offset = max(abs(vmax - vmedian), abs(vmedian - vmin)) min_delta = abs(vmax * 0.005) if (vmax - vmin) < min_delta: offset = min_delta vmin = vmedian - offset vmax = vmedian + offset # Remove an existing pcolormesh if it exists. This does not hurt in # any case but is useful for interactive use. if hasattr(m, "_depth_slice"): m._depth_slice.remove() del m._depth_slice im = m.pcolormesh(x, y, depth_data, cmap=tomo_colormap, vmin=vmin, vmax=vmax) m._depth_slice = im # Store what is currently plotted. m._plotted_depth_slice = (self.directory, depth_index, component, absolute_values) return { "depth": depth, "mesh": im, "data": depth_data } def get_depth_profile(self, component, latitude, longitude): """ Returns a depth profile of the model at the requested at the GLL points closest do latitude and longitude. :param component: The component of the model. :param latitude: The latitude. :param longitude: The longitude. """ # Need to rotate latitude and longitude. if hasattr(self.domain, "rotation_axis") and \ self.domain.rotation_axis and \ self.domain.rotation_angle_in_degree: latitude, longitude = rotations.rotate_lat_lon( latitude, longitude, self.domain.rotation_axis, -1.0 * self.domain.rotation_angle_in_degree) x_index = self.get_closest_gll_index("latitude", latitude) y_index = self.get_closest_gll_index("longitude", longitude) data = self.parsed_components[component] depths = self.collocation_points_depth values = data[x_index, y_index, :] lat = self.collocation_points_lats[::-1][x_index] lng = self.collocation_points_lngs[y_index] # Rotate back. if hasattr(self.domain, "rotation_axis") and \ self.domain.rotation_axis and \ self.domain.rotation_angle_in_degree: lat, lng = rotations.rotate_lat_lon( lat, lng, self.domain.rotation_axis, self.domain.rotation_angle_in_degree) return { "depths": depths, "values": values, "latitude": lat, "longitude": lng} def __str__(self): """ Eases interactive working. """ ret_str = "Raw SES3D Model (split in %i parts)\n" % \ (self.setup["total_cpu_count"]) ret_str += "\tSetup:\n" ret_str += "\t\tLatitude: {:.2f} - {:.2f}\n".format( *[rotations.colat2lat(_i) for _i in self.setup["physical_boundaries_x"][::-1]]) ret_str += "\t\tLongitude: %.2f - %.2f\n" % \ self.setup["physical_boundaries_y"] ret_str += "\t\tDepth in km: {:.2f} - {:.2f}\n".format( *[6371 - _i / 1000 for _i in self.setup["physical_boundaries_z"][::-1]]) ret_str += "\t\tTotal element count: %i\n" % \ self.setup["total_element_count"] ret_str += "\t\tTotal collocation point count: %i (without " \ "duplicates: %i)\n" % ( self.setup["total_point_count"], self.setup["total_point_count_without_duplicates"]) ret_str += "\tMemory requirement per component: %.1f MB\n" % \ ((self.setup["total_point_count_without_duplicates"] * 4) / (1024.0 ** 2)) ret_str += "\tAvailable components: %s\n" % (", ".join( sorted(self.components.keys()))) ret_str += "\tAvailable derived components: %s\n" % (", ".join( sorted(self.available_derived_components))) ret_str += "\tParsed components: %s" % (", ".join( sorted(self.parsed_components.keys()))) return ret_str def _read_boxfile(self): setup = {"subdomains": []} with open(self.boxfile, "rU") as fh: # The first 14 lines denote the header lines = fh.readlines()[14:] # Strip lines and remove empty lines. lines = [_i.strip() for _i in lines if _i.strip()] # The next 4 are the global CPU distribution. setup["total_cpu_count"] = int(lines.pop(0)) setup["cpu_count_in_x_direction"] = int(lines.pop(0)) setup["cpu_count_in_y_direction"] = int(lines.pop(0)) setup["cpu_count_in_z_direction"] = int(lines.pop(0)) if set(lines[0]) == set("-"): lines.pop(0) # Small sanity check. if setup["total_cpu_count"] != setup["cpu_count_in_x_direction"] *\ setup["cpu_count_in_y_direction"] * \ setup["cpu_count_in_z_direction"]: msg = ("Invalid boxfile. Total and individual processor " "counts do not match.") raise ValueError(msg) # Now parse the rest of file which contains the subdomains. def subdomain_generator(data): """ Simple generator looping over each defined box and yielding a dictionary for each. :param data: The text. """ while data: subdom = {} # Convert both indices to 0-based indices subdom["single_index"] = int(data.pop(0)) - 1 subdom["multi_index"] = map(lambda x: int(x) - 1, data.pop(0).split()) subdom["boundaries_x"] = map(int, data.pop(0).split()) subdom["boundaries_y"] = map(int, data.pop(0).split()) subdom["boundaries_z"] = map(int, data.pop(0).split()) # Convert radians to degree. subdom["physical_boundaries_x"] = map( lambda x: math.degrees(float(x)), data.pop(0).split()) subdom["physical_boundaries_y"] = map( lambda x: math.degrees(float(x)), data.pop(0).split()) # z is in meter. subdom["physical_boundaries_z"] = \ map(float, data.pop(0).split()) for component in ("x", "y", "z"): idx = "boundaries_%s" % component index_count = subdom[idx][1] - subdom[idx][0] + 1 subdom["index_%s_count" % component] = index_count # The boxfiles are slightly awkward in that the indices # are not really continuous. For example if one box # has 22 as the last index, the first index of the next # box will also be 22, even though it should be 23. The # next snippet attempts to fix this deficiency. offset = int(round(subdom[idx][0] / float(index_count - 1))) subdom[idx][0] += offset subdom[idx][1] += offset # Remove separator_line if existent. if set(lines[0]) == set("-"): lines.pop(0) yield subdom # Sort them after with the single index. setup["subdomains"] = sorted(list(subdomain_generator(lines)), key=lambda x: x["single_index"]) # Do some more sanity checks. if len(setup["subdomains"]) != setup["total_cpu_count"]: msg = ("Invalid boxfile. Number of processors and subdomains " "to not match.") raise ValueError(msg) for component in ("x", "y", "z"): idx = "index_%s_count" % component if len(set([_i[idx] for _i in setup["subdomains"]])) != 1: msg = ("Invalid boxfile. Unequal %s index count across " "subdomains.") % component raise ValueError(msg) # Now generate the absolute indices for the whole domains. for component in ("x", "y", "z"): setup["boundaries_%s" % component] = ( min([_i["boundaries_%s" % component][0] for _i in setup["subdomains"]]), max([_i["boundaries_%s" % component][1] for _i in setup["subdomains"]])) setup["physical_boundaries_%s" % component] = ( min([_i["physical_boundaries_%s" % component][0] for _i in setup["subdomains"]]), max([_i["physical_boundaries_%s" % component][1] for _i in setup["subdomains"]])) return setup