Python utils.FileClass类代码示例

    def calculate_environment(self):
        if not self.modelrun:  # no evaluate done
            self.economy = self.worst_econ
            self.environment = self.worst_env
            return
        rfile = self.modelout_dir + os.path.sep + self.bmps_info['ENVEVAL']

        if not FileClass.is_file_exists(rfile):
            time.sleep(5)  # sleep 5 seconds wait for the ouput
        if not FileClass.is_file_exists(rfile):
            print('WARNING: Although SEIMS model runs successfully, the desired output: %s'
                  ' cannot be found!' % rfile)
            self.economy = self.worst_econ
            self.environment = self.worst_env
            return

        base_amount = self.bmps_info['BASE_ENV']
        if StringClass.string_match(rfile.split('.')[-1], 'tif'):  # Raster data
            rr = RasterUtilClass.read_raster(rfile)
            soil_erosion_amount = rr.get_sum() / self.timerange  # unit: year
            # reduction rate of soil erosion
            self.environment = (base_amount - soil_erosion_amount) / base_amount
        elif StringClass.string_match(rfile.split('.')[-1], 'txt'):  # Time series data
            sed_sum = read_simulation_from_txt(self.modelout_dir)  # TODO, fix it later, lj
            self.environment = (base_amount - sed_sum) / base_amount
        else:
            self.economy = self.worst_econ
            self.environment = self.worst_env
            return

 def convert_code(in_file, out_file, in_alg='taudem', out_alg='arcgis', datatype=None):
     """
     convert D8 flow direction code from one algorithm to another.
     Args:
         in_file: input raster file path
         out_file: output raster file path
         in_alg: available algorithms are in FlowModelConst.d8_dirs. "taudem" is the default
         out_alg: same as in_alg. "arcgis" is the default
         datatype: default is None and use the datatype of the in_file
     """
     FileClass.check_file_exists(in_file)
     in_alg = in_alg.lower()
     out_alg = out_alg.lower()
     if in_alg not in FlowModelConst.d8_dirs or out_alg not in FlowModelConst.d8_dirs:
         raise RuntimeError('The input algorithm name should one of %s' %
                            ', '.join(list(FlowModelConst.d8_dirs.keys())))
     convert_dict = dict()
     in_code = FlowModelConst.d8_dirs.get(in_alg)
     out_code = FlowModelConst.d8_dirs.get(out_alg)
     assert len(in_code) == len(out_code)
     for i, tmp_in_code in enumerate(in_code):
         convert_dict[tmp_in_code] = out_code[i]
     if datatype is not None and datatype in GDALDataType:
         RasterUtilClass.raster_reclassify(in_file, convert_dict, out_file, datatype)
     else:
         RasterUtilClass.raster_reclassify(in_file, convert_dict, out_file)

    def add_group_field(shp_file, subbasin_field_name, group_metis_dict):
        """add group information to subbasin ESRI shapefile

        Args:
            shp_file: Subbasin Shapefile
            subbasin_field_name: field name of subbasin
            group_metis_dict: returned by func`metis_partition`
        """
        if not group_metis_dict:
            return
        ds_reach = ogr_Open(shp_file, update=True)
        layer_reach = ds_reach.GetLayer(0)
        layer_def = layer_reach.GetLayerDefn()
        icode = layer_def.GetFieldIndex(subbasin_field_name)
        igrp = layer_def.GetFieldIndex(ImportReaches2Mongo._GROUP)
        ikgrp = layer_def.GetFieldIndex(ImportReaches2Mongo._KMETIS)
        ipgrp = layer_def.GetFieldIndex(ImportReaches2Mongo._PMETIS)

        if igrp < 0:
            new_field = ogr_FieldDefn(ImportReaches2Mongo._GROUP, OFTInteger)
            layer_reach.CreateField(new_field)
        if ikgrp < 0:
            new_field = ogr_FieldDefn(ImportReaches2Mongo._KMETIS, OFTInteger)
            layer_reach.CreateField(new_field)
        if ipgrp < 0:
            new_field = ogr_FieldDefn(ImportReaches2Mongo._PMETIS, OFTInteger)
            layer_reach.CreateField(new_field)

        ftmap = dict()
        layer_reach.ResetReading()
        ft = layer_reach.GetNextFeature()
        while ft is not None:
            tmpid = ft.GetFieldAsInteger(icode)
            ftmap[tmpid] = ft
            ft = layer_reach.GetNextFeature()

        groups = group_metis_dict[1]['group']
        for i, n in enumerate(groups):
            for node, d in group_metis_dict.items():
                ftmap[node].SetField(ImportReaches2Mongo._GROUP, n)
                ftmap[node].SetField(ImportReaches2Mongo._KMETIS, d['kmetis'][i])
                ftmap[node].SetField(ImportReaches2Mongo._PMETIS, d['pmetis'][i])
                layer_reach.SetFeature(ftmap[node])
            # copy the reach file to new file
            prefix = os.path.splitext(shp_file)[0]
            dstfile = prefix + "_" + str(n) + ".shp"
            FileClass.copy_files(shp_file, dstfile)

        layer_reach.SyncToDisk()
        ds_reach.Destroy()
        del ds_reach

    def workflow(cfg, maindb, climdb):
        """
        This function mainly to import measurement data to MongoDB
        data type may include Q (discharge, m3/s), SED (mg/L), tn (mg/L), tp (mg/L), etc.
        the required parameters that defined in configuration file (*.ini)
        """
        if not cfg.use_observed:
            return False
        c_list = climdb.collection_names()
        if not StringClass.string_in_list(DBTableNames.observes, c_list):
            climdb.create_collection(DBTableNames.observes)
        else:
            climdb.drop_collection(DBTableNames.observes)
        if not StringClass.string_in_list(DBTableNames.sites, c_list):
            climdb.create_collection(DBTableNames.sites)
        if not StringClass.string_in_list(DBTableNames.var_desc, c_list):
            climdb.create_collection(DBTableNames.var_desc)

        file_list = FileClass.get_full_filename_by_suffixes(cfg.observe_dir, ['.txt'])
        meas_file_list = []
        site_loc = []
        for fl in file_list:
            if StringClass.is_substring('observed_', fl):
                meas_file_list.append(fl)
            else:
                site_loc.append(fl)
        ImportObservedData.data_from_txt(maindb, climdb, meas_file_list, site_loc,
                                         cfg.spatials.subbsn)
        return True

    def write_asc_file(filename, data, xsize, ysize, geotransform, nodata_value):
        """Output Raster to ASCII file.

        Args:
            filename: output ASCII filename.
            data: 2D array data.
            xsize: Col count.
            ysize: Row count.
            geotransform: geographic transformation.
            nodata_value: nodata_flow value.
        """
        UtilClass.mkdir(os.path.dirname(FileClass.get_file_fullpath(filename)))
        header = 'NCOLS %d\n' \
                 'NROWS %d\n' \
                 'XLLCENTER %f\n' \
                 'YLLCENTER %f\n' \
                 'CELLSIZE %f\n' \
                 'NODATA_VALUE %f' % (xsize, ysize, geotransform[0] + 0.5 * geotransform[1],
                                      geotransform[3] - (ysize - 0.5) * geotransform[1],
                                      geotransform[1], nodata_value)

        with open(filename, 'w', encoding='utf-8') as f:
            f.write(header)
            for i in range(0, ysize):
                for j in range(0, xsize):
                    f.write('%s\t' % repr(data[i][j]))
                f.write('\n')
        f.close()

    def write_gtiff_file(f_name, n_rows, n_cols, data, geotransform, srs, nodata_value,
                         gdal_type=GDT_Float32):
        """Output Raster to GeoTiff format file.

        Args:
            f_name: output gtiff file name.
            n_rows: Row count.
            n_cols: Col count.
            data: 2D array data.
            geotransform: geographic transformation.
            srs: coordinate system.
            nodata_value: nodata value.
            gdal_type (:obj:`pygeoc.raster.GDALDataType`): output raster data type,
                                                                  GDT_Float32 as default.
        """
        UtilClass.mkdir(os.path.dirname(FileClass.get_file_fullpath(f_name)))
        driver = gdal_GetDriverByName(str('GTiff'))
        try:
            ds = driver.Create(f_name, n_cols, n_rows, 1, gdal_type)
        except Exception:
            print('Cannot create output file %s' % f_name)
            return
        ds.SetGeoTransform(geotransform)
        try:
            ds.SetProjection(srs.ExportToWkt())
        except AttributeError or Exception:
            ds.SetProjection(srs)
        ds.GetRasterBand(1).SetNoDataValue(nodata_value)
        # if data contains numpy.nan, then replaced by nodata_value
        if isinstance(data, numpy.ndarray) and data.dtype in [numpy.dtype('int'),
                                                              numpy.dtype('float')]:
            data = numpy.where(numpy.isnan(data), nodata_value, data)
        ds.GetRasterBand(1).WriteArray(data)
        ds = None

