# set dataloader configtrain_dataloader_cfg={"dataset":"IterableNamedArrayDataset","iters_per_epoch":cfg.TRAIN.iters_per_epoch,}NPOINT_PDE=99**2NPOINT_TOP=101NPOINT_BOTTOM=101NPOINT_LEFT=99NPOINT_RIGHT=99
# set constraintpde=ppsci.constraint.InteriorConstraint(equation["NavierStokes"].equations,{"continuity":0,"momentum_x":0,"momentum_y":0},geom["rect"],{**train_dataloader_cfg,"batch_size":NPOINT_PDE},ppsci.loss.MSELoss("sum"),evenly=True,weight_dict=cfg.TRAIN.weight.pde,# (1)name="EQ",)
# set optimizerlr_scheduler=ppsci.optimizer.lr_scheduler.Cosine(**cfg.TRAIN.lr_scheduler,warmup_epoch=int(0.05*cfg.TRAIN.epochs),)()optimizer=ppsci.optimizer.Adam(lr_scheduler)(model)
# set validatorNPOINT_EVAL=NPOINT_PDEresidual_validator=ppsci.validate.GeometryValidator(equation["NavierStokes"].equations,{"momentum_x":0,"continuity":0,"momentum_y":0},geom["rect"],{"dataset":"NamedArrayDataset","total_size":NPOINT_EVAL,"batch_size":cfg.EVAL.batch_size.residual_validator,"sampler":{"name":"BatchSampler"},},ppsci.loss.MSELoss("sum"),evenly=True,metric={"MSE":ppsci.metric.MSE()},name="Residual",)validator={residual_validator.name:residual_validator}
# set visualizer(optional)# manually collate input data for visualization,NPOINT_BC=NPOINT_TOP+NPOINT_BOTTOM+NPOINT_LEFT+NPOINT_RIGHTvis_points=geom["rect"].sample_interior(NPOINT_PDE+NPOINT_BC,evenly=True)visualizer={"visualize_u_v":ppsci.visualize.VisualizerVtu(vis_points,{"u":lambdad:d["u"],"v":lambdad:d["v"],"p":lambdad:d["p"]},prefix="result_u_v",)}
# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.# Licensed under the Apache License, Version 2.0 (the "License");# you may not use this file except in compliance with the License.# You may obtain a copy of the License at# http://www.apache.org/licenses/LICENSE-2.0# Unless required by applicable law or agreed to in writing, software# distributed under the License is distributed on an "AS IS" BASIS,# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.# See the License for the specific language governing permissions and# limitations under the License.fromosimportpathasospimporthydraimportnumpyasnpfromomegaconfimportDictConfigimportppscifromppsci.utilsimportloggerdeftrain(cfg:DictConfig):# set random seed for reproducibilityppsci.utils.misc.set_random_seed(cfg.seed)# initialize loggerlogger.init_logger("ppsci",osp.join(cfg.output_dir,"train.log"),"info")# set modelmodel=ppsci.arch.MLP(**cfg.MODEL.model)# set equationequation={"NavierStokes":ppsci.equation.NavierStokes(cfg.NU,cfg.RHO,2,False)}# set geometrygeom={"rect":ppsci.geometry.Rectangle((-0.05,-0.05),(0.05,0.05))}# set dataloader configtrain_dataloader_cfg={"dataset":"IterableNamedArrayDataset","iters_per_epoch":cfg.TRAIN.iters_per_epoch,}NPOINT_PDE=99**2NPOINT_TOP=101NPOINT_BOTTOM=101NPOINT_LEFT=99NPOINT_RIGHT=99# set constraintpde=ppsci.constraint.InteriorConstraint(equation["NavierStokes"].equations,{"continuity":0,"momentum_x":0,"momentum_y":0},geom["rect"],{**train_dataloader_cfg,"batch_size":NPOINT_PDE},ppsci.loss.MSELoss("sum"),evenly=True,weight_dict=cfg.TRAIN.weight.pde,name="EQ",)bc_top=ppsci.constraint.BoundaryConstraint({"u":lambdaout:out["u"],"v":lambdaout:out["v"]},{"u":1,"v":0},geom["rect"],{**train_dataloader_cfg,"batch_size":NPOINT_TOP},ppsci.loss.MSELoss("sum"),criteria=lambdax,y:np.isclose(y,0.05),name="BC_top",)bc_bottom=ppsci.constraint.BoundaryConstraint({"u":lambdaout:out["u"],"v":lambdaout:out["v"]},{"u":0,"v":0},geom["rect"],{**train_dataloader_cfg,"batch_size":NPOINT_BOTTOM},ppsci.loss.MSELoss("sum"),criteria=lambdax,y:np.isclose(y,-0.05),name="BC_bottom",)bc_left=ppsci.constraint.BoundaryConstraint({"u":lambdaout:out["u"],"v":lambdaout:out["v"]},{"u":0,"v":0},geom["rect"],{**train_dataloader_cfg,"batch_size":NPOINT_LEFT},ppsci.loss.MSELoss("sum"),criteria=lambdax,y:np.isclose(x,-0.05),name="BC_left",)bc_right=ppsci.constraint.BoundaryConstraint({"u":lambdaout:out["u"],"v":lambdaout:out["v"]},{"u":0,"v":0},geom["rect"],{**train_dataloader_cfg,"batch_size":NPOINT_RIGHT},ppsci.loss.MSELoss("sum"),criteria=lambdax,y:np.isclose(x,0.05),name="BC_right",)# wrap constraints togetherconstraint={pde.name:pde,bc_top.name:bc_top,bc_bottom.name:bc_bottom,bc_left.name:bc_left,bc_right.name:bc_right,}# set optimizerlr_scheduler=ppsci.optimizer.lr_scheduler.Cosine(**cfg.TRAIN.lr_scheduler,warmup_epoch=int(0.05*cfg.TRAIN.epochs),)()optimizer=ppsci.optimizer.Adam(lr_scheduler)(model)# set validatorNPOINT_EVAL=NPOINT_PDEresidual_validator=ppsci.validate.GeometryValidator(equation["NavierStokes"].equations,{"momentum_x":0,"continuity":0,"momentum_y":0},geom["rect"],{"dataset":"NamedArrayDataset","total_size":NPOINT_EVAL,"batch_size":cfg.EVAL.batch_size.residual_validator,"sampler":{"name":"BatchSampler"},},ppsci.loss.MSELoss("sum"),evenly=True,metric={"MSE":ppsci.metric.MSE()},name="Residual",)validator={residual_validator.name:residual_validator}# set visualizer(optional)# manually collate input data for visualization,NPOINT_BC=NPOINT_TOP+NPOINT_BOTTOM+NPOINT_LEFT+NPOINT_RIGHTvis_points=geom["rect"].sample_interior(NPOINT_PDE+NPOINT_BC,evenly=True)visualizer={"visualize_u_v":ppsci.visualize.VisualizerVtu(vis_points,{"u":lambdad:d["u"],"v":lambdad:d["v"],"p":lambdad:d["p"]},prefix="result_u_v",)}# initialize solversolver=ppsci.solver.Solver(model,constraint,cfg.output_dir,optimizer,lr_scheduler,cfg.TRAIN.epochs,cfg.TRAIN.iters_per_epoch,eval_during_train=cfg.EVAL.pretrained_model_path,eval_freq=cfg.TRAIN.eval_freq,equation=equation,geom=geom,validator=validator,visualizer=visualizer,checkpoint_path=cfg.TRAIN.checkpoint_path,)# train modelsolver.train()# evaluate after finished trainingsolver.eval()# visualize prediction after finished trainingsolver.visualize()defevaluate(cfg:DictConfig):# set random seed for reproducibilityppsci.utils.misc.set_random_seed(cfg.seed)# initialize loggerlogger.init_logger("ppsci",osp.join(cfg.output_dir,"eval.log"),"info")# set modelmodel=ppsci.arch.MLP(**cfg.MODEL.model)# set equationequation={"NavierStokes":ppsci.equation.NavierStokes(cfg.NU,cfg.RHO,2,False)}# set geometrygeom={"rect":ppsci.geometry.Rectangle((-0.05,-0.05),(0.05,0.05))}NPOINT_PDE=99**2NPOINT_TOP=101NPOINT_BOTTOM=101NPOINT_LEFT=99NPOINT_RIGHT=99# set validatorNPOINT_EVAL=NPOINT_PDEresidual_validator=ppsci.validate.GeometryValidator(equation["NavierStokes"].equations,{"momentum_x":0,"continuity":0,"momentum_y":0},geom["rect"],{"dataset":"NamedArrayDataset","total_size":NPOINT_EVAL,"batch_size":cfg.EVAL.batch_size.residual_validator,"sampler":{"name":"BatchSampler"},},ppsci.loss.MSELoss("sum"),evenly=True,metric={"MSE":ppsci.metric.MSE()},name="Residual",)validator={residual_validator.name:residual_validator}# set visualizer(optional)# manually collate input data for visualization,NPOINT_BC=NPOINT_TOP+NPOINT_BOTTOM+NPOINT_LEFT+NPOINT_RIGHTvis_points=geom["rect"].sample_interior(NPOINT_PDE+NPOINT_BC,evenly=True)visualizer={"visualize_u_v":ppsci.visualize.VisualizerVtu(vis_points,{"u":lambdad:d["u"],"v":lambdad:d["v"],"p":lambdad:d["p"]},prefix="result_u_v",)}# initialize solversolver=ppsci.solver.Solver(model,output_dir=cfg.output_dir,equation=equation,geom=geom,validator=validator,visualizer=visualizer,pretrained_model_path=cfg.EVAL.pretrained_model_path,)solver.eval()# visualize prediction for pretrained model(optional)solver.visualize()@hydra.main(version_base=None,config_path="./conf",config_name="ldc2d_steady_Re10.yaml")defmain(cfg:DictConfig):ifcfg.mode=="train":train(cfg)elifcfg.mode=="eval":evaluate(cfg)else:raiseValueError(f"cfg.mode should in ['train', 'eval'], but got '{cfg.mode}'")if__name__=="__main__":main()