You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
72 lines
3.4 KiB
72 lines
3.4 KiB
# Copyright 2020 Huawei Technologies Co., Ltd
|
|
#
|
|
# 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.
|
|
# ============================================================================
|
|
"""train"""
|
|
|
|
import argparse
|
|
import numpy as np
|
|
import mindspore.context as context
|
|
from src.read_var import read_nc
|
|
from src.GOMO import GOMO_init, GOMO, read_init
|
|
|
|
parser = argparse.ArgumentParser(description='GOMO')
|
|
parser.add_argument('--file_path', type=str, default=None, help='file path')
|
|
parser.add_argument('--outputs_path', type=str, default=None, help='outputs path')
|
|
parser.add_argument('--im', type=int, default=65, help='im size')
|
|
parser.add_argument('--jm', type=int, default=49, help='jm size')
|
|
parser.add_argument('--kb', type=int, default=21, help='kb size')
|
|
parser.add_argument('--stencil_width', type=int, default=1, help='stencil width')
|
|
parser.add_argument('--step', type=int, default=10, help='time step')
|
|
args_gomo = parser.parse_args()
|
|
|
|
context.set_context(mode=context.GRAPH_MODE, device_target="GPU", save_graphs=False, enable_graph_kernel=True)
|
|
|
|
if __name__ == "__main__":
|
|
variable = read_nc(args_gomo.file_path)
|
|
im = args_gomo.im
|
|
jm = args_gomo.jm
|
|
kb = args_gomo.kb
|
|
stencil_width = args_gomo.stencil_width
|
|
|
|
# variable init
|
|
dx, dy, dz, uab, vab, elb, etb, sb, tb, ub, vb, dt, h, w, wubot, wvbot, vfluxb, utb, vtb, dhb, egb, vfluxf, z, zz, \
|
|
dzz, cor, fsm = read_init(
|
|
variable, im, jm, kb)
|
|
|
|
# define grid and init variable update
|
|
net_init = GOMO_init(im, jm, kb, stencil_width)
|
|
ua, va, el, et, etf, d, dt, l, q2b, q2lb, kh, km, kq, aam, w, q2, q2l, t, s, u, v, cbc, rmean, rho, x_d, y_d, z_d\
|
|
= net_init(dx, dy, dz, uab, vab, elb, etb, sb, tb, ub, vb, h, w, vfluxf, zz, fsm)
|
|
|
|
# define GOMO model
|
|
Model = GOMO(im=im, jm=jm, kb=kb, stencil_width=stencil_width, variable=variable, x_d=x_d, y_d=y_d, z_d=z_d,
|
|
q2b=q2b, q2lb=q2lb, aam=aam, cbc=cbc, rmean=rmean)
|
|
|
|
# time step of GOMO Model
|
|
for step in range(1, args_gomo.step+1):
|
|
elf, etf, ua, uab, va, vab, el, elb, d, u, v, w, kq, km, kh, q2, q2l, tb, t, sb, s, rho, wubot, wvbot, ub, vb, \
|
|
egb, etb, dt, dhb, utb, vtb, vfluxb, et, steps, vamax, q2b, q2lb = Model(
|
|
etf, ua, uab, va, vab, el, elb, d, u, v, w, kq, km, kh, q2, q2l, tb, t, sb, s, rho,
|
|
wubot, wvbot, ub, vb, egb, etb, dt, dhb, utb, vtb, vfluxb, et)
|
|
vars_list = etf, ua, uab, va, vab, el, elb, d, u, v, w, kq, km, kh, q2, q2l, tb, t, sb, s, rho, wubot, wvbot, \
|
|
ub, vb, egb, etb, dt, dhb, utb, vtb, vfluxb, et
|
|
for var in vars_list:
|
|
var.asnumpy()
|
|
# save output
|
|
if step % 5 == 0:
|
|
np.save(args_gomo.outputs_path + "u_"+str(step)+".npy", u.asnumpy())
|
|
np.save(args_gomo.outputs_path + "v_" + str(step) + ".npy", v.asnumpy())
|
|
np.save(args_gomo.outputs_path + "t_" + str(step) + ".npy", t.asnumpy())
|
|
np.save(args_gomo.outputs_path + "et_" + str(step) + ".npy", et.asnumpy())
|