def main():
    from preprocess.config import parse_ini_configuration

    seims_cfg = parse_ini_configuration()
    client = ConnectMongoDB(seims_cfg.hostname, seims_cfg.port)
    conn = client.get_conn()
    db_model = conn[seims_cfg.spatial_db]

    spatial_gfs = GridFS(db_model, DBTableNames.gridfs_spatial)

    csv_path = r'C:\z_data\zhongTianShe\model_data_seims\field_scale_params'
    csv_files = FileClass.get_full_filename_by_suffixes(csv_path, ['.csv'])
    field_count = 7419
    prefix = 9999
    # Create mask file
    mask_name = '%d_MASK' % prefix
    mask_array = [[1] * field_count]
    import_array_to_mongodb(spatial_gfs, mask_array, mask_name)

    # Create spatial parameters
    for csv_file in csv_files:
        print('Import %s...' % csv_file)
        param_arrays = read_field_arrays_from_csv(csv_file)
        for key, value in list(param_arrays.items()):
            import_array_to_mongodb(spatial_gfs, value, '%d_%s' % (prefix, key))

def get_psa_config():
    """Parse arguments.
    Returns:
        cf: ConfigParse object of *.ini file
        mtd: Parameters sensitivity method name, currently, 'morris' and 'fast' are supported.
    """
    # define input arguments
    parser = argparse.ArgumentParser(description="Execute parameters sensitivity analysis.")
    parser.add_argument('-ini', type=str, help="Full path of configuration file")
    # add mutually group
    psa_group = parser.add_mutually_exclusive_group()
    psa_group.add_argument('-morris', action='store_true', help='Run Morris Screening method')
    psa_group.add_argument('-fast', action='store_true', help='Run FAST variant-based method')
    # parse arguments
    args = parser.parse_args()
    ini_file = args.ini
    psa_mtd = 'morris'  # Default
    if args.fast:
        psa_mtd = 'fast'
    elif args.morris:
        psa_mtd = 'morris'
    if not FileClass.is_file_exists(ini_file):
        raise ImportError('Configuration file is not existed: %s' % ini_file)
    cf = ConfigParser()
    cf.read(ini_file)
    return cf, psa_mtd

 def output_hillslope(method_id):
     """Output hillslope according different stream cell value method."""
     for (tmp_row, tmp_col) in stream_coors:
         tmp_hillslp_ids = DelineateHillslope.cal_hs_codes(max_id,
                                                           stream_data[tmp_row][tmp_col])
         if 0 < method_id < 3:
             hillslope_mtx[tmp_row][tmp_col] = tmp_hillslp_ids[method_id]
             # is head stream cell?
             if (tmp_row, tmp_col) in headstream_coors:
                 hillslope_mtx[tmp_row][tmp_col] = tmp_hillslp_ids[0]
         elif method_id == 3:
             hillslope_mtx[tmp_row][tmp_col] = DEFAULT_NODATA
     # Output to raster file
     hillslope_out_new = hillslope_out
     dirpath = os.path.dirname(hillslope_out_new) + os.path.sep
     corename = FileClass.get_core_name_without_suffix(hillslope_out_new)
     if method_id == 1:
         hillslope_out_new = dirpath + corename + '_right.tif'
     elif method_id == 2:
         hillslope_out_new = dirpath + corename + '_left.tif'
     elif method_id == 3:
         hillslope_out_new = dirpath + corename + '_nodata.tif'
     RasterUtilClass.write_gtiff_file(hillslope_out_new, nrows, ncols,
                                      hillslope_mtx,
                                      geotrans, srs, DEFAULT_NODATA, datatype)

 def gridnet(np, pfile, plenfile, tlenfile, gordfile, outlet=None, workingdir=None,
             mpiexedir=None, exedir=None, log_file=None, runtime_file=None, hostfile=None):
     """Run gridnet"""
     fname = TauDEM.func_name('gridnet')
     return TauDEM.run(FileClass.get_executable_fullpath(fname, exedir),
                       {'-p': pfile, '-o': outlet}, workingdir,
                       None,
                       {'-plen': plenfile, '-tlen': tlenfile, '-gord': gordfile},
                       {'mpipath': mpiexedir, 'hostfile': hostfile, 'n': np},
                       {'logfile': log_file, 'runtimefile': runtime_file})

 def pitremove(np, dem, filleddem, workingdir=None, mpiexedir=None, exedir=None, log_file=None,
               runtime_file=None, hostfile=None):
     """Run pit remove using the flooding approach """
     fname = TauDEM.func_name('pitremove')
     return TauDEM.run(FileClass.get_executable_fullpath(fname, exedir),
                       {'-z': dem}, workingdir,
                       None,
                       {'-fel': filleddem},
                       {'mpipath': mpiexedir, 'hostfile': hostfile, 'n': np},
                       {'logfile': log_file, 'runtimefile': runtime_file})

 def threshold(np, acc, stream_raster, threshold=100., workingdir=None,
               mpiexedir=None, exedir=None, log_file=None, runtime_file=None, hostfile=None):
     """Run threshold for stream raster"""
     fname = TauDEM.func_name('threshold')
     return TauDEM.run(FileClass.get_executable_fullpath(fname, exedir),
                       {'-ssa': acc}, workingdir,
                       {'-thresh': threshold},
                       {'-src': stream_raster},
                       {'mpipath': mpiexedir, 'hostfile': hostfile, 'n': np},
                       {'logfile': log_file, 'runtimefile': runtime_file})

 def peukerdouglas(np, fel, streamSkeleton, workingdir=None, mpiexedir=None, exedir=None,
                   log_file=None, runtime_file=None, hostfile=None):
     """Run peuker-douglas function"""
     fname = TauDEM.func_name('peukerdouglas')
     return TauDEM.run(FileClass.get_executable_fullpath(fname, exedir),
                       {'-fel': fel}, workingdir,
                       None,
                       {'-ss': streamSkeleton},
                       {'mpipath': mpiexedir, 'hostfile': hostfile, 'n': np},
                       {'logfile': log_file, 'runtimefile': runtime_file})

 def dinfflowdir(np, filleddem, flowangle, slope, workingdir=None, mpiexedir=None, exedir=None,
                 log_file=None, runtime_file=None, hostfile=None):
     """Run Dinf flow direction"""
     fname = TauDEM.func_name('dinfflowdir')
     return TauDEM.run(FileClass.get_executable_fullpath(fname, exedir),
                       {'-fel': filleddem}, workingdir,
                       None,
                       {'-ang': flowangle, '-slp': slope},
                       {'mpipath': mpiexedir, 'hostfile': hostfile, 'n': np},
                       {'logfile': log_file, 'runtimefile': runtime_file})

 def calculate_sensitivity(self):
     """Calculate Morris elementary effects.
        It is worth to be noticed that evaluate_models() allows to return
        several output variables, hence we should calculate each of them separately.
     """
     if not self.psa_si:
         if FileClass.is_file_exists(self.cfg.outfiles.psa_si_json):
             with open(self.cfg.outfiles.psa_si_json, 'r') as f:
                 self.psa_si = UtilClass.decode_strs_in_dict(json.load(f))
                 return
     if not self.objnames:
         if FileClass.is_file_exists('%s/objnames.pickle' % self.cfg.psa_outpath):
             with open('%s/objnames.pickle' % self.cfg.psa_outpath, 'r') as f:
                 self.objnames = pickle.load(f)
     if self.output_values is None or len(self.output_values) == 0:
         self.evaluate_models()
     if self.param_values is None or len(self.param_values) == 0:
         self.generate_samples()
     if not self.param_defs:
         self.read_param_ranges()
     row, col = self.output_values.shape
     assert (row == self.run_count)
     for i in range(col):
         print(self.objnames[i])
         if self.cfg.method == 'morris':
             tmp_Si = morris_alz(self.param_defs,
                                 self.param_values,
                                 self.output_values[:, i],
                                 conf_level=0.95, print_to_console=True,
                                 num_levels=self.cfg.morris.num_levels,
                                 grid_jump=self.cfg.morris.grid_jump)
         elif self.cfg.method == 'fast':
             tmp_Si = fast_alz(self.param_defs, self.output_values[:, i],
                               print_to_console=True)
         else:
             raise ValueError('%s method is not supported now!' % self.cfg.method)
         self.psa_si[i] = tmp_Si
     # print(self.psa_si)
     # Save as json, which can be loaded by json.load()
     json_data = json.dumps(self.psa_si, indent=4, cls=SpecialJsonEncoder)
     with open(self.cfg.outfiles.psa_si_json, 'w') as f:
         f.write(json_data)
     self.output_psa_si()