From edc3512c3402b5669d34adcc05fbe523f5413871 Mon Sep 17 00:00:00 2001 From: Dubhe-Chang Date: Wed, 25 Jun 2025 13:43:19 +0800 Subject: [PATCH 01/32] establish example chem --- docs/zh/examples/chem.md | 142 +++++++++++++++++++++++++++++++ docs/zh/install_setup.md | 13 +-- examples/chem/Chem.py | 151 +++++++++++++++++++++++++++++++++ examples/chem/chem.png | Bin 0 -> 31491 bytes examples/chem/config/chem.yaml | 61 +++++++++++++ examples/chem/data_set.xlsx | Bin 0 -> 219784 bytes examples/chem/requirements.txt | 3 + mkdocs.yml | 1 + ppsci/arch/__init__.py | 2 + ppsci/arch/chem.py | 99 +++++++++++++++++++++ ppsci/externals/__init__.py | 1 - ppsci/externals/paddle_sparse | 1 - 12 files changed, 461 insertions(+), 13 deletions(-) create mode 100644 docs/zh/examples/chem.md create mode 100644 examples/chem/Chem.py create mode 100644 examples/chem/chem.png create mode 100644 examples/chem/config/chem.yaml create mode 100644 examples/chem/data_set.xlsx create mode 100644 examples/chem/requirements.txt create mode 100644 ppsci/arch/chem.py delete mode 160000 ppsci/externals/paddle_sparse diff --git a/docs/zh/examples/chem.md b/docs/zh/examples/chem.md new file mode 100644 index 0000000000..7200871784 --- /dev/null +++ b/docs/zh/examples/chem.md @@ -0,0 +1,142 @@ +# Suzuki-Miyaura 交叉偶联反应产率预测 + +!!! note + + 1. 开始训练、评估前,数据文件data_set.xlsx的存在,并对应修改 yaml 配置文件中的 `data_dir` 为数据文件路径。 + 2. 如果需要使用预训练模型进行评估,请先下载预训练模型[chem_model.pdparams](https://paddle-org.bj.bcebos.com/paddlescience/models/TADF/Est/Est_pretrained.pdparams), 并对应修改 yaml 配置文件中的 `load_model_path` 为模型参数路径。 + 3. 开始训练、评估前,请安装 `rdkit` 等,相关依赖请执行`pip install -r requirements.txt`安装。 + +=== "模型训练命令" + + ``` sh + # 训练: + python Chem.py mode=train + ``` + +=== "模型评估命令" + + ``` sh + # 评估: + python Chem.py mode=eval + ``` + +## 1. 背景简介 + +Suzuki-Miyaura 交叉偶联反应表达式如下所示。 + +$$ +\mathrm{Ar{-}X} + \mathrm{Ar'{-}B(OH)_2} \xrightarrow[\text{Base}]{\mathrm{Pd}^0} \mathrm{Ar{-}Ar'} + \mathrm{HX} +$$ + +在零价钯配合物催化下,芳基或烯基硼酸或硼酸酯与氯、溴、碘代芳烃或烯烃发生交叉偶联。该反应具有反应条件温和、转化率高的优点,在材料合成、药物研发等领域具有重要作用,但存在开发周期长,试错成本高的问题。本研究通过使用高通量实验数据分析反应底物(包括亲电试剂和亲核试剂),催化配体,碱基,溶剂对偶联反应产率的影响,从而建立预测模型。 + + +## 2. Suzuki-Miyaura 交叉偶联反应产率预测模型的实现 + +本节将讲解如何基于PaddleScience代码,实现对于 Suzuki-Miyaura 交叉偶联反应产率预测模型的构建、训练、测试和评估。案例的目录结构如下。 +``` log +chem/ +├──config/ +│ └── chem.yaml +├── Chem.py +├── data_set.xlsx +└── requirements.txt +``` + +### 2.1 数据集构建和载入 + +本样例使用的数据来自参考文献[1]提供的开源数据,仅考虑试剂本身对于实验结果的影响,从中筛选了各分量均有试剂参与的部分反应数据,保存在文件 `./data_set.xlsx` 中。该工作开发了一套基于流动化学(flow chemistry)的自动化平台,该平台在氩气保护的手套箱中组装,使用改良的高效液相色谱(HPLC)系统,结合自动化取样装置,从192个储液瓶中按设定程序吸取反应组分(亲电试剂、亲核试剂、催化剂、配体和碱),并注入流动载液中。每个反应段在温控反应盘管中以设定的流速、压力、时间进行反应,反应液通过UPLC-MS进行实时检测。通过调控亲电试剂、亲核试剂、11种配体、7种碱和4种溶剂的组合,最终实现了5760个反应条件的系统性筛选。接下来以其中一条数据为例结合代码说明数据集的构建与载入流程。 + +``` +ClC=1C=C2C=CC=NC2=CC1 | CC=1C(=C2C=NN(C2=CC1)C1OCCCC1)B(O)O | C(C)(C)(C)P(C(C)(C)C)C(C)(C)C | [OH-].[Na+] | C(C)#N | 4.76 +``` +其中用SMILES依次表示亲电试剂、亲核试剂、催化配体、碱、溶剂和实验产率。 + +首先从表格文件中将实验材料信息和反应产率进行导入,并划分训练集和测试集, + +``` py linenums="24" title="examples/chem/Chem.py" +--8<-- +examples/chem/Chem.py:24:30 +--8<-- +``` + +应用 `rdkit.Chem.rdFingerprintGenerator` 将亲电试剂、亲核试剂、催化配体、碱和溶剂的SMILES描述转换为 Morgan 指纹。Morgan指纹是一种分子结构的向量化描述,通过局部拓扑被编码为 hash 值,映射到2048位指纹位上。用 PaddleScience 代码表示如下 + +``` py linenums="32" title="examples/chem/Chem.py" +--8<-- +examples/chem/Chem.py:32:54 +--8<-- +``` + +### 2.2 约束构建 + +本案例采用监督学习,按照 PaddleScience 的API结构说明,采用内置的 `SupervisedConstraint` 构建监督约束。用 PaddleScience 代码表示如下 + +``` py linenums="60" title="examples/chem/Chem.py" +--8<-- +examples/chem/Chem.py:60:76 +--8<-- +``` +`SupervisedConstraint` 的第二个参数表示采用均方误差 `MSELoss` 作为损失函数,第三个参数表示约束条件的名字,方便后续对其索引。 + +### 2.3 模型构建 + +本案例设计了五条独立的子网络(全连接层+ReLU激活),每条子网络分别提取对应化学物质的特征。随后,这五个特征向量通过可训练的权重参数进行加权平均,实现不同化学成分对反应产率预测影响的自适应学习。最后,将融合后的特征输入到一个全连接层进行进一步映射,输出反应产率预测值。整个网络结构体现了对反应中各组成成分信息的独立提取与有权重的融合,符合反应机理特性。用 PaddleScience 代码表示如下 + +``` py linenums="5" title="ppsci/arch/chem.py" +--8<-- +ppsci/arch/chem.py:5:99 +--8<-- +``` + +模型依据配置文件信息进行实例化 + +``` py linenums="78" title="examples/chem/Chem.py" +--8<-- +examples/chem/Chem.py:78:80 +--8<-- +``` + +参数通过配置文件进行设置如下 + +``` py linenums="31" title="examples/chem/config/chem.yaml" +--8<-- +examples/chem/config/chem.yaml:31:38 +--8<-- +``` + +### 2.4 优化器构建 + +训练器采用Adam优化器,学习率设置由配置文件给出。用 PaddleScience 代码表示如下 + +``` py linenums="82" title="examples/chem/Chem.py" +--8<-- +examples/chem/Chem.py:82:83 +--8<-- +``` + +### 2.5 模型训练 + +完成上述设置之后,只需要将上述实例化的对象按顺序传递给`ppsci.solver.Solver`,然后启动训练即可。用PaddleScience 代码表示如下 + +``` py linenums="85" title="examples/chem/Chem.py" +--8<-- +examples/chem/Chem.py:85:98 +--8<-- +``` + +## 3. 完整代码 + +``` py linenums="1" title="examples/chem/Chem.py" +--8<-- +examples/chem/Chem.py +--8<-- +``` + +## 4. 结果展示 + +下图展示对 Suzuki-Miyaura 交叉偶联反应产率的模型预测结果。 + +## 5. 参考文献 + +[1] Perera D, Tucker J W, Brahmbhatt S, et al. A platform for automated nanomole-scale reaction screening and micromole-scale synthesis in flow[J]. Science, 2018, 359(6374): 429-434. \ No newline at end of file diff --git a/docs/zh/install_setup.md b/docs/zh/install_setup.md index dac9112aba..293c84329a 100644 --- a/docs/zh/install_setup.md +++ b/docs/zh/install_setup.md @@ -265,16 +265,7 @@ PaddleScience 提供了多种第三方库供用户在开发时使用,这些库 cd PaddleScience git submodule update --init ppsci/externals/paddle_scatter # install from source(recommended) - python -m pip install ppsci/externals/paddle_scatter - ``` - - === "paddle_sparse" - - ``` sh - cd PaddleScience - git submodule update --init ppsci/externals/paddle_sparse - # install from source(recommended) - python -m pip install ppsci/externals/paddle_sparse + python -m pip install -e ppsci/externals/paddle_scatter ``` === "tensorly" @@ -308,7 +299,7 @@ PaddleScience 提供了多种第三方库供用户在开发时使用,这些库 ``` python >>> from ppsci import externals >>> print(externals.__all__) - ['deepali', 'open3d', 'paddle_harmonics', 'paddle_scatter', 'paddle_sparse', 'tensorly', 'warp'] + ['deepali', 'open3d', 'paddle_harmonics', 'paddle_scatter', 'tensorly', 'warp'] >>> tl = externals.tensorly >>> tl.set_backend("paddle") diff --git a/examples/chem/Chem.py b/examples/chem/Chem.py new file mode 100644 index 0000000000..34d4d126bd --- /dev/null +++ b/examples/chem/Chem.py @@ -0,0 +1,151 @@ +import matplotlib.pyplot as plt +import numpy as np +import os +import paddle +import ppsci +import rdkit.Chem as Chem +from rdkit.Chem import rdFingerprintGenerator +from sklearn.metrics import r2_score +from sklearn.model_selection import train_test_split +import hydra +from omegaconf import DictConfig +import pandas as pd + +os.environ['HYDRA_FULL_ERROR'] = '1' +os.environ["KMP_DUPLICATE_LIB_OK"] = "True" +plt.rcParams["axes.unicode_minus"] = False +plt.rcParams['font.sans-serif'] = ['DejaVu Sans'] + +x_train = None +x_test = None +y_train = None +y_test = None + +def load_data(cfg: DictConfig): + data_dir = cfg.data_dir + dataset = pd.read_excel(data_dir,skiprows=1) + x = dataset.iloc[:, 1:6] + y = dataset.iloc[:, 6] + x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.2, random_state=42) + return x_train, x_test, y_train, y_test + +def data_processed(x, y): + x = build_dataset(x) + y = paddle.to_tensor(y.to_numpy(dtype=np.float32)) + y = paddle.unsqueeze(y, axis=1) + return x, y + +def build_dataset(data): + r1 = paddle.to_tensor(np.array(cal_print(data.iloc[:, 0])), dtype=paddle.float32) + r2 = paddle.to_tensor(np.array(cal_print(data.iloc[:, 1])), dtype=paddle.float32) + ligand = paddle.to_tensor(np.array(cal_print(data.iloc[:, 2])), dtype=paddle.float32) + base = paddle.to_tensor(np.array(cal_print(data.iloc[:, 3])), dtype=paddle.float32) + solvent = paddle.to_tensor(np.array(cal_print(data.iloc[:, 4])), dtype=paddle.float32) + return paddle.concat([r1, r2, ligand, base, solvent], axis=1) + +def cal_print(smiles): + vectors = [] + for smi in smiles: + mol = Chem.MolFromSmiles(smi) + generator = rdFingerprintGenerator.GetMorganGenerator(radius=2, fpSize=2048) + fp = generator.GetFingerprint(mol) + _input = np.array(list(map(float, fp.ToBitString()))) + vectors.append(_input) + return vectors + +def train(cfg: DictConfig): + global x_train, y_train + x_train, y_train = data_processed(x_train, y_train) + + # 构建约束 + bc_sup = ppsci.constraint.SupervisedConstraint( + dataloader_cfg={ + "dataset": { + "input": {"v": x_train}, + "label": {"u": y_train}, + # "weight": {"W": param}, + "name": "IterableNamedArrayDataset", + }, + "batch_size": cfg.TRAIN.batch_size, + }, + loss=ppsci.loss.MSELoss("mean"), + name="bc_sup", + ) + constraint = { + "bc_sup": bc_sup, + } + + model = ppsci.arch.ChemMultimodalMLP( + **cfg.MODEL + ) + + optimizer = ppsci.optimizer.optimizer.Adam(cfg.TRAIN.learning_rate + )(model) + + # 构建Solver + solver = ppsci.solver.Solver( + model, + constraint=constraint, + optimizer=optimizer, + epochs=cfg.TRAIN.epochs, + eval_during_train=False, + iters_per_epoch=cfg.TRAIN.iters_per_epoch, + ) + try: + solver.train() + except Exception as ex: + print(ex) + paddle.save(model.state_dict(), cfg.TRAIN.save_model_path) + + +# 进行测试 +def eval(cfg: DictConfig): + global x_test, y_test + x_test, y_test = data_processed(x_test, y_test) + # 重新划分数据集 + x_test = {"v": x_test} + y_test = {"u": y_test} + model = ppsci.arch.ChemMultimodalMLP(**cfg.MODEL) + model.set_state_dict(paddle.load(cfg.EVAL.load_model_path)) + ypred = model(x_test) + + # 计算损失 + loss = ppsci.metric.MAE() + MAE = loss(ypred, y_test).get("u").numpy() + loss = ppsci.metric.RMSE() + RMSE = loss(ypred, y_test).get("u").numpy() + ypred = ypred.get("u").numpy() + ytest = y_test.get("u").numpy() + R2 = r2_score(ytest, ypred) + print("MAE", MAE) + print("RMSE", RMSE) + print("R2", R2) + + # 可视化 + plt.scatter(ytest, ypred, s=15, color='royalblue', marker='s', linewidth=1) + plt.plot([ytest.min(), ytest.max()], [ytest.min(), ytest.max()], 'r-', lw=1) + plt.legend(title="R²={:.3f}\n\nMAE={:.3f}".format(R2, MAE)) + plt.xlabel('Test Yield(%)') + plt.ylabel('Predicted Yield(%)') + save_path = "chem.png" + plt.savefig(save_path) + print(f"图片已保存至:{save_path}") + plt.show() + + +@hydra.main(version_base=None, config_path="./config", config_name="chem.yaml") +def main(cfg: DictConfig): + global x_train, x_test, y_train, y_test + + x_train, x_test, y_train, y_test = load_data(cfg) + + if cfg.mode == "train": + train(cfg) + elif cfg.mode == "eval": + eval(cfg) + else: + raise ValueError(f"cfg.mode should in ['train', 'eval'], but got '{cfg.mode}'") + + +if __name__ == "__main__": + main() diff --git a/examples/chem/chem.png b/examples/chem/chem.png new file mode 100644 index 0000000000000000000000000000000000000000..9d243069fc42c220690c2f048e5011365d6d7eac GIT binary patch literal 31491 zcmeGE^;^_k_XZ3P2vQ;q3K9~6k`e<-D>a0qbcjj}Fd))3lm$qGl!SDHw6vtsFmwzE z3Q`W;@$T`y^uE8(_5K0R50A&;!7zMg$J%SJIM=x*NaKMbDKR}U1Og$2DamO;Ah>=I z2+kbgMevTu;PiL!1II~A@gAhO@7gl>1$oNPP`Eu3P#ywWwuf@bNO0?#w#9>T$tL&q$Wnm3v{(r{^h45}4yY z@UJJutN%QCBH@%D7c2t;c{$(IgZKND#;gqF_j}Zr2xYME_yyqL5Mp0w2(UnKurG)& z;;~?V%6#^N{QeNv6>>l98xZ^pxa8QMX^;NDLH=JU69Jq8y^`z%A*-R5mn6$|n8v`* zj|8q{$ha| z>Vnv{&Jbd}N4Vobwh_ z|C;KTJ~s}Md5Z@WI)X4{RvrfoYJRz`<%ME)%|>W@U~IQyqo5KW|56}N^-;iZUpatpJY)g!#^d3KnCOQ>=PvJ90n6F z?09nWAnEW4oAF-UiC|ZHAodNU1lvDHBmO@d|l7 zH9NHRqq~>2aBi#b?kbPXR28p?1N!~Z?y92GV$bdJ*v|#qPkzT6(hPv zgUN1axTW?6RFr_pJr|=QRQ0 zpAd%BpZ1+d46K;oy=up#LvZ5#1PG$=h)85;ADH><()et)(i?6I?W@41p!F=3fB*v;Rhipo*X0a}$r z54BTbQw$R>A5s&Ykt1Mq)_4iICXJL3fe*bnJ6V!0h#kmK4A17Ro{mDEAJ4`mWB#l! zd?z8j-`2Z6h3M~ju@h0!`uQC1@v1BgeqwZ+{>Zc5#a+S^)RMcgCj?fL&6(<57Smyr z^(=YTd%L%+-qd|~GryqtmdX@7C;8`AH$T6M*ZJxG2-o!xeQk)T>x4zccAt3M58su% zoZ&a~>GC06LnbV{;JE#m{UMja{wNBVSQ~;o`1X4E;}gY6MD@(@tnuQYBD0z=I57z; z3eu~Ba%w+bO%#+g=5KA~GLJcY4#2mta9$l%3!TRQ%;v{~Zampi5otdjQsK{oPwSa_ zWW~I$uKt$yje<=fw@D4heNdim(9~=F;Zf365u2E?=0V6pQ_dih*+;2Gu2*}@a)jeH}CM8za1{I^(OI^$b?0l z$vs7yLyy*laB0(#pj7qmf=r%2Ka@%2B)ifldvATs0;@PU+3qjclQ{aBVdG>;8!df) z*#A)Kfm7)4EVuwjJX`^>|IufC$k{wylSV*jYTQCoLqHe4BmUg(n3=ky=Z=PU71a&Z zG$V(GhoO)yLZ;S!+%2adv|dR9Zp5d}4u%n@P#)(b5nF{o0y5o0Cf}3IGHV;o_iTcMOWnbdF(`_GI%d0Sal5u7@tm6qM$vgx8c!NzJ|Jt>t5L{rfTCOQkvS&k{y4%F0 zUKrM{LCPiIj)utWR0y_WI2Z2{@Zt_4JohH7x`#Zv;;OUE-3JMz%YTYUExj|c{{Tm2 z$0jemq^sBn{w|Cj8LlvGF1)-} zPkIT{-MMHH!dsPQ-uf!`Y*ACfl&B{}M1MeXqxpi(?k}L*nyx8yO*iLOi|fLq4(2X^ zBMBeQe@vTU$)u-*wKH>y!8~X2r3WU|ED$R4MVTm;O)9@j(wt+(CgO%=PlG(TXuW;V zgr29nV|xAH7+uHBZSeX@s{Q3tUWKQ+&j$^Ean)fZb>SD8nPgK{`Kye_KCS8OW=@${N7(l@@C{ z$8m73F#T%;0=!4=e+5MtWhd)<0piC{`T_4dr`}d=g;{^XzM{+z^RT{8s z&5VXDjYntu!ZvGBQK25^$F=0a|6ZkiME09i^L(+t*7@(pNrzY|5FVyHl>gqb`nZrE zxnDcvx97l(M>JOtw$oNo{5-covT=Z~COuEE}KQaeN(JjvYN zwN6W5vNRu4`MsFuHzBD$dj_>9Tix*SQu9XKjY->*hw1ihZzzP#Tkuv;Ib?{kuQx2_ z=bA5&Wv*Axrl37ndFh2LFS%gQ$zf7QKg*u8w_N+9r{rPmeVKz|N!wWxJc7OOh2Ijn z$A_Orrw!Tq$sB%S;(NNARq>L9P79T-w*0Y^N!zzvV{RMmv6QTpvXtT#MlWnFQ*PON zD+~6&%TKyFU21j0?r6mZ;fPG2!K&DBGUWYTmwaWcgX!K{Z`emWgHiHkF?ARz8!-69 zW-Pa!R_iO$k=)1R-EsVSed&nBo&^8Pe@-}ph{`=mCI(HO-hCU@Knownl=r{slZ>-| z=`~1*qyHLr16Wr*c)3+SU$@)$NW~#dMaC>7v#>S!?$5!DH8B*p`LzpS;yvFfFA!YT zu}E56lKFG0Pn9yu zYvp+$+OoObt80XicSC#RBoqa3rnm=GjXx$4S#Mmgni^_4+btqtx!XYC?=x^B@M!z$xr@tiCO-XlihX5kz^wcqxd z=NAY-@@WfrOlm7)`VK$Y&(yjUkoy7OFPQ9mb~p^Yc(4^sCB*yNOS+%HBi|}#M!wlM zm)ITFhL`t3GvdsU2C=^TGj0V^^mm?TCRE~)JnM4;N}RW6=Gaa%T;urUXrq-(#lwzv zxfD{rl3&m@`SI!3s3AwH->XfOqj;%+H8Am0bYXpz`JcJ^dZ)P@=rLhF~kq1FPOvdfT*k4tu)gcUVk4p1~8 zyYoGwtYXG>sz`VUUas`o9LIgTS2XVI-EUu-|B`)3^*Sv+tm0(?t&QRPNn-Kk#8KR5 zz4WXp*C$<#E-%V89rjKxoe8o4OZd{5ScdjKkMsD!MYJGUF~+JC~alM%i){&5G~%)!MZhm>WgfFb1hj0dq!Yg75gZG4t+!(a8=eJlGp zsYt`no{#ZEy$pz z3^MGneDTg4XZ*u`0&M)dUfO!uMlMnj9{A6ugg$={tPo)9IlHqWZei&e;fUtCcRQwf z_Qs^!Y-Ni>s?VE_FcGIrWqsgXfc28#_$t~b%p1QWWXZ$g{xnE_Boh|h)ioEIX0Y)w z4L*`jvx}LPMunzTLf(ryq~tHWy($#5DdcG8d+PEMG7VzepC!Jh1xLv*2QD($r5(v; z;Lmxk7U^AI5tbN{u;_}>FsiWDyeTcc?S8V^S@4`e#70gnQTW4aCeK{wX|j}GK0AYo z-9z?Gr>-`$jSWXLi&xbaI~W{v))#I*M*fPxKYdAOs_z>1pQ~V)$pOiVY_~~3ulBn= zjaEEx&kTP#3!NgD-0qo-Y<`OkhOfurmg@eFF066B?5vK3pY7yG=RNFr&5W#3-X@kw z*|Wfby8W6SI;RX>BY=*gnY>D)6 z2VN-3d?M42Il7JwV{}3)@X~bbdG%p!-TGib(d zQ1zfKWy6w9CqGXkS_Yge7R371p=NbUKPpfUGR)4l;w<2%?$qJ4sI@Gs{uG)JnHJjR zgU>B7@|csu?cV#plZ`Y(rF%t8_>W3ZIfLhGUW~l9V+o^j`V-u0wsn;L{|z75W#?)9 zF2#N0RV0xy{npcuUXGOFzi{yq7w;vasf#yqFH-_L7;Z?^(?7441Px_!pTpUGo94i+Xgf{!jsIQi zV`3R9UWkWqufu}XD?Nej-NVEB($KV4h0&Hrn;-z>%pZi&BI`d1h0A&BF_$at@adYn zvsoiuBVKZod`SyU3#47HKmJu`vw8*1t>GZ!(6G~g-sBPPC+?1{uN)z-=}mnG*(W(mB;9`cE$6f7c<96cF`<$ zOGpvQpS@4Gr!#9O-?35rKjRiZPdWNK!r|LG{0^-x#}`(F;j@ig(2W$2tl8s~11Zwe zW!O&Z-O%OLMaE$v-;=dekIP9#&jKMmmw#v9f8vhaSL=kS2;|qZ&9#}nV z`wP~N(eElU9Z724#zGz@CiL4@wWv*wCx+Ky?c#s(_dg?ebX@yde)W@(e5|nNjRa@O z$}+d6Rk-PPn9Z=mPxWI|iMn?QYieqO3u%C&K(A9nN2kfZf-t{&O1JNQ)DL0!kJ1i^ z^(TxX`P+I1yI3-z@rW=Ltl9cc_V7FuddAP z=qL$gLGZA6ZMQ&Civ?1C;eTUWnIe{;D|*@Z$+KUQQxAhUwF?uuy-KF`eRIc_wYj^t z65ZNA-sH<*-Ay)QnE&@44P_~Bj|iruiwwGimY0onqzB7rM7v7F1^ef%Cs(Fd^Uq?}SkC#M-H}giX{sS{9fdE>E zs!Gp3OjhsZtS8|fS*%S^G}}Pk=L*^1Ksj(YOx0yt!dmpF9Hb~yr+e@Tn z|DIWOF1Qb88O9ql(OA^lpI*2o((CZybxOX-wkIiW zsVl(D00vcLSgA}D<bz-FYiV z-8~xysiV4J;S$*HOGu#!Nak4X46IF*tFT5W=DIr*FQTykRidB;ac}(XK#puMl5-NZ5`Cfdj8F!A71lp%C-MzAbimTs_mvW!^|I9ZtW^2c=0$2y~)nEUuH3gx{ z*o_D8mAG|>7kv{Uev|hoMPxH{qr?^l1?U1aC=<{NNVWj9+n$2SL;>##Q=lDsk8?e+ zBJw}bnqSd{P3O`(Huz4q05^ozF+w2TLEMn6ZPgHNzl!FCrC<7Q;3(W~BYT77w{CmN z++HR|8cidtRhxAwSR@e({%yTx{~qw-t#b-o_Q#9sp@KHsVHdY%(XkB+;_H@`Wn-1? zaEHqOVuk;G{KWfO>VgPdXJC)MHmtn!LFT81aE%rR*=k<4qBR$U?5S=D3-UhX2JU45 zLLc6zG?2ZgaJPT7MJ{b!B)eWZgf`jJzA&}{o`AF;&|OpRzne|;7oLoe&V-M$;5_E; zV<(}Ef9o8I&UujTipcTf)Ad+8gq54fcJPN;3BsY1;qxkR6G@R!N^D=c7 z66rCO9Bx&M7KBRggBw4EXSNyeb}U5+RR8#ez)}L@4PaR zkqT`mcn`Gl9eY_aYscda%p4LbOUuS=C2*}N^}e)A_S5`DS%XzY0^CK4ufEWsV_0Z7 zRdu((xcP)>wK#0vFd9<1z+5ufWs1^I^HjvSjT>(Uv=2gn3!zT1!E;9@lHgg}A`l66Vo+5_JDFCY(ej?}SCg(}V$i_!QpmQndc|s`J;+WbbtM zKte^Kmhjk=wF*+44abRp`9=lzkD~aovJu5(vZlSme>AOm$H@ZYb?z=7yCh0ayHB6} z@JaWTs?2?9HS(tG5Gow31R15jwe&}mGB+p@k(D@lQfo%&K#oR2Ls5ZYn_p5AvimnF zAHM$|Cs@+tv7Fpc%V&v<%bv2y&Ovzh8sh)eG_FpCp3*O6*Nd&q9ifkR(t~$07is*f za7!J|H+R0fa+dg5Ih611Z2HbBL}w$-Zroa^LiN{}B3nj&h}Z$kRQ1mt)KJF@4Nk}~ zGan}=O%g9(kxYm?&ba>~XNqT~Z%)*L5v- zYiDI7rX{s+5onrrB?(D9fF)4Y|sy^g0z3?04R&?rsP`|`|EH(|Y#VErHYuW1Q=PEb* z0cNUFWK^9Ez}SHoL{z!zsZteRZzK(rSpX3cW9>>VoiV;s-uv7i@F8*$&CvEFy_%0) z^A)}`f{Ez$ve!$=4Q#jicFToR`pd5-JXs)3-JQQC^?J3L#I1gL@}!NKZH8(}bHV1f z1U4e>U;YsN;BFpIky#VcdF^9o{;bDNxZ+uH)7jC`=Kp{Z+wW;E0n@3$uf*D>TMGLj z$E*It$evc05{k;EVDxc@TVQ;doy{`c_Ew}KUg*%%R-OlLNLl@?G^F*tH1EbpBEy6; zZJiNnJ(LMs!E7#cY9H^d!@opaU%}*Q=j11ypB-c3lB@U8UMK)~=)KD}X!QQ{oE9Nu z)yGp0YBVbsV2>tRuRJiNM-M{~%I@@ZzXlh|^d zqWbJW>WEmhwjA_>6h6VJ@!Ykk{T`)~D;s!m7+lib?L^z+n~VT(85`+MmGp9*@K9v( zU}tQvD?1i=HYei}2A6jQ=j_8mRLB(&`VwEAGtt7GH~%1Fg|1!y8GfkoGTdbQ8nmIz z?P0;F?d-&HToXC30e5CkJv6(-D%%QPwSPUg$*#Xdgarc5d`tM(D+c>qBy!_3!4631 zwWq2O%by9DbvhRaW!^$?AuBYnwl<<2KKTo&R--rbC&D}v#^z%i#(ufo2ZBAJ>fiFQ zDw*hGxxQRPR#Gg23RIPjw#w$bU7RCF->8ZBJxSR?9RN4FhKd`HY$hLv@D5`fv;EXS z;mgmnMZK+!#;*9byur0{3USqC!*hd)V}mK)ZrwcvzRt+yWAkjzkEK=Cz7cX(`d2H7 z|8rjJbpXQY8geJbX7YMXtMV+O)6ecNi;3gX`h(ik`t{Iz6Y)hF-@J1eH&y4lOsAe3 z4^P+{gh3SYyfm&>{yz0)st02Jwa&|0lNT)kh*t6WBEx_m&IM%;cZPTC!y!9q?Lo{XwWlL`_Xx}u}N8~FS7 z4yMaHW5M=n`x#{hSu^C}kBZ0b;bHlr0lx-51kxWagr@E0n=@cbe-V60)R1F4b?Opt zj_MbmZMBImw1Y!^!A>8#*@=~0USlGDXD!XXU|3Ifo>$o?uq)T%_hq%j!G;S%Y+b@n zjg#&yh^K!J3^I%>S#wdNmbG{?l%!dq5ql({G@y{jkM~IGPk6Dm$3LVT-Tu{yESC2(Rwa+d60J5?Ynm(JJ|XB!!_ zHUBYu>#yVVV{}wVPk%#hE2O)}clZ0V^qbl}D&m~V$A@n473ddzye)J<_INa2YS8e$ zXYd_m03-eJ6_7`-*|`pCf}9`dY#U9iy3JISUJZ(V@mQ;z#!gGpt~>AHUq1m6XTQC~ z*_P>KxBl}<=K#sf@d1d&>+H!*aXQ%$o-46AkMD%~uE57x{5ikn_sY^%H#kmd7yR{S zkXY+)b3O`cp4q+X@$5ABeU8I(IIG1Ikoq~deP#+JO`7^EX)s{KV&Od4MJbii5TD)Y}YZKnswnn@@~6QuqTJ`9dg|~geqmE z9icbhmx9=8!GrYDiFKb7r=u&L+tqm!4X)U{@7h2e4oCCJKvVf9L5 zx4&_gj1Lw2)14|-#|sv;8T=#EezHm}89lhVdx=K1KfN+?Y`ME&Uk?loq;2Dd9d-`2 zq_bSkZHd|(2%F`a{beX%^e?soiuChFT(>1}uiDrmas8&?#OA#9&ddNwJm(`K2SOoA zc6|tZ^!9LYmdAO+X@Yax;ArZ^$uHjq|9^JCMXyC38K99;N~2Lp&y(dvW4Y*4aWh1* zwWlIYF!1P|-d@;beF;X~%XY+g)|dAvU8*T}R?#l;*M6GRzR(yVU;d9&%dQe|p zt;@T1rr1ydy(b>F>>0l-OFNqFiSthorXifkk{+0p@T)RyPR~{ev0HVA2)-Yu(>SKz zb{IsR)EI%6R&6Jlyj)!-`ZKxVz1Hd6lmvzx&i4Eh$zz9n zOl7@AX`Ja=V-xPoIJFh1kd*#_!&5DUIxoyF3bK=EI#E}DpUx1KOVKO+O_lz1x+2)y zw0yHD)AB*a^I2F%w(P!Qe!;BA&>cr!o7ux4l;?j)70+1y8T|y;jF+{t^`Teh>Ua!) zD70Nf(#DGIm)+LJw)ch09=R1|ZiClfQ>mpCPCi1BX)&RkPYuRl$9o?ScPpLMnw;Xr zmS9V5B1w%XpNDvxoquj%@GELLJKRvSxdM4R`G-S(%jr#-f@K0Q$WhVSZk3u%AT6gQ z5TVx-uM{&gm8FX_zDaxWO1}ppu~BRY_^f;YgEN zC^N0pIw&&Zwce#;q7R~`XZMUJRuM@%{V$PH>MQ(BWr$|+_y=f=)Nt=SATv+st;rxf(cei_l;2m$T zF6N1lUL*8V-lBYKJSwC#?eQ|wps{FTxK($qMv#`3zI;8JlQ+{N?s7VFVl^$mBv_T= zI6|!_^r>;a+#AhKJ&owPT$Z+)a0)-tB6sXZu}exjfPJ$-+{VobCSAr%miML{asWM} zMM4!+nQ}n2BjJJcsq^HcK+>N3=gAa%)I+8!eqDG3B-X zTBpkGI4O5ugxk)$1}$^aCj&VxeC<)_H2wj@XM1duj1*g!O1oWX+_TbKob51mo~mnxoTWO08s{@hzD zbh@L-9#7GljEqlykzMjqF#+b>ci*XxBb(MfSg~`A!&U%0uhThD2 zWZ2+tBv_?98_wu%$)9{@v1R zXQrfeH^L^DR*5m&C$D_vKn*l2wjqxBVE+s7zg4X#oqCiHfklzprk`x-+q{320_FXQ z9HETDeE&eI>B!i$t@IFYjz-T|7O*BnOaPGX{28B41C;ZZ-#hi)E+3I7X3TxVd=C;u zDO$>TSdX1w?khj9A7Or6uSh_j2vL3@klsp%+MByB>37D^njg!?Mt4qwZsoT0RQWi5 zt5u#K7ye3Si}TfO=+~Id++E~Sq$=K=IHq2EZyy}g1C_WlF~-{fhsIVb4N@fLCjEmC zuH&%~MY$7_G{X8~r3y?@kB3#s1mQc}KfRi>AIL5x#_#Te-0!%f7ygaIm(>uIzlQ2x z6-gd zJ}Grd_o2-+=)`c)KBgSb+!|mt%rHBssE3yCo^<;#Ldq~mr+wvLI;nAXyO@l|+!C z4dD=6lX=bP8c*r((OG!`W+ny47lK4-6;WVkiK@`cFYis_JL!sB`EFHv4fE-w&lGOq zpwXm^*e0Sa$u}??cB=TLfgg>|W^9TTpL=Zfg`j!)~Ew zHckps=|?Sg8c;-%?_zS}*P9@)4`VnXLAcT%zH%Th>Sqo9clzm;^5@cJ6{xY>_uHR%f z8dLCRjv4mxlZ?1b6IpKe)3-U?Id4J)QrJrxi zEivAFc7$(4r^IIJNu524SEG~4etI*xpx1bUVXw~=6;SNpwQeZT1g$BQ#(RVaJV=g& z*a6m@T7sbbX&=cwP~BQ6TL5f^jBdWVLLOeSW`93YEsaeQIgboOM8zljQCwq6H|Tf; zH{r4|xPf!0_B;*cZUQb(NPYNB&>84$&`(Va9=ex)@2 zu$6221gCxIRJ*K{ZH9*B&$__G2$?q6W@lN*-cv?s3%LeO#wWDLL(v2SmLksw8aAk^ z>@1PX^s>7bD{~mFI>j6Q=;tIzL)m%AUxqSiU#U~U1D)DhZ;h9W5=6ApuVI3 zY#BMqy7iRROUqJiH&K92Hv4<)1)o7J4xSMTO{Y?qi0lYsQ~?Gk zV=2|Q(XL4rofNTu{o3uvcxv$ivXTWg0G7$U2?zOQ^TnH2CV~SDJ6*rC6Kf?veYk&^ z-+uItxf(rGc~C>Db;o?MXn*)NRgvD_X#T`r0$twud<>j5Xs4#Yk0X`9{ zYEXRnjjmt|=Puiz+G^s8j#OmITUqV|oYvTK>!dqQmfsi&$0K8{m6L=6o@&f0aVxtB zT)JyECZF#Vzid0+a@Z@F9$)M<{2b(SW5PCa&U+JbPD{avCASMn@DTx@ez5dRtNjXo zr4XGHg`b(Lin}kxhVL;312*49Tn!=1o-22k-aAgI8I=CE1vsCCR@9Y#+qw`ZH zoL~zbp?P3s7H~lA;ahJ!{;R5Bm ztHjNP%cDgMSYo9G9NcXJmtWHcr%YDFSJ?d0eu7!FJj`i#4UtM?`=)G4N-(b>Qs&d- zE-Rt-yb>!m%PxSyOlMV;jTyZ%(1VuvMNKk~>0g?vN>Cf?Roxewk)oQ~4(y8P+ZXQ? zsnmE4=)3xzxEr9RUaXJT(vP>IDceB!9iU%aj8=-xtmKe@AgZ*M7jJY-;}1=`X?+lI zZZaQI23gvtp5o+5Ri0DSI*K5-vZ|OUd~CrVZllbtQ%M*z*>2&;m>H>ss(^gon&^SX zXZy6M?w6+>+j}83t_~d{muj+US2U=9$7BFV4XOO}9Wb)86+-C;(_{mOjg*)Gr)Bs3 zb!p$+Z-__kcRjyzxrhu5x7GE&SI{O zjqgxByD;x|5jVo@c(HP^HwnuK9Zn~Qo+YF7-9IgYqPuRA^mZ;9us02=h~oztWZixa z$o}{ins$ibr~=vS;2V1OtdpTH5u#@CVnG3+?v9nW;zDO1m>a#aPZAB)+e23#TSzyS zyTyf2maK_jw69fINB{cJ_>Nt;*OkCPV+L^~T6k&s2 z0X`2;)+Po^s_j--+`1=uhrH3*k5n@-lRAcDGp(2ZQD^&N;^Jx%?hKdsr1%l+5S zi}m*U;DdMeumZ3%qs(j^M19FlXl!|wSP^j}_m`Gr{h}M<;%s31hm?6YV?i3fB>URD z$*A_HAndkr7H}92xd8JgGD9}2#IW4Nw@4ZkgEplx>xP#;8KLxhyJN?+-2jKa0h z<32Qt(9{w~A?8GBNc6}w&-L6&T@D7hICB4p3FY8TbeEYr#4~6X_C*I$T#fcznUJ7D zK4f6-X4stv(Sy5wDTL(`!)zIazvaNY0H1gmZeZ_@M;@bjM}Z%ro5*ltb+TsXBJijK&G3~wkKp*>&ncI8hQW|&@6I%3^NE??Lm z_IgBOlPbJrMNH?AswBl=Ipi{NLdz%DP?hR8vUI~>VwE@AH%^Y)3`LuVK$m6K8}sLt zWYv~0H;RH>ZBGb7i5*{q21lPCd&*|tHE{{wC`4XJ`tBy1jh~fZZ_A^s|9(@6MV-KQ zq3JaVCZh5aG>)cX?$J<1R*>GFXHMHEZ`jl26WM#yMv8d|O}sR|S=If-B1pwe&RjLf zfhPZ^I8h!+GYNnFz~q)fgP2C7yi#Q(Zj$KSD1+8iV2=WMTO~yRWO%+sl}u>&(4hR) z&)M^A8l&^!uU)2mLduevBb`S=e7gmD*0A~B>=O3P{O-t=A69&g)pl}1UpRlurxr2| zyho8;)9Id&dXdMbQ`*^a5n z@-CD_rbX4P6Bt)am+9nlmGnrqu`s;5D}h`~ z%BSgydlg(ELpidv-?wkPS6BaT{&dWi%mfwXrV8t(W-lp^^oDf|_JK zZPrY-SXA*7Dc?@Zcjkf4jB&J!_v!Pbpca0U_?zNSa1MG!X9qyVrliXM+qzlygXH|7 z<9&=eor!Z~Lc^l}1st{MM-=<@o}Vq*9-&%}$5L#Yk+z7=B z;KJJqdG$uWx{!e4pPWBYBrRjJFO5+6ymH#8a=K{c`Lbws>E11VpzW#Z4v7HbPxI zy6tq+EW!2%F=yMW3q8;4JJ@nyla!W|`#n1>B-;q8rmB_IPaF6~e|I|npG=4)r15Js zXB=`2cLQ|Md;r-U{U=-IvCHg5&L|whVJ1S_&4{Dkj-EzRs;hPsMqjKMF zi^&Wqo8`@_rf%NPFG1AFBK|C+X=^>0yI7;C9P%EDpf56QNDPzMQo4Jx(b~Hv5^hTn z)fj&$ma9h2Gk<+IDRctQ>BiSeeQU04bXIsX;OX2olc zc!=H6hf*#aGHKiZ0a4{7xxG{azfe4b& z@ntrh{NT444p+njT+ru3J~_|4-TWm2!*uqtWl_p$;|$^ zRK^VugcTT8euI;w7ZkG^&M!J~RT_*u*RhFE5!v z;(zD#PB{ANVDOtVvJhOl@r&~{Z8y#Ca@qpGeQeEF8vyPB3;isKLY;-K?5cwtK4=#y zY=~jn+Or^H)@N|~oF!7Akv_``hqwfvYX~DeZmdIY#bTEFqcBOR zPm`IxiRt>|U1)*wO@=9{&BQ&J(2Ol~EG*>g{rFp*$+>)9Ewjl{0%`{Ap#uYG)f2Dj zIl^$xWa+_1N^q-I491`*%hjLBD3i4nSlr60?w~y{@_?*m&(M6)#|pEAfM#@))yLT* zzcbL^_Kdt`MwwtKe)2|J-Ne4KiOLvhX0&}(JM)FlE&iWGJq)?#_9V*HWg!*PeJw5WqwX9NF`2QnLgq#*xY+}!hb z4o$1|i0qn+?5jW09?ksGl&1eX!%P5N^n2?#;{|&ByL@w=kjpXBnKFfeC8AesqfHwg z7By4p%$6Rt-Qm1+lk51ZnR;teQfJ_nXV7yK!1UNaAQev^T+xf{_g;7~UF-4j=KWc! zl;?X7{8#b095HH=^OL@^n_$fv5W(%mUhNEj#AL+>r+!In#?&I%jTOUn{Kv~4B?9FX zxW$op8oD@|D|*t~#Wrz$WUEHnmvndx5Bj^%H>D@;?j!?zC2Ga8{H|H8`j9l;WLa%$ zWw1wy2~kr;7k(wlyj#a_2E`7bN=MQIsh9Jpn zPSNm3z~seE@%uOjpucVKf%v){AphTP*cnn;{AP?x(8W9y*RZ~Ai_{IBv*W&I7Z5)z zYHl$;+A7RjW%p7vyy}pMtDFzX?d8Efz|-F`&Utp{*2y}RgC6Vey(8aC5J2`|$merXO-yx%zkH8=(iuq5wm~g7E;o>r(BLEN0~T zN_gp6AoZn|9^ zJKcQ#`&#-+)M*e2%~(#sM>8$a4tZu7`s0$hZkfopQ=;%#5}V)G@+m81m1gYn3;y&K zHlnLddQq?z*52bABV(g>*-crl@u>GupXfF|sU2pJ!XqY^x_Xe$JxsKXX^HZfdy&-~ zMbBjs3&5pO0JWds}Vnrf=52e_Tsf+ ziRlWlf$8man7Drvp`DBxhhnO~Ggn!@vt`|N8$Q+x$>ey39f$A#j6)5PamkH6?)a=s z$tjz?`csfx`Sh_ZpmBKF0Mdoyafj^@1q4{G#nX3NQ;7z}LC%cJ-jLuISU6OR&Tslt z)jLsRaSe*Iz!F89n;Xh5L}kJ&SE~6YTXgrfTEcQ25J8F3UoxN1+p#lRJsk`QJs>m6 zyCgIC3=6HV;FbawN(E?`1Za*@tG{trv!7iAuJ^nB2l*-yaK{g){q4ur3&hRLdV8)n zi`IO$BEM$UKAUltboSxcZ0$HSmWaOMaw}n#D@ft=jwq^E0@?OyN%I#QWLN1QtEU*w zWKu@J?5(j*#cPDt^Mw0iOXG8T#q5R$*Nm>48%e_k+|`=k zY%5)+p=vWOX7S6W<%~hwdmj-3#>v{*awAK-fWLgiG1n zK(U2O+{4n%6O6Y(;%aP%=Wl!R;(EAh=#>Vgpetz6{=(O(9_@i<35FYT6gQYjeF7IT zRjHInlB!jh!}+QFM=8%#X#|->LmF%Nn6p=}vrxE>gSN!P_v%lxP>+*{VvLm~Lv#3x zL*ZOzkl@q$fqnrdwG@d7)of7F%6kI;T#~-P-~*^2aYOHTjTEesU1%#2Q|~~J#{(=- ze#dDsCYcj(;T$6aT?dg_SCyio8$rO4PF*kS@+^u??E4`D97I7=V5* zwBKk_V`}osR0|<%S{NeCOie&QY$LttSxb??zUo8vehvomH6A<`GRWHeDC`bU9yqp@ z8ZdB*RUs*k9ao1dZS_1s=Q13mzZHPIj$-m!({wEc&F5|Hy5u$Yi0w30$9l6#kpVB) zyMvtGAkV$hZ>U*Y)P*@OpvT$_Cv54}Ig-Y8jnGq8{K~(Z=a#R3t9uctVzIGX&l)(@+QK~F z%Jh3)1$tPAPc9$1-$Qc0lE$N3j9zOWu`TaMpwIi^WGk7p>XVMQx63nd5xc%6Kkj1ljvTdH}(Y@vctMp5%`h?Of0o zEbOwP6N>~lW=NKR@u_8aFX+i12}J-d*4G0EbWEbJFDNQLR>OIMM_$9;&F>5=R9d7f z43O(}2=_`pE9o@1eK|{OL>_uYw!$*>3b{$rM7Zqk?qb}QyV{C&d4CjeBj-J<{y&xW zXZK8R^I_t;J^d?XAeKD8tH7XD^Yt)>*RX<%s2NcE!tXl&eDpo6AN0lj1XAHC2boR+ zsF#TC)|uk91n3Ak&ddF{G-8SuC9DLQdtE`XV^Q6or?hRTQMz@s8bf3@Se;`kQ;;P5 z;o8u*>9N~w{$UE25^{%DlK zNgqQ3u?*8`l9W&F(F`}_kTaYxJIIYV9t#MzVX*yQQ_UMO!ba4_;^d8UAukJL7C3TM zVG!lTqG*rMuhEU#pd+B#{{C?Nox#-e9+i#mm3Adro-a(yyUCPBcJ*7MRrOr5n$<8$ z7t?569L`VSvhhA8yrH)KPiuD6<1G9ufBS1~oZpoLUV|&3yH4|SQPpJEmvvfg$VzD1 zn3=DS&1hkQVdTk(o>^=%+qX&ojqCiEp^v3)O5?iT+SW%jW9W}lZzmT|oG+HeHk3=f zH_+rX!DO=PB-u9=^sB@mQ6#+1jOl1Ljcdk-oC{?4lm$`dZ z&HwOwC~IMQeB-O(*Hx)>dGEz6a9IwU4zFImzomI#;rzP}mj^K}wi{ncV3@5 z%!y<^_mI<|EJ9|NT#scQHl@nDmL!OxCsHjijvy`m#hm)&C_U6H;L$%*4i|zi=a4!( zctS7kTIgw{XbN(7`J;%sqsN6Jc2gWPwqNTT9~tRJoR7c90AArfaRE*34V8C1nw`H>D_l(tNI%2wkp*p* z2@n`h@W7UMr^@&tKvdFv?R{?*>E{sn)R~HJCO2rVv}gD_rjYEIw(356ZLZy4Dv0PO z<8CF(n@&3>TA`wA1Zmel^SF!}yk6`}w!MJ@x=^&Y)C~&FT)}={gm&r2(<@ZExDDQy z1ssj^V=jxiP|i$jH^DWdbt@_~g=&4{La_9>yf0RU?>#?yq^ieczYLtDpvy}c+RfZn zD?h!a@}q<19R#wEO}JShp8`EFlE_8S*(;^(YnvqP+h~YC4ce-=D-tKV(=4-iE1<41zrU)|ufe+mYd0FUIcB z>_2}UA36nrhun1Sf&3!QGa0g(;#w@Kb2g!M%{~y?F7*KD^Gvlm+@d6Pd61L$TCDrd z+j*GR_k_h{OVa#Q=-7TJ(H2D>wY| zbOTfFpqC@swZ$IAeI<>yY4FCVN_U7?Wzu*7=lk6gd zLY6EcYuQTHA`}TBdyAzZRLDMtR9dVNQH|_NmTZ-5$rcHNEEQ#q$i5B3_dMzKs?X>5 zdH=rO@9p;WA7kb@&+EFL=Q`Ipk8?lfk_g}KnyLGg-@q%#?JIYxHhK$~4Cx0>Pi6Or zh%~&}?r64`Qh%z*-h488oJrvl@vE??n(zz3fr5@YyQYxb%BLp9jyo}SQd}>aF9*a= zG86`c(c$?QD@r9rudZpaxTl4Gwtgqii>PfKeb{hNftWhho}Tc2e`r5V6zVy5R2ls2 zI9hcgnUUas&Dn>dLHA{`PJ1k~@|v{4OBW3WE<(`Yse45QvCU#gp-UwODlMEraH#tt zo+7jV##;RpFNJ(%Oa7!>3U*113;7nlDO6y?S-rWQ8Y!R^|6cY&4Gn;4?tx1RXOPco z5xX3pDPua+_qP{gX3nL3-kmpLRkEQhr3P>fmvd2MJe75^V@$y=8$W|o()v=R!e(?z)<5Ws@XthjmiG_s%l|@aSUvKAp%p!Smd z*eC^ZdP{S`hFfF&mK>&d!kM}qioos4sQD_X!or6k`f|oKP|Z=l`_n#_sC*gCL|Mnw zxz7}<&+-24{LQ)wp;}{U6)KM^j!f(}Z%-u4>3&Y_PpW#Lv0>&dwO}Tvw|X?nfxFyp zn0Cx9>Sc>E;$5|6c+{Gq6uG0ZAFY#iW_Vx#F3>-JBmT zS~2N5*!ZHR{b96b0wvJUjHMJ~}ax{a1 zQMzoP5pPD#b-8Akmy!GE1BJ7YzA}|RoImiSoLE}2aQ~uGbJnVELeZgpxrZJcKREc1 z9ue6*kLp+cz5w0{^J92X6*OtKNOF z%yGbnSy@skuRf{4l@X6>?a0cr2yO zH1^fMRLz-qQZD0=<6P)nVfN73w>bYnP12%rCb!?7GyGq7sNyBa)$x-_i_{H_=(mO# zoxajcP1F;R=b+leh2u;T#8Qu+X@q4jbUNh&m8UdF)t+pJQY%88ijdMDSGy&pZrqt(>tpi#gqN~fhuI3|Lofg^4&NnYPyo+gCa$r6aa=Ps<|EyWrojJ8od{+`l#deY~ERZvy@+Z&I z`JImVT!c9$aQwdX17l7!yZKG6B(VQvhm?y&plA$N*ZcyS(D-~I?Q^`R-PCRF0QOe} z3Jo7rI=Vl-_szTfYbS7RDt%VlPeSI0W1co+-vU$RgQ!eosQGCCT7iVhQQ=@>aQz*F zLPn(eWeO|q0Y6%;!@_r~q8menh#imtI$uo!-I7I(qU@<6+6UIysA((I*$|m<65Z zs_O={SzcR<wo>4{06IfQ}kPD~v z#`j`{h1|$qT^l|DAy8I8cUTo&YU>?KGAp!4Su@Sg<@A*goqmkl_vGT7qD)>7-Mymb zyqpq>&L{mnM>4vai;6Cqnt#o|X6G7l6`!v|uITdQ#<hs5Ryc^sRTTI%wC1Y!j`_RMHI#{Bp^_E8CZ~hlD>Y z_oym2|J)T(?!3e`Id9R^%*oM84hny&z_-1AKVZ z@{q?W{)%uo@j3ph7hF)fohF#?#CXhIdR5kTCrC*PX`pEJyT3a0 z>-(J8N7{orNF8@hM~ZhiA7YoXuEk3pA-U;h!dfq1N6!PB(T3~N047to)};bIwXS8p z5bd%VKjllMc*X{-vkWLIikA3@wHSQN@E!j3?v%Fb6ms?1VnOzLeXEU~hCi2ABD3N(5e38DaSFg6vq zXdVo$TI&TZLT6p8#pM^2%}0x zn)4~m_?(+L6U~hueC@<|l0s@6i6`aqtn`w;@FF;i-v`1yz$8NV7DvuQ&bjCT7>Wx! zFFb%lZ>BQM0=sOyUDLxhGlwUsCFz*Hvcn62FaZ*h(L>?wU5vI|-Od>|6ZaPji@hOM z+V_bi_=+t+fsCd+W)M@zDER`j13_N>vm4lK5W6}G>`sn%WbWj? zQbp$8eg1aLqmw>Mzw~w}O(Q6VMHbY}ip$?_WLms*DFURA;nV-dGdv}6q6dgCjg7RB z8oMU9pO4abD%!fEz&Uc9ujkwCB>I;rV3ntSczc70*IJazqxf1ecP+Tse0 zyiv?38V`KTp(}CZ^ck{<;}N0ew;khO7kCbR>H*9H{})!0`haD;NmsJ3y${PtBjlh# zFWkIN=J4_pHaIFf9<*y@^(aIWx(f>7#F>72c5#uW`s84RukJ=>=1iDF5!q@Yp;>Vc zkK;TgjuH++V4?kNk8exIh&*LqQ!r-kCjy0=U|< z*DU*lVlAV_K1unJjxin%hWdsU!(KEB4X_*0ruvD0LVRW~{8p|Fi}%`pd_2|xMvGmJ z#}Y;YckC}aOR`HUyDnMqLRCE%$meLQ zvU5xC0)C2n;R~-+-B}Q$W(I7m+n02Fb07ggN;#yG7lZWlus!6pfnEeQ!-J)zSJ3Pn%aKo_UAyM!9xJ0~fRy z_xoC?i?BNFy|m^8R(9JO71#}RI|DV7z|v^bE$dM5U6k<= z=-k`jsE6cQ^}Nw@6!UF|%J2qX=6HGqUFH$PBAo@DaXeIovTd3y4J3Y*M7eo-ki6kx znZn0wTg;E3OLD<~Za;wl@6C|-AeC}3DMUWjjto(g+pLk;N1l&{mzZ!f63sE?ghDq= z91l8l;<{cscEL}a1`5OxBxB#7%v0mRDk?qd`mb;F?-}XTl*(K8tadY|Z#dXZQDp*d z0efKeWaW;y)sw2GvGWU8d^RdF{5eK;N%}Y1nE~lo54C9I0MY1Tj5fvI>Q+^j-VqUv%q>xtEen~B*c7$U`F)4~S|A+4t zMgCm0LifAfJ*XZmWILuC^^?L(&SbHdj!WmPTmK`BHGL`GQ1w=)zM{fieVs}ABfF5_ zd&$e@m4ACfGk!j`8&_|zA!I5A{v=gK-)@|u%edt7U* z*Zzr+>b+~>WC&k4dOicuHW_T>AIa|@daEc#)=uJvpDU6@)rey4Fl)y zz7ni9D%~)cI)ETU41b-I0()8pD1IL5YO71uGSIj!7{4Zd4^03qU5&=xBWRfk6bppu z!(2Z_j}bHg=!Bh-HI5{HDnCr1QCFs4jai;ERVy#WINvOHBJ@J3mBBCl$E#7SpAC*? za(8ANK@i)`qx{w_4(P_lzwI@l@S3ZA?GoYLFO|Ji-;Pmlc{&zTZd3e4x-h03TerGQ zSuPM8K-Jmb@?n*QDmG_{ZJu)p@7-Uf@%nb(wL_KON-kgj``ydn`l=*DCp4LN>L|#S z6gxuATOr?ffrV^5WsKbgqmq9A!ow|2eMIe1q7{t$_@)o=`>ZG+LwM>;SSd*gj%+;P zD9pQJX>7VkeAI)FKjeS@bC1pQH@4~d>7af==#)+OVzFKi)~@HVaLjkQ`y%AX>|Yaq zE1c0r@NKPSP1sKJ}*Q+Px*~0 zRr)?*kJ6`D5QER6|FbeRImn`eazn?;M?bjdJ$id&qO=%Jb$1^*FYr$m3Ulw+dEQTa z45*qxG7QyzK@!hxfLh4;*YnUQgEMUPu(O|De3J`hx8dbehTtafEIMoUnU_d^AL??+ zQf4~BXtCf!S3)99ez!g!-O9> zoo^jk@ct!1mp z%uL8&Qgz!G<8V>vyR;)~TiB5-N}iKD#*a;H_s!#Xl?zR;7kb8OyTFeK;c+w+2yQd+ zvSlwqG!r}_)ntytB(U4rmxfZ(3dByztSZn&*`QQ#NN1QbmZ@C8%u9@aBazG7?s29H}7p# z4Ewrv$tU@4AM+8{H`PRm4@rHzT$h9I0Iy}-SYI1oDUGK_hW=XPI%Fa)d^=dEMvnN_ zGx+Kc$Cwm3($j^!OqHe-FUPz~?Fq}EH0xL$qA`rH*dEss_6{ktb25fBYmvV^<)*}tDv3Y) zaTKWY`3h*4eQZUbOxk{>gq@YCo#~Q6n&Ja^vuqN(EtAEH#%wlU0DSu$f1XesvIf@O zgc4z)z13pAEi1}7j(f@q6v8MeHM+=KzofU6;0NMgS`n;o=bEksjqhDCj;P!Q$f>_J zJf!BH*x86HDmrT~^q!VM;d5-fTykc*B^tj$KPSPTT4>N=J%}Pk_F(#|NG{Vt$;nb- zb$?`*`j;G1`}WTS?NFK=={XFZM67i{DzE$+&=H%Ta$7J_K4LddgA_9%LFRH=MJkTR z7Btmo^dopv0>szWImnQqzs%R`v&uigQF`Uco20KQ5i@l(&dyHqxJ3&3E=|E;b<(-H za7u)4zfFg43SH!sgmFt!3O&i2Ni z`yQaFnvE$>h|ZaiyLm^@-YfVt{a(z>4cn386;~PkxJRR=J268=567W-h1#b6@VAY~ zi0@)l9^2;Qep^NfM0=$d?D7uB*HVOr^o$hyhegzik9LZALDJXU+)AvI8*TS6d-Km# zb4y*VZ@!YA#I^5hpY*j<@0WAdJIR%!wPSt4uRo@I&EXZC1<<|6IzhqOgR&!8rlq{t<$2~QSVn>U6`3I1?$d>*P zE?PS-`i{uQX!w;h4%RO5`lf^VjPe_ZB~GCP&(*0)JQ~JO1afR9dDwF3G&Fa zvr{GYj-S6D^qcHB6w|)3^zFtD zG^TPCV%YJ3)A{(A7hg;dad{np->l&sD4ccm^Ww+SRzjn6M#ua~W$eb6b-JcyDxJ=G zNW%DKppS1#D`4*zK3Sk)2S-&k9%es~hI|kYghBTp#ku$5+|CnG zhvs&M8bs!PcaY45^B=PL8Pnv_v~bqyP1lc)U)ZTFK3#218o4d_aV>Wz+ZZQyWS|tC zuRl-Vt#Iozdwt5TFXgYWVK6j;m?3HHL&wR`2UV^3FPBcR%u6%p(TLoxYTXBmg&u|$ zr0wYc;OPI_9q;g~{5l%d!9vKMXt@@g}QwPh^CujsZ7nj<|S9~>d!46Qp( ztNvxS=*i-0<)Y4o(U~~8v*1#yV#zYB{(N}!%|lHCoIp&wZ@+CzT>0B)cw)pV||#ef}`;nf}&e2H%i9`a_{uFoyQP0MCDa zwZQv|(zu%@lR(c66$n55`*lF@m_(;U%Vj|OC1~@7%_K1#k|K6dUz6A$xShBM4wtcI zHB++x(3JkiPAyjwuSrkRA_(7(-x!d2fwRl7>iDGo{auORa91U7<|_O(k^MLF@CO_@ zVZjVh{)WatBMEd7T5rCr!O0lib(3}T@tr-2W?;6aci4L`B?VuG&@3^W72)IhYrQxy z$2iQ;jngmlBWiznP2TLd?Z>UuBqEktm=QkzKW_-GqEx#!eV824{<9`&0_kx>-`psW zh(AvP6ZtOW{o7F*gzqn>PN|Gh?Hc;^Hv_FOI%=GZ!8chE!EGDbWQZVIj*e1Oukay` zJJ$p#5t@JSE9kHU{6kBDUxs$nOewPAAV9Z`VYf7{4|qPIs{%wEK?*xVPaIXQ~!~TuXH}; zyPDT^m_tN6R(8O%FXJTHk9v@f=eR6F!FBHXIuwF;^h25PDbW^?uxmjWC}X$i_Uhto z=fSgw=#2{M%8A}b_CS{s7EhvmKNN3ur0FFTC_#If6m;v!KmV?gAcqp}N1dSLH_eE8 zAX3F>-a9lL?dikE!Gxn^#-a7}D92oRDxKpu!CsvAUxK}D?Z1(|&835z4@92tBs*l= zz~tCh=KbErDCcypSQ#ViliV(Q5NuYQ#^%uA_u?=mOl(JbL)sc#P zpaaAcjJ@hVj3f93Rsz~$iExEamkb@(u}!VrmN`C)_BlefZfh$Pl8=Jkeox+(vO5St zJ9ZH^gZ*3Y+I$4nh<`mK=wpg#^Zs$dp>_uTGWM6!{5d(*-yZ{(CVvMX{ADW@Lp(R!udXesd}QA9h`<;v6dZ>;iCw-Vl^$06n9`Lb;E8 z`asTq1%TeIz|YWyWMd-aW()ubh#9E4W9R*!^xVT_ZM%U~&|Cc-o(ZYcXZWxG0wai5 z(Jmn1p+7`hXAou}Jlq;%YOMnm7-#Ci(baq)tS>3HEw~T`s|?%~Lt9{(_`gCkR>r1Q zsYH_-*syvae;x&Fz6cma9e=IuyeefAukzINh)34R{O%yhoS%162J1 z!H;GK4HM*@FF^^Fy~#NpfXFTgkO^4#Ti5~Q^rso?htWDc<&Cwud2r=p zW-Xx%XTM0;?gK&lbmP?giSy6&@0XGFC%U64pH9`tP44Ip z;$a?u$cS43MN#&0x3mq}rS?2%;>G6-#jx|AEL_{-4qVLeUHat?xvSWiZPj1}4}fZR zqOe7h)I23R@#S?PryiM09AuV=gzp&fRW{PDuWIbw=QpR<@~DbxT?*P^ zDWn@DA8!M+h!G3!hl7IUwf90rUDm#j-V1rn2Wy=|;iHC}OUf-GNz2{m8 zr}1pHz$c~I#zg$rb98D6SQ5HHc=s7h{g>t_f#^(xu!)(Rh6X2M>s*TK%Y&T(YE=wi zf(>kVc(KdRhzVNB%Tt@W^Ff1Zv`k%`p&`H0`pR3%anC$NL-o!1?3@xql5Y*OR)!WQ zG%TVng=?IsPlrNR;tKFeq>i{wl`o}1$F4ka28wqAI3_!$K}JRhDL_0(`|?j={nBo;Ee4^mck1T)^z%^X{D%;oDMP#~{hOgtv1m7Dd;6u9=QyVi)c%ywt|O zV42$GwfY344W~$9*z4Q!++}wKPRj)PoH?zk5dbQf=rFV!%h|_F%si7aS`WPF;&vT6 z@Xuuz_?h~2)ZU4Ge?cS233Ev0y6s)q$M)?LL+FW&s88*iT08ndrsa$U-F+VA_zL29 zb`+%5i=uC*-}1H{Mg`_J`}gQJW8|5*r&y(LU{5i+zIfM?G9$ZN>e4{w0)U-j*oCFb zLa72=@9#NmsAn_V7ds-)$@dr_s_LY4sLQnA6_u%%EAmpp#9G%b+FcVwJ5yt0!+N?K{GOaT&I&2 zM63a_Vl1tUmS|{GyqDfwUizw}BJ<1dHpE?C4?nca#&MKKy4+^uKEu87)T7D`($svD&3b{j#C~W2B#gF< zvOkC*A!xAm;5)yC2NW-ckgr0eA5AV#=>whWz<>H#BNFf!RJf_T{xqIZV49 zG*;z9Ta-3zXE*I#Cw5h4d3yP@N0l|TvzlSQTx*i(C~>xcI^*l2lTfx}G_{8WsGp|U zUKTB6eUNt}QKF9kfR7bDel3xT#$JupeXAwPuOozC4|R)%{WiN$c@Dt zT8r>cs)|Dq$2$7)SfVJ}%~AH+koLmMe%5I)`zVpfa77;2zluv=+%WH>gi;67kjF~i z2T9O2(bVG$umw9Bo{AzXMDZ`OfUPk0_OgxOkUob1$HHuU36$(lu#1KGwk+f96D#Fq zvyEaMDLOmNOWiZQ!MP+41pn_0r4ho2wl2Jf4zU!3xb#RvZKY@t8#n9*KQo=Wave(_ zq#relQ-;A9?*?Kq)=l%u&9aG(--lzpQ07@R%<;3(=exV>EH z*~IuQ+}?>8O=c;=H)aKArO@6Af8UtVub%zJ0zqZFgKGKiNk|z`@U?@nl*{^2%Zd|k z@bc-jNC(4{sRZpz?=%jA@Yr&2U&3WWlYLU(ray(^Rfj73!R;yV+gEle754w}*V;Uq zwt@9VwUfQY$+SE7MGp_(lma`?DM=&l#1S7WrOU0wTIPM(~> z%yy$whc%xxiFwKyP$sTS@LZ%*!}3JdrkM%Lm!D9`LWbr5TX}l#IlaZNo;F~-oJZ!o zvtqwf@&-U}%f{K+xjSgp+F#chVU-TQFcY0Fu>msd7-=G`4MAI^l^BUb>+|o z97im+29y<7Mji)U9<$k>uA^I9N1>0;>dy7|Ei|)@Rw^G`MrZDxv-WyB-GMS1GuZ?4 zFGb}uO zYqtbA=%$6uWLc@ywS1r1Z_V=voPwnjqbd+89)EP0Cy>6Es_tGWVte-Ti^PlHYr95_ zUjvM2Kes~D`Kyhe>#lHlwcYBc`3Ao0^WapT&biso{p%$cp4E51xY+3~tGaCkPWG=S zx>5cg^vIMv2}e9ySpG-``IQKtjPcX literal 0 HcmV?d00001 diff --git a/examples/chem/config/chem.yaml b/examples/chem/config/chem.yaml new file mode 100644 index 0000000000..979183b2e1 --- /dev/null +++ b/examples/chem/config/chem.yaml @@ -0,0 +1,61 @@ +defaults: # + - ppsci_default # + - TRAIN: train_default # + - EVAL: eval_default # +# - INFER: infer_default # + - _self_ # + +hydra: + run: + # dynamic output directory according to running time and override name + dir: ./outputs/${now:%Y-%m-%d}/${now:%H-%M-%S}/${hydra.job.override_dirname} # + job: + name: ${mode} # name of logfile + chdir: false # keep current working directory unchanged + callbacks: + init_callback: # + _target_: ppsci.utils.callbacks.InitCallback # + sweep: + # output directory for multirun + dir: ${hydra.run.dir} + subdir: ./ + +# general settings +mode: train # running mode: train/eval # +seed: 42 # +output_dir: ${hydra:run.dir} # +log_freq: 20 # +use_tbd: false # +data_dir: "./data_set.xlsx" # + +# model settings +MODEL: # + input_dim : 2048 # Assuming x_train is your DataFrame + output_dim : 1 + hidden_dim : 512 + hidden_dim2 : 1024 + hidden_dim3 : 2048 + hidden_dim4 : 1024 + +# training settings +TRAIN: # + epochs: 1500 # + iters_per_epoch: 20 # + # save_freq: 100 # + # eval_during_train: False # + batch_size: 8 # + learning_rate: 0.0001 + save_model_path: './chem_model.pdparams' + # weight_decay: 1e-5 + # pretrained_model_path: null # + # checkpoint_path: null # + # k: 9 + # i: 2 + + +# evaluation settings +EVAL: + test_size: 0.1 + load_model_path: './chem_model.pdparams' + seed: 20 + diff --git a/examples/chem/data_set.xlsx b/examples/chem/data_set.xlsx new file mode 100644 index 0000000000000000000000000000000000000000..d556bd35e6256955b2bda30ccf4b48aabe30a7e3 GIT binary patch literal 219784 zcmeEt^;eW{*Dob7(nB{xmxOfJ&|T6EN=Qjd=M0@nNlPOoCEY3Ap_I}H(h_oR_&(2j zo_C$I&Oh)Tei+w+eeY}6=d<^b;8HE4=6#*Rq0Re(w`76uT6cGU-0R;g89|0Z7 zNXFU8-Nwn?OvlH?#?6Go+wnC`J_-_ZE&>uD|Nnpf2S=bHWkjuq3s?41;a29yGo_VU zNp!x$kO45Org&#>a(|hXQMR4keRkvxu5=!uy+93ke8u-fv0B7I80Ye8dY0?`*gQHvma!GDdGqc#2nFwjFjXhet&~+)L3Nl+lBJrTjtfhe$ zc&q(iFv|=}ga((p*3@2n#}gbboyOgq#GkRZWAQaT-F&8q5>4y7LgZ;$r~%Q}X0g=z z^5!~Kx!sMqiD%Z_3h6snL_!<(x3N{!mUaOuRbH~%(x|gC!~O&ZwhljIzqcLCanWZQ z*UkKdwr0V(ky!w%O5 zcSU>;CVm=){`&oXH}LsiR0M?kdt?NS{|3u?T`syapw?6X)L{UyG<#$7+KrRr;q(8) z@qaJ||HJi)BvthuF3gA{gA{Bmss`UJuPPB|rL_+~MN~Hm>$jtn?dv4LScqrmJl$WLIsO{YIG9$ISY5=P5L^1OLwN znc9azEm~gi)$?)r0fyZ3K*KuT-y>N*y#k<~^1yc^Mb0URa9Dyqw?AlwQZIU#<9#*NlG`uvuT>Df~hJ>Mj<4la~8@(=s{N zdEU(nPT>e2RDe8DoI=8-fltO(1^uq;(Q#eItL%8To$w3Slb#9V{c zAY$uJ%+t|H6CW_P2x6EK{~D`w*&CCeR}aD=8KiAj+4waEOY0$ggxV`hy z0OQkx$bR)euP+UOlLy-Yi$RBKUOkI1Bv|^ev$ozxt2p8Ruq;1<^{h8!wf;;WlCq!V z3eszSMNPW;qmx8F(b+!_X@p%0E#@UfBQ}Z8ICH%(tA>NGe4Y4}{iOkBX5-;9 zedsT%mWbv-7o*<0SENRvmCyI75q$WUq@HC;rxcF}HD;#;H@u+XcVJVB@+8H;(flDg zk|riQ(PVV@ZRiKe028%Z_xC}G(`-oSYizwI%x?=Rmg}u=y3jA#mFLow9hay?yCAZH zXH3X}1kC|U(JL;3Pgg7!FjnjYF{W%2#->gO?R&G8f8Q@;89DUdAdlC6jeI^IH6c>H zPTU|4^INF?#1rH^-%!zZr?M8J-W}qXHA)q4yVEd;j{U_x9-JlSG zc@co@|LmU%Emh@SF5EVZzc4PZOix0*MR!ix1KoWb^r1SIWkv{3z}YG}YqN>|oH8e} zbA->)=z!-XANB?^{&^Q$Nfa36E&l5xktfchr=!T|mWQ9HLgb&|gAexh%MUTIGu-gM zh$juQXA=l+T`@C|!U~FLp=V#FB6#h`xICWQyG>9I0|Y%%_ao>!O#9$dNL^ziO$N)$ zpZ8y6mM>6qLVTzvPf0MWE1K7#(l2!NmMh#w^FGN z;jYh0E^EKRMorR5%y-B0j=@=#*d!f4x-643vZ$C*jkCkGq$jJySgd$+dwc47f=IWo zuU%~@Ks$hD@%wJ$uH}zfX8ory7j*R4QP-cwnRu2y3rh|mmrfQtwJFA;w(Xkvu429L458!I7q#_ts=tj~qSDWcvRpnHO0wHL z`_?7HttIMf$hTB&gXMZ;*Dbh2?oD^Kq1hjA#>_Gl+tBQen4oxcYIDW&oGyl1G4=Oq zSRRPmlC1?LTVrzBIDeXsO}J}ApnHn0LWCpP(K>NcV*Oa9bE*=gJ}P)g%()dFhUujV zp&OxWDEdU>|4+v3z^z(Ga#qcWZbOad7(<37tz|li9DT8d=kj-J=+{wQF9SL@q z45Gwyap@s7y`5ltr$J>iyiNE8kf6sAlrzH+o|DJK)U!5=41ire)r>H;yzb)s$KR+)=Mgm(_R$g3P2Cl9C z-8PFb`+J)M{C>P6v+vTJ+y7?Tio1yD{wtBT{cpX7H0@~W*EJuBBo-}Ip9NqLo$jYP z#rNePE{9V7<#IkKLojLR2rzQ8t(CaCJjP6um^dKieeG&Jff9ZAbMtGiMKt$~$Tf>f z?dr7=*W!`Hz{z5x4aJz9==<}P{?q+eOoE6?0&1y_y(~BT_f*-lE|$?0ObsNT=Wo#& zotZ(*7D>p2DzqA#ydRc%zwUpD-+b=(Y#F7b5V>T8AL-p;$kDn17IVX3uw^v&>g%qL zPgA&tf}b)Xmh66}I^}j=7N%SlWZ#~rRSW8-9&O9#SVq^m;JyOR2?EYx{?|FWxnYXf zI$V)#_6A$7)D0fri#_yVReFq)>sd~@kB_L{{@)(~Nyg^53B~)9Mpzk-@9qmQaIODv z>v3O{OKZoDxFLCX${aX|@!tof0S9Swfw)GgEAo+ZPVNN9-O^#4nCnq zyZuIC%Pw_G8Tb?On+DDtAF@;*lj{C4k7u$NR_eI{q zg)QTPzTNg-*;B;%EDSglI`+5{jrG(pS$D)z_nzp_04F%W4!3^BqB<&J4;Qo7-LtI3 z)b)N(ytyTv_ZQ*!`!(PyUNrSZ@FMJ2<8Ners^q0l#CiI(R$`iFKNdJR^->*0b4i!*p`VV7Imc~blQoSX23*V%Jg zwv+hZTl0IgQfRjSiS)BIy|9pVuZib*4(>Y&Q2!v&j(jIt za_gjC4zIWJUk}$$d{la~#YEGPl><_&^hw`76pY9tXjc>a{nzpKBl79-Z1x;3K4Xed zTfjus>UUE(ReC0?Wp{)n#^*@08rQTr55pq24`YsO{+GqsLHGOtxCJV1B5ML@&VuON z7t#bnD}dQ3w_5{)0^K!k3zfO|O2w9;MBr1C^gmwq;u51B>Ty3>=|sjEp@|2o9=)QZ z*IrQpD5atiq?3K(>wy>s-K9(`9kDrt?_oC=b+Q$VC_}IJyoJE$TTI!>17Ygo z&%>vmQyY~YPvak|AF14>3{+m!Md^tw32-{Im)!OQ-=^ICyd!YB!oL`skI%VV@C>`= z(JH6Yf``|w&)^@d_Kzct$msbRK8d08xekjom#9Ks%uISB2Nw$QKgHk0vzN?G-1Jmo zdP5YBvp{(GR_2Z+tRFIybJ97R{o?nAVD_;kJ}cT{=AQQx31@Osi!R2b0)N3v3#xHD z@&PxF$)QrFnMWUrjSuYzYdlLhTXBV+Ka{Y>>a9M_g+XU@k(D$BKlQdgw1R0N=X=-R zlvk#YqSgN;YQB2wXVVvPs6TFzCU7hBn{5vNK>f%(uxO$_YhD0tR#L_K;?ywlODWwJ z`o-wDOgv*R{&q9O&9!3iWdCSk<-t#NMG+j6Lp#SM1J!& zLn~X4)LK&cNz$y@>W8|`a;EP&GMJPb^Fr}vk8x{tHvMk<14S`S#pg7)dVM_Litl$W z1*+fq_NL%Ap~tWm7rz;ImB1R1mUw-t_DAdonR&0v31M#m(C>6UibAo=nbb#&s0)oO zK&~J#JVUuXZLxbb)&mrZYz5^p1}8M}E8}8HT|Wh4$Dhc|H{ZMnmM;8-C+sz^wl1MJ zry@&cCksG-!J16|rc;pD5>0A>oRW; z`Snhw>g`T}l{u7|pTcWda<`!4toOR{cPNR+ZG)A+n}(4^ci6&XVtTK)1}=WwN~Myt zU*6PHUv>+;-iyU<7YBGiRIJlTV&56mHY-+IN?{;SBnq zRiEk$5H!m}1%Pv$0M6;9NNoSZ{a4IS!@_9WQJTV2?I+P~wtk?`z_oIAX+yDY?@|qr z!WAHea3M?oQ!6un_xY^Tn6H1^ac&IJd){X&0QB)fpKAA!%KojkwY&T~|Ltrb>>58q zOM_dQ!HYV=I zS`?L9ohN3;?;4J^ANht{+>&`O-?1YBZ1jjS0}FSDPY)+7J;i5)0@875bPEC!iB#*| z%)>Njm86t?roHt{3qTWWJ(&QCPfuTDf!GvIk*-R~`Cti>sfd#NX1j`lW=<2CK`Tk^ z1yd(gaCnrN zN+tzn)V|QXd&s+2Q%7LLonXSUnN-iDr5Wi!Z!JtW3t?fnU zmTz(npe-~ea>rV;MArwmTNbu1EaCrJ?Z>esXm98N5QEhZZOSnde=Y!2H(++;mpgo~ z`uCficMVV4su;lS=P0+2u>J#yQ0a3d3P_K6V1J6m5OwR* z1Wz@B)R7)+>%Br+C1=z|ryW$b>|CP!I zS*ygpsJcEE6Gu}Q>ef02U(370Z+)M)Yud#=b?Np65LLGpA>bNCO-vkm=3;B(m4`K zeB=^&fJ+*B!xX7@G4NTbK;+V2Uza78FZMo*Y*`hlol$@X2N2WVwJ}XSpv%@zUt~%E z?SX~iq#*MCk;Dfa`kM)CaVI$S3t}W3Yx%r6dZg=`2fBV>Rlvy&x$X!&H-Lb*EmoOH zELt(!uLcZR8x9{$SbI0{>2Y$IA~yAw%i7r;|9cUTk#imAce&4*^W;4KuDL)Y*Tapo zvN~HKFp%luQWW5805gdX%&M?({L~E!Q2UtTjP!Bez%w+8lr&zAVfdqRiur}aN9};w z00ziG1KJwRuO6tb`-bt*TAHTx|D!`xDoDi+@C`L>B;< zw=Q)7CjU@a#5Ij;1(21v@OG53&w2{a&DqtrR%7CuZ$3cQv=@e-`G@N%4*^DH@7H|$ z-RPBRk_;F1*Jt0)?0Y%_kDXsyDfW_*W^XD*2CG@7XhsH71H}3OcNo#Xa(MvclJa59 zS%=~BkSo&$B+N;S*(^*r1k2TEjdxD|3kbFrfc@NbG{9)Slh{@oM{;A)8xl*IvWaHF zw#7;}6xt^Aj{oPQG#>ZLUsc#bVB>)9;L>lv4NPDhP550-QfqFNkVPUdmghk|AOF%A zn?-sK|2H*6I*!Nmbgi3858?p3?Z(ig|6D4=Xe-r_AIKxYJ&D}PB~BeQim?=~47s1g zKUwWJD`$`AIB+Gb_52D9f6Bm|Xno6SYSYzEBO;YNQ4_DMz;rT8!c(wWQ^2&<5!*pC zf{f{F99}_rhQTR1ZG>|6$A{&X%epJx|8CMDn20LIyA8E=ed`6JQ}Aub-{mWe+E#3* z`0|AMf`IjKZvggsq#aQZ41S^>HEI=-P6RbuT*<|J+*4{X?t91EC;HiHbpyJcOHv)j z(!7ADL4c>#ULxa&tnKUAKZSCS>fO~ZcOw5dZ*sPtxaW&iZp>m|1LUI!j1SRfeco-_ zw-b1dPuKAPIMU)XcP&&SM<*8oYat z4Y2<%kLN|$UE^P00T>(FJ9x^Jl5par6W8{TtV>Hnh&ml5svp$s>H#nBvbkF}cgdNB7@F7PlK#p`D zaunCJdMkPWMr40ZLXs4H&t|$H_Wx4+*?QWYwpJ1)+H>qp_~k+54G_P<6Pme)G8wo7 z%7g(Zlm4oW@q3Y9t7yhvM%X)7G|Vy>29c~;Vy+K~F`(*S+GMppYyw(q8PKyIo%2~> zkP9|vWoo}05JRaxezqTZJMQIx^J~(^j{31k{$OyjhxL~FiE~pXcC@icdIoepO9gp3 zs9BSAr9b_I@=E=vG6z+)>llq|{|73WA>C1V0Tl6u*a>n4~ z)UDqX7>{2FQ6oT~V2DCDI63>x56vgXR#VH)&q4X(mr?(gG%IJ4nMNRo+I`)`3$5zje4UWi2zZl_Lt)H6vFzsy+3SIwqLWblBDo=#>zqKrTHNb`* z6p1|Ujv=_Kr{cEyfs)-pciH={JR!7JI>0vBUbBIt=J}+lFSm&=6p@Sau__D_NaM*=*_ym zXH$8k^czq!6;W`sHpraoheWfl*ws>aKx3;r^rmvsWQro}s+PNU{f zngfE*`&Yu-0kTblFU^ql(?0xdRnV9WWw&Dl zc)sod;aHr)MF?MeQf8dRy(y4sb}(pPd~2DuZ6{>X5N)F$nG`m_+;{3SYV`axv}yR= zn(^E)eyvLWx84F54o|_46!RTXYg)}a6Et^-#F{(0 zBjzda5nfs-U8e@0YJg4&3`khdCu*|>@Gh|`xd$0`noFRp_afS&2Y3iX{Fs!FvdC$9 zry7yaC*CF!Xwu*mG}nGwS1s&4K(q`+<>YCBJcj-Pm%9o=S~!c4dBm6u16XS9$*97- zBbG6Le=+uXH#oRW*7eS;IZP>G1_{g?ms@F)tPuc+n*idP@8MO^?62jVD8J#5aG~5j zZK+o#Cb{c$ekvD&@!RKJn{Zw5JLp?_&ptBF#@YEsLu zHhZKt&!+a?{O)jw4$G3S?!v??wFNFjZTl|F55@qZd_B8oVJ^bh(doLx=t zi#eb_e{0JFkLUo8Jjq(aThu#{cL#W-hwk8&5#)7Sze8mcBvP%&O6Eq7<@xRgUStE; zNl~4Jkp1@NtY(tNb@vxC0@jgB)7O&Q>v)m9T>AGhQ~MDCyM&*`qXh1A%NFV_l9Ny( z1yG!0t&?%_ww!@OzPxtwGe8!`iC3tLXk;r{=KrYDPg?xCh7hwxnb;i^ug@aH-W%pK zE94z^+(V$Eh<(}!7zCcQ7|a}NLWT3%cU_ya#&N`g*X}*LMmGM>VH< z7N>6UC};Juj(fWl^TkCiXY*pocw#Ce{EfYsYvwurtaDFNQzwJ?tux zxYK3$&iE5;rM=?*mC9pRYm|-FQg%VZ3xzTi6*>9$N}m|9SvGbG+V@bN3_2-W&bodJ z*%baAjpBZ^JT`sN#an`J6<+0PWyGTH#dgyygFlVfFk#KWx~@UrufzFNb2kB(|7&*& zGk2Acx@@koic0UqBhHrUA0FP%3!<9q9Co zS6R{wMbQHOBhSCd8{g?*iWPg#j9NGfM){_1QvObcsRX5l%h%vd>vMN1)r;B9Ex);7Xa#fARi;J5)kN^uB~_ zQzpSJ5LW5xZAIsppAyR>d`Tx<7PueAByNH4x@jc~!IaVDJPf5|H4W#RFLV@T+~`l! zR_GLZn@Mrq|66P*gATihAkJfj0D}EJ&cCbatwK2w>Jq~D%EJs^fy6lI@!X&Q zNXjZg4Hsd$(*0P9NV?wb*_EUT`K?hFZ)8o5b$EYKnHha9 zMis_HEVH(ZqjTpIt{Dt&g{qYI`$>|h_=>p&*3~jkgsw|Cg(ruhDax%X0FJS?c3B!t zg@k@g=5?X{)iCuiVSqz1Nwc)vGlacEaYeRtnUrv$U&81RvtR=B8G{37QY#Tgm4XEC zYbuQ48@ls$8ONwendsjQi%SboAic?7E;UF%JEE-+JbAf;=xDR~ESVLsTy49;7pF+` zOP>sM`T}rvl`stYJnr(srwp1@ddv_*B5afL>|}LA{VDlHXL!I_(6jxRi@z2<6k$ki zJ76WbUo?N_G_Tln;!2Ny;wr}-DW*MmtRT6a)<|IXYT^*SK2Z}X&4>DDLrF1m$IFhf zz?VNb&QO`6c4$)mJWZo}2-+N$2v8jb_|Q6mAo`S<+ zP2nmbd64ujleP|142lcY_XH(oEB*C4wZS~P%AVhbAB`Jhs_&IkQ)oQx*L#G% z8!Ne|2cPQDfE8_4{z|KXOy=}G^xw)V+UR({MM{S1kTPH|RenYAZmGqp`eLsA%ok+w z(@&J(g!%T-)-G$?)(XTSw~3errR1s*LeNf`z2$;K3}v zfB5ZHz0E}6Xyy2RCml{G5XO0DG~+z1YHmU{qQJ&9;i|Kdizf|DPyLXK*V`F3RjcoI zpeRWkQxS-*jIAWv^gcj3(2jAa}8Y1M4KvErsr9aL2`YG zX*52H=S?BtNv>4%V^vR;7^&>+ELKO7ocbz!cGQPaxOrO@>=dG;M+SprUw*IdP7eC@ zZKy6Z*rI5QFx7lon*O8C#U9v3F8sKDuk+ZpM;IQaVFJAJVBx_>#-I9%B>Yzvsrj-E z{f8j8KK?VmotGgrQg}_3MtG5er!9EAl|u4N$+%`O*LpvpTbfJf86XT~qqrO)q4f8u zPs}V~uYgSxE7?5d+#`7lq|+*dhFyf+GYaCbA&~As*0-S2mxz?U;E6V6u=yQYw{T2A zS6e`@$CJ6h5yF;I2iiHzF65;OY%)l1n$KwD1s!lgjRC?#qkVaLvIxPA=2qp2=vXUr z0X0>U<3+ya2{T3)lqXM4uz9(5Rf~M2>k6S=Mv~B+_^4B+@8V3x9Y}lAZR(yw%-bBB{+?U@CZ z+bz~B50K)sB*#<8cVmOsC*AlWHD+hVrsUFX||;6-Mio_)babGIL3Bg2NB{Fz|b zi8*+KL@C>it-8Ha@K&5z_tl1(bZz=WgD&U5^ORGmrH#jKQ~DU;>ffU8bvWhOFO@7Q zFC+|UR-y0CJEB-kt`G(eC~4p8T}E+NB*uSflouBY#fPK3gb}1esbPX;HDD~ zyiX(*9%$7oJX0a$;<1+;09rZ&JyMT=i=ZUSkt7a@;KF?Lw6jU-cb@mINR+xalu6^0 z^8P(3U329^iblcjcBaDQKJf*|m4wLQ6p!rxP_M!(`q5Wa` zC@fJylpzz7@|2~=+4hqo-`^p~7Hs$7fVX!uMa(zW!|>LcW}-ma^L&3d@+NeH6s+^O zH4?F(pn174%WPgFQc|e2gbDJFFA`(}Lk?1Jy@+N(zY65opP($7lELJ9f@J>#yYSN> zO7P^QyvHOI!6@1_YlC?5FAa8EfTle)d z#R>|~=4DU}!-^Lrcr%G4td7_%DyF#-rI}bFaUe)OF=D2K>ct{wb;LBGiEuCx18`({ z99@8uC_ybCjpN0o{(}3~d;-GHn!maa26;N!&3bq|XPBC_-5Bfez&_m3TsSVnUkSEN zO^7D~0wY696T66p2F^%BiLmGUfgP2nkQ&1K(wd<>4#^xA;eMuR=8=ygWy1-fdzO3Z zf_7^C4Ns5iu8H22E0*+pZ$s>igTbzj;vZ|00Ba^dza;BMOkd!r^>Ts-px#;MuLE*B zn~$0d;eD{%Ew;m~q&p7TUGBt`>=Ou)p1k&-utQcE`i&|4k=KGVg^Abp$^I5SpJ3AFhCn;U-YWH)_0XjSXYH~$bM}?F z)$xGHO+7X^kS96v_yCLRc+!g(#@x?~#?u_aMevll1?Cy}#Tf0UcN{*R&UP_`(eoh_ z-|85o;eOcK-QBR+khv@Z+b*^0G)vLH*cluB$zRPeYNzm{eH?x(;1@q${<~Vn)9{Z4 zQBg{<-9nK!xlF39`sNF^r9ZBR7;UhtwRg?eEgy>qpuTS{)6rp9Kmf&z-`(=MaurI9 zcKX1f`I7#szgi+>X<>0q*gor8jBUyxWW!&FG74-LAkvinGkN&X5G6AuSVrC{ZX`7F z?M-ec&Xh!u+<^u&3pkR4#ok>Ph&Jy5Zoy-;X{8D|3*|>A9#kt}t@FBr%!o|{zUK50 z{ynMPw|}VwEro0(QiHK|_cRwKAbIYkkzp(r_Wff$^kBO^*V_;JZ`e>ex%}Fnri#C@rwbJTSL8-`q;bIJpjOm1g19fL79H-4g-%ALoeGG^3>Yb;xyo>478w-hXFE&Zl!tK?{U5hj8@OWOd;*~rA#Uq&l-E{L zddgqS-Q~QbQgLAnI02!JMD7&Br)BuPkZt4NHds3py^yz0Ck}ILsJYtrGI%6Ln}+Yt zH|U9D&`c&62+_f_z$`gY-;$a_UFQ3!|2wW&pDrwFJyM9gi30{ztm%($IN{il4x)@e z-*k=`gY!G1J6Eh4Xz}xIa*+W$tmBBCA?!#*26Tl&teX&|F*dc)xfMda8=UztfrKK| z<{6giQBZ2_`XQ^gj2hrHstY!7hQS?TBTp5mJB=?s^G!lJr$*&2k>7<(>KVr~pJVdY z`UZzg0cHvMbHZ4%2_bKpGlG4bAPgRyiRGvh;$JAlA!cWX?w^_;PLMmE3#s(Sm;RGi zrpIe9xxI!pnM?`ZM$_Kw)br{TFkh3)6S8ccExs$~SjYNUiq~47yQz>ZP_c$3^fQXn z;qig#)50DVNDN}=$zd~Ap-Yo;DcNw6bGK|18L;w@6HF<4M@?%yS%Uq+n;g~?sdav0 zFo)bG*wh&)`S!-Y^qClab4nT0w130&J8Q_ssj&N3u}DB*=BBO>$XfPx^6+)`vC^&u#7l8n1;ud?&5K%-w(M5Um)N} zExwF0SHZP7h@sNfrv0-?Qq`{wzd|@5;Ex!1t$gC{(%~7zSu|juZ*1)kZBT(;X?6rH z^<5>%eOYLYNG;D?&|_BzZJ)`kR*U^$qBKt2KApivCJQ!Mg7^U@w=mw2@GCSEXWqpJ z+A|-PcapK6j~A)N*N$3eLdn~mr!B^kwb#1$kOd!GXufu=MI!o7RO_2baTjCwGE{wpT=U7gfz+3JLkB z&3DUZvw1q7Zn@C3m1Ze2yaCSA@yA-}bO~4)ZM`>3Z}();RcFL@l%iy@@7+Gv-KGIc z|IQFf(<*zpBMR*nUF7Db)fSLZJ{5lEa$Y`o+tOR+s{FIfLQ#`Wyh?f-xBNvw$tf4v ztU!e~Ia+U|{(0o`&UTc8puY3OymXote6P8;O!Jpo?#PYOLz1ys^XQ9C$9k@6QQhO5 zigs8knza_om9xl3Z$Mf;Y*0T_`C-ja-0o>;b4rK@5yc%)1r;@yPxf5p@j=YUb0l%P zB`1Vi6w*7@Fd_Pl8Nr%U<}rvG<@W^zR168@E3h2zm~NY@B=L@{zSRiYv@^`X{d_h# zU{saUal8@?%~y-v%m8XR>#9s9>n)hy_)4pLx^+)<1)^DO%Ee~cayG$bgM7Li<;Fh| zyIQ}O1n!mFHHQg%cw0Hly7rFO0e8|0py?~jlzJ)ThrV9dZ|W}$ZFx`PkSfgg_A*(W z@<}JQ$KFoCuzol-hac$SQYu2fH%Uu%{wWQM6v9cP;cC(a9%m2e-h%U9Aqk;Pb6RaeSWK)q)J z%q(sSb)}pP!~{mJGT6_+J1U%&%^Z&QXG85B?!W^%dZEI3M5+@kj`mS|DF?h+Ro!rz zWN7AF8DA7(BhEj$gr6AtFlL)ZCO+?h9qrvIw^xPpP*ln1@>D4UberF%yR$6XMF2P1 zuiEfEHOvgy2h?A6mW(WwQUsRG95(itbET&aGGRAE{==hkjt8in@`jUn(OR*LTBXGWcxWH__hQ)WGnUB zxjC$a`P~VHGIR9qZ{MK#$`Yq81#m_pO#KBV z3=$U=iktp&X935H>?y<^F>4JA6YzoaTSzNK(Mf@RGHJPf(_sVz%oza>aKg2*5fZ12!p04H`ifgV5Qj;;yZdQWvh zn)pAT7reEl6?L^uLf51t$NWb0bwXm2&n2>eFmHbYE0TcE@cXGQQ&`S!n1_PqrsdO5 zDwS>^7JE~V$2)9WSXabyt!3N<-0!%LDr&;RgOfpP#+mAs5Eud(oe)yO?ZyF2k_Alq z8qEKiR#uFb5|b1;h9;}4J-RKzM0zHfg@54m9Xt$hc#Q%oeYI{X19%XG-=>VY(t;G6 zev9l>V><8jQstUs`{zGJa&597=P;%j?4qDKaM>zh+1r%@_7L9@5iWD%Iwcr`+YOKg z<9+0r1HI%TgsrNXOwDu&U-(-vll$*b=tvLhP5)+n-E_ckR&M9ya@&UmtPEXE+7NwA zxJt;F*jWX`y_=LX$*|1^_KkShIfSuSd^&_n!VOt$$HhZ}R7>mf(M_@zhIg0-g13Ixrpp7rz}Xb+)h{cqwc8ZAi%*-h-veslxh!lz0RfBJQOH z@q3=S%`=N50Z@65oiL`j@X^6`w`H_JsVCrDLWoJ^xmT%{dL--lZ4~%0$fZO}p4wv# zC6eZ^4W|PFQxp!tM0zn3y8(B+yKJ1BtntZ<0A+_6ld>5S;%thC6Lf$Rgxx8wdF~(H z(VZs1q3_OuJ8y#E^CR zC_wX-#7t~nlxSL2W4QB}g)hJE^Wfu7_x?MWrdfp^()YZi%Amg zwW>%$e02?&cpELJUU19K7b;=BIvULwC|cb-wO2ENE1qUx9F^8`o+i_zs;BEoD15+i ze#%Xm^JK%9`LpE?qM>UPVKs-Q{Z~HKgaQS9Tgg)=4-J4<7gIF~ixM+_6VET|l_|w8 zCd?0yYAP)iS97$K>FKKsB&bcC;{dMCrRl4frg>UXMhm+5gC3 z1nWHj5-AktAv{RbbpUOlS^aTG!rzN6c#A(!p@`Ahwd8vOMb^|yubp+`ast7DXe_lTzN{s zEpGz-^vpeN?OF&Gy#N&zZ*ar&5Ez0O&p)K1Rc}$Ks$#3!|AhTXo%jg&Y8&QL%Snlf zZic+&s<4Jc=S2Ko%e}Ggvxsae;OAY)+=l|9b&W!~Kpsl4@jbX~253=%@&vglart5U z2xE3VLtz9ew5jq;+(Z$VR=g^w5}OKwQe66%;k9jd>$Mcpm*eQyLsH$iu`u$N#TG>x zSYIlOilFbwkhAP2g)%ZIYG;}V`yiC(G#Q%Grl*k^!w>7a8Qx0VqKFOVqC_YBTGgL8 zCi6CMXgb*!?T|6eXcuPf&8iQ1bf`rsrKFEecs6aIj`6Wj#>(IDTbe5Vj@h9FM?OCn zb>%&v0b|?jDJ}q$+3*bJ!xTC^!@Rk+G0N~4E<&&G9*$#@%E9Nv#y~AR{11ffC5wT( zWH5_~J(Ak47JjUPN2#WpEVCDj|N5^f2iK3Vhk9crcV1Ia1ctM0Tx?w+2O_xdSb^Ii z1PD)LSXF}_=9}jCl-5|v6&=d&-YV&ozM7hXkMBEI8gHv5R`D5i$at#QbC+>E_HQ# z{8>(P64N2~9~nv*mC=&)Co?2<%G;r#l7c@=p?O7Wt49j*ZAET(3iUn^a+LPQ%#3eIZz)KSK}{8f7OyzMX~ZM)2X;7Kius0x z#i&Q5cF#AF4_1m;5LQ(WbeN!r-#EGHFEO2shpN4(~KQc=X9N9`VS&aZ@`$Nd^|;L8X7zzak`KiP-C7J{D*zFHDb zVTG3GpN0?2JQU^e0Y5*MGQoK3YXLVORIAHt)2d=^M14`At8~uwL&w^+jQaJ$=HZ6S zTf-YBlR+jhwX_S)6INXW7pkagEB1|~1u_iJ(DgrGSL)^Fv_(aJ!bUeQOAk!7XERg9 zV^Z`KOEYQF8zq4(uJ*v&L;A1NGV~K)d=554+9}|_3DqMxT=!V;in9&Kr4^*6S=Moj zffe+IMcIRVh+3pJ{yQM>HG3CUEtU4PRC0&)lJw2JrQFp&{N*ZAh&|!j!1%=@4alJZ zwSf5n;DC*prkyC?S&WE6oTmFHzstZ!N8sXq*DT40{q2Xx9qTVyjXUK?ZQ$)EZ6aJ- zhT+Q4)s`%RO3`37S2=?$;|%0NfF<(rp&`43$Z6c?B^A9~6Xm5%Eb7mYxt~XR>S`Rt zQP{Sm3sp0 zj|?^GiF|N1u(q4->Vkw(Q$4SOmIAPVajAP&o1XsBJQK!7c~>v8YXT6b?FUKuSO#nZ z;40LvCv@`aQL!)%3^>zOf+U=LP?L^vvF2)+TTu@}n^>4_pwR)8mSCCHy(TjN~w|$JA4JL0LFMS^Gdkv)tq6|EN}f2B_)BT33Ef zVpH|RV?rseQu#fZiKg?8K4CXR`C-im3hwfRb%D=o{C0UJ=Iwwu|Lc`jK*=8sQqUR&1?}6~!5?HD!JR+AiXz|eX8FSsL;X6Zf6gsp-F;)G$&CY$l$#?Mg4c!#ELuU9w z8s{<`rqO@Go|+FUAZ+IXrkr2I=ODVdP0|I&92eT^D_|7xFF!GnrdKeb^LCJ$X9N%9 znA~g0Aqk@;!n3ofKq^Gp3D6$3wwJ}kU{#WpssE6Duo)i80^Dsw6`<;zfw^$S6f9RR8~R)~+{m3Fes--rbDYh0y7E8I z*{hnfeeni)J8(c#^FSYhL0d@jZgJ6)r1S2?K<>sd+aoccW$#SV1>Si<}A{dwz6Jku{^BUssyQWg|09SEy+k3ztJ|zH@rrjoC|>oG!(r)P3TXs^w7~ z@Zwuy@!rO%04tBoK$qAvk>s)f5li6zu~j_rTz5R`2=lWZb^h&zc7Sb+)KfUTsb6!z zLAs*lK2@kI#$-46>!aB-Ak^l z#25HpH67RQgUjXonHOt)E*<>7caBK`Loy!^H;=tJ@%K4>udcz1gKI>@&`ON{Kag3W z_q^L%9Bq5%clEDMUzkjF4S+9V>fe z@4d43k$EV}mX(l~eT-v|Y{yEZ>`m4=4x&OF5>ZM09!Kxb_xGRM&E@g99^<+m*LA<{ z*VCzRqU+mKbzk-)o~x_M$X34>zWc^?fD%%h8h}uFX00Ybt}Zy88=<$roeZ*mIQ9Qb#+&S z+3Gd@v&O5jahh|FOY1WQzzl=03PQqFW?2s(p4)7t)tRT}q~@0XN;T66(&x#TpWegJ zqMr>HtPg(YwY5vYb@9|T{iNR$5MeY1_)sjZKSfLRENB4>oZl3!e_(m!GW3quCK`%3 z%Tw%su`yJy^Y&Ufl311^Ij+7na#C3^_68{1?pWgwuWw&Z@N1J?D%IPk#n%Z1y=}Wh zguCmB)R9i-zCw+yh^twQOELm+8BG;6IV?Snv`Y^K;vPO)U>Dgc*xpjxByY`25}dMm zvy4%rbc|ceTmDI zvBRy8X3eCl~N@NSL!$aPR&CNzAzeVG2;;JN_sA zzV3+CDJM|Xlfd6RR8t-4Chz=9%=&MWjr&SMkBENe>kzxLewU)n)6{01N}RRB7AJ%D zq-<{hQT1gSIEQiiuvLRaDLd(0=@nL79j7+$Uf0JuT7JH;tXZ7)-8?!%YWdwCNGU}} z6}QwcVdl)@Tsp^fi$^W?&0dnT{2hgsY|N>A(}gYEGc8ZHN^o^o>E>STvqdj_aCePf z_?h#|<9*+xNsT7)%gZiB)AZd1tRj&6A@7oc3q|yHWoBCKNk>_($zmEL$+gCMirhp9m`vtFD=^tsXA!%K3XDPQdqH*o6nmDQ8LIqX3t zS&ft{C(ur#6~py6=^_J*$|#A%U8j-~+=LKp1Q zZpXWc9OpCM7TLJ2y<4Ab!Q-^3GdUTB=s@)yB4RAa?yWQly#GI)ZYO6E*pet&$W=I% zco`|&>~_%SUVXF*iK*Xnb3nQQS5EnsoD zgl|nwyYr^MXTLu4j8(~Tec}u99`8Z@L?@F@?wM|lD1W6><^X{X%M?M1fq;RRmvMhq zc^@q@yXQ8@m5WR)$ygW1-uUdI@)g{!4-B72PHu@HZL`M7+_*CseHFg( zjmUYc${e8etBvDK9=?x%9G(AqsXx32(-1nLF=yRRwR`Nm^u*=YBus+!|XS|j7K+}WhQW4u5 zgzB-ijFU=*S_2EPJTU|O@OijKhaudxCL8^5c%uC#)EXSM=dE`&GGdGx;lmP#>xJs& zk0zEL0<(M6qcluDKS0~ltjXsr(#;af7AwcNJ9vd31YBE7+#OZVDs?`|14 z?dw5iB_7Po{ghYi1a%IJ>OyoMcH@F%r(b?vZUK9f*-))nk$ItrW@ZD;F}9?eeUrY# zlt#!#mpI7tzB7eh^S-p$NgBn|AHm@>XVZ=<74QCuziZf%6Z`Ww6X9LICmraz^^Izs z^L~az6e_~Bw1)J3{Losm;5Xa2lM3coTr?rA>8f-YbN-=a2mE8$>}ZV2=A?`qNiW|L zqYr;-ld4{FvJ>gZ>EuiUucrCsxGNz|2$r%pNlLwg@D&B40zK8F+2o%56p4Y}iagOJ zeOu?Zy$c6G1SQFf-dK|%K%qr$2oVvSjqr}rFC0^R$c!!gbLvYGc}_{-V_V$f?L)9M z2q(94HutGv++r2a2$%->#h zM=dL!^*-PkAtDV6`@2r}BKq4y6SJn{v4<66OTPvCo9P#qIO`{}^xg@Ls|k&;k)5sd zwVZ$S)-rdmX-4{3>05|7-Bwo4`0i^|p|s?8GY0qW4wS&#xm&HSKfRIn!AyGNJeV=%MYCjaEGfW|wt}XGj-yGxWqORoF@bTK3 zk}ym(sl-EXrsuRE+R?=Ist=VWY9}aVM>=BfIPKjZv!pnddAioQ_B7bntP~({swZd)b_K z(MR5=B82bY^qwi>giRN#j9e1CJIbQd2~8~eUD@JX>9_QoCAC(cj05+xFM!XKXO(G|?tv#&*D=%1)wv;*XG?*XRUtISLM~CbeZ{pSXoP)mrq^1r|bNSJ; z>J&CR6R)RJDkF?_Np2D`5u@-= zq~)Mj^uYlKgGo~o!C*Ws_pS(it!7%2;HlY=WdV znf@8fIUp-fXe27`}1eHQXR+JtKXz$o5=^ZwC>eBo0 zR|~{Hez_|`IP5Fh#pptv_7XLgkVG*2C(Yo{l97?^nEThwAZ)jw5;SPcLC*gf)e?G1 zgzF|7myMH4DqeUvNokmmKhlZ9ccoK?&piI=ii~?D^YZ&nVoaj95A*UMJP{%wnGQek zOY>>7(=aF2JQ?KGj&eSVxXE=TkM#KBKBW0ba8GU3McG{3`dm%k+)|fXoJ2FXVRO;6 zp@0-=l=hgM@ue<*oS-(8U9`W8-pX-%0SxUB_tqtuCospS4M713G$*z*LhYRqir0$a z{J%f?gN5+ri`2e>2YC8?1=8epj#rlgk~b{Tv+HxZF1k!EB=v_xZOPOOK~gt*dKS%3 zbTah2LkKge$}~_u74`jnwe#vTiMKV}_-3h#;wxc3?Y;F2X0IX~QvM4>g@!y`8P9RQ z8#nrSj)f7i!pv%lvI?U7g+(5KA)9$n_cv%dtpz6g(%-z z!|QIiC38x~yoZz4zLbMrZTWz_%6Nu47WlYfAMTTst;cx3v7lohJ`VhnaNJjvB17e4 zPVTU!$PGavv3zK9IpvoYB??PIQ#kKHHUmO|pGb`L`PnOHs%`v<3bKjY{Cbq5JP2pe zqD@$S@avI-lmgCHMOhVRANIs?<{^eet0bI8wF@{gO^T=Ej52P(CR};cmE06%oOcxL zn=#5H#*!Fpi%ABzQp^3?jsemZ#EttuXScK_I}KxEuajTD-q^8fYP9M-mW9fF=@)~vZz55j^421c5 zo$=uZ5+xLG7Zyy+`DZiLCDW}!a(`F$O;lWx2}qF$(M&$B`_SJLNF_Ef!S@Az*L^J@ z*ga>-4a=*-5T7UhAynaM=1V*t|nh(+S+$ne6TjQ_5hL^dmBwunOB!T|YQU;)tC^aOS@r9s3f#p5FewC88#@`UUBt zrcQ*I`0sfVMPeFD5UL3=>EhE)Q?pCmMM*uHEXzUBrbBJzKBiyNvLnN>fLK|PF>EVS z=HZ`VUfq3@d=o#EMtEVN$#|iL~-lxtiT2oBP*M!Bf|){2yV~T1C$bZGWFe{Ml9fd!YFDJmSx3?=#%E z?gs0<%#&m>nt{7(NGn-n5bts}H{yPXvbn&{TZiU}uKzu|Z)c&Xx%S-P5K6NZ@*}@X z!J>jwRx913`7|BX2(7-&19R?FtLisw4u0!(6?}Xx@4gEqlN!l+N_BV4 zH)TO>PraG$h&*mjdiN_8CPUIUgebArQgyTp8O@gbvc8rV)_e=%JXR{((gVT_tUkkpGp=p`7E8xiZ#enPU=NnH*{OdwFyF!d2+h4hV9-V zI7SyX+g)XNI#;=@FCV(xe0HEs}7u+(7*5NOr4ougL zKmGy9n4>pctU6!@(|@br0jt9_Al~G+K56^uxAjm0?}RG7Z6N+0mbk(7oAw7`nbpT zS*5lmk^*Oy1fBEkmZLi0V0|peSy3-)Z3@5ZY%Qh_{z52~EXA!`t;#9l(=eQ`&rjdd z2R{dyS!?c>t{aPi>V3XyZc*#O$qzdadCgcOGe3mTT6b*C(tVj`gqtve#44XZXCE^8 zW2k3bdT>Z{fOH(;2jw+1RPbO~k{7CIL?4&QkTrku`_DStjdG(li*@tl_ph8gIG3^LI%bs9#&+bwqpJ8^_u(5YOREDXTbZ@oPzE2(Gx# z8-VP6}G@Df{V2#Xg5&c?nXPq|2a@vImk(y?h{W$3Z5wDT^uk%%-uRrZkRMsSI zMRrV|#wzs3miy-ICll2oW-`{pvMB#eg3^?DJD+7fz0aT9vcX&(YOFeVN2ycw;e()_ z-G**KDhX&INVnZbxDe$xd)m|xu?82QK}!vR&VvzSkp4@>z+Du&UwyA^*868-X+ zy1XVg&(XE7o1n~-%jKOPoK0rlx@Ol<{6zB!QR=HSDGdenzjZlI-n~RNVKQ4%zvm0@ zV?^F8_|mgyXZl#FPGl{d?lbLxs13dQa^>2Qc)4lV<5%6hEyscR{zgRBVYKE03YcMa zt>QK;KKZaYTUni(aPbMzjzkYX;Q#dExSQb}{2)+jU}4aJ(#G6&lhUSJFF??p-3Itv zXK#F@gq$?R$}OZ-CiKky$Ql(%9X zqs{tsE$+t2Wb^JDq2clTUufK9NU$OCanls^H=0*1uH->2XP^trE;!dxozwflbrA*6 z-lO>Rev9Yhk*i(#bPQ2-RAb6;9}QDxFI2tU4I#BOndX0qqE&wI>1XTeuEvSMJ%r#g?+u}G&Nc)s1NShS)xU7L?tH`hUSQk?M6siD*Z9%$q6z9sQ&QlV z7v4nN($M<+v&^|S@LPlAPHVd^+QqWDN@!C|+^`l|u#U^bhwalt3Yf29*;g6j>t05W zr?*`M*2CzKU$?RbA?3d5H%c+-*@J&2QCD1GvhfAIAuP=&e@2xHhIDxZlVQl+z=eq* z>(75txnI-?!gEQ~NWP;FW`!GXx>XLN{W4Mu4^{7OQd)t7+F@Bt6d}oFx$*1(J7#K1 z+<4N}pWPXFG@?h*_uRK=U=d;wcKtkTO{6aAGf zTbeb=M|A9#T_yQ)>F)==noXo`RNKyWpe?+=2Ha1SP6EZ{1xDr1eQ0p1Rb&pZUy3Dv zcY@8j4zoNAM8Eq|)J=@g`-Ifn`{0Gklw=xFWD=JbF49l}?@eT1W~aYl=4TGmCnp(Wjni3_u?nb%^hqsq|Fzq0hDz-Q&uTE`4ZN#;)_i zFhx38?Eek!u@|)??Y1MVW03jh-`IMB9d-W(w+qoQPv8l3fL){#jZiAaZPmXYy4W2E zdW_EYv6*#X>Gh@q@LY(%r3rJ)4&IzB!Hr1=8us`dKCDzFm?W1XN(82O8WdDjA`G*0cSzrZ4%~Z zAHvcI&;R%>Y^x7VcjH1EzRN<9Z6miuX4EMj3hXXxKTQHyO56fzv!889bS*fdT@fQ| z$Esx@4jTTFwNI1Bo39L@pI@gVa!B*1RTA+~oRm~wctX+)U7T&+qZ`2QjH`Q6+f=X% z4CzN-jL)i>K~ z>l*S^XVgZ}r_i2W^J~<>*ltl%S>c<&QNq)X6+FJkG3V6XbERqiU@uEkV|{ z4}KdG3vbi`Cpe8J5bEF32Oh^7T;v}2XLV@;{~MjVxUEy;;p_nOLCylMoap$u6m3_j zv9ws=5z90>()(jrcA`1652_=C`lWEjcnRLTtbn#fxUEC3m@TURa-s{L+tANvkK4uKZS+Rqx~!7=H8wobJ#|cJZ`@>rVrt-C?zm#kVHd}YK;xE zS=dlz&p@JTdyqp4WDwswe)a5Y|+8PkYCyY}{pvF9s`sN4Wwk&~z#&>kk{HKOJt`@fY zznd}a$9aTp3HJ*=c)hqMvj9o073PwV99|MB6(q9ApFE8f*nT|qclxBhi15z9yhd?s zF#i+(co3U(D!4tXNV}T*FdyV7Jhti)FnCW6MR5ti#tJXURMSi&kaKLI&dmp#tH@W2M2L-$D_VE?n!kMyg`;B z0c*aI|A$X%eI`XfPZV+d_s<5S>drgnmq zJ@hK=PqgPGlC>;rOic(}hND2ySL$V!%BnBCR0qrD8|psU)wN>dJ~FOU5;psW z}8pSFTejO%&IoPN1gP>krh*&-!L|rKtbygUkmeb zH*OBX;$pn8Th`V}*y-lyv*j5&RjXURLWV{^=gh`mmnYnMJI`k;#n=KaPxQE7+cTW* zo*Jy3=GWX}OrC+ne#f6%Mcsu}Kj_Lq`c5oFeEjB=i?&MB4}>pyEgpEaKJ!L9l^DM*r28;Gt_;Tmj>TZO=Z98yJrPr8Lj< znM!Hci4@9z9nJK9FF<=!VGmp94PcExV=xr|XT$Ju%oo14o3<;2E*vr#npn5AKT5V?AR$^!;|C; z{p*c{ReG;1%Bhf|IB_N4IPD;j%TA4T*eiJ2Fz2VzM z{elfjL%h>rVnUCWM#} zFrwW_VttHW8U8O|KG{a}!4>G9?W0TKdxMC%ewOa<(McxKIEy~dIZ0E6U>0FVfu*i5 zy0b5~8)aww4OquqTh;nJHw*mKe}A0)Xy1(cC0yLgx$`@f%1h-3$B*>f8_@LAkA7%6 z=Gqui)L2zwRTlFZG09>f6)e=ik2#R5+j;2iN84V2okU zSCH5I`f)UtKFiJOd9uPQ6ET=CZGYDbo=8evu%j@!8+`iL-}PeDE@52XsYQCO-LI4D z3Y}&wzTLTtXfP3g3H^xv^p0z?xD~^*m?h~7;ePuCLz*>~kQh1pdQH1)X%ka{3rv~6 z*RY?Qre8!$SoSz5#Os?HX_?-=FNgA#&MrE&XR;m+ohg}QoG@$b(X*HvS`E)|PiB&w zD9X-jOusz3SLgAZQWJl?UrY%-G{?o7CuGuR=z=TOpd2gpM(TJE zwDP4aJt?T@{0!tX$0~d_j-7wlkCel5{ModxXy=PUp8x!i{EdrRgOy8@&|^v8!6^x; zhiQIWxxYkKohxFZ?AW)kaW!B3+Eqzr80_;--|C}2j!X2zVj{cSrC&L1p4xQ#I_*I; z@)d3av{c=PYOa`VeJr*VhYXE+HbMCqSc4_nXp@7g6Ut5OL(dJ4{bn6;5iauNV(hPXP zF_XoDYDp{!z%rG;R^w6wHXh%638`APs#$(Ns^w>n9CGb) zqUo+mUWEZqUr;4E|5Xu3z{ok+>L!9qCL9gYUlvV}eb8N6%&8=T8dKw^y9z_2wiorw za-sl`zSqM4tHr-t|Gu2ZtMe}}MB`<9zoa&iqt8ZbBv-ouF#25hPa;dHxOmJo z6g5VcBx}L|x!=)0NJ3z?d|o(pKPYH5U}&AVz7%9#0(1m00iNdjW=SlOsHTFST95Si zV)7uR?uSG>zFb=VR$ceIs?{nf>gvzTN+2xZPzdEiZmqggs%O7uWR0)@$f%s*Nk5z_A5@W6NJ05%#ZU z=w@1f+zW+M6jRyVS<**rq!efpv|^I`*}FBxFHdKkG<}9WBDqR{UJqT2ldU{cad|pc zx!QFJmUO52diUd?-3e4-e#Bpp4%f>`ZNJ)IVNQ2~8FnbhU?%^R9Tmc2{;1JgnA;v2 ztd|6B(g$he0>BH>e*|!bR$({*x-|?7oSk6EC>&iha>A4*DLU$N(JGM~XFhM)Yeiny%d+1ncf!E{N=K}o zzjIn9PykI*6)+%vv?)r^SFUd*LPqiC;a&@saHYs=>q)YtK-;<=m_uhQPN(xx}JHSoL8WLq1t5e|>MMsmzZ)k=Yw~Q5_e#w%k1C8v~x@c9ghfU!F_H z=63!aM}>BV5TY{9HKm{|Vu2M-yWlf53#YmXddd)A70-xtS@=+YB9l72<58dC1GP}g zAVdYfEx$o7(CC8^&`4@edU-^T$NvKVQwv8S z(3GOwKz#B`8z|w_YJ|CMpU~hJyay&$k90SArbiLBQ>r=385?Qe%q(nF;lo56*Le42 zStAUH!|AnL>Zq;(HzYuG_sqZaTE3fHP)NBT=ZAs97zvUoKG%Z|caqODg)}F1T|c_q zpA}2r8wA)9-ns~;_-NKIjUOq0BbOpvc?9nbk^H_?Mnr-^JOFqkxLX;>*8;W*iOhjt zePrAxbOV0Ae{%Z3r7oyQK+j{AZ%?9^H#FF-ZCA zOCq#Lzm>hfPRlS102S3~JC=8MxAf`del`1_s?BZL7jpAptL0fCaSr$!v85KBkK;4G z18^yj)(8+uoxL_z^S?RGPeCZ6**uobf_U9=<>Q{2G_Xij3@;?E#M|3rJiX<6ST?ya zalna2E;qSAW8pC0$GnQL7zy>s3O%c(maV!y0D!yGC-tYds6Uel0H_ogTjuOLpJLWY zzsp@Nk1I@Vo5=9o8Bm#p#QELbrqAdQ(@$%`s$F9X{D?1Z8mR?9r3P;gsJzMetJb}- zkDG(Y`vYtKDOs4 z;Q)ZI5DS8tIj0r?WW*y$X_BuJf!5Dx-ckIC+|qwAq@S2Hn(doKZ{+B=I4G%~Kru9- zaJtW_{;~TqX{8UJKG?-pc|xIBWQV)}bvV41Q0aIwR3IvU0=r<9=Gw6g3JRUeu<>%v zK()QO`$yMLB^X?SD;~kx|4f(VbZz`s$Zyk2=9=5@c2FA8A`q(%O3M!{-!3o;dk=bq zj#Fk-G@bA#awf0%bUD*zs*lf##6v0Ul%WF?gA)oIkbdO0M7#ytt&sgyV7#}o$&97FP{ZPq|H_tk&VK%R*mvWNWW@){ z*LJ7K$4Pm56#I;+8@JC6tmgh0A8S(J;07R`l-$IWPZH7GVl9htS28o0y8lDLr=4XxukulAljeeg#5D~zi0WGfrbnoUsY!ZyhF#h0*Yd*OmVg}e&fP&l9oFZU z*LVVKJt(J>Z|T#vXpvA<(9@JF0DxqI^A6$*ftfRDE=VaM13R*W7+ee-^LVMHEcEqA zs1Q!>%rQyO1zJr4LQ0rL5pbK+pcqW`d|o*w*!^}K+V^@Zrrx$qu{xc^_c9g1USJ4T zk5_@j1-ce}L;|#%7)WIk0Cj+htMYDX%WhZ}y%Lb<&IjMKMA_0kxVgwra%I}{Rf{fH z%jrz7lmOA;1N{3NEmAu9@xzb5pqikmTD#Reb}tkZ!@Eg95l%5ETf zdrmP=R5J_!rKQO#iZ)Q<4Z-@Ky%W@7pVC&1+n4WyNaX9pImmrCO1To@OEJyYi9+a? zC5>~ofWjllHE=P13K&{%KJsbB3fB8iDd!q9Fx;|hK*xP}qi)eJsZRjhL(LGS7{+q) zH&5G2Q-M>-A@vBqZOe;SE&h{mM0d3qxGsUUy!`4KX3QL@LWpd-ki^?@C`wXM3l^gL zgy{3n;v#uh!rR4_2Jq8Ngp=C~$ww$yRC(vhu=>bti&T?8|3{K6InGGe>Of1>mvu&LY{-6n=^#o)~tO_&3M<+bV z(Z0)&RsXZ)-5>9?BmxN;ikEAr1%N3q8ImTvyz=1mNOaRVHTUs8>_J32dcR8OYf2}s zj4$%!+B({Ecvj}->-C*uVQmMu&A)i8+zncVYjXD)tm2k*fbD$RS*`FwO&*7h&UQX8 zkGPCbJGYEj)~ik;nBdz46eQcwLqp14>uiJ(1Kssn`nWVaho**X+9yL^aDp%6*;8Y& zDy(s)sEBOuKK!8jr@#h_d-3wIqCuS;d+8_VU=Slp)Nre`bpDJL)wt6~K^12%A+HW4 znW*Sgm%jN4tq|!^i5) z<;^cs-O4s%Ak00te`5YlQ9gO8NqrJ-uq zGDn0CkUGNPJ}O4<6GaAcSZ=TjVnnIkU;8dO=Q#1Q5bRS=Fen-XsVkCy8hfEHHxQNG z!=WAHHGnKNNiNp>+U)6GJQQpCoW%hJphd&?xn2qGgxYWNHd8Ejo)XcNZx)6d;Rzmd2EhM-?ad%%?m?DiJsa4Bj?^9bu$=pK>JIMf_CXG{0;@AnX7=SvQ z#tDD{OGJaR*<0sGmn83RAI6%pYL*La!2a6{As_V;bGqB(gFBu`K+t&8`NE##VuExq zfmb(%X6uqfiZx{4K;(4@AiQw4JQr_=o4fFU^PR|+2K6YwN~oLMGcqxl$Bx{(?> zu!T%=r^uZkS%0H0vq}nZ<+Np3hR)WcZW*CO%(yu|v{KMka%1`?`Q{--@dE4LsqTDP zsQQFGc-fY8YnCTNzXq5yr?wR^q~;0l7{ufGHM9&-dFLC~*7LRB$mk>-uD+ZtM7oJa z<@I1cHp{C{fULRUqRPhNh2MMA_qiwXI_3V$HLc;lA}6CAf4k=nCP#{f$v&PW^*uRc z?0PBp!n^E(&mnczI`@Z9Zs&}IJbG_@(krGu805p~{imTDr&UX!19Gbt9@)j8c&KW0pU{oxs2zD;dBWciJOUFM&7 z!Fh=hD|=|-NHm4R+zHLHppg;z9BgXznSXpiPmVIkH9&>*^YHRT-R0(^B5ZH9GB($h~~(XOuB77#R={3 zv-km5@u<6IV|qjU^5lYgwqOO>`rZ(`%;Ccs`FI-khXUv7Qsw+avSB7=_ec>FCI2lo z%4(=OQzJk$swPMzd-gLhY!$lC)4BU`wWa8@Ox&bBkt}P&J!l;{`F8xga+l+&c=j9d zdy`qDuc$x2tbAu=kS4AXCYOQGUgIWe!R~E{LS-KlUR_#VzdR7cBA@Qy{(=rcmMRm9 zkG%Q0@1DKM8Yn5>mcoNu zscxB~h9Msf4QZ#_u$hHqU-xj~R)#?bKVuHr5$@~AxILEFsPFjLk-GDe0s&IIM ze9bI2pgnc+$Mukj6>x^S1MR=Z2FOz$SqJ+N6;70^t0ezVA)>pi>^kU?Bl7-)tY#ZU z76*!=9^!E5L$j>|JJKTHNrQ~o6l_{^-nu8_Rtz#)xYt`o+HVwS&TOc>g`dDq>vH_JDqcS+`fV#`~1aN+a7$kiM-!m2DIm|z-8S# z^Fkpnp2gkCVG8n8naF?CVWl%I zZ}kj@5lfz6pG$i)Ix4wD9!JZRC_uXGSOi2mm$*wFI)3AXt@Ch=ooy@s_Y9}AGTj1A zr}?)bbnvjK0l`=(T4YIoKbh($;1;I;cG+GEh}$4T&(6%9bKk-KN%o>1-K>1SpGmjr z)VT^`UNHoGpGZ0X8kdZVWUc{uG_q6ots? z)l5Cd?IpJvP}rbOy5I$|x%C(11uCK;ge5)l^Wu{(3cOd* zwzFTl=#@WT{Z}U^gAublBw(xm+9KGsThaFRZGH#Y(!F*!j-jQYLXhxWpE9rt@tuB% z?G(>=xkO4L9)?1kYi%{yv6}Y64~j5bsU|}aGcI;}^8CX9j`{>Cz*Zh$6)?r4OH{^~ zYDn5_6Yzp?a-$iy8kEhGj>QUHR&LR3LH5(wh7Xya;D2A)U%ETRA^`H#vFaBMS&|tW zD?6_*z#u&Nst}^)Ty$(+MpxcC{v|flG8K>V`_Zjqb}^A!^5`UDB?}Qi$&PLPFzYiZ_#eH1Ok3Q1pMgN^e?(ao{5S z?8R|`e7sPa%Sy4|C+a@Lt%g^hbwCjwf@tnaAmQ$LIb`4Uw~qd%ImKWnc^b_)OT;@(*Va76o9g_&}0P^8$=UABi8 zY}-a0KVWUe_;M-c7ZOaV%ygv_&X1K__DX;GtwjrX9=>&rQ*d;XS66I}SS!rTg|r+` z(uO&7KQ`=%L~7b=i9a<5;Y58$lo_Kd(WT( zIiHTVeh*P+G}x(m&VAbNfn1}J(V=P8Imkq)c)pdC^IgUj+{6?QpeB$#Gz+o5= zXk!&T)>>lSsx=<^UkR@4I{&j{-(gVFDi}?18sLO?u#*0DS2 zQZ;pWZQtlbQ9fa3ZW*$b83{11eY%?dl2-L}mP}L=_76ZBPn5rZJ|daTQ4k;~;ke2a z77}8@iEAUO+(k4^%cGB?KGEg%_o42nxha7f#~G-E$$3$HVf5_!iB#mgF4vDaZ@`*X zaYwdl)n~AK+lFGB05yZqvz9}W6BJ&~0T9ma%%DE@#FEtV88qr$d)TRg2`7xLu+drF zl=;th6IgbFnRFUTK;jZ0=$br4)_|d|JfMBTXoQ=BKc!ktcfKDRWNkv6cET5%G#>@Z zMdW&p%*z2$8`5b02lVaI0QPF2hIh!E7*b1Pjk1 zfm_OGlk9c}Q~C1avynwsXk@G6L6LFzY4P*xAt4a5UI3tf)% zoPa9e=Xi^-Sr#~Pz>8Nx7co_hGLeKHtMaxqUZJq8R033{HQS_ifJnhWG!Ag#vd5E_ zKD}-Z-6>_D_IC}nCUB+rIHrH#!CceimM=xQo=M)OGM5jq6ST58u@AOINba^09?A4A z@@#Z}p5YYs*YNSF$a^UIlI_P5-ICVD$k~e0Oh9_i_*WYZJF!(+zh>uX=nJs{gY-%;}#&ANd|a0M~S2WI&clpxYH?A-z>DK z@EhaVguvwp5$T7O)SI%Ig!PCkFmfgaU>v41#0R(r+?V$&+`2GtG8bDMzG5bS#P=D%=wDSqSa4)5d2H~ z$ITHdYC>+8^jY;ZOs2NZ)l<@mqMj5ULiF-2ZNv~i22(3uR_Me6K5oqjWq{2iB&sMJ zw>F>qMAfCNC_EuOJxN7px&Ar$pP!^{AiK60q^^94(%A+bT=Bz{`XIk6y~(YArd18K zYw38iImw>j5;CTLv10P{t(N`|AUpzg@YXY#Pf{9RfUE}v$k?g2Mn;d_D8ijWz3e1B z=fafM4}I^O*V&3|kUCLiPo6kX=HK8Tc~sM2EtR4=*HNd)Ty~e@b)cx!`w0VQ#mY*7 z`W51qR?tFSDV#UHj<2?Mfjmf+VVlC2P&+&FyVhl#Z}b^NpY=qk5iAv`9tJJBd9l?0 z8@(I1EpKm)G~BeWxzt5&2#Vy^a~!ws^)mcKyQ(k*XOxaEnpM(Qb4!W=LVQXM&^}Y0 z$fewaahsh`;+h$}7j9C+z25r#J1Nc^z#ilUP~7d;5oLbK$W;UdDOhk#RYUFl?f_#$ z5t|H*PNI^O+KIYxpyA0?%iclbjc zq)e^a6KTnrOi*$)6z`Ijz$8^T;P29!UB~joXOEuzJw0EHs865G2km z^}JxgPPbNb-`o=ydNvaCD6nW&hav~**#IVF90VOk>;-qqe9{?GD}-f|fmvL%bPM+^f5|#S;?=qb#}e zmPUKUdtpgcJ)}q|^)&dn!q@o#bDjh=PJL*iA!D4j%@v!g{AG;<3D>GBemr9_>eTgz z}r!8=}U5Bz>L^yR$PdK|`i^+hX5)zH^%_Ktbop0lF(s zCh-L7{v)eaK!qP(?=$i}2#@o4byvN@m3$ZVA*uW zbQ$Eo?nxt7%Yl41806ZVQu0`&Z`sCJhQ6G@e|NDy`~%6~e>x2D;r%s9?hUVtGrrnp zQP3TIucR>A3ua9m6tZVZYRYRMh0)4F0g0H5;YbO3Pul-qdc^wUApIZgv(@I?yt}f4 zCTg&n?C1P(k$_qErVDV2Rj~_X|C8R(ea|_wWYn3G!sXizZxzneoo1q^?gy# zw`c>{R^|DVB%8hU`7T#_fADJX-I6WD!`fw4pp(?|=sz2-(pNvP=S2s#`y6&0?e6FN zK2`{E{d*-kq+ed^Yb$640HiNq=;@HC%4asZEHo1-Qr`xLH+Kd@D?-r;8=G7Ezj@!n zbC3BYHQ6CA{shUA{PBHoMI3+KjMWO^)zJ8xF=0fTF}9hCz55@Ly5))_g-}kTpR2AL z5c0y4FQI6<>w$OwCl0k6TIkm*4_nZXv67100(N__%qd8G)C81%azKY)z^M3zm4lkH zXW-9w`sb1U_tF`zJrq@ds=w!Z&3aM?dW-K1qai4#g!RkXJx@#Y}9o@2waYt}HbCuNlyw)t1rsJ?Eh3cIT3dLtV`9d{hmbmtQAvSZc z-`js(5|+DMgE?SEo_rET-Ogu~|7F2!i^2g*v}*Z@@D#)>|5oc}vx#|}vt{kh#n^XT zE?OKxC{`;~Z+f5*SXYCDq%N1S)opF6h4Y~@`MVFV&2zFIb$AGd&+QOBUsQQr z|KNe1!6ue=TL){UNYei{*uHUiU4K6r)1{||!y*@6F#*WWjr=2D39f@oJvf@U|8J9m zVuKgs9Y4gihvTnw)W9)5H=uU-l=m5RY4}`;h{Km`{T7i*phk3PUpcfw3#`CjHsb~f3BkAOZZe`6(n;d;%1*i6wli5SFaAr4?kq|qC7|9ky1e6Vxd0`% z)!8^-<(zvI78q<+;18Wcj}<luKn{clFpB1vJh`;Q;V*XD^GD5l+Kj7nv>Tmbr2 z$pyYq$)qUnrVt;fiJ&OWQ`>e_kHsjf$oI9HD59bx`UHQBRZ=NK(SGs z|Mooi1t@xM^I7*j>PKf`9nt9n{x1aI9?2?vauUD&=Bg_1VNZ|1fdT#7buf~^1WImZ zBL10xn^LmYncFx`)kRjJhN3H^;S`i(8Z4%DS3l;zCs#mc%Izk2W)4>JPRarPVNvq{ zXZ;3J(LKhPW*6!?A|@^|A{UFfioji{)*F$YNC6dZb-Jts|3=A8K0$ehPe*uBN$3`; z|9Jad^lo}Gqs=o9x9BjXk^jc-rC#m-eUXpsF94_#7@Pqp5%aFgsTxe_E6s2+qbdlc zcwFUL)4eLk{{O2(&^s?mwDCQqrgWxo6@TNoSq=sxi?|y9cf7OJYhWysXmbI}5%AKAN42Xa>rUK9+2XDj7Sl4*@;p*|qeoeOdpB3fch6^gJ z#OqXn}00ZX+7SHkcA z4IeUA{I{I`4V68UABle31|2c$-Y=)Mf5Sg(by9DhfQFa}W&kDI<(dJ0Kp&X7>@B|K~G*9se=Hj;ZZy5K#Kc`2>;f&>!N`8}ATACVnFKY^_tDS6 z4C~eL$zL{+CdMHbrlA81G6JHTJdta&fL#S`WZ+Bx?^qd6%R_i$V@K=d3*u!Pn=t#xqEd542A5w@3~=#+`sp-Uy}c{iQQr<<-pHU_wkBYinlO+{ay*tZk>Vwt0F zblxi;)vpDA^0!y8i=P=ZkoX*tS{1J$&EAX0zAOa#jmwKau9=YFME)M5wq`C!WY{wR zQ+BgIV?K4wh~u~0cg`%r&$;PXK!_BUv2H`PX0avHpov=g*|9lu*-4BygP405_30ff zY++rCg|6nCG1VH=R@kOZ3<*TkVH4dBoVMGkmuM=(@-&CP$shS4F6yGh&{J>GV7u1@ z(&YcCrtQc;N`<1JaiPLz_ZWKjm4Enw#Z$=qwPx>-vIg60FEDaz#vNlF1}HXZiTbM8d8ne7x_aI0U` zi@9!zID}3zRM?p_miuNzslP&$1jZ2 z1D|)WawS0|(I6MXO3y{->$o1tc;+8QEj&)nJ*Yo&Ue|437(yxWHLL}fT@xzEUY+fg z3VoaPjL?h*@zwdbsPdVCvF#e|NabGK(F92xpJW@Nr3_D&*7nJLED!J;FU(ycu<2l` zqjqof(VQtrr*cNjSuxeG9*}hX3>lhj9K_EH%&U@j2r5MxO)z7QOXCWSQn z;BLXUPH;6|%$8t=*W)K@?!?gzQPfVJ7(}(qgm4lxu$6!C|3ui1UM^Gs#vNbI&ikWt z?9%HZBI9+nBNf`AljXnT}AMDH+*=d;UeyZ>}>hm8BNVya~3=OY}_zu7n? zP_JMLWu4k2|Fwyb3MxPgGiZkS)wdrTYiCkOjdW@JTkgXy1KKz@h3_{t?^(N>3N0O{ zI&&gVI1`XfFCTq3c$_|@j+G`rO5#V?{H|4^YvhCb0b4YW?DNM{V4u(=X`Jbwu0|4H z%+C?(0Pa*Nq8xs4%9)NB&qm; zR18@((Hh|@H8;xaK_7`?X4K%o;TH9$o0%Jbf*eNM96Tqzh72>C3Ix9opby=W{0e{l zyq~q1f&BpP9l$ny(PSY1I!jOd5VU|Yz5!;Hd7wQ1Thth9bk{Vq;x|9b9m~z<%!QuV z1(?X}Xf>%csYRO1hN*jk=&x#d3P`hzDz`dXapSq?yJCe-m2S0o_9)`=-ApkS- zd|A5pV<4PJ%dO*TIKdYt;T<&?q9x_Pq?{E~!FPT5WbC^YCW%7K!c7C5MB&N+I`g|6 zJRE^s`)uSjH`uVJ#1PInK&_~U-feL4SWp)!-qR>uoOcIL^M7%xmArwWx~6aL6%<@w zb$LbI*!L{YZ@n=fG4uk5RSFL4s*urQ0pTw>IDKobKu9B$fpMgB0gHu1KvQ(M_nStgZ*d)t5LHx-k>OnK75 z^Jk!pmsCIJ1|DFNhbbMiBKu-_NbhO{%A?@WG1>KbX4l|FTDvO%hR56;jFl zNQU*_YJCYlc>H15^9g>blDa)m`by!zo`@vH;OH;}G^Yt{@ zwrb;XBeqswm zp>T+J&3_(N?FnOx&O$o2m5GA#A!J)hKYd7pkr&1&%hWl)m7&Y%W}eNo*=Dv)-1|Zf zlfI|RIu#$e8l3&;oP-so+4vp-T4N=rcG^mEl4mn?%|+3OOM!pn3qzk;D00w7ES57) zFh>Iq>1AwfrY)LY=A*z1m2H`mXF z_`<-MCJ}$2f<(lkK)bA3eUj;{+9d85-lPxN3%(1=bE0=51jN={O(D-zBND1|G^F2t zvPLt~D`RdA(M@!QUnonou2jbxT2aK=KzYErOLR_W7GWh;uBmer3^N}X= z0H;ub%`JFI8(V~cSy4~aT%M4|?-XO}ChTyYaQ%8J>3%OtVL}R+MJat%&bBi4VyI7> zp&!5OR0N#DkA;+#GTIa3%ElT?mul*mB-7`ZbT9hF6tm!|?|n+0{@V9dFI4m=UKlE9YRAF( zM$?=ZHa7UX5%NLG)eJw0%PkkK`kbrVYk6}QFJa+f)82^G(3_vX_XF91Rz9VmPS{tJ zA_@IE-zD^I3fR_MoFub|LWk4rDdF8V2X(ktibY|dB$jBdCgv|i4Dd;jcPNTcW(cjw zd;XtZzt9`Nd%1kB!Ucu2ag~ep`>}Tt=c?Cz4kyd`%9Fxs`4#KupFM6c@8I*|uksct zF*Wd6ASkw`S_9LQSB`r17xBaPySz5vTaxHGHgkMxO`{9P`lJ@U9u=6I%Z3!fX-LK^ z@73D=c-La(XBF__(os+I41N+h{fx6(%N)aR@O0YjHLv_%uS#iFmdsDHxu+*8 zlSqLMQ9EpkoFvbBEf@HgV8%t;xBqThX*?@uEn$WE)%R2vJQI6Xm!3raS3%_WTVLf1p zLSf3bDN*r|&?u-5k3D8Zj)VHRZ07UwLR89bQogky!$K!$UwbrwC#!dsef4U$t+4dn z&1-sJ#9wYa?i;n}Y2rFS-_zwV`+f93eonK?bciG^Sz-NN%x(ppA|-4Jq<0>Relb@w z+iGaj=?AOjrm)stnP>e+eVN7``q+&CnS6Y{Af^=Bc5CY9kDpa6us|eBw~5?Q8b{$u zpCoEA3Y2TlI0yL&kWIRbPju_#>dfvYR%sl-Ri_=iNJtaaQTjn*Kk_YqxSf$1N9SMT z;~^7F+4%8**80p)g2CNNaV`?fAJrhb>6&C#J))&9Irz}w@|EJ5K;NWXZrE)OlbZau z6)DPS10nT>^P)E_$*2IM) z8dKQkWpA&62gZ!8-S%{yOF0m!DrU?Qzilx(fBpEgjD9uDhoZb;3A_-rP64AI*ewV~ zmq>Z*6|JtYN$*NWV3yugXYLl?!*Xb>_haGrj>eMA%&}gAetc39a~B?LUi9A$ITfro6Qyms z)t;gdQi*6QJKcXVc~eP5W_1sy)-m8DQYWWKx>vC058C@K!HYy>(f!Sz&FF6t^8J<8 zd96fPe{m4z&ytR$3;|Z60~K3DAgrA$0K$gXS@&VRp6`&Iv{obSr>mWr<^DKW$OIxW zek=!=8PV0aCTA2i#M$1>Qdv0dpv|lR^IWJwXC&FQkKZLWZsBTFg|42h2!~4b1B@~Z zP5Y=3+hDim^2-;L%KFJy5GxE|O8OQ*xfODXJGzmq$l8zkb!u?Gb%s!Z)_O)f@Mhfd z9xPTuabKj6olz3ZGCrT2xn_i;jYL-C@Oj{BndAnb(M;~Lb@dnc2NdfI{@WJV zPa2fo9Iol&pU!sTmVtIH`}v3~$4ZVYx-wn+n7f^dKcX$mI#H`6FGvFjuPa6!Pyrkn z8v@&)H3#4Ppbh-jF9 zt9U!cblUYG=vEwaAW^}|a+hnTgGmShX`PVDrTg(B3B~RYK{{7h+C0Uk=o;DP>6k~> zf9UzpKB}cIB)05g-NFk$cKPc=L(fM$dy(Dmqw8qxRH;A$&i=9Zvi9mv%Bekp9Nx*h zG}D~b5l1=AN~7OFNfl+psyy&!&n=fe{2YT@fPFXAreoHA)&&Y*5@ zgx%)Y$|0m8xHIQ%ahro=~ME8kA39v4GTTyVus$D=G`4N;#ZYTgWUR zg6^wACII_jtvK8ZKH=}CH?Ggsm*N`?wKMB)>~Y-dJn6jk^H&q-8gp4~c@*C7+;U4@ z){bnb=9%4ivdrK5PY<%v)yaYjuWxA#VNn`-GAZ6&Ds@Y)RdK0rnnRlDl?Y~4zlMeE zJ&Xu8uQ*Y1>T7YW5p($79$WvO7KyDPmeD_=ugX3s)RjFf`!S#-9;yEJJJ`yVrQdaVvFB8)bq7Q1ar^zu&#>SIrmtK4s6XBubEzd7B`Zeq5 z8G;&@W{27rIGOy0H(JMNvLf=ADoBnAzVPzp6>^uT!B(Ur9ayrW8;!ppiH^Hpv=n=y zx6vBgevcQK+E3z)fuf(FxHA@(kaJS9B^zV|>6N74!c)#feci_rt7Eot;Wo05xH~>F z81rJrkb<-eVQP}W_Zy>_^fRKJ&0D=}X4yjl8bp^1lbH`e1$KMO8D{N~XXb+=IR}r6{A4^3ET|IZC=vnes9C;$l z{l4>CI)M*r0h5x;OI{i;kk4*4xMcSZ+%=}vHyYHCu z{5T(n&hG5YyRB-@)AKgCM(TuD1YjsYZFa5P^LB_#5{^)Kwc9>_vFN|{3ml#g7T)|6kx;N-3LkR*}a zeeZB;AoW(L_O#`_3wUnG!$3|CyscKpo=bwAIBEO?v!&iKR#0|+=MKSaxEcp68FnPkw5r38l}-(eR^Eu%_J@4~a8rKGWzF~GZsbx`i=%Ol8l4l= zTU29@v@AWQy1Xe})Vb!0mL98L6>%hZ2EV=y6RFdf1`lDE%IDAOD%fZ&%A`#SL>-FWe z2`9Z+o<%p_V|Mt&K8c?11A7P0oZf!QX3siT%2i6~k34S94sdTu6Tp4K~RSzn)A z-XU|8yc*9WBEzVTJ7E+UpKPGa=Ie}fb$2K;b$560FqC^E(+b|+d8M{RnXmI<*y>RL zLh))_;Zt&4F#3z<0-E-2_Qbllkqc?}Ce|C~(BuzcJhZWAsUo75N$_aVJ*A2_pAKQm zyI;$d@}4++^Kj4lZYLsLK1^+-mz|zO7W~;CTfl9=Ds12vLG{YZ-=3?Qx|-kZAQCbM zs}Dq$_uIqjc_W@0?c7@4ZCBc$ztxCZ37WQ%9NxdpynASx5FTq(?|w5WXW?=>EKxK^ zOp1BG+Y{pcCw3P)7df=3a24}QqfNJbsn$qt?6}{|=Ud38k8li+?sHVj8zQFf(Q|oq zV$X;ciin1;1P+6hXf%lxB0u=f*giY zp|Y(g(2LJNb!+*dtA)MLB_|EZ@5BB*f3^~eWfGfSf0n5l6ElD~C6eLSa!HZ=fo+&#MyfSV`Oza8&G`(6Q^i-1L zM|MA}%AYop&EHvqOkC`7t?Agj`M_ z=B^PWmYK`YFDGGtr^xhQ4%{DagcL3Dn+*!_FmPd7T^(Jy8IkQCb_Gt~TKjC4wPCL@ zL7k1}*gDfJbMP&NG|bSFj1pl{BzwytJ(|^a=*PI(0p{oSgSajd+e|k6Deu=%Q-Y;_ z@B21rhTUSGstw79QgW;r20eR_~yYqhjzAEfzk6HB2gOrmEdGtHz}pdNAc|SvO$b2#M9L% z5qtCJn(3t)rg^`PJ5NE0@$%DjIm8fG4u%ad+mg$Us;%0V5seO6R{Z3@+$G@3;160!SKYo_U=G4OZC$ZTeI9Dw}axl*|Zw+Wn5TNe*&{hq7xZdj-%dUSN zGyZqKTIzcWCJ(DtEHocl3|`l;8LhhCEXoBReaM`|^Ym3VZh}9Wv%K2|1d0{$9F@^- z>^Pjm+sKm!;xvaG`2yTgp>f&v-uR9_7i8}x%^RR42(?2OaX_Q^A373-Y z>B5om!x-;cNijUIif-??dkl6!D;&D&sLf3pk#9Rv=COr5DRn*j;+FI4Yh!8W@%yAF zPcK*422BX=g<3tQa*A$$uq+VfjoDz0Tiw%+JzKDU3^y6Mnp zS1xZ`b35AB`DElB++mPdg^*(l?U~oe403{$y*loEJNAYg1XzR=bFZr`^#fUN`=`Sk zF4vu|hag6W?yKJF#_ikO69hGJCqx&x18Za++W##N{@M}H$#?8p_^tBZvZKcYNpPm` z+&g7KY1N_HcfP$EbEQIWe!EZLikprH z*uMW9p*JDH%3fSItL>dq;uo_tC4D)k@Ra6z_Z-E^h@McfH@X-|=~I9U1Jc_)5F_w5 zQuC;vf(19-~Sye8*2A#F4w>xw&m zpJ2}W7`%u!-4+sZ{2Z>~O72$AoU2y_qx2D1H@7uX<4OBzNaHnipvorAEe7|-i92QW zoU6F;DD~nU4BURjutMtO`kyQt6?KYO^57&!w3VKn$c^*~DjcxZvEd2w1Ou~>0HHpx z-1ga|bfGJxeDEIFR|f9#Vwi=tZDJA$$%S+NKpAof#zr9coawph#wK1A;KnltO2$MQ z7tssiE<*C)-myb#W=MB(iPSa*Nm&!B#Fad_XDkuBMRGi-8Op)AW+qnGDkn0A_kOvV zn}Msm_$!0tNC`uY5)|iWjNyhAPVd;9_KFwL6D#?ZvEkpH8n5zNM-!XlKvY)nw%kj7 zg*N$zWk1#J?8|3l8q|(AR5A;%{IW1y1UVL}H6-(dA-6XK016UgTXxV!-ny^yjPg28 zX&5yzoQvt=)_b8>39Nv z9?V802>&qU4Ome_voa3Zkh&bji~b*}_*g53<9w`)y$LcO|AzGEF6osGaye<4JQvb< zWRVEJY~g(jf3)ckBjxAUPnQ&`R#VEi%I7Azd1v5Mt~Rzt<{^ywFL#0$rp8mp^A{$F z`%j9vn&T55(CTGb&0CL;esA_b2vau+@qBN#Q3~*9LPpL~ehEHk^28ZXwlH_kS?RQT z2e{gd&O5~s)kC0Ex{5ph*x1*lCD7bXmgqZm3FNF89>2Xm|~&z#hZx1>I>Lq!%1vcq9Y`78${b!a$T-G)^)GdPyLR@Hw7w0&aZKcRE)Z? zIl*-@-NH-vRd#NHKhm;_AE_E2xfk2DZ5UK(Z+3_6;6l^~l2B5*VppuwWLVcTXDd<< zT?vEtE3;6~=Hz9J~aY86y_&<^k)rJ5xp`|oiPro^p!q3KGc@@P`Y$a%&rR@qjsx5%zf77eSFv3cRGd&b z;m+D+No5BKO~@R73cZ#vE{P(~AtWytrOpCzE+P5cboHqCuaD|cFS}!B%>ogrlOm{g zTJC0SQUp}qF#A&s;~*r%2?+H0KV`$r_T82WC#)YK7CH-c88pzn+J245Ca)#$jo%O@ z9s*%#u{!<)*YN0i3kbSh2oF~@A^beiA4^&KF!MgNc?}xzdaQcW`F?BC&Pq=rzb`J> z;vJV*oHKQVUhhqKv7%IHU>KgP6nISTyn3~$V!?x%o0*0cmk3#k!OF{we&ms34Z;2a z^leM_mb=D{3Y|6dcDFx?S)tF9*a?z6U%igAh1D?PX)@8x>)+_6{NU5yX#Kr4tl2uv z_FzNBjz1hzKk!4wUw68&S!HEEyQVp8L^f*e9}eRWA5Ubh1IvN4uud!IZ^JmtXM704 z0el_R`GAsp@a5gmlkqc@!?!N)%|=`+`ILX44ZZgir^9%gv@>Vv-IP1G69-M)53qiG z&xCoa8{OMAg$QQ+IUxH<<>Wh^Dqx4k6_U9?21&kpT0iDx4YB0q9c#kdu%r+DI~hWj zeXLJHj%YfQ$ZW|+{NrC7&CoyoccH~WZ9$JRxaCs@1i!j5b=eoOA!J=#0s=SD-Yv#` zxs2{71i*CTitl4AZQN*}xFrtiL8Dg(eFKdw+|7@(T@3;8GePkGxJ~lepCV+S&!lf7 zn0f(7CDAY%Yx?n)KfxOP7WYOL>;O%gE+Y{@PK)tZ+yJ4o0tnrIls{V)X+@I=rg|p` zrVnNB(l8sT2=nZ4hY{C8sh+F$_`|*B^s`cRePLbR_g$tq=gdalDf+Dcq-o~kPOx|Q zRJ{)g|GBFDH%=prWSUHWfbG6g6NHg84h>LIm*{miV9Y4b- zXXZL8!?Ek*C)`v3(wqR0#-w#Zx9+t58ya*8`nkjY&_Fxl?ACLmK_9PSU)EnaFmFCz z7m2klr=19ZKk(oU>ouTLazhZ}C6Qp1$53RM-5!*e@>W z+iV{1@II7Wr-PC2`Lrj_tn#TKGE;BF|;G)#VoxVyCbS0neWa+ z8+ftm`ITG)lOd$_XJJZ$1HCUsyLpDYyPak>r8G#{f(c*1^}B%Alu{~zTqwuypi(WY zpWJE?{5>3lx#`!Wl94`W*evT~AJ-*>7|WvDY>(}_m*>ce*l!Pcz~oiqEWD;+rIF|O zA^m7^n9KFS@2blRg&^6Vs&`Rb#{Wo z-8r<|Qmw-;A0M3o3Figyw>YuneEh6?7IgBKta@i;7bB%AfesAiSInb{s(0s*0C@B$ zJ(www7*BKok+h7aAhC6t;C(#-ILhp2MID=_AOv@_Mogy+q(4T>Smn!<)1y$ne`3U# zd!U}MSBiM-hurMTHVi~4-;aI2yE88tUlp}85l$ACyZO0Fw{Ulq>-oflm+FjfaM13} zN&st-tPXE=POrmXIa{1LWqP)QcDGaoJ>Es{`@!pK@J|b(0|JG-)MEM)6feV%*pRnl zSA0_NCX*FY4RJOV%f`Q+nBtfvg*AnGIca?Ma*lm>^h>Knr-kq0)MShJ z7DysW0td6Mr5ugTXS2w>yQspDee`Y89KPZ@6vv(ppF@#p-jZR{mj&bKu`|Nd-HWaWAF zYR-ECyxv^#%;&f9sB6LsVWW|P$IYmmcg>(dj){HlH{Yx!nVar!trKv|vB@rP@IxX_ zdYVpma0H;oIF~Z{Mf5yAJnd}8XDpa)ueh@%8o*y-65+wvcT~k0(5{DKc<$AB#F!)Q z7@uAK!8q>0BDF-+A-j}3fCv9_MH~!olHDM|-q%Jj4Hr^>(c`zeyXcVBB!=`se_?gu zeUx-v>=h)5{A{>tkFzkQpLE`z%-UB|Gv^VtAKv4U%kgM=qD8zu;%Ai+tkqhgo+o

uF}i;Q?G2Vhg%=JlPClCEHGc|&TyhlO8q96oHSi%wpV5RTO$b`M@hEWf`-fw zF`Xw0n%Ot=a{6b=@99X)Uj3+8ds`=efHew`zq<=yBiBXW#zmrItkwlm_+i*pyt!f^^v%&C|FB)YNmev2V&P_y>;foNQ!a`}s| zuN(ET6v}eo-(cl4qt9WEG zwWL9a8uKfaJP|HysJ7Z;L4&iYa;31~t+B062vh!gE&SRg%-+LL%23F<$-1r;_*np! zPdd_*9(+-B<1*mTXYiPTYfL)I6+~vU^d>fbe=a9b=@6w!U9VSjCl^Kp6Jl@7V<{wL)TNuj_=to zLtS~yu8OrjF{Iu@-W6p5|BcrJn+I;NzP#qqf0=JR6$0Z~;bx=@?QHt*3eA!{cpmX^49e}75 z@dN%jd21j`M(6m}f_TgPKan?E^ARnWaNU)aN2Bt~_mDfOv;zVRym)qc3005Km;gxc zUm43-3tAldnWYPkS>lps!1>SF>;;J-n+w_@=8?$xhjLOTn$j=7{wUKz(n7eGRXfQj zf?;DE4+B_FqHeiY4c1q8WfXV@(l=0HCffrWEw9$LoLmjDy1R_dAt~mjSKZ{|uYQ}K z@I{x@BR#PCNgAo7T{d6E?o%w)9*xoqH{dUG?T#9^&3VToxh`VfBX5|@^!h{jSI(rP zjGM|`aF(y579B@(ymERxL9B?`Lor=07m5Ul-o02`o6n#6thUB7o(}RR9_Sm1vq9Ek zhKC%_Ffku2jb`J305J+Lq6eH3WYH1rK}yq3BgFbkD;e9%cna0)moFgU$h-c;@kn;Q zS~^#|jPx$gCYEq7@bJn1{z6&jN$KG_`)6Kh!*^fKx-%7wWGn^win5ksTX{tt{nsfy z=%z<`+u{)XyCA1sj$cxj>EGfSeuH5dcKT@ZQwE%NXQNT%@jxIbMBQ3CY@YM}h^=m8 znE?kr&_kU|3*qqWP`?v`hsCS^cQkELq{@z=F7uJqk~cf`C{(!=cVIt_b^_2p`A5I6 zUNcQh{(?kH5!NU&mrmd_?9Jz2qmUGB=DD{(pVSHPWtab>J(oQYq$b`F#`f+oy!Xdl zZ;gu>oS{{Th5~NtIuOXKfjW!=q-P&Ua*L1!I^uH2}7c+cEwG_J*2i z1y{zH{af6>nJH#yQMX8XcC0H@~hVT_(mhB8K*1mWTX0py=< zw=TyePs^n9@ye-sD1&Ei&Z*PY*ur>4+9_=~{(#ZPWXkuWyyq$@{C#_t2Zra90b7_$ zdCkA|lD>NY>7T2ZLpMDt_V|bCk6RUYF?H{vvb<>84-DCmB%iN{z2U#qhdUt#_^l|S z;qIeW)rcGLc#mVjtT&|v;XC-GojUQIkAK)shA)fjHQF5+4MFDlM8puJB1x^z?T$$n z@e3fAWR23c?aKZk8Qr! zen6VRupA#E*GUvEJ>rr*@aF}90;8v!1_KJv;(ZFfcm(QG@8SRIQv_{vEMGHBkNl<)gPG+EnSW%p zK1sLB?~8wUD&TxOw1Su}N4Kb2LmFlazSwPzt-sfMpyTpbX8>1tV8O?zmEjIxrWbZJ zIlSG*;JKHtEDwyZ2ar2R@uumC+9V_8*>g+cpxtcEMywG(flFk9^Mje4@d`)= zzvz`rf0TWjOB4i;MVWNc_GlQQ$Wy%)n~ZF67RT&E00x!#AVzOn7IP86GuixZKvE#U z)l_Wb=g521UvXS)@lp3oz#@9_I;!`;3}L?fj?&xY%WCVx42YpzbEEpRlNtW`^@TP1q^{$r+% zRvKK_(eQ6RduXI6Zl6y;J=xC|u3XNvXMLIkLLWF;TE@|3=30%peiiY-0oyKY7Gn)$k}EFnl11&4546 zgxiD%x*9_vi!3ZNDSXMNQKJWd5dpVCdW_qr^JW!9G7V5O^r;@x^R~69VKjztbG#7S zNF9&sO1H}XtB>X4LT)@buwL_fk{Rvm0j!9#zy>s?xirl-7Q6he%vdr-0*s_qTnwR1 z*ybG`A((Ruy;12*T2B$_k9^>ST!60af_)l2fwkID-})i%Ka5?fVTE4b5+6UF%whg= zKP}seY$*R3ucNvp{vwx#%^3Y=8xjQ*qzS#Gz*H^8#;FTr;=KblYb-dHjPMnW9ME24 zas`sv8KtiI4?%eCv3hyC_$vL8`BN^WC7d0`W-4f*?Xo%apcbvSL@hvtQ~#6tF;L;a zi?@?5M*(iOP@VaXfLwauZ)&Kx>-K?L=Y8WO)=F0y^wdmtoY6+3@m{l|CYJ?@=M947t5NEy5=Ptly~ylOy#zfQ^NN#LU)sdj^h0t^ zg?U70+yUFxF_15rjc?_iO1%yCp340ocxRtkJ9a!7lN5V@-Yip%9DlJ2{7t>s*9ePQ z*;}5F$#g0;e6HTR{b!=BxiB_k!I@75CtH{BshU1-%OpSBu}7)s^8u04Qz`r|M{x?7 zL6Q!Q`wHp|!m;tK8JM0k(FEBnyWJQ-1EWfCvG>!8ptsTL)e>`S-}2FM6JbutJ2`sD z*anIwiu_Qc>p&IDbLNBdjj@n1xI7?hjnPx7^3G7vdKNL-kE;w9+_~t~sfIOgzT=H< zYdb3W*vfI!sy(%u>b-uxxZ4&6YS@|||JCczj?bq597=_qChTM(vst&y{w#cd$B;Sa zZB)|ZTA(APvHbtHm@o;>$Wt2K3VHd_r4!e1zvEcpJiBwa!W{I2y^6KBUU^i;mK4#;m*`+5;5}%pSTnRh~tOBJj=SmFm>m991+nB_c&4+_di~HZ zj?UaDjR9fdalwZwFElIipq2w1N92B5h+G}I+{lYro4Iwz2!`^&sXGD6TM_iPw)|v3 z00)mLf?~SS#Zbyv$h@f(^50_m+O|ZQvVuF>(tNCf`PpjWK2rshwW}2)XX<2LCK3A_ z+z<-SIbHKb*$*Oh6}{nC<(nP3VZK-89ozfqM|n!NjD{Gb3miK;few}dng#*R8yEZm zMtrUtZ|vV_(9iuTpLzAy%jlWw24C(s=Lr7J*de^pottolNkVFhT$_{Xk%uaxit+w^ zgHdKvCd#h_#y<0_C>1^P)LT}OX_!(AA}_n9EA(w3wkpC_je^D@E()CfQTN|3iEl%2 z!uRQ+0XBhi*S~Jgos6_Hr;JT`AH+`_CHrz%3VFVToEv4@{HY%K>)5-!65sfl;=-9z zs*)|q`KsJmWg>K=iVM+;w9+EgVaK82#ZBiYsK!3){}qd0Hu8{L^ySm5=p@{M7JNH} zu$4~REh+yzxAjI_^6sMN8GLs_5(X3XC_ZnW-)DN8d;avLWNnty-|}yg>M;fu;$!TB zMfk3Q&s+{u3Vh3(5|7<-Vt zYb#44o=bDvXe_6(D#@7={Ix~@5kJ0LVUGqu)Tc^~VATv0>t0LwvGJuMmUu3uXx&o- z%s+vRJ7~e9!c(+&`a9T8rJ2vZj%kn_@^jJaKg0<&H{uf-zfaKojLGA}Re3P+*zGvt z+jxJ*=5Aa$7`6suopl$fxF%FHk~s*m@iq6|TZ|syOK{$&ZGL9)G*>btlXXgTQkZX$ z)OKI8>}!%AFX>bOd{b2R4gv)$xbd`MQ2OBCXU$@{6-^gPHKwBz=UE^=DSVW{O4rfS z{)CoZXo6NyWRCYwpPOkhf4vuZ`5QhprxRViieGU*&c7oV4X$x-{{E$EQJYIbpNKC* zr=vvJ;^-N>F!8RMZwJkvFyV&RWh-M*Axouq2O@#hc zMw4Dz$j@^b69l3xXlx`pp;F-|dX4V_&hq789(_MJt+t~1EJ$*Lqeu2{t|Jq&D7J7m zX@|G};y7J5CjQz7&m^|IUV^m|KbsbycJlhDC6+N`zcuXr=>EZy_G8sLH&cab8Ta!| zT;~8D&CSB=*3cgp_uX8Xx_g$kFKtvbZRZC1Ax>G=whiO2+TpWBO!^ZH`7iV`56bDo-NTaAl(HTbp<-BrkkB)KJ0abNWY9 z+9Q2WCaFFYPh9=u?koLqwI?Fa$KyEyR1T0OcU7J?jxGr2^%;FF+RhyrcILo$!|3@H ze?p?yju(vo5p}D7`D5@IX=t`}>Gw;Or!lcR6=>ykk9XTpujceBw>|jYOV-++=+^KHgy)}%jrF{$bqHzu_r(N{u9zSqQuS7C#LVt{n_rg;*lU|3VgHB%i}dv@>P_c zCLX3+VZ}&CSM2^{2fKCCPm}YZajm@Njq1;ypT?@^8VW%<2~PChS^4jFB}wU-rv;um zpGd7JJ@z+?ujT&mD*N|D?DDGcVX3yz@MI3fxOxep3+a6Nn8qKMfDy7^f(`Wd68!sq z6n;?Wh|q9k%*{^}VJVqaXwcfDIp>#X881O?9=$vBt-i=a%qo~Ct=-vHe}Ipo!LwPh zmnqcNJf@^+0jg1@srMT*@>Nhfj^Yv>u?WGqK?6S+?UAkN*z!O4w5${BxzT(x@{!-B z^D?>n>Tt`}(HU=J zT&vesbbyAu=J*}8tm0R_-^>3Z_?U{jKE3lTR4QdHfT56Dl5s5+ugI^-ZRN91Rva6{ z1zyMswUyb(e(wu*9+4!v&UzA!^FUM_*M<{Ru}bE(@H5g~DbhPTGxkDT{o%I5j2C!S zD=oamT9}{z99`t=M2>qNbOq@KTt<51=*CZ&*~&VzT}6MRnj#pABi1sa_-2*3H6F!Q zMKrjL*U^?#YxS1aE9$FVQ&09+Zdber4X{?J6GFI5RS*;YS^UOR&?TFREm~yGI!=5lzjQs zs***_p5!_{m)MD?A=Sq_Iwt;9N|x9!-rX@PgNQxYW8C62{lz`Q=#4k{9%pr;cSLwr zE)od|7Kb}c322$$A?8zF9SU!>JO{R!dZ$Y?&=soOhEsR^-97bn>w(Ij>CTh26nU|) zxE|qOMVYil+2K+6e$}$s@qU2*2lin98`FCCiwxd}HxpqDiM-zZ|sbUyx^?hl_6%vCuS?Y``d&Ixt%py>Azh|+Qo`ej|t!3&`q+Ds1 z_K}4rwfkmoirRg~Uo1@7syks_-NrJ43}crfF5G2Ax<52ChOGLW$CZOJ^77Naox=W9 z5f*IbQ?ZnUDbN;3%s=_)^n@gKQpQI@#`Tj*FdMB8;rI1ISL;n{n;kqN|0Xy8D+?uP zh{{V@L2jj?lrGs%e`C)RmBf-a)>N_F{uc7HXtKX(4_}HbT0a0>@U6JHWhk@DB>T?I z(vEuOdCd6)CjK|OMJl;{7h+faltZUcB z5f_~xg$JIwAliQHuxi~SRQJJKDr8AV9q)BVx|uSCOha|cd3rS;WUQBE##z7B)AUau~l{Y5I=tI%Ajv{ zmnt^RFXcOv*srj)em4CW+F=WN9_Uw!5n9bKARMZ;hQ!OVK4$;8=0@THj^xdKYL8>z z_6NB!T=Cj;0muA|0H33YUNnW$*pVw4QF2KWu@-F~zOd$$`~$^Fjc z$6Y^-^W4jAMU5A~r**vXh<`TCW&@9PzvS}f%ckcuO7WLkf+|Wwsa;w|;u3f)#(B(V zn)Ojr{F7qiXdgM3ut{{5oiSX&c(^MxHAp(=hyMkwg&Y_)D&d~O!#p=n?g?ihZd`Ik zDr5$qLkQeiN{-p@k3Lt?)L`yGyx0TGd?QlZBO(7fjCJ^s++@YZ=w2F+xoF@Ofyv&$ zq|F-me^a&Hoos8Y7|U~NXSfib7B*SZ-Wqg%8$S~Hjzf$b1T5+LW%kGHZ)hUr=HO9f zThWr`R4gB1Tw4D4?SPx|+qgeGI>&h;)_Bh>^J|kov>p_nl0O%(A(nnj@l#rCoUPAv z5{MEZ^LyU^xBW2JT<4r~#&H}oySug`zaxFm04MhB?7Uj@&d%&g=5C>hj=P%k z=us{DnV3_g|2`Ffetvy}_J`{^p9h;p-%@A?LVk}r&*In|z!q9Q6bVuMdOeNelzhaC z3|y!aKvL+^(KcZVY?3U$uJ6yG-=fI?pT1fv@vrxv+I=9E{!ZQLdUBuYxVrP0vnD0C zED|6O5A+J&ocsqf&7Vup)}OQI;3|JBFH(QGu0n4fDe)HaQM#fQ+v0I6KHz+ab`2eg z`XZzXjT$XXl*DrYQ9%&*#$jny#YZ)Q_UP&j*#Vs~_t}(Zh~FLx2cLnxn7@F>!XSkIF^Whm<|pdpP5i zq@9opCSQzXQ)1x)0D+5W4}b1u{_Vc}$@t^Zz&B6pd(=1X-n=$VKF8};L28@(+2Y>J zne*MPfLlW0p0!H8HY#*Kt@RS=+8j|21p=oyqapI#O~TzXb~J-6W!E`7A|82QHY9J@tSq6|-h+ z{)*M~ZHskqAt4px#x)^q!HCa*$Zj=g^XfaFcC}#Z&#(vtkpOb?1 z5!3H!(UQMYQ=5-)j=5R{F>e5@pFKaTMVEMl;xvv|t0mml&WZ2Mcfoo;JTR9uET$w| z`3bT3a?SFToqEadqq)i|e@}GV*o|Gs4MuX&!!z?BlXQUv&>6U|67&+&^Ez@cGfQWe zZNNLGDJdFQpS8Hp9^K8Z!)|}RpqceO1D!oOZGb15$ou{dN&Qu@thrtNd0ZA|aByyp zycr!K`(#A!H{fqgkSLi(khJ?$-6XXCJV3fN5$rofDaa#_^dxb;c<78Ii;jSY!k8Wa zx+DB}Sfj|5E zoIR;5{IleD6bM}q1NYNbfAAfmW|Vo8vj4BD%j?4wUCI8{$BL`MyRt)XV+aMn%j zKI%mqV*K;NcZd_bm8Qp!25G!vZYMN%e)UJ22>ebWeg@t!62{huPLym(@Z`O3%Q)BU z@U!98bUOM#>Cd$VQQvNnkVFukE>``AY}b-5d(k}d`7yEGlvsZr#~fg*nEU)~Rl8X- zf!VaVzwdkNi<^iFTt;+LvK6*(? zm7CVUc_{9$O9}1#Bllxc?uL|;{wD*0oXAz{y`2?!E3w8`5Lj|ZZ+{^7mB0R z3*v+G?OxC}Zs|Csj}*VRV9~YSVmB7s#t%N|cD;^(XSbj7-MDtYX@JJ*F6|FF6-56| z^npFhrch^o`IZzx_{n!aAxX$%MyuOHMdY*(O9yUL@n;2WFy6KB!oBlcKj;3Rqe|~m z7FK<(&%0sixNwO9zQvUMl3rFH{`#~4iwGw&V6fKEQ1GoAz&*coIZjDPx+s_=eD(zarmVfEAu3Qyx z!MloPf|gVanpiQ#wIj9j*?Hb@Fs?jF7vUp)OnQq#NBcdPUc%5zVI~P_b2h^+Mt0p0 z5!ox55MaOpRjW>;ahN)ty4Ae)t7%)#+EuA0&A+aD*GXYOOE)(a8jowD42-$Tc{m5^cG5Z%*jTlLQgS9FVd1x<8eou;%g2+_#8g&&YEwp|EYV5(xy<>l z-Lz}{JE*+s(nnn~kvj5ks_vZ4`TQ?yn;N^4 z$a1<1(XhB~H^S?+As#RopP7Aqa}!D)r^ zO<*XAIrCSv38tgR5hY`g37w2*2*)qXNIZEFRG?R}U3@isF#~J6SIA>;&x9a@3}}aL z{@R)zT{uL-WWmqIt*kFnRafvgx%Gkw} z{|>d$Dd>}~yxQTg^QUjVHX0hqnU{F2{{t$|_$yc^x20a}7lYwZti_fePagKP=*^r< zUr#10ZP$+ELw)KZ+&n5gWn1~>7Qrpg{f=sZtLNn(tQn;Pz6@dw!ijHvgpwZg;!V!B ztc(Aw#_9aagcPmS!vA%4$_nk^_&eAc9juG2 z)tO&0L1jfv-=UCcXXoRmKS=4!DwU*lq@Zo#YQ5)^mPhT%WYh05jM{FA68EM5Yqv^o zD>AL#1wxzN_ca>l+BV3&r!Loz7DIj+zRpYX*)iCpA`1lmGemTNy{88TXujT$*&>ke zz79teB=A}Ilp!jUT2TW%`|pKP?A(*J%ci?8rtYGbl`0?e#v#$Zo_on#iR_ihE%=5< zUs`3wX#yvyP078=XVjQs?j89jxK#I;7vnP7p_mUtkw5(|m~szl+IxPNO5YJ8;cZ2A&*7srW2w)iZUu?5C8oiKNl+$^ z^SbJf^ZsDrFuIfRoZqH?{ZrlQ-eN;hZV~o;`pd{qo#{^r+H!Z-FiPD0R7KYQA;Lu_ z+)k~nKXh)^Ua#_4toWJ>)vS8%nc%cU_*B5%3lBc;^Bh&sO8(UO^?sy}wEL^fov_F& zhydqoiCk-SudIl&EDzPnc0G7VA%4Gi(@|_%y#WGJGrPMs?fIO*nVy`T1)}zTW%Jun zgTdl5a`>yYiDzUPNYB&L9BA|7mjE2;!4ZX`65fj*S~uD_*k0VtyGy5BqD>-Io?4%F zCh>WAjbw10dwyXKgWg?}i!1Wyzp8n7?09i3a&^P>PlXwGS9@+6MP|US^iGOa(MrKw zYyX|s`}Ec35JIxoO-Ed>_s=XgREXVECqeOxH7-fFVjAciem>}M^y$%)lNpb+!LN;8 z5+~j|MSrM&Bywc0V6bK6r6h1}ul$xwIeG676dR_Yfi9>e;goGwoSU(_`_Jg+qCd0R z5O{9|Yx1{_hGz~}u_foi>n8yc8==Obc{fVWcY$;sqtS+bnA zZDl((JfxW8$w~WhY&LdZnrB(pr5qd;zVuF6|7Oa$KhfwPA<#(N;m@aBUP563~mSKW)SI3L&5?AtlYB4h^$Z_JWDI`Aec`=d9+2qC0N85Axu;Q-FZphto=aOOic*A_)H88YX`Y-5C>pS=YNPkJR6rRB z6;&Hb%el&0-#(`OW{2tNgC4P7#F+NXvJA2B$4M4r1nLy&%WdwfG-Ue5prc4s^F;V` zSEj_{4rcdN>9bI=B{5C!#kXFh7F1F1zUV_sWDSC$O~3AkVX=d=T`|U#)zPhkgBXlh+&GPkesGHRXm-IJ%{Zwqs zt?mYBC>TBllkMXbO((CQA&px`QdF=HS|S4l;}iaP zmK=jjWTGU*o=+N-8+Lnj=az^y2CrWa|9A>a z@u#q*u-;LFLCusWZp|f&*?K?D#uHE&lZ`f&A8Y2C>Q)+)bdiKO{9}uQdQW0<{eQybq4}$m-DR>Rct%suZaZxWj4deHGNRSGFg@v}Il-e`LMe35l9*vNLacn_{No~A(3NNSP) zbf99>?|ld@XXtOA*&ye4CC3@am-~FT!;AE$5my9UbW&z4o;Yjaxf1nL#O2iiq|(x= zB$GgzO_QSj<5vXAfQ5UpQ@YOdRfLtO^?R$@ zrF+g0LTE{9eU${1>{H(OT$`VQlqdd-$pJFHLh&7Is|du_+fR?Vfeg5l09r!wvllsx z{7~WFPkkiK9n6uuLB+D(|Lwz(tyvWd=tSIo2X)^+%+nIXjT$MxYn)31+IicrWh(EJ}B|97lXBJFqi&#|(jab*`J zdGXFVx1vz9Zg&V`x?QAB-B~*<{vL3+2NwfH0c-SgY0iT(*0BCf-Wn4F#5GV zS*w=EM;3=(+q=?~De8s@yN~kvDPNF!+iyJNq{>^xiB% zE{j8v2|*9L$tRS7Qjw_b$(_w-v2;RY$D=pcNQAB>$9ku3d6|2YOjMC5|9l3La2HTR5Sv^I{sV>mz6Q452rE*g_nj^p^Y7z zlq=RS($wO!B{c{3E#}Q<#0FVoz6s!(Rz+@4zd!@W(?K5cHun3g_vW;lv9QYA`#QS4 z5TCR$Y5mWVaoopl$`{=<01twy;hTmh&?W*gKiOD1fF`_zh^CWP=ra-4ut#2D;^A@eRSM~E$ndq8KwE7;Np55%|5EDv(@8i;QOVtA6)#( z5o0Qw^f`ILoCmLi=R($P<)W=Y%;*HK|0u>pU}FgT%qjQPSx5KX@g2hP8l}WAIjWv2 zv?C~rSNg5>(H=M7tD;(^VI!v16eZU{rCN!AwksUD&;P8vPvXx?b)YpGtlvE+2K7PL>lISt82!UE2glg8` zc>8`pDZi|U!~8@P1N(OqYIZwZY=;S!SmuIBrEU*P+G}l4FdgqIQJMxoP}2F?`oWgI zK24f6`!Zi-&bb0+5{vtc2#qs8e^4xbw26?So~1Y6s@b>GF`X&T3pOt!H%BvZ$Lx*b zJvBsM_OoJ%5@ExulnWd`6^Kciy2)0PDJramVoQ7(}tCxp$9zjr*&>->=yu{#u~>Qi>(~ zDg4>m9?rcl_C451eh4)`N%(WgD9yK;_am7M^X})Y2SSr@-p?ejr4*la6O9Y?3AT=5 z1CITwz#~MOY;crr<@^0B~-G|64j@djBG(dY7r~uc9Tf~}2)Ts=(k;ZC} zewO{8C0gQ^jjEHn0?=Dk@a zU3f#8X0TAMo&QeFzBuW1V&xia*Yk#9yB>(mH+g<~**(DE$+|i*3NhXG(Ls{?KEoj6 zbDO(17Y_@cl7l8gPI?^>6?JT6V$s9shqWN1b*Sk34Lg=HfL+<`Z&4gDwRx^%9sl7C z?UT@FzSy|OHkj2{!wSoR_oqwCA^~Rsd9H}HKJbR~BckV|Usk%*aX?rK41-ES6%K%9 z0}1ILlE%XD2Y+ul)ab~OE)h0^;e<#~a6@4>zrOQp5{&yZG+$U6EuYJHiYd}D%_y62 zKezgEb+KLRyL=Sk^$Zlg^<#s@?Dd#4*KtFek=X3@+H^|@VKcd5`nw8%>f{Z&;#DbS zdB|r_zZOn^2PBJHSCJbtnfVWcMo*cHL2)DmjPzjf+@pIu2H0n?A#o%~YL@*8|7GEQ z6SKt9{cEk%?EqpCL}-+7J{g}5Mve5K10Zb(*|C3?ZHKoNQs%o#mP_-NEK{|?t%`i< zMKfp}BKjtXkHpHbvAdAXzk;0qH~mASPkA1S_8)%q=_6PDOt=AFj~LH-M=ymTL9$1c?*f8`w8Q)2mDMHb{h{rB>dKgJs5bbS9 z>pMJ?)tBq1n+;VPP19drP#u*^c7UOgcb41t~9K%}LR{ zoB)7@o?C6*EFU|d87MlqzV~kH%6l<$?>4++E84lW%WdWcx&aHersDw53YaOnvpd-E zQ)xkqJaQw4XXp*L2KINizqEoiADj$~wJ`Fz`067o!bp3_s0Mwdgk2*`!0#tjBJ=A7$;Bigc6?0SWK`^S z(`_=^RO9cVMKr1VbJne5;45sUr_ZAd1K#I@Y{{3C(P@yqYWnRTFt3x6FF_@mt&;%v zb}Ay#@rcrb)~$A#Kx>c8!d$U5)cO@XO5HHT4n_y77%rF)f?v@*&;O=d^K{;s||qN%x0^?9J+ye9AN zeG~qtZ@mgv1t+?3qg`oVmy)COJ;KC7V;H3*NJ9sO@(Wmn3ZCf+t430o)Bv&?lJ92b4~@woBDAo*^XeTnt$^g++e^~BmJ1Y9 zUp>Rb!ebtVqd3|-hRqMug+I6Fhp(G3Lnnu?&IYLnxmp<4pZr}2@Je$f-4O7Pc2e`d z4UL)vQWbSO(uHdEVJ%QWu;P;-zEj??w3ZQ3Ur-ZqOTY)4q@DmEYv~jLS&NB!n+r%_ ztCtcO(4OeRbp~TlEWR>HMH05r9Zog}$=CXqbk$gw4G>hNrT)VG#F`2v$wkIiOnU7~ z2kzNVlGj#!&6&hCJ_mZNI zAv9dukm@r#y&+=k!6Zw>SE9E-h_eKgIjC4xucoNc{V*`E?l%`@P0Nu2Q_qjF!O=w- zVzB_nm^TJ7<@$wizanD@6PMx1W4)AoUkjf|UAqn!A_|RR=L`vUyT*jjQ)iK|eCf(| zY3}=l9QuJDw>ulO+%HK(agJy!haBx1N<3Ft*wwg4Vz2Ngk{; z2qPE|qc^+2w&{G(rUjuyPa19+nL3zUy#;TX6bG=l_^f*FV5B>cCYX>dDp9-=9SM)Z zy8n`EE}70KEpp3C#1G@W-v=#V8Ohnms0e>Gs367tEi0;(@Tq3=KJ)isDWF3G&Nk9O z&?(>(yv=v}4XGlXo<#8+0NUDqQqsguW3?^!2dI<2AS!r6`xbWzTL=GvJK1?m-4mUsistOP?pZ?4ptTYrh0HUzh2P{+ z&d*`A$SO7&U76J)DY|t_Qvu~EsQ|DGBDfUx>;%@37*=e~$7tR+l(-zS&Zt=r1iCcD zC3`QHbVUVVX@j}9Y|82woc=k1jPbG=9Vk`Jpj5@|L_#&^1zMG`xRAE~nXl4_q@y(; z&yowNxXKdERi3iDkF`X2Mdc@~oRpAQq&QHIA`F{WWqe+lck46YoQ0xKPViK#IqC-_FjwaAuA8^ieFLvmuO(^q zCUy!YIXiYu-am2hN?NuSSu60bFgHNjg%GrI%fPoeEi7M2wkyp|}Oh_O*_hL6O|&=v@8Y zEeN)HuRaWMKs$$D{W=8w>=q8JNuT8p#piuO%%pzktfN)Pm;){!Ux|jX&OJ5PqKhT} zi|LDl#tRzo>GZ(=9FOQ#R6)YB_xG1K3w&x@io)beQm=cF(h}B?Lm!3Y)i=L90ENwX zot6bXWZC+Op1EXO{t#|(%#W?28z%v9pC^D6c;7DL3%Pf6lp}(^Rx6g59}YDEHsJru3k&L z%Jw#Dl@VZDTS?C*xK1%prI9SR?rBgma0^HwiZxT$2mTIA`a(Pizke(_Qx`vt@M zEE{K_8RfrK6p-xXa_}ypzW$b!1G)QTd1bZnN9rpxz5Cjun@T+}cE`q2YA$snu%Y;NOI$JG$tEt>9Yk(M>@ z9qIJ$cYhie+MFoD(7Hn1y#wHgYW$hkC7JgiM}Ca9$|dq9DK%#`ren3jN8Tg~uQC2_ zL={a`lpDe#H2p(39xSR{4~NXM_z)ah>n#<0*5D?i*|9Plq+zXjYY8XxPNonhp;_8p z_p)+T79cDJT+D@2lZN>z-@}tGF1Zr^%Da26zjxMJS za?BLr%Zn%_|BlY!jk%HpLX+Ixneadz|M#A%`g9mh$-%1PfRoxx+AL)wHLUL`npbw_ zoL7T#2tqEsgUU2O^Qz6>)m>B9{^Q-Z-gPCManIeA{GhDLi?W6Ctz*oZ%|}s2O`Y2Z zU~b(&bTPSp*}oV68`o zgsjY3ZjAQNf6n8~%(nHCYj?9mWG`KZsW|S>2fMc95;IdCs`h2}jE#I(X#@mbQ;+I61 zN}trKui2!#B3=mN>Z4w18Y0R$@|0QMEJCrKOc^-asrW?l&9G0xzN4o_aZGjG7YC;P z^lnTI{)|^Swek=QG4lNskW`Mis)!FYxf45kep*;JPHkD1R2#XOn$zkI2`#yRQ3KZZ`J6#@Dg(<#v{FsHiaepchN%Z9=+a0 zkEqiHmZ=5zUvS3dMf`=}XS}P8zGij8*Bga#{k1T;Qprh?1qlyeZ7k(Y&OU9`njrXA zLJ9nIe}xM9F+XCUHu^~FYsW7*q9&mz3(r6E=kAvIfUt&NTLEse(J3x6tw#KxP3|z& zpn<$%B*-AG)3>_}p#`*UyibNzOSdeA8x+m3zAV;sf+3ocLT9U1J`QyJD8(2JGl#U* z;$yr6z(y~jwd?%-c<~b@yno+_ZhawFD=w2~nvZ>F-d=;3x%ANBF}# zSRK3U%lbxk*AyG4i=8YClyTuSPWr94$sU|- zWec&t?p=0_##*CQKh_W2qg5Zhh1RIo=K#(NsT7Ftou6^p|W_L+wzOGl($Uy6=HhE-utUEg+s)^k=IYS>3*7iXH~Lt9p==s z)-c(BuixM4E;eAx(m5 z#r@pi&3njZC2o0kpXn*rXZ~HoHrGbf2<^-slqL6hpP6oz9z^R$R-)GGzZJ*6Bx}Gu zz`mk|+z!0rR*ta1@P`n&4E)GBT5Ijs>C*kfNSO1TP!65D)hLt!BCuS31H9lm67p*{ zc++fHIaxqiWGJjbDvdU5oi2Ik1Vn0i9lWj}coSp9T|zyhtmB`g>r#RbEl=;gd6Z3< zuz!@Ty1ZqO*Grd6C|qzrKq+B|z!8af_k|_~T}o~z;Bu$`$htw0QaK;vWB6{zlH`uQ z(iao^EUhpRWugUk{@oPF$~b9q+`cQ~E4~@a9w?56VX~=-!S~47YA*N3p9zusH!D5mwr(4Eo;|?SS{5%XU|s)qoKX15Apxmx=>kvl zI=QW;((WvrEE_`NEnw!z>38&kBkpm8;i#77v97j*lblilWdANBGI`9K0Y(ufMS`Jy z35_|b%0%hPYZ@Rr)W0feP=XqjgvuNVPP3h>t}qf<)QErw=ZQH^#I_Iz0kiK}EHlE)EcdjAmpWQzIl4x7r`R=i%V5hB32%Ht;h!+TQS+CM zD#Q3Na8&RvE8fu<%@+Hv-74IK-7g0+B@v*OnecWfNJp zQp&uc!8?ne&O<^uvW^ttn@i|-HJY+JvnJdS<6u6S^UL06+%+7& zAo9P?7-bPAauXT(<+iLwoH@MTq#)|wq0aQUr?W~G44Go_LF9wpK^phD zwJToc5dN!eNA>)OV=upN&FIF|X85UD2{v&8 z?yC+?%U;Lj^f7G>zx}-znY`xBAg_&(@^Fo;LXG*C0AQfFoX6f4ll4t``+kuR_oG zwNdpL4F7JN_mU(Uo&!mj2Vsy8i*ZD(_4omKD+n+es z#IuqtD@Tot9M)*g$+uc;gV(!;Mev5OsEcfGo{idva~=0j?) zMq~r-iI5rCuqk*$SCHhBFM>E&oYhz*|L#>ogIh7$W6 zWD#>=MFLYdX}Czbs9Goaia-s3oDc6DHk(i1jv&MkI{#^0y#+%?OA)It2rNn5C3kLA zBa5Frb&~U!BEC-vZG{;TP=Z(VNi8U68Zh%t#)pUvaSOf(BEusl^eM3MAQ_(3Xi~o2%=x02Px+xMO@}uB5RV3GKdJQ0&ex;1gime z(-^v2m7ddR_ukwy_KReat@0s(#k-nk+eEr?4Ov~--~A(R7IV&N27KfGry=rOpm_e+ zYmngW^QLA83DXE|2Ht}&xxCZ1GZoa(A(+u6(xg2>V`m=%QXWo~35S7?ATvC(=TA27 zo;#Og;NQJk%qeo{@k%ny0H zgz4qimCgmfur3^P6gA_i7ww*BxqLko6|8))mwY;k#~^BcntyEAIlOheIrn{&bPD`6 z!0TIU*4gS$_jLk54Rk zditrg-$@C(y)v9PtnbrdD!@ge>ml!@PT-;N-H#d*1#7R_w7)%)|I%eRwQg(wO2R4W zFNIzIk};Z`gir(Y+OmIJ$Q;<~T2yL8XGpVQo1fsXH+ZS;Dw!ghhuAiP&=U2`v|}Vv z#_nFsMC*hZptP6vL&y8oId(m`x0O04I?R6|<(ntgXKj^Zo5AIf`QFxL|3(+6S>P$Q zqdg9w1aFw+)lVwGr~I^}fI_@jGqZ@Cd;S@;3bY79}b3MN}rlFuj6@{Y>0q&I1 zIeDF15$tusD5&3pH>OJGhXYG@LBK^lKP0LD0)U4@A@JUE`ca1lWdkl)h=glVB%>?= z6?bCtF}+CUK!cQfZNk)P33cQ7g8ZaXj3l`V|q zcgcGwEtS}$h}R+F=DIQljqKduzv&j-?3<$Z$pv+i}_iMRLOnB?;8=o+V$9bqn zKi!c(F z#2--xoRr!PxJO^H)7wNGjcNKw@Ey=h;)b)8W_rD+o#u8e!>}$ymvi83XjNcNSt*SS zFcnf-%;VEMfKh_>5`whZ9uOUU$WGTe`J6d=BTkjIp@CXe`&?Rpl6mlxp%gCBg;x zYIJhIhUhije1?;XF}&I_#z z7oZ=!$-)&#tZBI%S+*#~N0*H63dHcaZBQX01Mr5PL-|Ga$l2DR+C#$^BI$!XmBS(Z zZ^Y-tpJYtgR@@E}jdv9AJ+|O%do|8YzL=Zeszv!HF5y@R3qezVUcjjh5ML8esfhWE zAL-j#$ItW#48P5s)e-clbHjR#&_11$lxso&o&Oh31Drk@dT!#lX3=^#K~h=u^sk0} zCb$Kb_Ep^Cdg-5GYuY$bz=#t7e8PS@5?rj@zCRw!UKG2)WU@Wyk!Hb@qdtA$8vD(h z-j}nDOb;$a^ou+(-B$oKd84l3ts706?UuE_wA9v@UP5|w!xJB#1(-)xbc1UYkRAZc z%KmTWu-g<)q~ew3kwowH=7g{m0E|V4R{+d^%9;W*AXvQAE&yGK{F*T-8{U3}dnNi) z>762sDocYOR}%U(03KNp(ov7^601AW+pNSdg@xG-^05|rp>Vr0*-!!$o)Np5$~pb| zvMLKUnO}g0v5skA;e98ZXCY8HQu@i_zMi) zn{bbR@0hrS+?K3C9{f#}AhB{bfLF9=7PpTH!CvV@P_36DI*u~Awp<+o5=RzNiUV*U zoIwV_VzVz53_%`1nlTID=j~6{<+fwD>gN=ePuPT|N~nEv7U-bXs`Zj`B}3!z20v{n zb4vHS`~JD=S3b#YncS<;*!vOiIJRBZCxzTroFD&n4$TT%tPeWN{uRg88t*1Ai>9px}k7H8OY!?vv@r1 z^(p-^^-Xm)bAM>Iv#WdT<3k^uP2XK1OORNV4o$Kxw|G2>^f9yV&Ga1-aPpk|@Nx zBnzcm_cR!VB^VhTK8SHleq@D4oZMKnO^i@U_b4MzKMIAQp^dSt;Wr}!?fEF1oS)d#S^Tq665)5ZRFih_1I0yYs zy-IeAYxeuCh|6*)IZh6Fxgk4ay1q>A{9if(M&mhq{Uw%p^5bn_Ek*qyvRcIXkboF8 zOdL_2Y!tYaTmf8o3S4N;ISOwemLu;gqPkH(;CFnqCtizaQO{UTRHk&w466?Ya!>B8 z<3ETZ+^n$O0FnN&Fiyz1`W~u1HaG*R$Gc;P=D+U5s*2~|_`#+E zshAw4D^3xD$^Z!d8Q9chOn%f^9kdh}L zH(&vY__(vfLQVTemBdl>53!b|(f_;xeTC%KOO*DdPI+tupStwsW8fb}M-59d@?}z6 ztrSV5i?AVfZ*39VBBUS5L$NtdFng(KfLhM2RmW|UUKk546>QgkySj!QOL;k#au21! z9RuVQ4WLkHM4IX3&L)=U>&TVqR#kk*gb0sp8Kr=_`gq!DS-VkX(@r2_O(&9G*3BTX zNJtX?XtR($T;@-}i~Yafb;0F1>iIgpI)@^+V}ZVn;Zutm&fL88Pwk9MZEwc81xYXA z+|a*@z-XxB@kTm(IjOzfP~`VAD~jm!k-oGF`9IXzYD#x%j9UJj=SU4@#ixZViy zQ5*z+_+?e*xC9x0>r}S{T=x=*n+9j2g`5iuF!J;3D&p$NqCm1LOGux8PNv7QbQ%9( zi|u)|Cgq#hL%xb59tcb^`%byd01f0GPuEctPK7(jhxh(_920%~ZX1~LeDFwMaJ3&r@yz@EHIdknDlm`pDZN^F?)K(7f`p~H5Hgsngf zJuMAs_zh_I4{D0+-e0pDtC8y2e$A)^)aU%!V*0P89}magA4Xc##)O7S-7>2ZIRmm4 zgRK{Al=|}gL%#WxT zY0{V+SqVmdg{wyKT{;P-X!2~43`6cLAa-miMFfpGpwI&{5U95Mu`OR&j;xief5;7L zOEwq_KgDL5GufC3+x_$cFa%qg{~hd6`I4)`Cb^KBjY){dkCxjJ1T;sDJ<)Xf|r58KAM*j zSLZJ{bXLI`w`X8TeOk&fp1KzlPDJOy+aESC?C}~nT>>X0=dnNaPhlTZOUXTV&ucv> zr{ie|zR>he=r9R=vsFQh{MUMLwF}5{m7rIgb3%b2G6|~fd5oh>%EqLuW2nemt5OV=A;j}@VUW3frcZs6xoB!^u z+dtgKLpleR^;OV}8yaH105Arc00D9(f(-VxbUpGpJBHN>`Fq^v%xYq^lmOc^5S-8c zl#if~uoSoYJC5aGVlPPif14O~0nU}BwF_goO;fB!zy0|XoJi(Uu+{*+`j0gr@&B!& znm(EtuX1?Qw|OE%nQr-7r7;XDg&xKI-|#_Za(M}CkpXT5)SsFz%MOT{EEWMVo(6FyH6?$RbF{LAi+p2-bo5vkZnG1o?!HR?v0 zSUqL~9X){q@N1SDij)g9*N~GPQ>`|VkQy1D2rYq^^p4AWV=Q~vG((_yYFkw1WUvZU z8zlVH9q55rd1$mGT&r!hZJu7%@qK^eFL=ZdSVQYVtZvyh{p~qIFB@BM>eTu_4*+?R z)6nv-0A#{pwliq$CzA;)Z@OH+cKm}_`8ATAa)V6X4t?%ZECr@o8Sq`_aGcFu581Ej zyY*J7i5BxoOMq~#67`GG?H8$=XT+LGz6MaB$iHFvZzfR03NRvP^N-a8o#P1F*FSkz zQRn7$zTO*xyqLzGj1szQHn3d%cWyc)pyP4s77LAbqvk$}F9X+wF;E<=J7?AwQO$tyakAY2cDG6K3F&Z2w8neBZnL0Z0;!p;{Tf z9TKNtqM#c={4?<(m8xRHbYHZFbnPP*P0F$TlGUP<<0lufnv}2KNkPeJP7H6bxUN%& zF8;;cGHkP$`NNVF)HfKXxO_AYuE6PS^K<|Oe_U!t>Y|0Gm1&2`KmOa_6sIG8kGLOi zbgq`ssj0s8q@1^B*#QM2Y4^KqGqCK38ua1|498W}P`OA_Tz~Ibv|znK|fz zNd(~f_y2%%b>ut!ukBba@lOhX16R;5h+Oh|j`~Kns|CFCR2S}MVGG;`Konw`2V-@- z!Zq~Zii-cyo3JM+71U0#%zWd|4)LzD&5^S!*)MsNlp70qL?r{%Q*NPAz!7z0?K~hJXF5Kj6Fa#%mE3S;NYA z_Kz&#Jw$skTN)Ma6p86oS-x!$!ryBou0UL~I&n7R2vx{D$6yb1O8CYY&T9A-*Lt#O zYw4&)&Dm~ic2t*x;wi*Nq*5+f`&6m*UpWpaGd`nHH!wM|;jALf_vSAYj(k>Pb2aWW z9QoCX-*YB=y(@51>A!D>X&Z;xj^m{^Za`%;&kAYXlYh!&77=g2%6E=tVerTqdRGD1 z@wFCO*JK#DehIE$Es-_;NUwGwy^)ABuMhvt#3W7;A3_(gE=Ja^c9*3s*^45ih!K&) z>@Q7w&GBnH6D6yxG)E~<=3ff2h{wjbp;&CNZ*|yUwj5;2;Ur6&e_c#&yAsiSNgad$ zbr?2vSR{7kLXOYtAXXQ#(~2D`4|#vMV0TvAli+@(E= z4MiA^06BNF5Lqbbz`KAX0j2LBB3pu}6vrn|)|)++%Oo)$s@CWqjN=FZ-&)27aKei- zFATEl62*a)mb#X7&#^T#YfIyC4YIV^IY(6_L;+C{_U);yvthS-MHVbVnbrUDCIrx} zK&jdRKOH2kE_N}gS$=zxH<{n(%mi#QIQ`-#eZ4#Dp{eJr2z59q5v__SuyB;yWoz~ASfD~gaCWoDL}4Y%3T7AZGQ)_K2po)% zYc-GRddveSQ4!2s0u@X?z*%a~SN5B1Z zbkxb_=JwkedTW83M|0wl{S|#^{r^W_Diz&o%j@l3e+ws)5fO`L_MeiUVb?=1ac>q; zEHD;$WRxZT{||#lmE@(UC*#;_o^1{XiD0%SCuFhB+Zr(|GH<60PUD|hh^H+O{YGU=-$&13JPENQojNUoxUfRp# z{UTLY(iLDl&;N0>_k_6?2`t@}Sq2R=G?@JuVb$-=>F{XCh^`X8W4W^R=GuFWThk2c zeU7)AolZ8e|qN?G@XihQR0Ud4C^U3NN^7;0WWid^zI>`=jPn>t4cev zL0;Q;C#8xMdkUIxD`%b#WLp+1838_LiNxnvBQ5B-8W0-!)&n16=*;J*k&EAejcKmL zPOJ>fJbejb)hoPq@;|g#ESO)#bOcS=+&N5nw39(!A~{JhYzMX|dS;WwS!M3|@NKo7 z4h@AJKt=)H0;D_eBE7`=Z)>zgBvHN-D2%)Kj;H6V(vP1G1hMh6Fw3`sWz%?LP+u9s z^r7uPf?U_5WA+Py|H;Bfb)^^)77gEb_V$9w!GO;cXh&klBY>VBuX^d&DyiKHVr_&< zD3{7jCkEc&3DFm?HQ-_)pz*G?4^r5?-6tGiGvCmxzWCF6Zx3w`5;z@08`U7M6o1e4 z%hwX{ZhqNmnBk6e@p;VkM<5{SRu8G7%$-6*kQ zlt8MUiJ+Umuj6GSac;bYx&ebcVYhBR{GjEm$VKK456nIM1FcZ~FsaA*Ea}U~u0-_* zQY27J(o7xe@e}TvH{0(@6GTVgS8aKxZE}2Oq^xnnbINOFtO)(Zjj-%5nh9g_#g7x) zD&U>>KZ)Kcq&O_5Pl<0uxNnn1)uFWknN>Fi*QjsY@Bv;QW;7C(^<&Y)UsNFXy z3LRehO-b&+7M@vY>EZ6JnTDMKr&42jV!d62fchP-si0!1-gNs98W+B}#g8g+dU`2` zkFX^Y1zHkr2b4$>e_ar>5Tkmn&;MjpWUBP~HPiv~lrg?N~Q zt3DUtJg(1bG$`{Vk6LS*zVeV7O*+-8`uM|?t1aIJON$zwu41~6tqVNKaBjTuxNRBT z8Hcb-X-k(B2s?Rzx7YjuGuFEVOdC9UL@V)zM$ zNzG}6$Kf?X$`sb+=#6_i)*)@Iof|`h%9x}FdJr!ZUn>3nf-|46Hc!qK?4osIS|Qkw z(0SOitAZ5weKHr#^Pg>6uz1uIVe~*_5-e9*WK_-t^`O3upfZ2c&I+CyGl@Np?^knJ z^!>3~o$JXyd@V{xHqIgrggZxxc9&yK!|3L@8!2Oj6qqPm_$0Rl2keY3lIL;|8bt)- zz6-A|+`Gz!3LwtqL$=gzgxfho)}2@W_aM$^*frQ`O1s6=`#cgMF5|IZmOrH}8F~KT z7Iu_Shc(LU3kTiCl)^27*FV}n$=JlbyIXF&E*hiFaQfkK#iS9B!CDirSkuu*J=Tfs zf@(%0QEO?&KWM{}ty19U3Ib20BpHo@~%!I(SfR=!fl zhDe8ozDf%76W%U%3v|Hx>1vrRPgZeuCa$D)Es)!LBS$&#a3@Cmeea?y5y6&(kDFjG zXiRDA0E?555Rl9XhiKzVWy<+Zfz}0}hF)|Y4ZMt5rNrwR{C(@-{rkP*35d*5^hjD*t%?QY%xON3sAu2Ir4ijp;vxTF2kxX;X8EdK9fO}7`tGy@m>~ zz`9lh6AHa~aty&ns-WWsfBD@JM9{a4*WJ^#((QT~+r2j81_df(M0%i16ZDWS{svlRQujN$!^_G6gmE}#Z~v0E5p>$*fL?y$vZ zln(yroCw*}1HFCe74g*SNY%sR5=ckDFJ7qeFB-{S%VH zCST#Y+VXZDeEFC%b|fPaEqHTaZ%RL#)dDur;oaLM#YqGb`+s%Mh!}@z6(+TprKAT}2#hil5 zuTl3`Z=AxXC$a+@$OtGo4+5LuQg(aN@O6j&+Op$XUz>YUMvM!ifBV3)FEke6zbs6} zPIF+o_+Fzdt#g5|^3({ml=BqO;wRMKai1N^dIA2yp!Bz4f}T0#ePwV0W2U_eCCb)X zQ_NowIRiP9P8O;`$yg@gYylO~uhBv)`B%s{((=pw;GJd}fRNnPmTsplSo~cFgqs{y ztlvqLvUuEQzM8^FpQ_Z@cLq>}WYW8YK}jeNjJ(3b?tn-OOpF}-6(y944IkinFCHEu z=k(-T_Vyb?vPs}g*RX_oWX&4LIsaW;d;kAx^v9I zyuH8>h`YU{cxkh;fsBM9P)a^h=4MWjfh~EK~}(X@E{a-zJh3g6H|D}TDdf)m<0DfR-bR= zX5`LOku@Rw<3@czL6eqfL%E0pAu?WW_cE{p(Jbrb>;xmFnJR-^8Lkfb%}Fc4ctaU~ zKK31L2j+?kQ^K+LRgj-Pm283!+j;6Ee^wej)$E9j?@V|<_cQ3-c?xxZx~onscYL2o z@Wgj#e2I?Fyql=%3H zI>op3l1!P28UH?n(^n8lwAosi?j65UuhwlP`#mSZkrIDt13;nHbA^;2cx>#j%d{Uy zZ7wMw6-_Q%NYcxaUzFkDznv*c^Wu0y*RsIeBzWX9&is-8&+0+K`svL3Mvqvtm=kr~ z9svOwSvh{vD=c69Ii9EYBNhD=e{^Ff2*f#Rj?0`dA+d8*8dyePtQb-xj;%L5 zbUJQA_=c@Ln;JIYgK*v5BFc~u+z!Pf;PulW=E8GAv3_BD!;wcat7rGdiIeh$&z;z+ zB}35art!k&mY!FXkvi;?ehgdOH?~0-C9+%K0>$JKp0Ji*bcr@dn-_MsVd=>DLX$u` zvdqMf6;~{Zr`=1T0ED28klpVw-qDoYmOGVP{NTFqoGbCj*XYvwa#ASB*-@|9!^Dpc)5BJg}a0U-)C$c_~=xQ4hjBU?j z7Ea2U^4v9+(RHQAcbUI6oA?g0^Br|qgAW60d!xna6-98AJ$&hJFD&=8 zr#0j=LCFZZzY%kU1|H@VrN6uADE!YAYr7ObGncA#uOP?9$8}RO1+|}z4xqCHl?foa zU4Kv3MN4#3+_M5khc_8lccnakJ+#+fZLiv6(%XHri0=fbnZXrW;mAszjAwOtBfPto z{?vFk9F-BWy0G(=J*uL0;1n_SD^*9fl5I&{=@7%@!K$#X7>Ro`W5$~ zZJ|uZ9U5^yZ_}AxJS?|*wZ*JZC))+=wTcOLyCQ@fH~`|bYm^W( z*2SFenxbg-^yZump-+kC3IY*YI%gcfL2lIql@(sI#Ks&9@{m|%9cD# z3)oUcCDVMNpuPv+%5gl*1zEH1D)szq%fnKk^-_vZl|k)Ri)%E8IOxB8R7h#BP?ZU= z2rbTS{@BAvuJe^{tX&fTBjM{&ben@sPQ<3E)Z-%}R$~`6G--cu)1lRkaW?Jkp@Yvg z2}VhCj;ih5jnt2rNZ<(65Soj=o6#({I=E-&2TU4;9%oCUb}zCOm2cDsy9yOF{a6){ zJ$l5ckYZI70<$1(n{|g@+K9|2qN}wK$}R!VDs*pbgq45kiMKTOl7*eQDl8?5GP!6t zdQ77D_a_M50Z-G1rNOQsC$m?qgpXi5^u!Y^Em`b_T}b=4RD-qJt2Y=JyP;ZY?VJrq|LxfAzO)-g8ovFTQl+b(0dard2A{SFGIvTM`b|OC?kAah!=P#a$1# zE8eHl()+PIg;N?|FNf=p-3*r*L3HuGz_zrF@1^KT(+|^Ib_nuks4iSYwTh)LwcNW0R+Yy!#rMah35&Jfn}D!u)`nZC6Vz z|eLI#v@ih+@u%)k=FFplON|g>GExp9*`Zi;`ADA z&KCgyT+nrME765VuaWB}-$;qQQ4mC}!DsotBv}0CgbrU(* zHg`3{g|Ny)d){bo2^h!M(B#yvHmkjHhlA@>h)cHxqn&?vL6G09em-Yro`C?y!PCr1Z7uTl1sRib^}Mh;McUboh6%=17jxwys0RHf zJgj0lUj7O)F(*Rn@Dyj>rhM5n0I|VU%|mH^Pc=>b&wR%D%{{}G zogi@Q?yVbv&kMZ3A3R7lTuZLXEWD~-f<84)9vOW>aUZ-a^hNo`Q$f;*wt+o?wX8wP zVTTd}3i4mFDCj6*CX-9s-%)jT&2s=UO|o04__6PbDr@E5+BV9}2A0UrjJO`fWj;Pv zQD*9yKs4<_;a0vC!?R7yl`BRbCIVV20X&E#ag0PvY4_}Mg{VFK5_ccft3MEvZHQh4 z+?RF&&&d18B*Pb$hy<7?C~0S2f&*$5-*tQ?MRFK{&)xIj$k91m`+7|VLSS#+AfUNxvJoeuaC52O$Ok}?&EkM8mcR+YLaO_%)`00_K-uw*GOBw zfBV1oevN+9SgTG*+EeFybFJZ>p2oBktosnEWM$hpTw_?d3Se9!0wOGET+AXYg}sk(stT4ckv=SbDk z0O$8E7n4YkddH3$s#2{T23s((Z4IE3p9E??u{2*AOcMwv|7kA!upJ_bDE9+4Asa2tY({xqpSm!NBaN^LPX}- zpoTJHWZv52nm2E}q5qNQOJ$=)({K+d>-+AgqJZOne5F%+XOsO{k6|_cma192_(MuI z@%>2G+c`Q3HM(&i27OCW*LD3-nwOo(NB>{z55g9mc568;z18{J=1V|jctw6QcwFrS z?_c|(3hVQRfC?XLr;>x77m(t{T7Ei*?pq^!%_7QZ-;?d8`@S~4j$8T`5n+F2`eN2+ zG^PfmokML9MOM?Apm{N~k?LYYnAD3(6$|w9t(ZJ_E!>h6W-;|!*ryNf1S9-if+52Z zo`S=embMgZNpA#~J3cp*KQr`2_t}^uwIE7|buZY3wJq?pcsp}=xOkP7itwDDr3qptU`FY5n(?rh!a5~N@h%ALHO)sf{as0$8n~I+ zXpqCCneCqK`qVbeJWEH92r8jGEOj#ytOqv*)$8BjP?V~yXv{;e_~c1tsz#c!0rVzK zX3@u7!+W4g2E;YRL)K7LZ#M40{UE!vE;d1+vb)q*5n4%MLPU?>1uL6`QKn!esQ=#! z>j`pM@CTp8p^meQj-{~Wr=2*;QM2t6QLUOD&}c$btfa01lwBg$k9jd{r5n!h^@QaF zVw%Hsvln*bu;40{C$Z(4j)z`Mnjvr~q$J|zi;*W1f1S^Uch6LZ#L@4}xMw_yCl3)@ z$<%$YRuo^2Mt&d};^?#~=d#kPV;91#GF=k5|6 zu0u^JpyZd^{m%Q|@L&H{px&e~xBa+{yjv`$;;lHJ4@*`Y|-V*Q0Qv=e=ie~H0&Gwm|$dw%kh=3%wM{+I5Ytg*>N*sk&dy$Cm`SWCw zdb}hQ1D9vw@NzlAY-{})oh8n-|JTXAV0ZY_qkY)|g94QI!z%djkc-oEgjzPwgB1Ng zLYWjMXpIo0<{+lDHEt?Ey$dB03b)XjcO znw;y?xpgT?ZF&yx1;~N^xRPn(y)rlNm8jaf{6ea3Dl5&I>8l@c`mZ4CnnyRAh7SSj z!h?9Qk)U$KEw;eeZ)n9l-6BgY{GNy8|7pE(;?aEA3RO?UD|IM16 z+Tp!tTiPu2Wlzy!$XB#S|0kYg3kI1 zivbyLXVy5Q{A&ZwzIzwAzrDY=EE_nszcbLdzw=nOx_Nmft}2;Tn1XHH#+Aj*g#U2{ z8$wopzbEbY6e|aXXtVXU87?rbeS2}zd2d7pm65kyp(-)L1eLN+Zua&6$PMW=W`Nyj z;p}Pb1RhV}3V0_@cH_5YG{@SL(wj1sDl497FGxV$Br;e15KEXS1K@^wN!(Cym1DzN z|n;*OJ8)@^4ucz;3gs+X;C*cII z8@=Yd&gc}Th5*!8$H$)d&H>(ach>M_om75$-phq#cerM+!PDF=*_W;%(h|N^XMZV) z=NT`NWD>$(&rSrlPb+ofu`vYvo-Gkni|0@Aq*?d*$e!(dJ_f^R$B!nU!>dvi{BQOA zv<-X{QH)yzS1em=#ukm0Tz1gXS8&bH{`GF7o14rt<>~GjGdyv*4mAjZIFfkz{`w+v zY7X+;%AFm?hm~*(5Il5f@BQ`!x4rDEIC;YF z_}0A+df;s%>=p`(@-e5sZK^pO#^0nEp2l7IV3H*WYB0@n&scN;q2PjSK+nD0#sQ}= zS;04U;c-;TXRT)x=3btPJf0#&Z!$u237me@H`MvnsTsei_BW)~BWMR@m$zOpR{o|% zsTDl8+%{|-s3#DBR$=uyy&IIlwI{f`DW48e3K73r{=MW$dx8(+UD;UETLDAp`e1(c zjOy}x6m`s)O%S=gR-$Si zpj&ByT|gx>_zhER@F7Y5$LJhG$L-ah#(u(;OD5(I?%dBj@y5XLz0}4~fQ$%$noVz? z0t;<|xkp9sN_-=Zno_S+EAVw);2>yR#)85)lj;Ka_lN#7vP-)<>ROVk zQW&;`MoA_FYLnrK1k@1*r9t?`R)AvXx=Y0slqNTD3JQbS`_d-h%s&&Os$xU;yxB!u z^(O6IL2L0mk!A+MSm6y!0F@%_)QI%wgY1;7^}oN8<$b|B?e4Icc=s?KKB?*-P)4es zFy(1K(|Wk}JF4ycBwz!hoG?}~&1jq3T@sQXgcF?K82u zflHT+v4jLg*Ye(o1`A9CcW>AORel%QBn$AiA&2d2zU`%Jukw=SZ&?g)oQ4{GToafp{HLZH1Lj*0=OyOm5c(vJojX+QfmahAT1-Sv zy^n~{C{teir@IZ&VH>Cr@IL(J6V^{;yYW?9D2?p;FfROacL1)KX0k%DWxp7SG|GX_ zfEs`{B+X^kYATx6m8F3zxaAihsAds>tyw_MdP%M#XTWhY3l_M8%!0E&Y^mAg>LeE%_u8(6EKTa; zCn?O8(;v`9#ScRL?0MCK0*)dv)rFKH>V&+nj8r$eCb-MkrvQY_1)dAQ+_u#E?&082 z(W|UGpPbx&d{0p<2K(x;2KypVF+Ij*Amn#kwgIk$dFI{S+5B+xkb}zMYzg%li=XG- z@-fI3XPC?`h00eFO2rMJ+Af(ZLG{oC0MsiJULJ~O=1;H`gN;k~EGDT}g%ZGesJ6ph zq%`(Px=S!(fB|H-Nvaqi82TE&KBiDV9J5Z#Qcb-gEnnOWAW%i$gm>Jcu+7S^VdH-j zFjUJ9rqLS(t_K~7d1STjT70Snbdmte9@=y=+8Agx{-Bv)h>j5(C%FNH2#mF|Mj{p^tgnb`NQ%1|n9-P#jUAA+?3A%r z%jy2L9un9<-wA8T(5F0c5v4MR7BbTWi-BdfbLgYsmwYP^R)2+D?BL#hLp=ibBDXn4 zf)&?MKS%V00T5o9SWlFqp?|5v8ip;!h?fp=ShIg3snQJ#1S-oL}m=XogF(+nc&F^?&-2vd55_yYr^nexvH}VxqgzYx)%y z%*xSE>^s#e{T7F%ta0^&YZf+~qIfuNt}PCUD@v)vYMbIKFcGZ=i9 zej_us`^%QC?n37T1muFNl>R|xFoc^x@5N>;+`TxIVLPdSPZqS?O^ARR$T=xlSNIa3 z4?uc>A`c=Oz_IyOwfaffpfnle8K%D)|0&(VlCQYTe0|n+Q0E`8)SgjAyQZcr9tDI} zYy+}>Vp!xQxHn{KFN&2uQo2-!`V-l?Pzg+ z#jG>BA4wPmB*fOiod49E5WWO#>x`0vGQ-^lVMZ{^R-J`~O#m4O(XkMGaZKh`+O0=s~6<80S84P-@#1djTW2>!1o0eu63c-1V^n)xqbiN^w zQB0(joWBi;eK;v*Xz#-&WJSoPU!{7O>b--?{BZSHH~A64Ai&O?CRqTaE`iSI6x{;- z{bhoGAp7om6*8)Veo_~@z=XJjSr}4Y9x?5K;Q7bkq#gD@_#7r6GwsH9?uPqe*Ap>7 zzJk{CJ7`$QGuantgUtB{YLQwk>bGxd1>9ip3dCA8fhiNs2nazk$`y2?BTt~oOpne01iTjICRNts}G7p@(L8yC} z_GHK%kM2zveV>vWx_mhq0+RpmoKVjsqqPe4qA52_&NO#g)&tNMa7ezD09sLBk7W9| zfNt`l7g8VE$^+b@H4e-sqFQ>EBHVr1I#(oK3I5$T^CzZq|#$amAM%=NZy z^tx#w=_#)zcwa`@Lnf+cyODlyu7zK~EA{eT+^X+cJhvTsGMMbIA~= zeRU5{LYN2eRM%BNP4&uF+so~^lI@DJdN9`T(q&S|g`uVSM96~W7r8G0t}IKakrj)a zj>p^a0($HZs|^ulygfz;>j#uUL#(h1FD`2PKD9XgX);36a7)_{@zFXR`R zn+EL3u}H5z69%^)q3PGrki|$OoLL$e>lLw!7T}4?pE2mb*3`SSNg=??b7wB*6h~l36rc>07TUz zyWe#^^pB%HP10nd{=neXfpdZ!7MMv^m+-SK^Xv2ReDH%>guTn;)CxduPX9GY&aDq` zXjiO5oH5feDb@)KgAk3!TG6!0CyFMIQuUz8Q9lz^l`-N{7cTzvvdJ zDq)-=`QEoI^**=TJ6f}**O5SG_Yx3Ne6m$(6Oium)rA1o1p4n<5?g5V@17E0dpKXD z9fE$KZ9ESi|Ysu6M`80bM_+q=O{VuI%7>M*esA7efuzCoR4fqA@xz+e&%gREb3+P_{X`x;_FidC>38Yj=4^6sL=aSO(p%qD0G9F58|x;V|l;Q`MW z945GDqqRW{1(I1%PH2=a+;{^~xwZf&kR+`F(c31|{M_@C?^c#)r}pwq{Q^|f%>q=A zYZ>S9zNK`jdP)#WWeW<8F1``Ll0jDdy^!?hTk>Y7u01{nQ?GA_EdKGdB(AvI{2^oxWR z!SVANX0^>I(=LUi&@N|(dEcKx|2@;_MyFgQ>;EswD~<*sIQ(x82=tK=7m zbNpPMo4cmrs}jHz0kk9)s#AqHyN9JBQ$*y%(uV1kGu4Knuq4b-b`OMhIV;NWf{YkV5f&nOyA<}cFMtW$MpTZHCr&YCLYx&8 z0)W!CHT;UV>`(4b?{=7sUR^+NI*iw}%DM12CC-n8S^IDNw8=%}^nBjn<-yebZ1bXg&(|9P z2`JnqPJZ#r_!8TA&irzlNpnbWfsAWLrjxVaWd?e+)IgVn{vURdmC!HrJ_nBgilH-l zB0QJN(F3Q2(B3T)S|WfESr3cnE0F-g_V=pKLUC-IY`sRZQoUEaz{jYZAQ2@{x4qW? z8PSknM182W{Jq`0NJy@4EvO6)`&RqDv8VEcl|tDwxaNvQ)FEe*?D7nzpnorYnO*1Y zn80KH+mD8GpGPsdJ3xI42<43s80JehB{?5-@nGn_q*=>JYtysf)kHn2jj03-w+z~K zTY}`R=M*3U5NRCBKqUI1wknO7A9}&TtDgKRL|n3d4c9|U3q`BX)la>>Gf9Wm;cIa) zkp@g&!fNT5BEib^{dB@#@^9{9vcmj{-LnM=CShb& zIDHvOv2{s~;E$ZQHJXtRF>0WIDN>wxcZu)%7QIJ56hA*!4dGZQJ@t|zzX2S1qh;d5 zEGudh7XUR005W+^cN~qpk$hauUT#Kt&Z0+8=ueTEf)D*#Tb{^E>kfkMCwi4X5nWbj zB|pRXre(|HmC4+hpc%O%#<}hxmwz~)w`F{Z=sp=gh6t#z_NwK&)KN%56B%8l6 zTs>eUV9oMQ!sVF|i3uPjsbc{;Zx|w$iBD8Esr@PW$%o$J=bpm9^)U!3q0&RyTue}UZ7?|@IZozU~xy>B45m8Oh&Y6EBU$!bZOC_>{z&v zoGu0WCNDF>xJ#B`Mi)(IzzplV!TAP`N(~%6-5k%BkSiA2T}@{^>CQh8#96 zz!dIckC*8aVuycjG|hi~&{dAUvLmJbf2&At1%k(9X&&p@kq!`@pDMh`VkH#sP2%3UXEsE?`EOov3Lm3e#N( z`ky#8!AtB}nfTTdc$(tuDZ%JHh{n8lr*UCg)hv9v{u|-n`a7=-(zA~g(Q5~;(DqcW zO<2E7+b;mOb|}74tAwO)p~?A-&M|jf_Plv$_1*F23%j=h0oFFzabq;lQa(4&Q}8^a-g=syhhD%UmqD)MXp{luhp-c+NQ1KauYb2LldaAVx`0xa(SoA?0lZ zB@3!Dt$|TICAgUJRr9aEQ#G^7KC|0UXI?k{d%goB;8=5-X!;U6J`1##J%U*I9?)9| z1!|ZdLo0%kgzXD^Z)+_IYp}@5ZL}4l#8}HT(VcQm(c}5#bnpJa z)9M1qah-ZO+o}b?(`F>*)5jFKZ`p?9HXiR99dBN>1^=f}@*)*I-@*v0Gh2Nha`aH3 z(8#aW+)O!12bifpL||tyCGqT6yri`)fJ`J<)>%eJKIoI#QcW zdZE0B;vkI=NwFI0(>qZT>y1*Y-FBakC)qv_jhMau<{P0jePX5CD5)IKOi<6@fH;!c z^37m>beo!$FUqHV+S^PIG6f)_=xa8*sPG4&{j2DG87O2kEqC}@mPmH3tVbsB#Plkf z2jGLL@|98b-awBDDFqHr3@8Og2^4Nyk#!0kF?o!;?>9)Bc@73ruSnbFIFr;{{|XRx zA0pN0@Q{hYb`THDl*ozKl=0g`5SRiH)K`1o-Vkx0(390)z6Od}1aaHAzWhDl{%6Lmu4SU>a(hN4G4OzZccsXYYycy`t3-NI(DW7D=_#&1k z9z6{yHv)r%)yprIao;2Mz$khwpqXf~*Mb91ARF8mGwxR&80kdXUw{-cL=pIdEyU+x zc3AF!#NsQ80@eK~bjBRO+hyxp`05C+={;fobob>38e z2BnJKpuhZPuKn`AJ$StIt8?X9e}t_q`s5^#^0ph}E>McHh(>|%YKvY0b5)CAAhMWl zlLnKOkwo*()o`CJb+q1Sk_G{Q>1T9bGg~f5_h>LR^dkV;F&gwF6{%gh%$unAg~l#x zvd=6`S3kjmE(Of8&I!Qdh?~VgO?#I#XbJAyYwXvXAtn07b@m+TkLzEC6GkRL+n~;4}7%cJ zLUdqD8#K)D7-e)vaD;)*aB^emf-*nfrLuZ=s1=g_h0Bz3->YKXh_~k~yE6!>;igR& z_rrw%Ag_eZ^l@2#4EQ$`F;bHeTfXaF2lONZAjuSOXhNd7$Kc(3A-x-Fs3&qUNNv8% zu7SgC=YHV!(VBu=z>Go7of}eFgx7uC;4yomNTNj27Mkd!6PSN2YcSTB%W4!@_j z6a~D`#(>W#eYt3wYktylK76vi)>IqPOMxq(kunI|bnUU!P5Ir^u{w1%B&ga<)8a_t8{akg5{ZTR=cOO|5YCS0jA))0Zc^a62Y`b)|JH9YAsXV zWiSw>$4JK)>*&08k| zqXZ72=q_4{-H#(8;f256Ka0ln8!NbC82Wfz{6 z=Yqp?okb*uy=;QzsB@+bph>L%CG4D{imr|=9g%UKlsbQ-BWyYZ<6ef^g0xYGE44~N ze6;%~KH>zI$}`CgeB!sb27QX2#kfP`G8O?UgkwX1-XPS8#}I+emJ~{%rJv$O*xy*0 z(Vg^Bq=8$ z@Ec5-IwWD#OP_uB_(eH>&;^6I$yYowqx4*inx{3mYNQcsrGa`P&;_;dOZ1tx z-~v~Om3q+9Y6pmEYea%&OOR#M?KlHIyBRv7*2{ev6X20@pm`u9J>dqv*K6-1$?uM0 zoTV|@A1^S%XqCV92YgO|YoqDCahn`C1Hbv51#W{C)8fRw*-E(&q1EtyUyx^@P;t-16$#4$JTqH9v<09Q z3cs9qgYI0AF!Io?jiHuNt$hEz1YY;mf(F+rRULB1m5@z|(x zm~qQJwXM2vGB6dZMQvLO&q!oLfDy<*(g-9_aq$}0^%@^?ucn!fb$c08Zm6)j!hFM8 zu9i;Gzsm**W(G}DiV_d>R?&#(MYMP+pJ|aVDe^a^@3FSfRk~8T+fNAtMUk!>G$@8@ za_$%|F=i&5%-flJeCCP#`nx>=cX`%3?I8`1dVSS|R!Sm+=6mU^%qQg>MG;!9fVw&k z<`Vz;4LhLeG7Gh+5Y?a}PYU#bn{3H<={5ur>%Bxn9{85RohJX8m4uUKCE-ggooqg~ zIq)_H@BDem1&r0XNqw}s$qTOUT|M!urLsT!9jB)GwGU)#U|e#>bFC>LAfP%`Jqb>8 zqh4A}YZJ(?5=I}V!0$@MZT(Ix4e$yD;;WCNlMWAX+h&WpvHcP3?_(&~hz-?*^3QX3 zf6OKUb%mhT8KC-*rp?w4yqS>-x)#v*0j8og!E0sCEz>+oz~1NF(GQK9`=rK&Olwrl zP6b`&dCJxGFiLs{TJN&O)&1C1qt6^tZdvQSr@wDX#GKvL_RkvJj~IOLM(+4~St`nY z!E0v&t|rFAhrT}sU)Gv;2nN5ffprqJ@58A@#)at9F8TW@8dPtP!F&g4`tN3aP6So6 z$y@qJS(8pp5N)^ZHP9Lp8JCKdig=u!`TFQZR90ke`?Z6QIekt`8zNYdIUu#qTnJpd z1Won)YIj)7bLKY_#00!p|E7L+BolnC0xwjE2OnL!#1ym+J{f=${C>JM?UfA_J^g(I z)tqVtP$y5D$==W0D9#)R#r97CZe`Y8mo^qE-~VjniO_pye=EfL0?LW3SD?poPp35d z1FS}^8@hG15;OgcH82f;My!zvm2#{nqN1eT+jox$1!>3=^a(u#Tf&+*h|yj zw8Q1a^ZFF4Km+jm>FKmrHQDzD+4l>!l3||3anJDVW7=E-BYnDigzrza-lbx^KO+)e z)F?3|bpDVw?eM6I6uR(0zb*~J*=zlpu>Eb{Jaet{1~l#rwDCOYrV0d5h#LEm4X=0p zfT4Ozo*9T49zI;we+;#S9wJoOC$#IBWtuK?QWioimLsWmPBRJd2NTc^HWsmR3L*u%01HX z;LxBwINX{q2HyxZdMwFhMdz6XKin2tcnJ<5{fkb50Y#1cqKwUQW%o^1j<@twa?k;5 z(Qk`j>;)qyYccaH?{lYp4Z8$Aqs!fkUm}2Yx?ho{*I}S{_{w^R>Yw%YQ$$5FZ`6m0 zZEYIQojSt_BK>DH_=?Q=i1Xvx5J39Ry_byZ?%bQzd2SrFeq+BobWg8%gN&l`ZNso& zhUC=I`|D+CI?ZH$hj&Z+DZr;#iiIpUs#Np8-#3^VRgv#1+O3OsE3WWUMCP29TX+Ep zXMU5dbnA=4z*j@ImjXWPJfT6HcV1uDlBjGO3vqCLT5V|3kO60{k%v|z>9@CT?7!v7 zmzge>WD%tjaMqu@I#qCTG2gH;Kh=QKGPB``;Ft1(Z1BaL7Hjt_jiA|)y;n6t-y>_{ zsgCSWXc;Brac*p17?Ol7B$wqF)_)lQDQ~*=CIDm7;~|<9Oo_uN8TK? z#mIY$@&%c`fu)-G5T`iZvs$JFXG#&V@=X`O>NwDA+w&B6|uLBenuuH{7$}y{tb~P~z{IquUER+k+_)avw3G%h4 zS#{QBj|3id1DuXD>EzFCCkr*u;Dp|Oc8}t@yIh>U^9#(`QvW6k@6s3jqHWsTP`44Q z%=_uC#};kW^Pe02<+F{zkV?T~4wfWiKiBzOMFyk+T?%_BE$Ex<{8G&h{b7(cpi&f^)XVVxPwc}s8K|1kbVgwh?lIK3<& z5Quot<9~hT-H9LF1$SU+51ugWg30{o^JrJF#;GYLM!4A_g#W&Q^^($qIi`=l$ zj<6WB_)sKv@VEtj#IZ}NU(R88UxT;4_>dM%9PzF7d##j)PxO=d>zwGZ^+Or#u7fZR z_+#h3%D8IW$5WrtD`eqk{UA;iZK13+V@)eyhBEGB^ciUX`5ThJAHiJp+-y!~RQh#C zFemr)#*riB`p|5}mw{lS4(%=R8>e)fbu(~EtskPg&U+T$@~h3XJatIv*GZtXIBX$z z=pBX>LhPd=?i;iP)1A~U-8c&ypy8zp2ebg6I4e~#H%uoC_M7ZjWZ}&_B5@B03y$EuuS*eUNSUuSlcVru_d1Vtfn?ua({Z-|ey`q;k6 z-WwmB+T7oFl5M{rqZYL`3O=xsDi^~4z0X6de4PHmuM)oXrnvqK#~vM`0$fMXe64+y za3?AzMnOMF_*g+z>RdZpWs9i)r2uW-xNkQyh#Pu^_8)-{P$z8PzHX;TO@0J?Mg&1u zv!F>aCcPSRJ%e2T0nw)4+%C>o-^S>OAfM;HmYFFqZ_Qs|xRSq#;!Nb5#-o_L&$OgK zHhlA{Uh=Fx8^4_Ewpl}Z5q(@&9=VpH%e}r(#U8W|rykJ$e;KShUR}-*eQDs!_ zJ)fXDI=ZfxL(en9@R^_onON=TIN!c1E&AK`U}*Hezro^bFcLi#+vR+xhZeiHM_yg9 zw|xX-)-}O$$vNv%{6RtoQLY`7Zrkgts-59~I@UB+Tgh-@bn z9>TX?X5TC=;0V6zqH{e){WYIFuTv-HLV}FK}O|Ebd#BG*9M8G zn_K$KSeZ;{0i(jCAXTO0N_W*VCmz782STlT;5mqE z_^-y&GztHYt+$SgvWwbA1*Abb1OY*55K$39=?0M)x)~${0coT|Y6c~wrMtURfdQpK zI;Fci2F@Oy_dVY^=Xd_&dEECs``)qkwXSuowZ{}5>mWi%bt4;>4mr$*;T}s$R zIi+vp&oc1RTWkVg61dGOzV=FvRYQp+Y)rpIFbF;Ta713nz| zdRHeH19){Z@(eR!F^});48o%R^xhmr!JR8~n66ff;eg8GiiG@WDJj2*j5zwT!2bX- zFw@(syi}VHQ+#6r6K9_QV0Mk010#d7ZO%!sqB`|1--25h_Tu@*+X??bu4{kffiQA zPbf7`zBarM8B&@zKdbS*3Rt;?lZKb%K}{m++n#8Z$=Q`DRsT-6X#zCcSs-y)vbbk{t?vfv!z`&TcN2V_A2ati zWJVoHA4Vls7%0&%Dt}Vkyf)E44}^OBVkvSTFJ@O2*E!{_z}7hhLlZTw`RMm4lM`cz zCIxN5XQMc5t-tj*TAIyiuB|x_+}L%*HA=lRcv78HosIbj!a`B>TG*-%fQQ< zq!-IoT@+w3;X|H1o4!M+xne#!4x3(25-XI%!2PmNOIP024lJTMB;D=n!6eYsSAL1w zQKFz^CJ+7M4;%>6OUufobLBt5u1VP+I6pCrzAj{+u~L<6$yoFtP>l+EEtuul;$76u z+xd?;Y6nfK)X%3LQV(u6bwS9?g&EM!F0d@VMJ5a(JM#~=^+K!oLI1&+fbs942mOcT z2|m4?99=+WP5v0$WRxp8_=3!s_Eg+7;3G%spqe{@<{LkMQoTPB-2|vG3$kS;t~Q#v z&39FcLHI>Oc9LFHfrNs%CDS1It%g15a1j)enWd zW3`yP=2}uK#F-B6Kw&vPJVx%1Z5Wohf5MdCXn5vHi9#2M)J^mBmo+36WJpdDYMVOj zP8E2o_akai3os?T*7H4h0(8h#p4H~2-axH3m)~26MEFI@AvTQJ!%zz@r1ZI+`bYaV z8iHI2&G6@m>@>ad+x@6_)7NQ|26<{6`)Ml+bKdQIR@)qu-XO77Sm!(H-B1qDk*TxW z*1XZ6^>q|w!sIErKBx4RCGpHBnUbo4qlFX8OG1j^+YZFxg-oS1=l;)6LDt*18S<+l ziV+ITO!;ir%+;zR(^Srm{c7~YCOL8P^3&H63nvk~x8DKfNUCj`CCy4WQbiM2nHcsC zv?a1yIN^A;ZW~K&@;%rCFsF8xNGm!>h(}oa_9K`Pri&x1jbWdZ(@OrLOy+r>voGVE zE?Lcu)>TsyoVSC;4xDZXHQ8z54)}kaLygYx8_BOmk!*$$xaaB%r7Si)nln0g|Im=u zV3R}M0-s^-(SMH#+PN1=Z7&Zo?}cCYGnTn?RBL|+rQlVP@9v_{172YvE zZ^u|B6_{i#yfeRa#i26LVJ<}TOLCEomk@s-Cqp?TKsL@t+V|KR_ zLbP^`aK0QiIoZ=wTWdOStBB`En8}HyxE$FGMpp$^1#2*BDruB`a;VvWS);?=Q`pAN zbgeF&GWoF_DKYFO zxXr8p_w-w{XL}pCjL)Nu8cGNex+>+G*cHXG1>_E$b-uRI*-D-8ALnQYVy9h7&}iE9 z6fg(_BF|ug1n;F{yCdp7Y^g?Oatp-aDim)QMP-sKtY`sFEg^Qay|+p$D53vb^RZ8Z zdJVb%Jq&FU*wn7&GQW3z`uCUXm)SOXcFfh0Pp#NavW)(FfXyryACz|gR0EF1v?mQb z*n<7t+fJUr*K$E6uzN&2u~~P!YcwlbTEqvMbbDR``Br#I+}7>a=63INvo-UKadf4c zw~y%}-7`8ab`uhiW?ru-8Q1!hV`r3t`0maOYI1E)d-f{#fd3KeaSXs)eBp=PRn$eC~>0o=82eDS2WU z{a5JT-~J18>tnT|G(aHaDpd$1W&qkF^)u1k@Zqh+ZFF#qmPssUyw*A7-7rGW%5UG0 zSzQ3V0v{0@sGK^OX~uy2Q2Uq8m%fYetts!nAAy&lkYiS&a$o9$mMsa4l3}Bh4r|8B z@?;3IN#8NcyT-zYg1Y*$y}%e`nh^m|s^16x)%U8wN;!~tc(0Ea#+=@dJ)_}1$Ix+a z(dbR`yY{AqH+2~h%abL1oQHW6U}tO$m`+*l@?`a}i&>{}%xnF>#b7;41?>Xk=PXp%NDcvp`xKlBZ+R&$ z?L(2O<=0Bu7th?YE%^v+m-w=W;A3+cYhAID3&6-j822xfr{xSE1EZEjRZ@cNPRL-v z<&onH#eT#QQ0j~a)=tX9hGv08+qJiU@y^kk&3~OQZcBm0a(PyR5@-C7BH1ME{;9e2 z!Kd#ztQL+rmev;%*z_xO-`Xrixe9>?u$%8lh1Ui(WRjI=iX{zvmp@cLEysHw^Z|m_ zR7vT+f8jn{Jn`V&7*xJ7*#luFTn-#jxfp}SYv(~{;y_CWjb__vjr^JrKJrV|Ql-!` zxklk+_hMU z$8Gjn5SaKz;RB&k@cy@1sLssjll>%LyUM#1Uk6vgHL)1*Of1kNs7?I;c_a>(U&c53 zN4R>1QlXqoaHusJK#);e?$P+txPCPUS{btDij4;R>S53xqHtFsP&w;%+<$ZA;%~Kh zkUF`5v2YuEbt>~mBl_j4-KNjvyRkk5Ao8{odgrg4tS(Wt-_if=PqVQ=yL*51Ea#8>AluvPMj;2O~?I zcg(6(s7QV?Evv)YhG$M6DEVA#sGq)p?~H7QJ&&Lbj`zCUuE@c}*K|C>yi2^anIVti z83FvCwId(Dow3vF(dBg0T>Hqur+n%=h1P4|>K`z$)57u|*8?`lg4_}RHimE=uuItm&Lb}sccQ_AH4h{3FlV(pYGf~s(u zaS$evQu3=kaX4HP4*`0czsyrwy{V~ew*y1@Q_G&H7ofd6)wWW5YACM61LAZ z#5bnucD1Ocr%%#__P?&HZ{$+s|ad4ggScbycMGBTaV6r~6 z3oxYF>=0&`FQN)3>MrU5<$CN3Z28kKOqpiTo*DKT!`;|^^~>!iAy^cxVTfP~m12Z| z{^#*WmSZ39c*kaUqoOLOCBf1WNBya4y^=pp7-c!ZJjmLv~!+rF9erLt(Gqx*1 zRn!bra&2JAaIB-#oRp@SykzbRo$+>!h4^H-K7KOSDRqcHTv*_GsVzwpXbZOzKaDKh z^h2;Ul_bdvDSF`Fo7yj2p>@x+~ z=`y^{+>z-XdNB4ZfVetaNODgrOK@R(Qnh-5S(p9#Ul)m(35I#kIWklLKdQKXJ;ga= zC23!)_kV`$mQY0v@V8pMXt>u0`EM+Zxev{?XsfdwZmRBG(XU9ZC(J<0Q^bDIpFCfm z4svcm!k>dkSJvM;?|d_jJOL-Irp$&s-((Z=u}Pgm3}9jz`Ps+YQ0bKa%h#eB-{iIy z6LeCA!s?x#U0>3^A&KQnamfUf^M(eBPw4CVOvwz8y>I4SmlXlMPL&nZthK zQJ*;%2esmebg`BDeuWEU{2!(}Wd}7t=QFeZ)_TiItY@2h1WeNOxHAxSo@=O}*s~$h zXsee&kRkN?Bgcw6gE%yVmw-IUmQ~2H^X8UUb8A&WI{(pOj~8*%DjeFyXLH^@3(M7i z#9a;QzC=)}mV^zX2)u%QZv2w^mzvM_! zDIn3%>q!!Z@qx0tTz%qMDN+yVE4WL3I0nPdej*Q*C-%d}Cip@wZF%f&Uh%WwwMVrJ z=^#uL)oNgZg14-Ypvl6rd;T6wD^r7TQgLNy9@FvdtB*QGfD(1{}PN4%X;P+5cLD z@Q*!1oxMiM{^XMD{=>~jH-7PDL z{XFrqmdG;Iw01+1A-b4owUDzbikX6Gk@Wgr{w(`2z<53yxpRQcf%1i`$74B}W@B3l zCR3zR)HRt7o;`j7X4;UmYu9R@e$^^dZO28V0w}t$macEVfQ0IDQrY9paAQHxM2=aD z2;Xb-H(Bo$!BE~69}<$hJvo@hvn=vyvvt9uoiH9*?_#GyyMitXcLg zm+{*cskDG#>_R!iPfn`RTs;kTI95i-*PdZ+cFvS*u$oDoSCg8gKg_Z(5q()j$|yy; z&X7Az7MSM<0$$7pFKo4!@i&;L21gWy8g;>BwKR1#RlXEQJpf)`D5;zsc0?4KUWnmI z9IU`m%Zm6jt3$dcET($N7)@owtc|Y#(d7T9mF7atNNoQqS-@)<$B{|pmc3P^#`eV6n{`c%GfcivRdAo7`0z%fJST z-6|mq(Le>h^-$H?q@h}w0&2F%4$s5`ahg*fz#2KyobJYUOniuGi9F}vWY0=zYshLp zz-K$jTqbGTj_N?h7T%G1$j!Se^I!Y%*u0c{#>@phmgQU#iLdwD(4tYxr{E{3DF)hi zA$QE3YwmmqF-^&|2R;_q$M%Nr{|>jfN52oJF|^^Rr7cE~cB8_GmEablB-J|2kk1>= zz9*Ey6lwAV|81CSyC#c;LZ-yi1~u+x@4M4hm06l~-jaDgY~`2fi;QG2MuLaK)czY$ zwH@`+GqswOWX7)UI2zHwh`xlbCP#|P*MdVax4kLPkXsgx6j7}MRAq@)*?8Wx;SD+W z2FYr1@6{E!b$RshPIU7+VlLumMAHxTh}b#HVdhL)VhjBm`Ya59Z2;t+P9wXg+h6yERcG!h=ZvX3 z%Vt0ZR>YWx)`(RnKYgnn?8_5(1}(TB2b)q-J4%{m`{M3H zNW6rG4^!#?qxGpngso;_(KFtWn6<$ByK7ZyaH$2befMADjQ#!gE=%VpRVEVf;^1qy zx}}(v(sJFD6!&7>8ZIaUr%UpiQIlIGxlTKH3~Y26QHHX-$GwucobekVg2bi|{- zL{g!hpWGzb2Sszdr%(dyRiFrQO<206;Q@sJyyzvl?P<3ngmIRqX>j()yQ73y{SXzr zs`l9<{5{$Rs7xQGPhWpuZCRC5&cnm4%c50=w-^nL$}Xy&mVDW z!6Gd0oMTj>tf1~zfy$al?+}lE&PKvtbs>+gJkXr&q#cz(Jy4c|650Ynm?0McD3Ie~ zD828lvY1DXOK;huEOG97_L0IDMesf;Icp~rJoy66BnhuXHx#}3I6?=vbA54^PS=Ed zf(;+SXfQO?loGVm`AfWRT%QH|3D`k05PYr^>r7U@>xS`~CRicU`+9)@2#d;iu;D52 z$0_Pg+;sMeSHq28=j&>v^$PE|U<*F2M+x1rfMY@a{|Z8fyqPN zXdgqZ5Fu-{30mbcRsu*RpoPwFL&?@gZl*h~_%X_qMW| z^x`|@b|8FpjaoEPdRyI>+7b|jdpe#P8>|duYaDk{ItSeLt@%6Ops~-$5d7gt4gnDb zqNtsDFD%1P72{p}aKSHM<`kY&1}qylB~H4a$mT$y4`_dAl}m$W-fH+ThR4JmRh0Y@ z{K7*T*hyOU1M6esh`WvrIbZcvu6p4C01hXan2&8mf1%D~up}pMsdD-1@v&-Bg^@yZ zTJH9S4$k}m8`syGXee>UC_0oo)mVuw8`dtmF03kY_ONG%x!?@XxWHMUKw3rUdl&Il zlVe$^7+2X9dNAlq5hebh5|Cd03JmOs(K^Hhu%7$_h!5+?BdebcaGY~F3EX(ipKIwA zkW>s@c~v!<)Do7TRDx3VU=)urI6;aHxRor!m8Xi>n?Wd8?xlu_e4qqJ0n1gZ&X!#F z{D=#KDsDMyhHjaDBjA#P6T+IUIUs$bjr<9R)Sk*cx6lQe4`p9Gbm9A7Y2R;ysNn}R z!>1c$sL1GtGM*MYE5R?r!(G3$T)Ruf3prYO3fq$_ebDa(PXc0k2!wvld?|pUp8I6E zL5i$&m@`j#HM<&;_l+QhqJRkmHNa2!-Jgo$lXw$(xPDN7M zeiBZyY$_XzA2&ZvNXb7XAe6`_HasLU)Kk4ID3sq9I3)fpe2bJ%N+KrrQO9#5UtKfg zjBs=2p1kr|sC$;9-J={D_4o~EdWq=u0C>Gx2%llI?rdy{YK~=PeLjx1{40ODB5&&~ z=D8JY^8PA;UT7fHuQs(L{M@9VLK;o?m(H7P-LBmJp@M8{q$Y=h7f?y>P$FZJYwuzv zn(<2Tis=n#rh0$GWMdiP?j5V9+|57dLJxu2XVUB5S?<(ET5Q@Eyel(a zjN702eZ^PsX!r}02Z{0Gn;Y;#kaozW}#r0Y*#&!hBO@@@T*vTu-(WG!3Ndut0{RD_LPU4>{R)I`c zPeko14s@bM{plU#H3J1dtCR4u)%=ay;oE?seHM|*HD%ngAg&^MF`1H-4F7lglphna z7)1Q9({khYXxfy`rlJLYXNcx(bgbG^b@%X-_fwEuJV#>>slnGq&=G2g-m&Hb*9uXk zy7tt)k<8#DS4N7}MDkG!!`F$Lr1iNT0_Bd#;N<*!HD+Bj*gy=!qSQ_HeVAy>9mFH2 zx=m*lp1RR_MY^a2mX~beqYUO>>$OBdoQJ;CI>-T~OEi3)oc!N!G~E4`{s0B~CiB~EO~3L>bOZY%Kad%ITCbGr)&7Dmao^&h(( z4HTu6cIb0SgMO_ZZh!g)c#AWu7Tkx_n9=Ch+IJQr>Mi_Ay5z{0-^hddApv}idZW2; z<}DXue^Qs;)!i~!=*Jt&H^~y@V9AU zylYrUH#|qxVPG+U!U>pYC$Ompy`BuYvXO-|duNbTF1nO)khJ(fW1T)HYj${7g3OzM zLL+W;CjlkT4ZA&r=icjn`13&jtV zh(VN63|tFyGO}adJF3Z0i8Z-y`jOgJ)9LK$S!J9@+tZGzhx$Z6;tLm^-R1lT%&C#A zVJH>6NT+pKNn$=l%1?8-BFrvRYh)fo<2ew?t4?XXAM4KlSl!b`gGGF}b6QbyA9W3i zp;wUWI&`opopgsu@Y!UKI*k`vYw_Ey`<1qL8I}EIuYTh_yxV1;d@|Jh%Kxd9_EiJH z{IM!a<)&&xk3Sw|C&|0$KyXlYSLe#Uc_Z#)BV@fwZi;3YjyjkSHh-T7uD9 z@xx6~ET!K(8nl><(MH@rk)u}h=7+l=(jS=j9@t>&d}~#tm|GS3=3ftF+Dfppo(b1z zCdfnV;&y(O<(X#y)ryuYgF=D)TZ2Wk9s-KEfXv$}RN*r9DYt^$VLRa!F5 zzSMIyVqwW3eq9Qpxt`SUjr#*{vB5|OW<4&ajei*eCD;6RLK z+Ol<6Oi@-)v#PMHFz~Cm+M*O&Ul1*`5eo{hb-rsT^rp{J+;m7WgS(*SE!Vr+zxlsE z*~(MjG-qS!YX3sCq{Fm*J6PsrCDj>oc;$!${7rYEj-@)hwpEqNZ>ZzMllExMSmrw@ z4gCu*E~SWtZHQ{|TXl3>;$*zaYOt9p1cDTz$s*2o11>&Y$g2OzvzyX5%;Pu zPwyC*3akB+Ib{z8O%E{- zP9JcKAF=8Sd+QY)e+!~&s2%<5g-}WV;E+S<0szhdLE*uaXH??*)!ErGYDa7A>0LcM zskJ0NNN3y&ejT_diP5b_$FK1=R1S7#{PoO>X3a<2&HM;B2giLV3OYTSdm1W-K2V(# zUCaArXX^{{L6Oh(IUsE2;h!X2dY02YQb@fyKB;2q1VDlB9|w0@0vmY+%Nk%rgx(gY zTl;p=UuEF)0G6Do&d22fDQJjW^vFnhpa&m55{D89Gy^tigZD$g1mg=z0e$AJT5kW1 zlLktu01&U?7<|Pl*SefM{e#I3cp#E{G{Tlm{C5%nH~k+aFzf_4g>V>TtMA;ROj=lf z(6fOAJnv|>xT9;YFVpJO!rrK*mqx81FB(Sw@`W}aJPK2{{(`7I*wCHnMIK@WGzJ0A z2^~dtJn~3zl{g1)i;uCAuHV_$i(*J+K;a#URBqOS8(>~g)X=qlb8<;+*oDsHf|J3e zwEbf}(|M|=0Ko_~fATLEDr!@+#Xs|_1Y>&Z;#SQRsw#!d<&ql+d0*^)VlJi#)|P7C zUZJ)Djnlx5bb!gcyq&Ph@mQCNK%;0$xU_mcVf~RVp^6?SaRm%y>D3T zIKmw6Zw4w3fQ|SzT$(H&8S*)toLmqvFoN5fTD^AI6_i`RT{~TXr7=FjYzJ+rv)O>B z{dRa*JDNS_U^s`=+kc(d^r}O(ql@Ffny%vo8r#%gn|CSzHCEobCQ*py(6qa+MPSyj z_c`I@F>N!A(%OYgz`u(^P=2@gQE0?0PUXo3m_cCR`k;oiu*hZ=24Cr@HC&lc)2Ab% zoOWN!w8EN=sZFd0s&b;^pj;=%KTdpAs0B7fR$myJ8h$xvd2gy&-ELkowFeqlTpOOwjg(_-Q^Zd-X zp*&1o9Y$|k%uu$<08XP4Ez}he*mMIz0k-f7=XWToyi^#aWZb6!;@MO97b)gewRT%! zQa9?3SL#Gw#@r{N%xrTO+}YIl+P~6pLYYV1CsF3W1o%4w{cw8DZrug-{#UjkiuZWd z`1+vdXVsJ@m5QyWyy{YbtXsPO^YZpXE5L^ilb02Y21McR@Zj!UgCevQTEbv&Jiq(c z1F#uGT>tQa^#jZ66|e^9Gd%hC8aBU;DvLT@e2q68@i!bF^pxp-UQUVM(+?7Tt-vvK zfU>M=$%tkeZ;kyEiQdQr$ie79BBOUPY!4VROsEnLNNPZcACWCU3+ASYLR+T|th?0` z5E`_>wpmx#1_{dSq(zp&tFN@xDdrC!>WDU5uTk4gLCb;MiXLu+xz(>RqxzG4?*XZCrw&5w=6E~VMB zRk`0|vWe$2Yk1sFm{S?>IQtcM1sKv;DP7g0$6pz%tD{cqJcWT~l&TW!S|;UKt^E%u zsp>GI`HGYg{Y*L2VeHQO=G#6a*C4VDm?Z<33{dU zx2!&M7pqXgQCnlc8a}Zdi4vvbq@v~&Sa^Z8Cx4GcXb>7|q!dnJvW_EPP!6vhZWq5! zNS=lV;8jpGC-^i_-4JT{^*S#{>b15e9M3@wM?}(&dl8BYsK`flWr&Z=*@8gTf2+Mi zDWmOsxW!6`;5V#d7R#wY8;)-IT0fY~a8ZJ#Wh{0#*L3EBT1xBVAJ@K|anlM6q3s9_ zVe;eOMd)R*9Z}N^G_POV{K+lfJ14ZJIHzlWIGx@SO>W1?d50vpTYD7^1bM(=;H`;Y zhwtq6@5zB!`g~h>Ls#o96#yEGK=L1@03OfcW6061JQdY>$$~n0r;4#3(jj>@%X2($ zj>R;!QJya-3+PpDjxRcq{L+zQ6?I;s)UQBvSCw`)(8EZ`_BnU^2>4beKcy*9$cy!f z*tIlD%=jB`lnGU%@F&%6b08yS93uTOyj1XGmLf%@98;YdPy)aWL(-!D9ZWpB-T&#% z>|hOC@Ic+Akm*ePj`!9YY5#5KQ@tTyzCXAA{{Pc@BO$YeGeU%dE1$=k{E3D!=zGmx z%gkVzD0b!cy54OxL_G!NCkMkltFY+^1I9cEopdOoPZ4mI>ymliDBfayg*E^EB3FFaUS6@xX|nZfe@YTnZ~p7e+%q39bzb zfZF2XtbaZ^9uP?h%l<~h?g*|Y+%+oz>+6v<8lRhJMBzfbY}ZNI#`i*||FqB5(YfMh z{i|tT6i6<}QcY%6!TG4XMfN1f_|Ux&$MWb4Qz;H`vn~{`uYP5atdR}$3nBw#mwy3+ z01nDEzWwZcsQ^KmSV51V+yxW?Rk^I}J&MZI8z>Bl5HP~5L(y&jd5`z(D=D2QErqkl z@yI{lh}soZA;3kcfx=w-fl!V)sL1yhFSNv#bHM6v-9;zo=fs{)={`(rDn#nxDcG3c z>4UC;Tm#*EUbHY!ZjA2&I9j^U5!GhC_~i4{*m37yscQpLw$pdWUP+%cP+*8bZ4p6@NV`)ThgJGltr7@=ECxN{Q%( zi0UvMqin}s!yQC#ei~bvWY(1$NtzuG&Gc(MxYY!u)$bYzN$Z1$*u3VZBF_Q%V?&}& z+S~SB{FN-N-rbMS8f+-Djykha)G5L7;CxQ0*lp$g+L}B2fLkLYB@Jxi|MMGN!@KMe zB2_6HW-9FEcj$tqt=hpL8NfTAH3vfTm3zF**rWk>1x^T*{|iL4$!>uIj)X$No;yU* z(xHoqxA)1SKvze$$UsKWQ~(e~i%*{9jFgz!|KAA;WK>nYF^sevb$=g}F|=pl{RYb$ z6zsYRq$7l(HS=1LdpsG=Dp5|G;B(cjN96$i8`GS7#`76XiGh9>d98uVm4rOumiBXM9=L0NXen z3q|bN2n~poq7=%2llTt+KOrh2!={7QxUWN3MU%FRObzTfg4YX7^5nd4O?5@5lAnUse{S`vAZBErEUKy@Uq+r3 z5mtUK^33~o^M9+AC16Phn?KcB>2~+uo*-sJay22Hua40Ac4S=`h1alDHIuZB2vrf)Lg16h;zCHeb z)xyh9SMz%2G~>~oRts{B8cD}2017|#p>?zstxEPwRBJ|U|NrGN3EK)1()iA_3d(Ti zqMT=cNqnKV)3hJ>@%j6fM-qd=O@`(xi^sS`Bo-7F_}WqjHM&S?k``(xVy27yqQDP_ zuwG7WFU%tsh4vGJXQfsn-}?DS z*5cGxn;IFad0&MsgJh56z~kO-ILE7&QUdotQVN~oR3jzg zY-*S%GW)X;IR)LG`LF#klBlU=cKuMM^lWX#n$0t&V$}n1dfa8ktF&mhPshM0GNMBl z3whCRxHW1#!1NQ4iBnWbWUe1-bKLTx$~8^FReZhoAk5p7>eFOV{s%dR75qVmWF%|o zsW^9;hu(0G_XCHW96~n4w_p++A)|@19BgQXEEx#xKi`TDSWlLt&sm}~6tPO{{oVtiTv*WIRc@dE05aZH#0NXG74hx^8inVLwtg~z&C8Pr2J z1hg)4|M?VhuNqnV{hv)PHjoukXH&zyn;5KLnB`{vNKrlb`_H24j8Jod#vD8%!K*f} zCx)O$w<-5S5+pK{y8@XYNIT-X)~AIT5)Lafmv-5{AX2UuYwN4D0->1kej)3+AdW<^ zmNf1bx!3CcL{T$fU2pt4F>HA*qPW2wc@N4(#H`R6HlcnO1-@z|qbYF{-Y+}=0sGo> z05P`mG4_`>ylu}a{#gnP#WX~oxJdC&{M)eNbni2Mc|xJf7iIB<;r-T}l@pHqNY&XL1&9oE;-i=uUazF}7h0OJ1v}_gFg12hZx$gpv9cppFd6!UG3WIs?0Tq&Y z>rjG2@Z7nMDo<;jF@Hu)u7$zH5fyEN3)zK4^5^Tw6(LZ^l-MLn?E?J8=C^voz9lY+ zB{Y|rVQjpm#v$gQISP!c97xwqk^>R=Ye5|}439gyl~KD%(~B@hZ$WQ{w-vW=$4_6&m^-X_G_0jQ9f5u^XL`1^}==L-36L5KQB(CQ;e zt3lKxZ@dEy)~kgQ;Xn16Fm4%@9jGDqhQN!xi|f;O+>J zMVm|MQH9}&XU*;zxky{T@e{<=nxmOHW-ZE9R%}4J8Se%9^(}bmK{WT=4721({?avh z-Q{o^Ci!1i%%7<2OMgfPmt?H~hwJ9nfYpJ_;(mE#)2bv``*P>>Oohz&AeBVc?v{E5e{{XfQQBZ4C5mEp*fq($)&l|vT zlQiS-A4SyZjR<)`s?r62A9Fe1I6W!mrBdkJt?A}$cKp63G(7KsciNP+ks~5y5#OyS0x`5t*h>9r3fh|b2d|h?wywQPX-|q8m1kZh!;1;y ziD+&(!g+B4ZA=1)i+P0>M*Jzq0$*+fDS|WI^&zB^xexUwS-Dq4H{pz?V@25;M4@e)#1KB&kRGy*&s(|0CGIsUUbIe*}K)f<- zTqSNJ?EnIg!)P`~2j*YyuSg&s0D^9!JWzz@RyTFCx);~4hlFo5PL$Z!fB;(*tkzXd zm9-UlKa^;v1Z2W^zNc8i*>m|dGQCg1Tv*KWm8YeEv?%i?&=`8VxNWUN#Q81nz1v*OHoMS^pGF{HY+Cm5ik`)5;RloKM9Q`7jZ*r1KoYIZ z$P8gitNb;-`4?fBf}M=4L2~V`!n+vg6NWSM5n(Y0LTl+(yU;s{9OWH zA<HCE3e|MF0 zXoCA?h~?6R{ctqG_bW#3N|wVOPFiB9WOVaw$U_nJPmcXEuwn<1dQ$(WM7QpDpH@D_ z3y!#dCm4ibihT|zEU`eA-yi2{o|5*l<$aBOB|J+v-+}kB#sYF9K8AP7a=27hrN^|P zul-5XJp3q7T%zc-*C;`IZ&z8h)CFxx_U*eWKHajif15rdp51HIHI!+6=BZKmsU%dowAY%?nSa3Id2DrXJcd@FRW! zh4Wp0PKIF3th~%6R$M-z1q+{R>UO?XXxUjc!6KaN_jKDcKxcl!O#I)MztITP7l{?2 z+cVAIQWWbugFcx*Zxx(QKgTlm_=fq_-RG0Tb9N>)P0{m@eEy#<7YJH=H|8E`shE-M zLQKM84+F!9)XzQ9+LH?^Q3YUBg z7&@p9X*A1J0Z@r3@aYgLY$w+QF9+$a7$Dsh42>l#XTs9?6T1^+8dLDTpk4+?ObGuM znhW>jO75nH9>KRNpnzqVBhbxV{JbtO!*PBqb+KKO6S&AGmYe33fzx2wyvkqByg=*2 z+lrNUUW+9cCOkx5>UZ!?m#z%$?k_x?ilWI!3zx@#4r`{D;mjmgz9H zGuNWD0*vmGywpM)0|VMlu2{z7pNpH!pA-c9&^(vkCGdu@VBr!ym`-ufP<6cPp$rix z$7~+ucB*~z{OO)d-oA#`4)pGtKaarxf!wLB zf}(a%67tA70~RzzB=`45X`+t#b^oIYrzd=O0wgzR=68bZHAj55_Rg9fq6@y=A!M3Y zXnDcp)O>};#D7pOd$&)3H&er)t8#*%<0bZDdz31#miwdAjr<5Qhl2rQ+&{Ac;hh@xKd1z^V)r*T{_jT)M;gX4~bNVRBgz!4- zqB;^TBa%48>$h~{EkNge_x2AtcI8;a6Q_~iA?)#_rKjEYX*WSXvyb$-md(~XCLF|> z{dU*4HCfHRz%mAYtVt{0FVeLBnpitQL$}hTnkIeO{Z~I(VZPTXtLwq>HD%r@K_LYd zO?~4|8~IT)a{M}_$h2%lC6XS7QI7fexZuHuC!U$%DgdjkLXOf7c${ zaAn?4pBLNIP@ikx|8$q|w6)M4B-|C#OoqQVgQA)IiKQi^Vr>r;SP^3K()VIo_Mh} z>YZJqJB53o%dTGB5c|leEyIxDv35y-kwl8au5^6H%<~U|JFWgMT`GM`5YN1;%Uo%pJ>dH!+TsV?<*n=Wh<)@I-e%C4s&` z){@t}0d=+_)!zSjE&GjjzLNEFkeOlk2w_xzzPNZZO5|qCST=OS)@QfFINyF)7&Z~8 zXxkV^%86xE;@8jhY!9oVOTex~eU_^J(`Rqf={cw6o71qv2g}J;!WYUI*8bo5h2MUl zFzLb(mZ{Tp_SarctmV00Ac+u^Y=+Gpksp4kXMW@m<=%R-CP`RD7(e@pH%sHek^yFl z|1>RY_OQIq^mEx5O>BdAzq)4PdtJ-PR?9gk&~F;FUJSoHb6adz(*Pr1mH5rhHZg99 z)7Q@11c#%^J>Zon#+-Jf<2QTn0`jvIvSYESz0I#=5I4{`y>B9vZl9>lf*BXp_y&+f@CVQyWh=9JmKe$SLiu#k3`n0Aqz%ocFep1|+H)FPe zXwo6no#v#Ce{k<3K@;eBv<`1}=$1pUJFp8-+~7-A3L|om0^0$9nN`+Al_>C2=&&yX zzv`qTbrYhPP%5uu7NNU&Y>3+(lcYd0O*ZKt(Zc z$*NlVUC*!&pY9K`2#>7&i1^8)JZⅆ}f$_$}WqPb)LB} zX-vZkijy=uW9fnyx(x&jvCoE$VKP33N5vHK@+*nd4xD&K69o5?;s!bTq9`4!)gQ5% zRmxXX2}I}y91X+kzYXzAMg66X7P~g>3*oJ?kLBO){WlPDzi8oY-ao4X@Cg>tW}BMQ z*lWoXWiNJkql!fkFx)7IccBv$0plOBfl zbd`cm>)&>|tEt*@sWnvNL$#&Nve!J{59hqdP2KeopqqnO1{lK`>s7ue>hcELJhMc^ ztNt3(_?j4}!r42@Gj;X(p}Iw0PGHCbh244IjqE%g==;g0IkLlTy) z3b5W4w!8`^!5F&!s*@J=y>3TS9xI}ECh07oel^VD{q$ew;&(l{6c{D^D(Tv$k5eS3 zt)eV6?*E}8KE?kE88D2eR@m%!*rxisKPgm0mZ3G_X=Ro2)9h0%83dX|y4qjiv5hVB z8f-`O#bzTKHrt$jbQ$|KVPyAtcW}x$Or>;sS?)(J2V){6>e6gO??B*p6R7>dUPC5f zg0x4gRb|Q7zpuOpRPipa`|WmO@uF>cCcE@MaQ5-USgtgx|BMGa(e8F90uxBpO*hP# zfs6C0%Y1(c)BKx*5~FBs5=Fz11_5ecA9tQ)hpsuqp|+MmUUkCSqx`gwUubD41wX_X zFp?3TIn1@^YSJHMH{E~W1@?Bh$+zlVBQd5B>p80@`)@ zl{}rNu;@kMsD+4XXeklIGV5CF)BOA>l4>ztJ&3KQqsZQ{y;p;mNm4O$M7VcU24Ms-&J6dcLs1(Y0)GDmjC$DQqKb`se zq<6S(@_Dx7Z1MAD;h?D&8@RRbmN?Fj(5AY-N;-RiGhemP&#n-9><5b5?4|+hl_78e zM3qZ+=;ZV2#M?vi0^HD_B&yz-}oc7Z(7KPM$QJs;Bo1S`lOstDh0_3e0d z8`f-6PPC8S4beyBgY#DGI$4}?tEDFjh~}Gn6~gAUMokkz>}k-q@pn_Sm-WCY9AT-{ zfG|&^iU`X3_2L+|XzZ zz4|u8SKssdW~-iISDBaUUolxGTL0UxCxg7bJTYPH(`YjPc%OAt-Jb_)G}#h&`19{r zl99Wq_-d7RVhi!<)w|HQmt&B{j}ir7%0j@Ynt9g4vD&V7UrThdg3Nfim=jX?8(&#@ z>}%y}_(@bxXHqUZ8pi*-ag-LW%|Gg`Cfq8F0WUNG4@@t`q2w7^rur7LRlHvudrdQac=>3XkvN|o92<00 zzk)~~3T%LfIm$>~Pl((!Mg#g8weSrpKYU-?X%}cxo!IXUu_d%4=%3wcSA!yaHja@U zT&gS>#mawSJkS7ubwJxnQ{^c~0iU9n^wa?_eAF;vg){bz@Sk*Gg= z_hXtI(kB7@G}o}{$@ItkgU|F-*n*GswcV2nRqfrDBnzg*w`K^m-LGULQS6&9*xXko zs{l|bn*4Uy^}yw9-ld~PIo4H(9hx99;qQHqDiIHk9Y(^cl9g;jgD;YkzbDec+$LlZ8jjm zZnQx}!W1TWOiwFt5^Nt*NtkA^yF%XKbvYF%Anae}@=ncn>14h-pMHf1gS5NgB$|k+ z$)kzb?jE$i6j4ib%s2XP9Uu=Vz(cV9r&~$*ipd>YZJre4K{eQd13ro*uT_WYXurC@ z>2ZK&J)52CZG=Q_Q}ldV1;|MsOoNICmN9A?)=uVy>&hGPd%}aT)fntch)lIsh)<>t zTFhh-!<7_ZAT;l6{SrEGLJlvi6d96sd0B>+~XTTr+x2GAwA$Dg?Vnlkg4pa_Y=u>j3 zbJ)C()K$+ewryYQ^&~77REu}aZnZ_J;htR_cU40A_T3D-b)l!*D+*RA;{js1Pvc}w ze!MkfMid+NaAceScMl&T&Rn{=`_|Y^G&fB=Uo>ZJantkO>i zD+^+rF3-m<(sbK*H(7j(tORh_7}SPUEnX*XJ-Z~Bhqb>CULLffCnlvvUl<}-UlpQolrMWo@1oOmQlfLFpnGD>F z&cR@w%1Qd?ny6mCYB(7RjU{1R4n}t1)7borXuiL3NTY^D{=D#yg-U^_Mr*YVq=nWf z8HBYPc$LY)I>~KDig8-vt&S;e6_*F2BJ;~d{;`msWZeL-`viH9AjNpL@Gm|VQGkDn zf1G3-b$b_&f7p&hPmm zW(n08Wh>&s33(o4pKfDvMz8M#hCg@Oc@vXh1%b^ zn5VW3{T6?hA2yP?+ZRki@9ue70DW)jV;l1X$Me=UGQ!N6fT8y1MlEtIO^l-a9tm zL?>6i22TIN94eM}PVsyIffmtT|6bshUH9!iz&CHO(-Y5!HRn^L2*FSLbFJCYA?GtX z?qWCPtmT!C`EFo9XE!0PAqT`5t=PTfmI)r3>`d{>$~VSZ#E2b7U%xWyA4|);gB7sr zi?zRunU+DC>$BamA;cV$P07*;>ihqUx7Ky3$-hCBZL#Nz?1>5 zlbcEzEH!ot*(B(FO}z5F_b^<7>^)`e^+ln*HgV(J<=s9}_xZEdGd)?&)KBTUBXvo+ z*^+*PzjuLC+dferj>8c6`e0?@Y_BBj4)^XuOsNptb(Ou|-`&V{3rx0B92i3HQgr|K z0!5+3P2X-Md%+lGVd}B7#paEvbBULb$FTCAXym%aIDAdUY+H%{wa;(PFU-CcS2EKb zqeIE`+@WlblPDDa!N>H*Ig#BNBC|a&HKv3)?h2<_Z0CTlOEJ>v8wdo2)0`yF$U?~A zU0y4uK7L{JylXcN^AeF|5Y7rjPElQRHK|oXY2nK=5-XmpsoMaDrb6g=gF474f^Sa* z;;x4rR*0IDbLo5ngVdm?brY-kDS$uU8xP#T4aJ0MPG)luxQlJLy)<`r3`=U8 zWR5;u;1B4OZ!3&wnXD5rBKhFjV8-uX%=uQbJNjp_wl>zH%()B(!!J34*=Pp4E9ZbZ z8eaMy|rJFz-111MF-ehF2jeX{CE4))7ZARSKd~t`%y0RPhoW5EdPFuRV z#P$8W6o;*dohynq2n7)g&$T^_=eE5y_+ruIVAif`13j5bfVZ*M*_zAKcXE%G>)~+I zgU;2ziU)q&#kLSeux7R*?^4<( zfordcH%~^cQ+(>Nw0tPAbJvMUbA={JK>NjAyG@w0 zh%STuIn9}PknW;n(&jZ$tCn#`Jc)hyg;vpA@9!_7hV@HPZ!?o_Z4FWMP-mDeh26#n zKJF#NeTp@ajN?tom@2>wWm;|jSlT2DIRol)&jSU18J?*m1s1VVk zjGp?B>7ohgD&nzS=hz(|Ic)^Q%CyP5A1r>q!;b*N`7j__lDlI^2`!Fv5GZfIH}G1$ z!v2_r%(`KFkl)B1!3V)QPRK3&;b*Ye7m{cK+pDX58oTxR0?-;Q7@seg>3Ov0=%ZL! zMFa89vq#1%LbH?c!%zRob-hZUa z{(rp>%o*#x_gE5&LAqJ91<)#7f3U=EK9C;i?6?JOCsfUb$R*nKX=}+bfUQMtr9C1~ z)h{N3*;9P`H4)nW@qKeVa%XZ?Q6GEk>QuoRcp=zYY6wIwje#2d)iiBZTTQ9yMF~tv z6w}P%29NG0f0M|7$um~XYn3GvbKmL1sp>y4;?f-Sw%{!a_9R<&wK}b5Ol_{aZ|`59 zi5pDItEP2wM{2P~IXfhVcFA-rTkJ=&0JgFPL*>G-WTrvKAOF?;zmX3^1$pY5DgaL| zUf~!ekt9ti?QzDrV$PF)Pgl=^;hD|1qft{I37Jr;T#-_{{){*@tV~@e4?BEKEjP@E zxnkbyAQhzp$SN+@Q8N$|sBh*Dx34todakT*%gyrM>#lr*c&z2c=FBr6Xfy>}o|!Jh z9J{*K_AD5U zIr`03z@SyRx5hQYNm5)Qc11t@lR~*m#LUo;HG^_a!q}*mBm_0&@q_oNc6y>kP?Aa; z*K7Lxfk}Ub^>0er3%w7U?JBRcr{ba9SDYD(u{hxk=(io(@6Nbn{4Y{Mv)9@}`37~k zr%@`ZO>?(R66^6z5jsxiFztB**g>0Zr~RYC(+GJZ zeTln#)#CE5ps}uKJ{UBTKs6Eqe||G9k?z&o)p-|xClrIhT}S)$wU?lZYD;`E%_IMz zXtd}9?BnBG!?=@)wDg+7d_~F2-c2`co!5(8J?qW~A8e%Zg@5zG90?!d-_-9kV09&P zf4)Ke!By77)P-{12GBf(lad`VQYS0Gp)>ECs^gnVd35~6KY^W%P{UR>CRV2 zdF7;CzVgpv$+b;8D#Z6BRlC$C?(6n4_>SVR49ybe9)tr3(Lo`PPBI-vGK6nkI?dZ8 zSA*CM$CA^gcyQyjc|U_6*0oSDHHGmk8pn{(w8b_j9I_ytEL{ z6%F9{l7q_V;`=4{@((oe+8k1IO7~;LLcvarH8H(TxLuHR|LN3nT!);co@(duv9rHvLlMfwQ}6Kg)O|DxKpR;{ zT#2qC=~y~*R@$P7wbT3iLS6+pgnjJmorU%(qi9)|FlnnqlQJ&nSQtoSNe)+s#yypy zbb2wN(V*mIOcMVo)Bi%U`-`V8iOXvU z>C@YBM{ioYxSJ>9zBttUDveamPUM24ok$!#44&I=6M;E$tcF)zA!*J*Xr)VZcBE! z3^G-f?FVOjKLDEp8^&iE-MRneT-o$B4bxCDK^(<&akxPf3FLS-FdwN8^oT81_9uSe zf(BG!5e%LNgG45;F4lGTIuaQBLfvDz{;1~VUS8a!b`ikJt$@^4<^ZsN@eH);RDRK& zT%=(XntFq%k9j5i$Tb+gXUY$Rcl?Nj05)s4$WekbH-jtWj0F(5nRoVyK3`PB#~*6S zxJKcO#>~3BOjw;7H7C~04S&j($Oh%BQxYP^pxES zQ{OnTQYg#extOIlrG(ja8{=%Zj~dat^7rx26-2wxX|Eh-kJ-PVr)nBGUrtmD#B!m` zY7f8<;`5LfJ&9AKh5%#%vRWku%LVV!QT~ABhXn(&XDRc2R6!|R1L{NCZ3%ogXVvKJ zW{_t)^EOK0Df|#Ke{1MAhJ9eoPy#cc{*s5O_P%pGz zJ7vV=!9Q)jBkWkTAMYXRNEQ!5qoS|OF2sYK5#=G8I~AzF+NF=&V8?UW1ld2Y%51$W z)R?7%g|#;MxYS_3uG0xczezNN&->7ThCr)+2k=kOGvX%ya922hT({h1tu}v_`umBV z&pP8x+sms?>{NG^G^bHsvFXjebVNO-33YCPI!E$nTG?Nw*3!9K1E)>hCOetx;lg1~?p0=UaA+ zP}ru0!02LK%n{Luq!6}k#^dI4M>?>s1Q0$SKW(|uyJya)pCt3JJIj>20nh>rQ6+$2 z_Q%J`_qs9I5DHWbg?Zx=?ijwf-G8X>h6YNY<&s`=EL#L92*v@^@dX(+=3Y+A;Lz$} zig13pj(nnxKapbk6bDN)SNFe)L&F5l%TzOmVj9R9D#S^r7mIN_;D*oIqWm`_-!>MR zB|k-@LD4@FR6@|2K#BvVMd64$0?I@ez((+KH%{J4)jr&)k)Tq)`o3tjp>pxfV4B@c z7J549Cuh}4*jYA-H~dtfsDT2dSIQNrMNb5JmPQ+w5;6eT#Q9ZD4adfDYh+Y`8LBkQ zn36FKO1BE})M;i1M~2*}XfccY?Hhlmx78S`7bGGRCHt7VRjO+jY4L3G>O{SO35!-o zV7wKKWT3Fx`r$s*dC5yBp6ES&nvKH=m^l1cP!Y@o%>GI;z9#C`f=iwu8FOfFE_=My zCMjOwHpLh|w>Q{^XQB9tytrV&OYv1<5x@EVxK7fv?p9>Ad?@FHGFtE@xhYoJ?@KrC z*5N=qH+t-Yg%va3g1>NF#WvrAG!G@?GL!Wu6GHz)P={-JWDzD(8a_-m=#x;|%h&?kR$&ObDp_6ajZU)Cp zbcPts=X-4oH~$oDS}8w#`x6f{Ihsn6l#BmNC;`+)-}AVv!MNdzz~6L{uZoyYQ0iGS z-oz@+S-_=aOcaN!NuOx12*G3mJFLru&_(vyjXAHjM6n=@Ih)Pkf4JG`zq!}2G%iNA z8@s90oA($F-HAZ11l*Gw%DZFTv80R%uuCFnJ=Gp-|Bm&`9zuwyZ(Oe^F9h4U&F%^E z&PlTJ0mZ%5k;zdp)^W7eFAHLTd2$Hd507Z4AFqdQP&pC3b>reRv*TM|Zl`Q;r(vX> zllbPdF!3Mv{q~yRk=(@%IWAqrD_R$4O&fQ2 zzx$~T-X=KC_)d!!n7ki+dPogCBvmg**{Twr${o-h$qh44Gw z$HRbOt4JgKaz|ySi^3{tyjS@r%kj~^87M-c~pNdp$>L-ZX zTa||Dnd;nAy=ccO0Zg9l;B%<$#i^d{hPxPyn{ZK-qMne!Ee7Z7j zyyC_*+0u@{AKSDjW&c5cX7W=JXJaQ6W$8{C`*hK1e2xFeic7gtWA`B(K}E)m9GEL zqt5^@C`P(DaH=Dd{ds@(T+=S`=AG@105j%{fmsuB-$IQw`}d+Vo5F0n-1!6)B3Xjx&jVNB8fTvOV|61){X%X-XShGlNL7F+dk- z=8m+kfqNyOc^B4dtHt<6D-He)+ZUgu-`lm5k$H^T)GquTVLr~aJyg6LDo&xL@;=UL zrUumng~qSW2LeTN@t|j}%-Sip*uwW^GRSq`Yz2tdYs|&hMltIOkl3fAuu_uQY`>?O z@wMeRyyakn{Igz6J8_h^l&_6h)L=BYxXB3IygH*ST-BYxokC6lCVFE;A+NAIPBJI9 z0^X2~M{_Vr#}O>IdbIrV7-23a;Z>~nDCS==)O>N)IVZaac}4AxD0F8ysY6~+BRjIJ z$t~>-qc$VM2gStlOuk0FRvM_EHShH(pTW1Z=IynBxcdhJ{wmtpVmxomF4#UM@BY)V zuKIaGjA+39iXhpW=HF38V-)9{n{B!BQ5j|UM3m=qsM-s_O!#EM4xqC zdVVfMm>&YD9R<2N16?=UDP)vO-Rw_*3UXq!0IE^^=n-=ZDUI#IrTv16FgEE+_lr-7 zBXqF)%Sy6XK&~_uh@Z3~`DQcd{+lUOWwjM$RPd!?6Z{C~Ge(p1w6lWKHR}|iQLA1D ze^tRfaEvXOL6FJZ95<4w1(a60qB$Ju=Pc~FTZ*EwUPX?JuWX&)?YA%6`Eye%@M9QQ z!TI}hN6G-D)!viR%3i!&5%tvD^pe?Dj=3n)EBBDcMm$eu_6=$|?6d_$hsNlGLLwcu z%hSlSM$GS{xq0~)>HQgeC5*M(^L!A;d4>ZqjH|H$89c8OR$e4Gp>^95j>#^bHJ#6M z?7qx`TjZSf&$jUFB;#k@MwlHILtXQ>11F6f{cLgAKP1cO7@H^QS-W`oIl~=&HG+`E zvx39vkmG4j@@YC;QLg(&u8Nu=nI*;i%%!v0MVvQ<&c%fU|Ay|<{ieSwGJMpK@Ouct zl_0Sz9`dYej z2E`m4^2hG-i-DD`-pZC##DbWoo*EHmX9lrm1vdOwtJ6RK>QSfZSfcy(t+0GfJ>ei<~p4VpzZ?b^7ObaMGx`08c;^Rd&C%lYc&=$G2~K>{#hcvN=AT`H~ZzB z$W3#Qg(x+~HH)_1FD3Z}8$RdH*un+#Rk241U=Z@hwEsFFzBZPAW~1!q%I6rcuiaeO(k8X(227bSy0BdS|x>!Hh|) zg5@5zkqV5y32v%}=l5YjZ+>fCBu$L5`Rx2-3^-h+`wbllAg*t=saht@0sVo&aQ^F& zJDl2@_)XVgbf35t$-wiSnk+9QY@t!5#6uzfVC7+RCHw|3kPVhPZG}>dXXjBK(Q8{} zGOyJ)<}XvMHlCf|Z3ajw)ad1&ICJJaJKyiTc;o}hgrCl!kk&%@c~dE4(x@1K6bUl~ zzB6qI#=QrWSZP&cEdXzx1$`RaRec{ERvt0z3V5#fD~%7qyzH6JOg>X8X?d`AEl=x) zeH~oO-Vz8cV>bs#>4=phJ7mGm+8H1Ao5wJhwe~IKi)cQY4zu(}4w@sZh+$}6-gb%y z0t4&Kerm9~`aN3>we)ke;P4_3SdwtP#~NMjG0>8$(-2nkQi_humA?nY+Bq>-wu-C= zj43n7d$a)h27s;<5Ug(q#}H5CLOR~A{Tk)T@;Snknj?QG?0i9E7G)7|(h(@i6&tLo z9Cw%OJqmqT{O72^TLT5W9as25tq<^I4C4qBML&Xrfj0>XpC&9ss z_{NgI+vmC$xckq~UL_K30H*EZ=&EpN>1%1|QK14+{DV#~ubt+*8*-?=U{v7C1bf)e zKkbC3{Lb*Ts?p3LZYwNeY3Z;8A5j-)UC<1T#hVl9A{XsUoPp#^kScW>Urw?dpDIg7 zpzBjgZUZ+5A4ve(60LW3A(w+5;#W{<-6ac+767iw>?h{X*)PJ^nWR~?+nL%sncolj z;d4NoHGCO{{p~SRSmd22v9W*9ma7noIkw+n+aGgG!q_<`{k5=+)_Z?vP7%oRjIAN)-Q`wKWEHnqTShJt2wZVa)Gc&A8Rj4^eVl^<&=p+O@YP;9Ks>``-Y%|rEbo9@{V z?JUbC!0vKD*yB7H4%PU`fjU9X_qsL}mzL4!W)AzK8>jGH*%VD8UOBzEbd7P>nPs?B z?+*iK)A?^zt*Y;V`$vDyHa?r&6neid>ii4opgqt|E-QUDL1B7)v(gJ5CW(^=s7*kf zHn`md39i{CsH8;?62ASeY4fRZDQKL9 zb8O=_*IPk#fc-W$c_o2jRDK4p$jER6Brh~pDkY*gt@z6s)G%mVfMw$=*G006YE|*G zjiBS34z#1v@sw<0d!~vghAaidWKk+M2F{AEVUkuLQbC&s{c5L`w{)!2?dmE3VzsU_ zN<0uob`p@|KzfG_go>BKB#Wi&4V**EVd97(@LX4~dsev3lH8UBc$SGb~pi^*+6b zEq-ul_m>>j^xjjM0Z|J37j^N>%<`a~?8M_HrJcJ*u6w^}^l$hNI7ts_&+U2o$W=QX zsNdL9d2Rne;x-bt>po$<$*x1)Y!`1HDb^9C6y{)bvWo#rth@aVd>9VciLMl*z~46u37pyxJpCxQQQa{0$9AA!r#C%tG9Zs#F2V$n3V_zfmcYCjwpkJa^aJ0UI{vmcoy6#Y#s2t?7U z2s$^LFbcS^BW$Tlr{vP-)%qykb?kI*T%L`dij39&OzYmAPkb?2O+rY^51$PQctY`u^#)-FG~4IqzN(Y zA_3?2fzw8+5SD|FM;2xuIi6U?in^I~=t<#-!c{3?Vj10_Le|;oa&g<1bjSF4`9Xh= zC^q>wpSdsu%2E_?N2{}jxMVyDMu7&@z`6Fy1Y?N|#|U1Fm+tc&`P`{$L4#AE zOaoKQH(p~8ltUD?M_{5;w@u;yU93bT-W2 z1ENAwue9dfqq*hyia#B@|kySs0(s7yrQjAl1QE z>#d2ctJwe#_{3}phnBNvLPP635W{(cdJ$@Ciag@Crxs2yc^;p>>})r1OWpHX?{N1J z7#eedse_bpyqSdZlqmm!SmyknjITJHA`k{>m~aGaToMgESo3}^h*iucU`^d2 z!NSo_L5{govX^^5(iFEjLjPP*8|@$hlO^{?U+&T4)M~~DMtAn{zS!d~#pCYs@y6q^ zf!O?h&}y9D!5HP#$gJseXJQozazS5z{qzmh6+af`kosd`RpR64SQMr}3_%4i zAc^k_qUVPxaNwUa4|V-lcynFSr?qk%d^~McrO!E%>^&}DLPT!CZaXQvJRbagyg9fd zCfwCe$*O)!XWfhzYE84Ep=*CGs#(KSz`&R1&9b;U6Hd4K_Yc496pYoyWzit0;)KI-)%9Q65 zJS3J?qWrw?r#$?Z?YzBksqGf)Qks5W+)3ztoLCicYq8;QwoB%G&U%vGgB-cZ-uJ(0 zr}ZZHJ-cMjLS4Ih6?ALQ*(xsEO=6(AN(HCln4#A#Y>X?vWVA z&IB);jrk-w+DKKvn(?{FiUh-q5He>A2bUURxB+cny-7vrbnvIjA>+X*J!dj2$0PT2 zwZ_zoz~U90@#>%t>>Ly5+o8Df`4F^i#T#z(hKbI2#e0r&6Y54mqUg6W0_KRkPWk=t!^I z)~j-STAG43ZdsW+i$5+X5apN5Qhm|#(4kPU-ov4-wY;#4mr3}p$&>TN*2Wf+Qnps{&7h=F&(CA(Plh6xI`RYRC4 zZjm7iPT3UCW2Nb+8`>fo5BODv}4e^ZJT8U0|~1X+=rfZP0GLnMB#}0~(Ig zgywe{qD)+DqdTCNV3qAibcL{0oN75g}rj0<;w**_0Bg zQyzFppUi%WrZUvv(DT;);98;6#nMb9KHh_H61d=R?hQ!5NdR_ z{K7@8H8Ugo4)JxKU+4|~)7^>^``3_%coh^o_3y-%JnHnfGxA*4{LG2{az}^o=9CA2jAs)< zs}_d8;PE!0r49LAIWBuTWP0RzG9Yn+@rMD4!ym-^@n_41+4bML60V-jPNGqF^~Tvi z49jY4aQ!WuaUce<3EZt~3o!udDSiVuEBuNJgWsD4xaKgQ5@-F>wjZ+F967yL-JaU^ zfLl=iwS8y|lv%fCyq7psZVL0sKj3B-?M@#BE%lj=D1*zh=G~r#&<%#r&OWS5XPge_ z4(bi6FzH@?0t3CD6lp*>^`RfXWyL&i58N^2f6|5Fnl0uUl6Rzit&W;Hx7ws@5j=G| zN|#!-U=Y>ndxc{xfy6fpT%X-5Yx*YA&4Jn>hJW7oeRr6SlPldhCJ!qj&YO-pe4><% z5FMK;g)_i=kw7J6=eN{sSn7s1iTJ`k|xc}y}(!hV6^ z4PravL3A5>olwjlbJs4Ff-+VeXM;dc^n%{?>KnaJKh+g45iYxTe$Hv_!h%iUXzE~1 z?&Mv=tQW;l;5P#&dN^P!Uc^<>faPTQ=yTi<0Tj~a+6U2upDubH*M4jpOCFyj;%4X! zd)6nkfd9r|jAn0I92Y%%x)BB|K|1U4)({85`~FLu+;}^0=z&u*CS$m{$7owTk@Mi0JC5iKUmHVp?7So1{GPY2~O0>r`Q9%k)v1@Sm7i^pL4KE5a+^?bXr2mF>@ zj>D%x)&r7o8`MszU*&OqhOV(aWn5hihyw66-~xDH47fgj$Z)9RiY352l1+2n$lY~M zA$VZ*;ZJoT-kZ%RbDJoaX!kH%kfkD+#Cq;qqbxe2{y0p}zjXRNsGaZfp!uHdATN3( z>0h@u*a2mbPX_?ZbTu67ii1yq@ZloO!}z)#tHz81^|Cts%SHr~H+>RHG?;5+iW$Ko z7I7gd8FK(N8Azng{Xr2BC#B_aZ>&kG`-P@mblWZ!qKcE$*DIX`MKLIL%)ZI{xNW;t zpzjLdKE8Crtd#trJaJk1ORcy_Ms!>0!e7Zo*l5&KML$GuWh~-cp?yDPuQ`~odbawX z*Mao{EMGHr+qhfg$c~!I%?0KCh28|e;%WR&m6_=0!hio5SJ!&`^${> z{|xelx<|^M!zt8yC*Kn{6eM4wb*%<)=aT~rd=v#%E2`1u2xjsH!C^WVCPR$zFekfi}op( z&fxEhDB{!K%B9|3_@f$VKbvnOrWGRX`Xpf2hH*0$8Gl$?*B%?hMuq~ujyj)g>Z`oOg!E@xNLQLwJpwB>*gqpo%!?rrIfth+rHa z*m5zU4BogU+JRmlGeF37u^mKa)jey$Cjmfseo)QI~jHFyubi0_ZJ8LDHvJB_& zLVk}WFCR|TE4k5=-|F_CcdXU!M%Pd_s2if#OLq~!Us2B>gsFA#=gv|LVamAL$+wma zMO74fGbU72?3M3j~W?+4bB5d9WacsNEv(YSEmg za|#P_Gg}tyNWqHVrLMin6V#kKNLx=;Y&5?NY$K}dQC}s5Fa7W15xrY{-n9<=V+x_C z%T;El%MCx1CgEu{VNcK9PGYYhb<&Qk`PeUKNyFLlYl;ZmFvep{_t`64HbpAeLB)DA z8^|E;I|rn9I!>2gE_xkOj#hBatE=IuW{>c{09XVJ7P-Uz5469pOT?{K91ng>DJuZK z?OOfarTazg&EyYMw(f$tJr6O|2+}tN2yq5%H842yBq&tqCPX9vfE z-}uw^lm24FRNs;p5@AIzHLSY{9%-f0x*}lFwA@Z&BEmXmx$v?U5y52qDL*A^i>jSY zraB~1n$706zvbQT6q2nsEUif&s#m5(d|_Y*`EvT;JG#zw^g>@K$WDNd?3O#tF6a{P zl0zh2v1%JfiV)%65pTvel=c##+r>BSh>FErX%0JWMaIXIlv;J{HYQF_>WH{{jomM% zjF!ozG+qcBXITPGAkDFur9wfxa3im^>`&x+GfTv#`$rEWx3V>?!(lnr43krfG)HC? zq{8~o6^iJkv4!D^BwGS=WbTnGd>?FB@(wk%%9ogsMAP3MlIG{T@^;5z9}eo7-xG@k zP6x{%hl7{M2&jE^Kof0yj;n?~BFG>oFl3CC5MG_OcFAkX@n9K9B_XZD6^X;p?FAnf zE*!3;MXILlbv69QUpJtlFXB@AGh|Q%S>bJ*t=FZDMLcf#{yUtTQ!`e7L`L!B zg&&S1J#8$)s??BMPHGNaSlET@6%EmXO-PK2?U^v;Koq2BW(kK2E0C;*+!bL(i~$vb zW*8@YXv3iXz73Zk%!Ve|+_HgYgSMB^*+XijGZ`-!6>90m!ETJevxd!+k z+OsrDuVXHg7|jWI+x$zs8+)z6t>$+X(%rqYcD_m|eyx8T7sMFgg_b2Uxgz)YYFk;F z%N1T(M|+d~++_$z+~e~*8l;VvotFE{RY>4wYjLD?*g|L!CB z=yvo#J(s57U*E=otALjlw+(xe3K~VCpOFeXGU)rhw=|bzNWN_<{$m$bAyeW0(FD^{ z%~4Du15Jw*aYS#k_9$4M^x@7_mJ_>4;?2;0@awkMUgU&Pl(PCx%K`i(&jrSc#XD)< zLjtRKr9G08>E(hKz9C(22g zs(K@}&!D^YMvt;l;x&oasw1bTqFcDNa?K@Ikb*x6?Tn8)?L_cuoDxb9w0i_i7n`zI z+Iq^F<6NpNA~5ocvPW$!-c75=#?pWF4iS)zNAxN!tS`~ORGY5r(tZwxE&^uYiUS>T zPUmGTtB<<_Z`((SZU^n;0Yxr;&t(O}Pxmd9vkK`BO6-Na$S(~`+ridD7K-w1{72Jh zRb04LwX%QpJTHH|M$T{*6K}w0-QiE?ZK}0vAHFzl4xb+hmMmkErt9Do$9Chp}(S{bpYiC^{M=#+gf=`7|G} zVc6OG(x1YJNGO{u17kCQ*O79>jEfm}VlPI#iD# zZX*~aa)nXxOKH@rN)4~|PY~pyu4C9M@GY%;bXu;?5N>H~LAeZ}dn#?jzD$y@dp)I? z437F~1~b&<(~A982*02pUd6ufocuUz$~80iTobae8vfbh$}`P^3f=i!Pd+o+8$?ff zcfoP3zsAcF!sLYy5k$z3PF`h8wzoU4KSjGpr*g^Tq*(~vI~przM*F9yGFGjKX-iQC zEoRB^%ExTh$>4CTrpCb*F5=F-S3KbJT#-~uzOV%Qr$7a{!5F*@ z#O~wf{NbRHNg9@mM!oDUODp?$rD2Q(`=Q8Ldi22S`L8jZC6y@Q@*T&^aOm?HhXS zMa)egY7kv#PtOI>s2LICB6E9LOG;_(A@24MY+vc0AncR%2}U~|njrj_DRi@BkjGWG zzGcpG#SE*$qx+}xH4+mjt%JBW@EwP&pB!sfk=Abq?DtggTWV>*+0WjE;n=hDB@eYb zTF6XkGH=me5sQW29Nk-AL*QrAhyLQyj`V!_=gTa6imCKM#L!TgqmUqokow*72Xt5h zFBH62rDM(y3f-5mLF^XjFF413J1B7!@?yQTD19ezEaIz|%NNqvvCruD@mb@WQPPss z$w@TD6X>F<`?QA26e^d3i1ZW^7<}~ts)lgo4<=t@hB?$Xcdk+PiD|wWQas}GT$WUq zDm-sM6*QS(Ko+V1dw2}|up9Wxw-L^d7DXqBd}gWX@O*ZY{wlfInwp6@!5XRm|E&@N z19ZKHl-tm(?yNn%yAMBGlhQe{Ius%{0*jjl%lt!+hg#eMy%F{tQwjQ48nY-lamA7d z8s#0!Kcdl}6W9fy>yI|sOoFSTU=UbOgIJSRRb_%Vq#D>{x0pU!uy5<%yT5O~a@F|! z*8xIgk;d;jXzSAvI`kJfr%CV;RaZ(G2CZKZ(;4Ixe==x@9r8TXUApg^S~4S5oV4&t zpPJA6V9*%PPbFBUs(MTs?|uO<a zOmTMlTmnIUg6>BILl?cPUnQjQN5|(4QTVdBh z-UZF?e3(~iWIW$h8r2+-E*}z@nx6TVxLcQKc`kPUBatDOZ}l8y?S;~1(2(Z<|WMvrd4(_qDa(+4rcQye+!PNgPpswb0fPPC#X z5|<(;{$@xQ!jd|(B#w=BF{wg)4E|ty{_Xf`x7DfVQ$?#CbN3y9p3jw_XIjAmMAV@J zlYpNQi)_7L!g1>Ll0LNmZT(t(Bo*_@RQAc9#D3@_BW}lLIFo4qJe$$n`e0zga8;#( zs?0Crs#HP~EJ12WW9s>n=|gF<*l z$R_R`>-0{)cXa-4Ixu@N+_$~2q}4jl-vx8y-WElq;kVpEXVQCf&Hh19K7suE$}N)* zwuvdq{MT@IA~e(Q4}a88CY!bZbXG;2cMHGcvoa}dBPxck62(H8NK^4$T!rWOz5*Zg zei=8{Mk%#_K!S;K>KgDGbJJ~Ht$vOMenvy9han=E* zHJb#9_bJX+6P~o;Mvcmp-ZV0Y#n9hBc^S#+jy>kixf!4Tli!eljRSS|So)QF_71ug zGn~VeVY(ls$#xaKYR7KM74)1_!9u$=s!qGALpY}s#U8^@%yYHtB)H19ss4e#w+^}9 zM~iAs@rDmQklnPDfHPqzOF8@r?WDRxllc>P5yV{$oL1(c?u)^#$EDpcL5f}xql_25 z`kFiWKbIeS$~`lJ+8(r)>fEZ3$Ll zL^-7Tq&+c+P(^=Fyoj5JwWxo5R#mG#hsJ3< zpq|a*U0k^JA&)(xOiAnSh(R{RzvL$fuQuUI=c&0BvK%bQOx8aEJcWPa#jsWDnFuqC z62GrBdWeCSyjaOQ1pkX4LlU-KM+A25Lw8;)$dJY(wy~-|w!Vj%2UVQNU}Zw&1hTA| z$z|`}dg_(yQ66x1eOS^QB1ArN$M|rYSE=x{)~4XoMtH%~5ipgpLYcn#XQNa;RKUnj zPwAGE`vtVK?VE{tL*?q}<|CT{l^M}OXw zs)lm4GgQvN{E}Z^dCB4Tf5>|Cc&OU{5BM&Lklj$o5*f_cDofd!C@siN)@)_TR`w#q zj7Vh($*5HJEnAktPEzl^JP_wHoZSG%(AW4-iB_OG$E*!U>d! zL})`U&VmQhWv>N#CFcz0Yl!is$LTusi$!gW_2OQI5Z*b1l z5cM`Ah2!UOvse$u3S#!Yvx4j%#$zY>8&fn2U+&uB2C_}FJ_VRw{I5IE(0d@tLTJOI zhXrChtXnZLQ=H8kkO`o(#|Ec6s%&83VtkJxiNw zp|1H6IIgMPz52tO2wU1-C}Fd!r{URsbjY^p-iD}B9=Ul(@5kw72NpXx;YjH}Y5zcN zGkQa<)FHsEfwhH-y!BuF64%RT%P%(!psEV?CTG9v6JvP>>xH!n$VM;!>hWYm`lgaV zy7&NbnW8r%!h#+d*w}$B!>cTMjv-U5`zJAat#2ds#B1E~e?P*3Cz=uF?X9*(7r$6& zpxHGV!fUvq+^asVt2=Wpoh&#lA<$Zgq9ZZUx{B<>rglFu8jeBJVJF`3wYM{e34DdU zM>uqUUbp?s4*Zl#LxIrGD3YRA!lMGe+Plo;6)^^lRA>-{Rm}XjZ6k|2Ii&-Kj|f+X zoD*4AxMa0k(e$Pn888}c-?~Xvs(8N+5glH0h6m*rH%Vn*E0Xp@a!jhp z7|pKowWzHm+_xHE&uj(V=Lg6D5dt;pR?ijx*(ZI7o9Uy6(z{_MFPCuj1uBg=$9t#h zr&>+QzzBV*Z5;9{IPA3U!htMfPlY_y9A z{U&|L2~ND=?sW1SU1P*3ioxvNrDzuO#I)@7fW4%s-0!8ItLQV5tGr~pn>F-GybfKU zQG8EGZJWLYut;TS(?jY;nKzto6YO7R;cG>tQa?7EE*3v)Q?Iqe*+%Bj#s{>lhs8a8 zcR%59z~BMeSF9~;9ozGP40xTj&kbd#jIvY%_tr&aSLLYEXEL6NaWM0)&+>A}sEeQ|JA1edcoA0_W{;LrO`<#{&$rE8*qL)0Z&2Ow!! zk7M(&sNioKCJWOs6wA^`Z)!#JM~LQ+bft9MQK9US+6-|V*9I@W+@rIOjI)B`%PaQ; zPlOKJmt&-=rW~_m7Lj+cSUF&jQn@_b#e4C7jDB~Lb8TbBt4x4tlN~JY|WuMs5 zXR01z1ywoC2ah;LTuflLffn%R^LUu|3p#rvNlm+cs(IxMU#!?B#RMVhrw`PBV|$mT z5fM7e&$Wu(hBy4dGCx3g6bT()Oeq|3eOZ1M?M#&nJ`$1@Sp1 z#xmwti(XW7EQ`Ko_S{5&8G3ZPLON_pXCCJh!Q3hn7pv1_c4zsnT;NX@-CT5x(%lx# z`+j7GYWBqIkU1S9zs19G{#5@usZ)(S;$Py{61*Y{Xst3n0s@3g%K3U`du&8wY4E5I z{#vc^Mran5iySG=^l$t)!Ubvs1EUVZrlQpd46kW?v$rn#R@}fEWAm*A<)W+_k*|xb z#=*x+|5zt7EpV9mUGLju&7ojLaCOqipw8P3-|vOT)28$2sSH=(&9oW?=l`USAd=}| zpXN$-eLsr?SE;(ytGyiSo?R6o=NCxa-C3K^JUECOksa9Tklx`NeqLCf(mS}o&i130 z(OH|@^MoJc^x7fctuAXCHr11SlPGR>M~y=lU-O<{XHEU@<=GiGv{C6d1jMx~SNCpI z9LUa^^I#BV9l4Y+l6vr$X=EEg`$c?l8q7sg&tc@-1!RyGz2$wjxT}`leISA8wpuibt1A2Px4HzBwYB|D$bC?=dQu-O+sxxE~tJlr3vWA3%R7^*o7$O zKGOLw!XA0okR)8WSYq6-USMg%C6)&*)8YZ~w+R~F;WgSDH4k`hSIrBtUhrwc@g(p7 z#B<&#r>O1NzF!;>0!4sBHjCu2!@nTRwG|i{hz&m)+~4S8V;@^ca=jeg3U&&2l?B-B zTx1s9cb+l zv{REm$Symp+}OHfE@dz$i5EAVgX!490D1X2I^Dv3Wj;OjV zZ|+i7%8X3jwixyP;fKmLSI23jro=xFoz0GOaYxT{*l6)|{TjZPMfZ%h;YooBjFvNTkA23Y$*w;GBt0=szxTPD&f*O@kO(siE4csQ z3C!O}gZP@#61;%KA(S96jvlbnp9%;PHSeKdOyghJI=807tgGv!zOWn%j54EK%z@e@ zZ$NsQvKn0Td`0_oXhm4F2LZfbvOPJVD!m*LF0(Xz4drg)lBsR=7sJvltg;bJ9%HKF zsGZ|{yOp&DDeM?l{o5IB1n>Hn002x1AS@51Ik{H*x*rGmBlkxVn&Ob9w{6*J`eb{S z0P6|Y{QKDK8UQZuwmS}rs-I<*?52Crf_atIMCl11TxH2Xunu>Y&#Sg-a>i_tMs;dC z7!So~&Mjkm1kQSoD&824_72!U=T29P@TZ<7=w4-!S)!{$(6oIk@VXHwrB}O*D1Xmu zJmz=hTlPsK$~*Bv5t{%9hwb?vO;C?uh91-lL8HIQ3BVG!qWgw_&?IZ8Ae$#<@8I` z%gVk5gg+Kv=d=5f5>OiGaysTg#=js9vhI{&^UTA$Pl|9pJpE%%^Cn%B@Z<=zr2gfM zb{Bw9@&BTQ6mXHHUVmvzHHOW2kjwPeTk(m*^=KxQo46D=dTDP2o+P0bh1h zd=S>kSt2B7K~M{}u1@ocM>W?ovAbT1VSY|oMAp)>jx9=agI(^?4XkyWN^iVV(7{VW z(aPs^Bs|O%FXvm_CxYwh3rQy2!aMi`UW+~~jU|%LDM;-l-PRZ#?hA`ditTZ}hhs@$yIWsM z%q`0Le6p{z$4i<6D%Fc_9$UB95Ni_y`x4>v@9$R{-iNG|bzMJKo@S7mr;7>GYtYJ)^MR7W7u zx&lLu_9RS>KWOPbf_C?+DHDQtzT0KoBolkN;g=PRpLc z@rbk&QJOh9MmRy^9{N5r5q_7I{@6}CI@v7}!FB{B0Nc3@*p41EX9#)YYzzIXmY8-^ z#ac<%;$LSk1oUKI%9*Nw-~Mt<*0^mI&CEI8^0^j)Q6dSdg{E&!w!iDq@1L;S&wuJ% zn(t1P7{|;SyNC0{ZO0OhvY#kR%$Z2SqP_V4{=k!Rm$iYL->>v=e=gYa$ieZfW;s8L zBt5A#@qiZ;OQQABt8p36H}^y1LyaeJe4g^ zf_d1(ovL$pRwPfM>8xt=Q+ACTG_F0C8m}3Va^-%TAfvr-|JYMQ0X{hn1$IZWXHp7j zqEky+-(9a}-_CQVYRDGlot%$wFBD#njTzFrHQr=jG0rXi*y4q5haWv1pdbTeb%o6( z8~5wCcL}3NxR23ePF?XvUZ}#BJ4cgY$SuRAahh+FI+&dH_0OKo^#brw0NLIJ9 z=;C7-Y7`fzof&QsdM6z5#-F}qCX=oI9K?~xXeOT|E4GLE2ZfsgJc}!ex%JEtyUrmc z;g3FF{Vw(dUedJH`30Dh*&c~rjW`N6&VQRSy`VM^bF!2xl6+X!`P5nd6h@Q2j3qo} zg)#CnW(tA=(6ri{tHK}uOJ9xDv%kW_xbIuzx<=pDuxBUFMVX{VAt-3LOEtE_TihLI z&MNeguXc)vrj56kfK;>3OHESoOaU*BW?`X1-rTb_+;izD3fInSE3rIA{1$knTct8t z^@Q0zE?+yj!+XdT6Noe)cS{EpTu?_nx5?2!#` z`aRvkc#ynKnoxpjuD3zy!I9KGId|}!;o7X;wSB&NH4Rc)X6SP&y?w0-^@1WGacUnNPyNai>~TTP1i zl^S$*JIg;%kX594ufFq9JzXmE>qMt&&Zy9$q|8f$EC?uB)XCYq#Mx`Z)xYeoTF$E< z%&BlOrbyQ%STo}#M!f_ug0cNa@ClBzWI%G`xX74g(j)3u$2ZUZ!*`AP^f4EpsIu0N z#Ny~0Zn5W{p|7M7gdq}kaY?oF)(l^&>WO5=MMu-mbxe%RR`O2@1;}EyDkmM=pB4x7 z#9UGZ(oekWGf(Boz309LFym3u<|~xfiqMLqdI9AgaZ6_kr(@jKV_P%c|6dhE!HU3st2U!T_Bma8!Os`retz zmGy6>mP7$=S;IweD$#gD7n{@BcF_6W%APB@5zq&`MJXDV6{4p?B1sP8Bmm1xr?rOs zDM&f66gb;`Ll@+e<*JI!Vo$zj%NVSY?@q8fwe@WAm?%4;4N5{MN@sT+h!cclI!4eg z*OLpf(b6P*GwgN*WW9m>MKI!dOJhbljUWmj@2urX@pEdsNH+k?J#+)ACai52a8wL_ zYo%2NIRMczIQH9)yWJ#S$-g)?(L~VHIk|6Tj8ln1j+Lj$X_DFGZ22W<4nL5*iEvW- z+aJALp4>sTQ_Df!1ey*>MwEW1(VlTNs8gBUZ1~%ouV|y6YrYtL(VFK@<8J(~4)<~V zMX446z`(7%!?+WBzGI(xvaFIg%p3U# z;J7r9&=gXkkuY828{>V_9qN= z!1P*7T~Ch`otUMbj!1|ka6HW})0GYuI<=Z${`woV^emjCIwykY5V{iz7#nYAt>^P4 z62p&4s#uZ4FE7u%7qRuolrY-dKb=n06w$l_yz!r~c)0JCDZ(<-7wVH!l9~OZPdfv! z2_dO6J{)Ro>qiJi+W?+RR{m@cZqLYb%54C0Q7b?qxn3i@CMkE;osSa@@$;h z1EDgPeo*KuC${2cN=yT;>{}LDhVA&waJ#m5wQoEDs6d2@QaC9ai${l^*E^Uh3G*OB60dh)K)1*-T2V@zOdttetj+}m5al+Zo!JQ)Q*L%u~y6c3?4A+ z^#y)6_>ha42&`#o^es&CN)4}Eu+u3q9xR#Hg$qV?mIFT{+7Fm~2^&6Cm;8>>NnfSX zU%5VE43AwX?}f?^DyGK<_?)^Y52ITWrR=4`Q!9@B{JB;^6d~Vzx=RTC;CZ>#12x^a zW>V5GXPrX4EBrCBdiD3cj_7Zgx;}VMi8!_EAF@3f(t>;Y^4tf=nhJ!Tnh!H`AN>)^a0L)_xYyV%*Z@;Vh+2+cBU4`FQu|l+z{%fLeOg2F_I`B zL2ub~lHW9L)Nf^30~GzF(4-6X+B1xwAs*Bd3xlynm-!FBWY~DOF9MKDVz3d!kZFi^{>k1R@KNLZ%XSYIUh|!JY7UuH@L) zaAh1>g6@Q}3+3TP)!LBXh-fE*J@8JHR7*Ndd;X-V_<>_`L$zGjYN~W|JRzPZ8KuPo zi|KR!)fQw-K0?R

w&=_1B;LoyHFJY#Vq|MRREErJqZjkfkv{iFWwUwZyOU?XRe zDf6D%^+iVNwnwP?xla$DW4^@B1!VNpge;rT!q^O>_5x+pE&|oU{M~!k9u(A_SDgw8 zDkVhiH)K{m67?k(L{;AAKRX8=g&Vb}0kjOlJR~#EY@E+l+txkVUz0y}7gM-Ako(n( zBMiYHAlUBSN{%x@L^r+zw#RXs0Ez?cNzPlx2-mv5IX(8c@iw+&QjF)nN1rsIg&Ic$ zzVK!6cE(>3T*u^W#8yKV&HmkJj_Rt^+=sBbh+4~ieplsR00(3-@B`2=Z5L(;%m#ly zv#{`9s*sS@A%-Al6WDjW1p>>{w{`CRx#t%vFsz<^Tp2+k9FJz}cB;smM!7DWH$QSZ z%~b8!!qRs7;x@EyTIR3$N0<_azRt(9?p%6kp-K_!JOGLZT?(Y1tat=@ZMD&9N5+Hp zZc3h;K$dr)AiHC~r@z$@z0o;>W%{>m-+_#mhi$*oL%aJe&<`z}?QyWsX6pM+$fnX6 z4UD6P!P#|6YG|9~;-d5C2Xz-}y62fY8rLKOkn;c()Rdk`&A|vI<#B;M{jt4gFMb+F z)JCLpBj4>4l}z?Q7*Gae0LNExO!W>V`E9b;xZP*N!?eRXs)ES88T9kx3J&$bxSzov?=XRu!OWU_8t`Y z6iE@*pM>l~W=CIkurW#hBXmwer&*tR6w>-*+|)4b1w=RklDrHgtKbG4S5H6Oqf!3z zdTN`~?c0hiitn}j?$@4??&UeNvtF1nOTu6?lb@M%Ay`k8<)y#tP`G~CsjsvA{!Hda zE-R(g?9h3bI7A7Z)hx-YVkiwndIpaUWxFpj$b*a-;f`USnmb};L3n!VL2&z|7|pZy z1@_BItl^df-4a0IxZtB@C6~Ak1ku>MOjy!1iCNq|L?wIUj62t{(wd>XjD-Wt=ISb{ zCAaYB!YwhA6HTYh;J0mD1}rZHFYrXtktSNq+uQIo%MlK&ievSG_XT@B)zY&0kz7|q zTLEgjd?>4?P+%!&3p9u)bq95=K71Kd-@YiuP}`!QW%`;=MX(DR{BFRg5iB#EqWsf? zVT@~7ZcwyFI)HocZO*kK*p!H{fF(oDAZaaTl7f~YTbQE8(HnHPiq;9vXx3-fl41$Q znP>3_^5CYMVogelPd;X>9)je9WN2W4J$2cQ$Sy>1C$)dMJ(eqEIVbTf)kpz}O=p&5 zIrT1()j=f%CFrZFPKw62v-dsoqZ!%d-jE$b*7E3vhb@$rwzG-_zY0Jfz;^$yVD7p4 zomJN_+p-L;6-k%$Apvr=8={E5N7-lpf3RHrVsL)W{oG8T& z2QRMIX8aWDa3@zlwAgNCs8u-9TBMF&r@2n6galzRZBGCoQ(be>)Fu9O;y&_!bQ5@6 z4hcJI9+eRITFj8eg9xAC=9ReyQ#d>u1z}$6x<-e~jupsZ$6t=8rh61Cvj5s3lj8Lc zC)fPdCbj9~rvft2FU-O+s{g(cN_s`srdQsq^QY-!vzy+GX<`kvS~__O(WjPy$8{kb z!?1wlxh`K}h4Mww^g%2d3`JYEDh+7b$KPXoyto<#9a%{Z%HFpYbVqOTSeKO_3;tJB zwSd4^V&d=8-Wh5vyd9cP8`nG57vv#%^Q(1|9gDFdKpLMt`$bFw+!$|9s3n=qAlsrO z>44-Az@T9z-d5+mhUL8dl$!?}ADbTeRZMN?+xggCII*v~|`_ z-o%BC@-{V3##B^yj$F}IS^hMso9fi!_lm0!om`-a0hQC1$C6cpJ_0g<@@RgCQH=)Ns5S_#LsgR~N+2zE} zM~r(Toh%w{vraP{@48+_+cDMglXn7OCEr&%Vdcx9M>XA~2(cOvp_lrFgG*O_KhUb9 zpCWuX%w_eEK#D52eg{8#Q9Q33&78P8Heph-jC_+)$+X-rhlX-?^U6IW2{hN_?WCj~ zpLp@DJn0_R*f_c zWb;H^eo4I|#T~)Y|5Bd8x0db*Gx!36Lwds!EoyOJ~EA5(Ss)6=i)R6-cR%HKQ-@(zG#RF ztQU$>Z?GhCy(uc;o0OsSY??-Jdu8{eXey4A%i%*^mDSM=Jf=j8?MTfOr|vqNHqaAw zeZ@7Hxt+OOo;+HF_W$~{zlB*OugE4k*zfcm?t;l3_0N0Miz-x8`RQPVqT@KPCa*@Ti zSH`>>Qx<5_40|RGvwxp*Nqn<%z5?{pKZQV}$8axLk%7A@NrtBd#3IOyjkDb%e|+0@ ze`!~>WzQ|<@+T(86b%0;6{dk;CEV=>6A!Tbb}Th6&t7e2ApacYY|2}&^SVO&VC;=( z=D{oJNn;rHXud+H{I|}dJe?(rwq1*SAK-`hV2W)Acx;zgf)@O=$8N*Y+l|kVN|}Xl z+_6v{QRZc6q=xTZ>Bh^(ZhxU%T_+jCaeslDt8?qrgBv5g}|7XlBQi?d`qX z%=})pR4B~&)p9I9zz$27Oppyb^$cOAE^Q5-AtaRR8VwI{rtnFks6{qD<(x}nvy}i4 zG&QMZ$td!8LxS#f{XhbtTe!&reRxE#gcVz^N=LGNp_H?A>Q4)HXN*Q$wI}fa5cdJtd!dpd)slcBO@U(pqBy+FRNrfvPvh@!(RsRK$Ldh0NRLVXc*9l{9nqzftathb54#eIk}Vi_Tm0{dYe3KwEx0UTG!GG^h|Zk zh_3xPbegoV3N>dU?p9oX5GcX=G|B&MrGC!hT$EI5N}c7Zr#bo?;kv1r-vH4XQ`%1% zTN@FSadIhQlInF5dH&ClmeSeh6!p&rzNZXg?0ib{(vgVS>ugrib6cATFccVGd%Zzm z{m3zA#o?r~EB)FDaVDRUm9i?NlcVYRQQgBs{{I*k(jzD0TyNW2ur&w$>PZzZ(iIJ! ziInX?fk5-*|3nrH8c^IfPaDivx=S7u%I`SEpyQTCTmK}h%7vTp<&SDa+XKuU)KcbB z$x{ZP3W4DTMp|O#(`d5>FS1wGm6I{PWMGTO~&iOZ~mZdBCanvY9!_q=tBwGecC=nQ-6?4 zkmZU^c5%gY4XCI-;W-@C=&c}CnnJ1^qReH$#GACcqsyL;^S+#1v)j>RZMrz^L_EK3 z(k|PJ%2+0QXU20zm0BJ}Q^w+Tx*#YCK6gWxCS&pnvR_ddt|`IJvnyPbeSl*m*}j$p z?MV)WLxAP9Q+ob;d&4?ml9#UV*NO!zf?3!VZpxyiBfyqsDMZ(3YYMQv4eKPsEhk%Db5f_Vty^wgC_8_klXyj@eJyWb6wDaa@Xfn9aQG(XV(~qb(cbBwFr6W zp-6KFUpgpS?XgiyD|Tso6AHSTH%ueF#qnhyPXD2eDDf236CD^(=gZ|`2#?n;bzb3& zx>aRc@qCDVdW>_XqPY(4E!yG`U=KE`ad=$a5)PJ@t*L937VCP6|Jbpp{Mzsm+jTct zpxYZ>bCB0(s<(0ely{A!SE|?s5SkO1l%T#>k0+=j1=?YDfxn~3j!gs~;#S@8)qD1A6MiPf}>8A0tK+nZr(2!qGjdET4- z_Lbf>!_TN2oQ`+DECAZhZXPpV_`?6eI>?`p70+0&sjlHd8SVc^{qkZ4=PP^-P`RyU zIh6+^^`~weD1@$>@}l~|3*!RBXYW=7|DfvbJZdaMUDt4gjjOb^Rx5*}Eab}-2uCDm zJVU0_**SeKS@AREn*k~(N_m-qIruvoQ3hFdVeUB5w2>mkh2cXxjlN3r?`@R@cDr%<93x{M z+96Q*s^bmSmt-%!h-bS;8K*p@3~DeDCAK{3k^z#q0L83LiwhYI0gDC^8gMyK6D2of zcp!NC#RzD`iU(MwuA;GNnr})(uL1LR=zhUt%X!xgK@35uWp$VbF_EblBre{v=xk3- zGi1c`>#FFuL@ndUG#_SWcN#PpUPQ{jkW_8eLp6Vl)h+o{%9_kcf9eYgK(&AEjes-C zJIOY-?D^4;i(m4-JpWE~f)@Co(-B7TdWtjdQP z*`leKq)G`_q=C?4QE5lT-MBXm7EU50;N<_odN`H&7<;XUk!i0#XO&lmg2o0wMpXbA zGY8H&4Nu6IcuKKTU@OD`CWAARW>rh;Y6)xNLeDq0Cm9pH4(JHC;-`|c>e;~bd0qu+?q9rPCQ`-XozV4I12!f>$ zRMzDlL6+pP^2QIZooHVvNgsEq7BQ-ZmxEoJD|8AsojaiwFGev{`Y>qmP#- z-g36h5L7WwC}XY36_lb99ra(hq#FQ0%B$35SqYJIn6x8=THX{j#l|M^I1@XOaxjz` z&iR+P`x{%pBWA2u34GXdE2656ezg2zS@z99iafo?^+w7OJY{63C~f{?2iw-e91+LB zmdlE(ax`64Ggv`NwKJ5CZ(8dASMvjml`n6EH+Pimi!?n2?B$aQOljG{@W>Skt<3|f zrl-PeUjsx1ONt0w5Mfi@!3!n^JDq!D10wCcOnL%i?X;uT3&p{#`SB}3ht^EtY|OlH zMCa77EwrD9>yce5SuO$RDwID=FXP}t?Z}!9jk>WrZEP|lL_Td zW3uy(fzgd43;;RrS&~_4zCl>%-F4Gz*k@aHx+h68jrAS8I)|>@Y(Yb8$n%LmY|_RU-YJKipRlZg_Z&9@iolNR1pMJpP~zrKHCsPrhi2@ zk)fWxem~cvLUiI9hW#F2i|LDgxS7fkd`6=!-(ROPjB#=#N%rm=es{N-y<+mG*9>Qs zju7B#bpOOH0TaoKj-Z|bIo$FA(*O@(Ji)ceVZ6K)A%ZbKWXh=2zUX<2gFulXYMPY3 zn{6WAsq!l6MAR&(fvK}r7cq!cGnKYwZI$|FlHuK=CC8!I%ZLUHo`251H4nY-m!VL^ z6o_x9DQFq^WWg6Ka4Lr7!qg>+h0g)*s43Cl1D;?nP841sg^ z8#8(LcuDI%r*O1%*a!SYg--N`0l#r1^6dDzGlPxw2%e};A9S#vM$KgEs{8uL*mC5I zh;fcTu%fSI=pCBP88qxs&8-0e;Ga6vw5VK}e5drXeSG|E!b^_C<3e~3%+Y$$r*k;| zn|=j9*-xX@S?|YzpQrAG3*fI&<-i4xTh~>bt{sxkm54dl03}_JCt&nRE|6n?o z>3fP%jV1c4Nv$mIz_H8q?h_g)6ca}8%A*4KQ@<^TTe61x>oaArWlKmW0g_uXmT%+A-$}ZyjUF*dU7DKvX5EDgG zEQX7{xtZp)&m)L|edZ*$9`%JKbpo#gZ{qN=^If+s(h;lSf?{{V$cP@wJFFBmHfr?! zw!<5Unh1a;u-_5&ST5^H3-4Q$MOR=Fz3?U9+dJGDv#!A3e#*4mL+&F1>G3tjjD{%NX%H;0r+%p*_ zhm)1Ldctm9V_vm584${!`Smz8%t;d@XRoQ>cqr4)R%EspZc{O8RWZ&dtG5L@xl4G|aKkn(4(a`BXhm^EK-(uNhK4qObiT^a; zNEA|kK?`o%KK}sWRp@Rq__5$L`uD&DJC7j#GPMuXP2NA>rj$Om9#)e;HC+P{)eDEI%=UF-UNzTwvwV-j7-iF>=AKp!(x+e`a1QjC&lo$xV}_WWL|S290@M#N8SGvk-M zPivXIw!FqC(7ym_g?P0|^?tXv9|$RX?EJ&8cQol3-f$HXyFt_( zLIpZ~Ce;5uU%_=F>`c(kEib?|8t|DMdJUP8|7b?6OLUjHi(%m-A!HNhKG(_%ync@6 z2pPza&78jHX8=+*XbK|9Y?o&`y)DiJ4P?#BU2LKO}xpJt2FE+QoJ)MWqz}#7l?1>u=6h=0rl!aV;y zzCaIAe`h$EvxWE`!?rM;CT|_cZCxpzkz;*ToAvLKzX8GN1rN0R94SNiFl9|LR$kf` zno6Ut-%mG4YxyyC8Sg*D>g5CezBTQV1B}>IWHWJ@B@0Btq``&nl@(;|IB46l7E<>U zKe;cmzT&3ru2qGdy%5HI;DDN_ADqDP5?U{$_(z0lA7Dt~xa3eZbY&rwXRQdQIeV-p zuX!%2N0AjJ?sd21#KUXd%>>O~z*PZ`#11+Y7o(jUjr5@z<7FyVuT9%hX0lcx_~Hso zh}?rV5xOwRfjS3KZRiCis(PSHw~r*{>{*Ond+2`tUQ$}SLZy1{y4|G`bvf%K0R|0k zc#@MaD|6w&LCY98zJxZGzZ88NE4#3TzQJS zPoNkS>{r8$f}~ZtPK0lzj5vytYMuE(bhU1qsUJwwb!?#I1|!N$fuGJQqwliQAU%7} ztv=NkhL=jaz8Ex$*4)VA7kqWHR*V@a^_=2VO_kUI8nJfJPa<5K$+OOKgmy#Fme~bT z{3<h=ol9go?@{lveM&zcHFM`PBGKL7n)v^^9_uZh?KjUxccsQW!6@jNvmDd zxc9^w3F8~`AwZ>r~S*M-D&48#Cg~?wk)gZ7$U}V8C15{-OHGVBA;Qn`(V-< zmIz`uXvf1xdI(fgLLm)DWm0biJ(g}MrI;q;3Qb;e4_1dGw?w*o1clH7L7`rein^W0?g%{(BR}@dIMfnqjUo8{%V!Al%T@q+9tghW`VJ|fOlB- z8ES+@V5eJ6*gA0Z^c^fPlrTh=RuB%kPuxT*5zua+lda_ByZ=<@JEHA78};AC-Tejy+w zpGAcCFf>C>7j$?*w&93z8XQYEIe9uQI~vH`xYEGcaao_APYpmy7m@MN_7k!J5l{Y# z0m~A*+<9%1L36Jx-&|DU9nnRto()hoAw+@B?i3up)QKAS1s*(#=9U)+D9L^6p%%p* zXQ7E^%$L+TuKqkbP$aHW4q?%(KsGC)lHskVW;|vLg>a*+xWOkIg?1C;I?>q zZ?(Z;90#>%HI9$s-Zuk0BQjaEvtIikEh>NDRBXe;Pa;x0qAeYzWZD&f4$5&eCX~bl`>h7^v1jhWRO`({aO4jFq$#aNt5y43i``c8uO=JKMdm7Qj z>;>*d-0`{iT*Nf0RcCAHCZ&C%BsnD>8VH-hD{0$^OcRWbY1dVxEpIkyvz|X#&;I)- z?jpzS*yq~0>%m1Pb%G}*Ng#o`oAEin43ZKfe)=88j>w^Bab(0t>nZhm_1a{~%QxU* zmEt)>=y6ob4iW;xuZ`2l8~UG}BmlsOFZc8=t`S6NbmujOHcnq|YwUZPA^s)O5x)*A z+eQos(NH-Okzh;POkqM`s1Cx;Syk<+d$zHS8Ru>E)@t<$*E?GC9L+k*d_G!_rR_6%c% zxoV-pem1(J&FKp)m8yp|X5LK?64G>AK;gH*rwfJj$$9NjZ<3QhVftsexKTp$`i7CU%27TZ) zJbI;&j479;vjH}ef%-~eyz&Bp4tQQ-eoa!Bn*KvgCzkJrTG|qlvI}&DcjZm;G*NbX z;hxp~tdv9TT76iMBXHm(tPEQFZqCg?Q1tw+{QU5^jwhI!1yNefwm7%5jJ9Ir}|5u$^S==KdsAL4lQYo6UMSk<@(YR6GW7FId+av-p?0FNaHQI7e^ z!Yr@cN~HR_c!nUt+E^j4V~9lh!-QY=59s3rEC6Hz#GnceHmuW*RCwLN+ENi@q>_G1 z@lX@}-sR&%2IsToK?50AOzVmw%H|`2(@`DvEvJ(fswSc|D_8c`uU=7N(DRhAS>Lfn ze_A#xg^@;zM+c++1c&YggBk)*ktoU)O1;AmOj5B%XmvR5`K~+C;mEPe6aD6vQV`paQ^4Ji^bCKu zeiL)A+{dD>FJWt(mxnEYMevaS%pu^a!Q%Jo2k7=S)Qnx;Q0y+=P=`k6LJ~Q{GrxKLYf$1X|X2G68zFkY~;eTvNm8L&mjxh z`|d$9=Oi8CJdo7;4d!dum;tig;6}6E*AVlmk^4x@7`BraFHC@L0UpKe=!*#;s4D?E zM79qJ*$8Q)+xTmT58F$hu}l7r99|W6)VR;3`<-sCF>^#Y09V6cS*1u8ez{)`Gg{?j zW8e>uoM;er&wB_+j#jDQW7n{2&=)rt0q8?Hh(N5cvv-8I9UoGa)c#K6@zTU9@f?^7 zcn-91@ChlXzlGos#HbGpdn(%qIS{itSOdsU@?DZzw)#}wFT5qYwmVwdy@5fKS*ua- zs%hBQs=Z(mhnsIYoeaw`tk*?)*O4C16ETka$iLprmA6)yb_&PCSM)>)v7qFxBQdjz zSV1tJkj#SL*?eN+k9`P{w9V~YFQm`9&p)$Iu5;QgOSS5yC%KJcC;zee2`4{Vm_9op zQ$)2{EU~rmr_zwhdXnnDaLP6;oxH=lI!DP~!$a zdP_vF<3Vn9l3gXgY4JnxmWU8fDKFW1sav7mc%Ecu{>}WVe2AFN+8EYJ&%UCl=64Ha zm_a@oWQgm_FxkXl)_+lZ-)#Hb_d)yeU<&B_(#%DWgP$hrCU;5vtR!jm> zC%Z?O-ik30(mWgk)ic}Gmufj&PqHdKCQu@( z9_W@lWq@!wc27xD{ombLXH?U9?cb4kwy^|hdzO`ESE_;XBn z!kW)lSt$m^t3QD!2UK8V0R#`ETQDXX*3Z9HGMZvae{w)c;Sy%XISqlV0Y?~L&i7*G z6fklBU3bSez4LQ!`{|tThzcBWV=fGGu38>RGKx8KNvhU9#IpaycwRGD(6q?5VXpIu z_afnBSMZ69_fPae(~8WD3l`3P8n5U$5;8{aOdprvoGeT_D#0E(=SUEiD7kSF{A%d_ zU2xlD1rUlQMiy}EX z_T_DFA6c3Dr4=}+Ly;{pC|q0mqJyGiNP^I{;3ev(Mho50O)S` z+Ro4kwa%NdCL7AppVG2GvyA9fG3M8+=%j!73=|UFzck<`G6Rt7@5>Z&7WBmhO;j4lCpt4nJw7|RradAS34>$AuKg}IVy z|4>eX1zQbToy7E$ymo&p(x>)=l(7d~<_Q{a4JL*Rnqyr9C8v0tFeAQ`qTG^v0 znws-@n?Bxt7kwAxo0NX>N34qc9bZF%&S&;8wZ*B;F_k>}p1{2>FnQ+B%#1@JSoc+3 zKHX-*rQ8sI;^$Ei<^-Ak+J8P$GViT7>A{@>5kfSQ$E2mN0@zcktBG{U6utp;4*$m$ zR*scmnkxJrH?kJBkhW!T@1HR%_{MY#yvV&@`f(GGakfTt_V4I319+nyOTAt%_0pAs zOlj_wV}?0da|;|7;?`rxJ_Y-CP1RR`{s5b{?e*LLIxS$@NU}T;NS58pl>W+Yrdj*9 z0dI2X67mP=NB!fVXM=Gw;d+P2pRn7177)TD*!Yt7%?5Jlyfbdj)!tZ&K$bHU_U{;4~^GesDa-U)gGVshu(!#_2txH{epEPLvj zIt=dS8_)9v*z+s7RN;6!rykvHRDbC7@gm&gMLpLi4*|+Mq7|q9cYQzjaZ})mn${%F z^JlN$cSgor7tCjT3Yo3#=2m`*h3fD9H!hZu0RL0ihk>0v9by{~xx#JFcmv?N(7ldM^rsGzk!z z1wkp&QIuYSG>IZzKu|h}QUikY-lc<7=|v!+cLAkKM-f5~5Xzm+dCz;l?|1Ki{>~Y8 zX3w5^W@bHWtw#awpHc;QNMF z(ki*oNHP4+SsFNpD!>fcRLx-T+?TuSdYgkR}C<~7S5n1ATsQ6G`2W8n#;Pn-nzkb zwZAAEj|{C==La~f3L$Xu0!PSyrxTs;qrw%Rnl$1oR5?~|vF>Mpb zvuZFN{`rCIpko)&L=9wJ6ytuULm=gYR9OrAj>(;l8VJ)i1foV*ZZP*d=-Qc%0>Xw` zk5VD8D)R5f>dwtNXb%(k6$tIGWg1~lgTda586|*iDorrv9Lhu+=AR3}t;-j_CkhAF z!5A|aZt}d-xd`$J>W&9j!5=W?*FZUR35U{yjD3(c4Xho|lXGZaNXoe2|D}J6=1pUY zBib)v(^GN+=`&Euhz4Y3{}`Lghdf>yrr^WxD9{{xN_2mV-O7HZbQjk<7Qrbv9D;>C zos%V9=6es;Tg*jbs4*}fuYn%@f$o5`W{=VOidjKoXdoW}n@K;^emu#*vVQ_-`~YKe zZCcRW^uj6fdd`7MP_dZMgxQKN6JkwP z|NO(|Nq(1FFNo;aDCT5PuW8YLX`Q)-OO0?xww~^mS)cA4`3g z5{-x4DUd1x3}O!HWcLmd&_b*PpbkCL^Bb-N^Iz4Yj?>u^xYmK_`yQURpXqpqk*^&H zCvm`#AQU=$K|0-kpS??(yd`_NljG27;L1^AcCZc^;6rmdh|JyHHz)podlJ$X^+4j6 z-;%P1)bUik10?5y4uT8+-=5%+dCNjnBF4IggI>aiW(mI+a!GgR@u73dp~ zk;R;G(*DnxFsX6 zS}}?K0Qj>sjGWV&#$5{Q=9Y^P#0)K7C@AghDj5nZ%-21a$=sw^<1Ww4A^|i_7Jfm+ zdP3NL=fHoHZh>>#n=+|W6d$Dx-(m0Kop=-wEL%({Ck2 zJEswp@;ok6U^q-Vl!5cz2W7Gm0|<|>2<8SI0f3MFY3{cbidYB9X=iK+ZR(jxEGOJRLKV$eorw+P_?r&B} ze7W+OPP?&;M)?Oiv|x6%a>f>Ge3OCuKgC!q7Fld|D7*-+Oz(kx_>H9*o2$|Ld+eHK z;dh>h|JCaOrTc0%1vh|F>`@1p!N=l{P}4MwXVjgHLnVG5h7o6dI?E9?#-3)`Vx~(? zX~)}3PI(DuBgvs{56r0$gAfgpOL+kJOQ@6Wq6Qh`ySF!wf88KhCKNcq(4-Ac!okU@23`do9d8y!3dVWsX zP&FZ~y_-|>k^tXx~t#Hg^NYP@Ak@(UVXQzR78P? z;%5ON_u^di3ZZ{*7^J?AIfCwJiNNOWr2?5W{(TN}eXnNM2A4`@Y{yD`>*BqK_j@-PZihtxay0UwM8NwP6 z^$BE87`?&A{$lo29DnAf54e@ux}tUII<59btOHJ`I_&Hr-&$aW`dqqCpv7!NO&}&a zYMcZ_G3vl|PxlDAU5?twa}67V=)tybkY2oP@(OG38&2g$ zrOQV}R7N1lU(ww-czusg&LrC>j+1{l9&3>Cq#~N8;U|sU3xGiuA6vY+^c3u1z^T*u zHa68HemoerP_3B+II{L098^lRTa|I6<+ToK>iD z1Dp3VScRQ$HUw`&>Do+VC@TU|3kqa<9UB}MRR~XUEwFW0l-5m9>?G88@p+^z5QhGf z3I24JFT4zG?8iO?Bb4*JA}Dgb^1F3NV4Nb;*q)(WrSMVVFCc)nm@to)1FI{j??AL2 zrT|~=C;FZk^Mkp?s+_(-LTfuL(N%Bl$m=m9+SH+mpH&xEm$79QG-LWEETgPUzt40; z3o}zL5sm)FGTG5r+@en0Q+VdNu1!QiN{Q)k_W-P@rdA-ll6^{HA_*!Vkd#ZMYo0S( z)Bh4!bSq$188~tmnK-({KQ|rSHU+7)wdJyr##R8+`c|+so&>VB$;k<9?!I^F{%cC5 zz51F7BBsm$W|Kd47cftsb4CEI8wot1m^A=a^h=NdCl)t>PfF?U;Bth8a>s;qRwY8p zLu+x6unh~VJ$-*uWFfy|lf@hY{xZw4jU%$h5k&5L?cl*M4Z+h4>-%?Ch`{j^3Dsf>|tj} zZod7dQq$bA9JrTIehxx(uS8=Ej$FMWk?^>yTJDXZ@J+5q=!k98+2kJTHt+_X0j(gv z*MguhN#13jl`;OnxW8K24={LhRO#*QrlYp5X%+^+!-CpQgSUk#GIh@uWCqeX@HP~z zy!&>vas_LU!rL~~qXt76-Uknoo_1p}GR^Ma-T=FtDHxi@6cIAkH+(ih+xq2HoDXe- zqQbDiYh`A^>n`iF>7RD4XFBJW>DW5D5EYT}o!i;A@rOY#1GiH2;TQ$~yd(FlYy zZhL!s>BoPYPF<)_Jo`(oGqc@6vR<){`2#2_WQr7DmWybp|GoO3$;FF&0&I<~GUweT z^d2+!na=(VgeY9bWk0#X0W%*c+8W)mKc+{;3C~2?Hk#KKD{Y|?-pPQTQO@`!HJ`E6L7x+xNb#h&?7}LN4kSX0f z12A+~RXC3P0V4yT;BGb&T%9N=G) zQW&rKhUISoR(LLmW|z-sq}qN!v1(xH+^#hxbn&-v(UNE6^g=Glo%t|{iAPguABCK5 zE1~#-@w&$uSgw;Uv>sG^)5R8_emZ+R7Q;LaG?nvp9CzmLqeMq^$G3sIYqTTCAy*C% z-PAap0=b9ipY)a^8b@iGsE2wF98#A8b;*%Q!X=rrNBgFg6_abUswpBdV!rDu&mI;A;l-y zD&#c6v*PBJe8hwqERff`6?e0_#EV_8FENoW&ebGsP7cX3xi-8XG}sl|A{ab#{p~xo z0O?5}PD+`%J}4L+ZFj{%!W~hoy^E6voEpL>QVgvxXMty7v3kO^sb|Y0Qy{AZ(YrZm zMwwFNsQ9*2&$K*8K##y;%r;Jp&h>3zk^1ZDU^xwhNKx{d8dD7@B!$-9HK8y&7rriQ7-JlZi;1%I<^w~ob`LvtHX;YI(7du)CZatR6L#jMB3`^ zZ4TXL4GaQG)rv^MT-#f%N}`itdq+Ran7H545R z%RdeXK)^&MTF+N#aGV#&+SywmQ1{95f6^#bBBO1;g>oc?=}tFu1a>T%teuwV_^pr>)z7=8@vmNJ zVDr?pi~_hNN>Nf9T24#Pl5C-Ou>t1v;`#S zjw@S!pB8N2`=`HQRIoC78`l;|vd-(7z@L;+Cf8h|U=wv+K*0dGHCO8$9oE=1fjBr;O?JlMLC{_Iu-)s1H`BIR`RRsbJX~9LWtCT7aw}7?%wA&YN2D)_M<#`u&WanyFwS7L0~M$OkPW=@7brOeC*hYlvKYCGF>$pHi({|CmbK!5^OAk6`QL&?r&rxNqTqj$6~-&SbXKx_-kQj{o@7tKEc zyTQ4KkZucj-_W|A{AF8_cy4jBbD%Kd1`P0SGVq(1UuK8sIsXxl03isFdwPmIe|ZUY z;r^vl8JKM!9#pOa7*i=ioY`qS;6F*=Gk+2=uN3G^!ekGwo1ye-w?!lPc2eGzWF5DU z*^Jl-yYXB!E#KOcTUGf4f;N%q_F{5x@1H0uIdck>fxD}a)D$HsO_mQb6%(jTZLES9 zZGH(%(jB6zcnd#S33$|9VS1m(ttlh}Hz0hMitmjpsj2eL$^O^v-p{$a z#3SxL>HC5{vuSQkrkDAajyws&;%aRiv9cKtp` z!YaFKoNkd9CR?hXtjv~=TcxDo*l*GzQ=wp+p~)i#E({HAkG74UG5sN?ZImxkvzX>5 z6WG)M3QFX@r-yy!N8`0+<^trnye2o4Q+$Cz{rsh1zknd-o#r>mXxD;)_G_RLT_U}R zhqQe(LXqLu@wul7*$#25M47A}VRDJTU=8ln0qm84a?rxFvrg~7#@Jt1&{CqG{H1yN z)86}Ev+u)$08*ag9QKpdS$@K{aA%zmb=JNQYLALh;H68{2}2(PlX8pOCyi3rrlvVf z$G*HNA?**?aEm&&sPQvfb%=f7JKoSTpveXVbYBW5p10nwxHjJ`8uAlDO8AH0gUjS9 zBI}Ki@PDkq)87;e(?utA|KL<&saR*+>Bvz0D%fG)b=p(>yr9~B&hW}_luqe+9m-^W zh%*s%Ne%pkI^D2RnoBnS>bB28sroa0`>+eqf0^kVoy>-713qQ+ew8Bx)|zGV%_26z zIRmwTd%EA1_NK|>`UTQ4s?XzWFeTL=j;gDy0xw%@UXNRc;TRujO@CSbGSNz`y!a+ z;GLiw0e4XV3y2~?D!Pta!=AfLtv7>$ekK&oe=jUrve(|mYWt+ z#rO;Pi9*^bn;molsO={rId7WONx^0ANmcY{!vhg*C_1lcJq0t0%Ir$TdiVK;^*HRw zlrbz;HP2D?H!FG7bFT3h2E0?GfO!b6loe|K{L@Umw!QiNJ+{8(A^RhR2dY6P^v+>C z;+dn8#-9`C^)07J0Zj_DN2~qwPqTqq6R!a$uXzWyop&iY-1eq$@pN~t$Iq{i+ZYf7 zne|zmC{+9BpJr6G0Vu)g^T!YL?N0+nN5+ojRF}Nc$vkGLY*#--)G$JZNf9+wt%gxO zU~Zrhm*Vleez>{nfwu6i7EcNLkSdFCpQ>K1V-i$=0^SUW>px3*V{&Tap1rH`y!%(2$0^MGXdYEb z?5w)+uKv%wWLH401{_N;2tzj5o(>_kc0N_uwrA;!?H~jTC^%P6Ogtq7VRPDY z!qV(c0BslqM?rqKR(>tI$3K=)uSa%lI_aM(rn^r*y}Abnq{>K0=MAibht)1m0diLb zxc9kpEHV5FU-u_sO8su)@Z|OD&wd$O?J2Md8(CSX-(&)VR6sl91ej*wK+ee0$U4J! zsbVhZ_uD@OjmS9uvJhFEbovWgn%eeHziP|fm{59f4y_8T7O1Pxb@TF6ruNsr_g{{B zim1%Lapq5qtTj>0!_QbUO^3e<_`FUj4VCh|?ge~X(sDh<-p|0^gjGEty-`b-)fger z*B_lUd@Rem-Mi^oraIkrIS_Kg4J1u26{=;r${uugNFA z3a9{`0cVE_!)h^B9B(N;n4TnOu8t11+$S7$G%)DPJE3lQ8a=ccb*STslZS=+F|fDM zMAWql(ux~KO*SO|&RBp#!_du+sd@W@xoN@?_l{O%TpK7{1UfE91=`4AsSFw zR1F{X6AHfiilun$g$Z{5@V{)63GYfD^Q!2efM8i!-}mzq<-cCUK~rqHax;62Te*eH zdrw+>3qisasI;BCGj6||5H^I(iL~{?Kjvdo9)RLeN^laL#7~pQXPHgQQY|CUGtmC& zr5e1g`@0gl!w@7XO-dr7XON+*8PGP4-zu-O;yBZdKt`8;Htp;U#&32_HL*LNwhfkf zn@lis=xlq6C>bi&;8O>_Qq%0de*nY~0Hw8=QtO0b5!&oioylO-htlFex#2NsX(#a~ zs*JZ?%)h)&aE<*Xe`VQ?+yODES4dId4lLOxu&#(HY;ad>uOmglbVI|ujlaZO*QCn+ z>5J_2{ib4#FAw5x$~dlk4Y&qLGpyoxyaE8Id)ZdpWAiG~`x(E-bKTn+r8}(ssuq(m zi|Q=F0=hr5#dkD9Dgt;-D*v1zGUKfj5cTANv}&@W-~DxFGL-~YsnpE}*45LM5%2)N zsszqn(}3cMS8BT--rg>*xLHLNpwl$v&ji0TZ_3h1RBu+xo}E(uY_#uk`ok(p;H+}l z3|M9Ls_FBRaYij}j0k&pvYO?R=3j7+9!cJF`y$?ef9b`-A9rm13Tlo~wu{>`v-au& z#W4Y}1uCANc+{Kl-%y_$b9`87r~Y_E|GJ|G3UiZcinT9>3PaV@8e5fe_P)=l9cQYr{y4ap@e(aSk7H06nJv~* zi|=cn7CiC#YQ@JTSHHX_17cAQ2Edc*x~V)pPcwIcY%__F0AnGMk}4oRf{)JM@Ki|T zYX4Vei(XgPNqg#LW@B@gO}!Z!1@=a zp?W=FRi^gz$~7WOv+XsRZPwy~2PxhzkKX0<4-_Q)`+P~7Cu#uRQUIEF*oeET8<>69uEOl8D2z;IB5a5KJ=s&DO;=j88$$5eOc&T~w zL+g1i+07#oHlvMS@2fG*kqi}}4EZ_PC-;evDkov;^%Q>&)WVm)@YaCZxllcWEv154 z3fZ)?%+dIJg4+v0&RmmXNjY+@X)7tM_FMwGaarO!cfrx}KH{RK$epG$XXVdXEt^>B z$Mk?!|AF-qlDfh4n?r)VUMX(>3B}i0cP=2HUl}2Ve7Y4hEK7$03-Zih`7oKIR#{ka zv;1_U&1zcx{Y_ax2Bq=Wwx$BBV?!6fZuCFP1{Tm-?C>J*PEK3k(e!JDLuI<9fRCAC z^7WSmHL-JGrTX(}u$B!Tj{v+embw86<)6Itc)Y7)VI$v{89!2OyV38P@-gI*&*)VE zi&NuXQ4R6h(98U-(RtR;@2QNTyE3Q`pR^7$>0}pCpAW( zoQ6HkH-)r&PP(_f2Rj`BU1ZXWs;8|)&*@vHioc~xA`I$grUq`y-A0VxkC{tN8Mt$ybJOA*<$ag|QH zLyzWYHiY*UquUVkj_2FMh+0 zlgm7I5;6>!>o6$1N1E`Cqux)~7%pmpTL}AdiK>^(b65Y5g{5?d`VrFNJ2TS0Izi#f z2S#fxDv))*xg~i<9J8vf7Y&$VrBEsq^Mcy!i1DKMpi2nfA4`I^-FHCF4Y8anh^weR zgPaTCVwuex7C|Ef?fMmjncEuTm}Dkny2+wt?GH?>rW7Mu{I4AmF00~OWeNv?9x}54 zZu0^l%X%^}6w)TV*SIgvqzvIUJi z95pg;KM=`QGga6023$i# z$`$xmZb_#WKozpWLB&#LvXn3`amY;UkmkSBS^0RQjbzAj!e77Kp!4lZtm+S_-qq$L zN96~G_!Ju8K7CD@*Erdw)FuZwa4Cx@sJPc0!WO7G)UT3a zzzTr5x~u%_c8)1GyKum7r@EUagm- zk~=ADa{P!4f5g~YJh0>&)9(uG)C1Mf$nFzLI<;i2_UJfJ@RR=@+Q8SXL12MD2(R+( zgz*_-{Pk4m$&aGL_llm+@UsWJRs&Vnm&Vb#JTo$ML7`bhBxowz*0IU}mb`bPzs zo=X66P)+)(o13TJ@*0{}w;c)QkTmhcY%3e6w%B-z2jmM3C*suUMj(C`R4~Pe0P}K$ zlX_y-npfyX^Z6W{sm{J+DEphaFlC^_V){xi4$u#%rre(#y7je`+{SLCa)IixF4-SU z&HL3g+~cbk5+d_3kHZ1%wtkAeZ}Bew5Oc*hX@^v?8LEQyGnk_OZ$sgy zFN{h+LZnQmwrSf0IaWM1;b=}55?op=%}~@)jHKKBd*C0pN|CATVbZ&&t9u!eCwPE0 z136y#zgG%UmZ%CaExSZOb=Ndd!BA#$(jQE2kc=~niEW6bZ-MF+L6sMqpm{_XMgyaG z6GOiJP`HFK;bLx5M)C{RVD2!0+H6I~+WMw_s;}br-v^n1M-W>W!dQX9vAYevZbRxb zOLS?^tLynt*Gv0e-zyKj9WJ{CkewEei|Ycbw2B4s;iw6wMAsm7mi?~1Z|+(WdmnD| zynLyqqN5`xb>6gl8J}&zzfYHi1D@0XXVcio91*7i-l?G69+^DcJW%`ro|J$#jzomK z)iB(TW0S~VH`6Q&FxGQBz;4qAc7#>{-U}hCk?+N)kx_eV%XlYhFJ7EZFV2NUlLJg0 zEn9X2V7&KKoEUhoba~278l%$>1?yrg708^U9#i|wWGPQQ&iPghVx!M20$F+w=%6OP z%cr|SQi~oBpoUBqTM)X3yIb*xDdTtHZJA$^1IcqkPw64iv@eF_pOUF9pfJiTLr=uN z`Av}MRx8*jA^AX7`9q4k@9g2qdaAxhVsW5K%hTR;0uUCJ&$MflwKg>ZI5Vltzly?8j5Gqyvtfjx z49s_RH9$Op#+y#yfVygYXC!F5F_V~RJ5;3e-1i^ZiAbGKpj3o1qOhlz&vzBbTL8rI zkCz!F3Jc_@jlaJ9?hlFaYGc!PVsRZAXdDqzd%B4cJAL$^w!tp^&p%QdNWS;uH;K8) zx95y>$!a+xizGRGkXVtpyFGgxXqqslxuBNgnRMnKH_-w#dCs z-kB+&2qY2yB&wjlcwsG!K481IpaN*+9yq*!>&{(jAp$Die^FLc`XQJt9|$=xFsC;6S}b1r6Fp>H^oG31*(_$8w0IuQ^_h*T*TC`CcI zKPf!NL8QG|mfn>|1)$IAu#%spf)HH^q_IZyXR65h+6#q277OHEpc0TT_)S0&%@Jq{ zLjdBxF{9W^EVG5k2oX|@KaEOODf?SnB)>d&UVB(+*SdP1u`%T|WIE%b6{+9nbN<&} ziKtJwDu+LOhdAmVC@TDhBPMS!5U)=2WF^ttC@(t+Ynjq?!zk9!={(*U%YED$VNi;w z^T#pH5)Px`aV2xfS0@4Q;14I6*2U%X`6Srmwv(@FiN4QCh zeVUbRW!}85E6G$u>Zq{*ZEf3JsK*yW797Ye0h(&VuP_9EE!eL$P3SBd*gN(@B6OGV zXu}=KCRLh2>57YN@sM$b64+J!iDXhC56K{V)_QtAQxs< zKJCyTQND@f=WtUzTsIrpS%QFTh11Xs^)f~8PVGN@`z)(bEq8@n1wAoISZ&;ki7=%u zl`rd8G=qAB8MmETuH_IfHe5<7p=uU>VkE=F)#%MgdA-mxKv16{II%YqN_Z%U@_D-2E}j*XB4c1_jR``LEQ%xSzICE6v;^v2Kr z=JJ=1v-7w9dAg;P$RkTI@(k9}(*qrT_ifi{nWN3BvwFtpb3s`bd3>#oq#CYVQk!{KQ58^d15Qd0lA?xg`9 z%}?Rs?XJEyO|vJ&_dcD_LJ?n`gxr%Pq#m1%ySdOVj-`Uf37a2WHU#c|@ z>8(6r9~PGW8Zpgt+ey){dn?zJw_KB8X2urOhfkuM7vdvVQN;LQ?YjMR4<^mLf9oi5 zCY4$m>tkv8UbTDTIi@cTxBc22V=Rp_pf2^D(#k~t&VHc{AKg=^mloV%O&8vjiwGkf z+dOz_*6>5;_U)rD$O=9J5uv_7Q9kh*n4zc5T!x|o{m2Y`|6S2P_PZNteBrDUJFVis z4@=6Xd92t8HJQu6Cl!hKvw3GNNgu*hDt z#Vy}8tzYgUv??)}4_F`C?`Xn9cJk}hU01BcaPKfB8+HDh)kZDTGb%svQhcSJB<@1D z!H)zqnBR%C$#%O~ITRcmnC`#EC0!cRV0O^(Gn2@Q*lkA1Y+HDoEPPpOz&+ixv&?11 z(|ssF$4701iN19@nSb6>R?ZG_u`k=Vewn)U@JYw*kQ;I@4(PUwExS1m_&C!z`CN`sWnCnQcQ)qE`n7b{haAGw3c}i*Q%H~^Bt;X-A zk&NMMnPTJdomqW(D{EcYyC>4x{=r;|xh=mlZ_+<6jGJFAy-CmC8O^$im5u2kd61>F zY9QP#H?h%bNXD}!#gQ0H-!$|*IYH_UtWrAKJ0=b_k@}Ub;^Ef$8%rlXlizX! zx}PijjFk?W*!bW(n>MBNf=0KiS-a5U;1G6~`a*@wiIJ>atdVQ!9mf_gZ@V?`HgUd` zls*CTSd$tBl5Pg!Rs4XGS-1Yr=w}CRg^w{Pqu-@_EcC^^V=={8$|o(#la!%i^{?W! z1}CCt)U@$15`*hfDEhh#Bxb8m3WQdqo?GE@=DtnMkpF&Kj3vzUK3WZKj@9MpSaP8M z@<6%zW7mxgd;av=@>ApuQ%ungqRzJndzqHnke@duGDvJu&n13QVkVvCn`WP<{GP_- z+zURGRmnz1@hLJJi7DPexc|OzQFYyZx)--O=M`1$-6Mjnn)NnN;Tqeajz=YWcWBdv za&zGigT+9 zLb*RD2LutG^NR}BxVwjfeIi&_%6T2K-#$_t8iR9OjvDrVGPhTo;&Nd~%#S(A(Nq*m zdCiA3)xABY8gq(yE%#uaX-BOC_u>Y^$CTpSa}6l&q*M&p#atPb>YHW+i&UoSvL5FZ zF;C6eo8He!?=O*$=w%T8QL=ytk@imd-hzX6$VdKJr@Q#R7K7WubN%yQC;L`* zZ-p*TIi`hemqeJPTU6Nxe^Rvj~FX^!9s?ioj4Q8&V6?W$Ddw?)slG320%p*E+cl?R+tcW$fknDQ2J(KVnvK9j za}!TKg9Dm+^!m8@e3!WiXq>T0Tkrlg1vA5~J*hb6W-K^-WqYR0`Bli~CcAlG855U#y13S+HI-X{`?1julaY`WX@6s7-t-F zs5tlZ@m~4qiQY{{daTs2>mDKq^#_aMyElEDT%&Aadt){lJD-j(5#lsGjsHH7qtDU) zk~GeXUZujF{%n7@1ivt_fncb-y@LseSW|n>`a1E}q`{I4PGZ4%>EOYYKjsR}xnX9e zlB%M#nS)eqY8&eqJ&j`{S%p~h8WSlA`OIhdCvDjKUh6%suU`osmX%yCA0t?DEPY`R zB9)4eZm!egBe~D5EpZs5D}di0>R%+`z~k>dJm}qA_t7}4VN#pqo`ymxX6;8`Z)(dF zk&WDGGx|t;Qj%z$@&#Y5VfxI6_Uz1r>~@xQFXj{Hg;16lDUPu;m1^(`T*q`+<6K0U z0E6{Jnn|h(zxQ?SCyW=2e^EYM8xHLiSXp~E6DA(sT1MEz=+s#aoA$izjdB(O=P<_h zI`2y3+#hPz-4m>Op6vk`@oC9Bzr1AHIh})`2}t^M;i+tvdD!ABv@ePqRNuvgF$-Z` z_1=77c3nO=-uH4>sQ>X1ET$Qb-X#r=X&l(WTB_0YsT|X!^eSx_YU~VYBI%P=i*rBS z9*XSGVPaX7vh;ddxkA#J*Y7?+i|J^00L*YI{Y3&Tu>z-dyhn7L61aZ|y*<9eEK z4!RzjSysce=Ks%M!1Wp1>%Ap}o>3lfh>|qdpd1yDZN4))A){I7`sGbiZo8~F+-$l@ zSag!p@C^6ToXFzPbVYrPx%(oKwFiQ*U+!bgZhi8pGx>AYCi(gqFVO}%S~xxquk{W8 zj%`Ic?x!?Sy|;(Z%wD}+`biw^ROj(umW>oe7cc7|rFEK(#`ew_$J@*;C&Y7}@}@f` zJ=zj44X2BE`Rr1##fu7iSl2R{eltUf&67P;)#Q6kp;s6i?d%L;!9nmB4VtAd3|Tf) z$~nR|#9uqsbfBa<<@!ai4?Ds144dR$0@E`O7NSz&{ob`0nmWN>uQ{G{MS6evx=e+u z4{^v)hf7`js@vCV87m0t96aY(mQQBIaa1qm67tbC zScjYRnN?J;)a=)zQQWD>QuyXk!v;y9;sG6xm*vS_DbxP&b>45gjv^@~DxbmGM@z(w zbawCJZ*ZyD;$^TFYEQ3?tR$@JPwZ>sG1y}HouVmOjFd2qRmMrGW+QhPTGoh%z)&n` zf&Z0p<0hQ#-UA6{w`Y@t1fp}pS0>8W?ZVB=2_H+AV>-1%77ao3QP5ob>dVr-n(7L# zvQ);xx0St-^peZkA^bDRkGRGq9u)I=@n(zZKDrovKrIhOmNXET*!rfmBK1YevO12E zvklA9cIz$i8?Z>VG{mA2t9c0A`I!N45v@%GbEToshER#SQ+rL-#gmf1E}?Hvg|B+( zi3}kQ0~ZEt!65o+1$n6ruKS_5bLK2?{dhluh`SiE3XcFb@wdJFzZz1IoAvI-7ux5M zyci^CuYlq{l*H6A^2HPe&rdl)L#)Uasu&d&%HT3#YH}H?p_<%rq=DHn? z^?SOxh?|0;F8$N8{-`bV-lMIDx4zBlIfS~XrCyU}oDVlp|AP^T`}kHzVZj_P9n})Sg4dm&i^Jge)rKt#(4?;j{+7_#i9%n8SNI-wdhiuR8} z@*E{Hlghk7bDfYeADUe+{Mc;_xDw%Wa{q~6!|qFS%<50)`86WHqtO124t>>3j9_EC z8cTyq$&|xXpFMjWJX@P(sT$`L*H9%DCDD>2VS~u3em%Xaq!8AyY{JU|b|&|Pcd=ZP zipjBVD7FwuH$3HyFBlE1gFr|7UXKAo!fq%t?@-z6S9k2pPo*e1Bg{zSjF*`5r6%-! zt%Iifk=J%L_+ku$F=g84*H+IxQ`f`Q55@Xgm(E(HxDU>i;;I5)%-(ZJelgak`KTK9 z*K@}U83ktkQ+SNg(uU&4e)h1Z9khkVFeSb|AFONGx~e!4{F%3sI>8KBwuaij@GHk? zSbD}`U54G}FZ|u(4C;;6OUFI}H>DPPe6RimaY94#8g*vp!Y$fABcSM4Kcjt<(KK4r_tMXApc??j zp*^X3iZ=}5dDUFVc$XU7>O|bgXHS5q5xq(y1{sfv`nLT<9>d&EQEU#ra1|39_`*vh zHu9AjDvjW#zOMF>EWV9{nn~P zP?Q6J97gi+*^^;dIx&X%Q)X8$^_=J8J|=O}>y5j6nuv%&!)eFg;n73;GZ%a@MEESOMC0rr z--ATt@SgHJW}I=1b@g*gqGYTwE!`8hr$F+<0{Pz(3mYaL-$_HUPrNxJ3+P2 zeXpUI*JcpVBSBh-$whXSGg(O`2Dg;UFPnpFno0@=Y0D?IJ{dSA^+{w@;^g6JN%VnK zdd4;!XDwk%YSWMqY>)(cub;794)#oN1qe3d^Ad`3tBjE=R+N`;l`FXk&4|98<2ojh zyAr&G`{tk*OwWn``}}xRShsBC*xe`RO_+i=dCmJunXY^N5FTCp8EE4zC=Nz{A6%y` z5wD?+FOMc)?03e?k|Jd(aLeM~kuLj2trkx&mKhf7? znqh(o9tiOCx9d&g(N=H646Z+HPtJ<67U}V*OsGOEkM_%oQ+uBjdLtsZzauk6sbiNH z9XWshxT@rc7SHoF3BtaiCZ}pr!?-jQJ=h+s_?as&PpyIgLkaD<7Ef)WuF+!|HD7zs zbRLT^7td1&@o$4en%&cjSqkTVW8POnhq?-M1-Kw-!{fo|C}HmVmG2Dg4{R|>T^dds zG87NQGD%01l%5Z4IyMIjC5*mFNyAxNMvlTFkDyh4)2G1K-Wb1QhoKGP8_c zz((aUULRb0BWz0%l^GM(cSzP&5L=#`x1W0;495WW<2J!zf0y8Kv942RB>*XV&ydG4 z=-Un@Y^k(`1>{Wj3%>i{KZALj0N|Ag2A?`=3%9H}1|BeXD&@D$o8t1$&eT-`w}pS@ zjS%q1Sg%J;nSppN*qG{caZ~ZzZxC+gLGa%;2e1G*?IOoqH@#+wg@(;v>lQO6hqZ&) zP|G9?Vna5`MKm~qlnDY{)B6n9G=j(ExzrUl7{R)I^gUi0=@J~FQG25bRCX6rKuZvj zhduvpAfX3O1EB2^mDdPr$Dpo0#5p$loQdAxyG>6)skxDibN#22mrmH!#H*~z7s@R% zL$}a8-pEfDk@<=W%gD^fa_^ib7t}FRhxXepuU(dmMK)=#x-D-70u#>r0EM|jNt^C_ zT&|&%6&*MsRI4G2wv7tvkh6Fi`QU~{7?0ZZg?f^t@icPA8J;UD?cjX_r%NwSeoXrj>q~RE!VAlDeppgdK2c z{&gJ^GWTmz3#?8)bf%u~EgW2+rB(r@^k_aBgBM16Tz90szr87NLN0o#fo=vj${u}k zm2i&%p5! zym@~6Kpq{TL!ItB3zmJL$n#9lyzA}KTR%C*2emQ-vY%pQGM|gum&DV~iF(+5R+}BO z`;s3fMqf&K@0257U1au$QgSd%$qzHBblfYiPEo|ny-~bcElEt&Ax_<(7xXDcAk902 zVQ)zaMk{g{{0rxno9T}mFwp&oqAd-(N!~=~|Mdj>uIbymkNtP)lzJ>z9O2Z1WrUCK z6e8&9`3@VT_-IR!I&OuY*4g5A1)|v+PT%ddzdn&$2anee@{JLA#)rfd=Bf{3cfkwL zmk&K(yH%>c__b`k4nD7YiEg?#2!Lp}Cb?dn*P`ZstKf$UEe`9FFZ>DMGduBZtE8yk zyBBUFd(t_-O7mn?@%P&gO%r5u=>A2K&`X-IMF$~D`6WlaJj?5IJx_;TC!XldNA8tn ztakGx+3wFQvGCsYlrIz4``^tQ#T{ckn;K!h)n3x`vG97w_{5;;FY+A##7PbUAegK| z{|_ZI@FR}^>;yM#d5wQi5tF z-s4sL`4k;;5t)UZ<)!p0erkJgCs;Pcd;z1d4MxG2z7C8+aEYt+u@0G~}V5SjBx4G8aQ!T&<} za#wL|>g|oJt;VU^abnh!D{Fej!!&fxyt3GC!A7di#`ZNekhA( zA+7eEDM!4bxkS;nce*`u+QlWYr1MJc>NjtnyB^Z&;&43g8pUhGq+>qmmc|F0w5f7-I>xh7%s-9H z1Y!$Om${EpK6brCcIG4Cv<@``)%6XPr^c=yL}E&={Q^s=506Ju;KqW}MuTl@UTDk1 zX7|AIGwsYp6dK8vb=!0O&OE*GYF+;qVEzLO8}zTA<>rs^3_A3K++8oRXV^W_o} z<2?>o#{PyTrRpba4qg?d!ku3VOK_V`7MV2E7+5%^{ZIF+JY~_}BtE+n3tduVYZd<^ zNuDxFUgb?TlfIH2rc4JHr^A;%hna3)+%ZnFM#HjxmF}^_RTE8Dc4(NQ!7)ix%JmRSHAzQBXq6tH0?n%b0$fZPjF|VIU{n4!I8C(bWEMovQ zW=M&Pzb5Lj$X@+Kq}(=m?R5jDgpPa-EAG1S;(kKa7q>dGxX8QO&lB$|)e#-8O9yPa z*mBZPt5ARsU@1vh$H)=Gg8DEuAI&x(S6wpOWh(yJN7DXPyc`5pZVR&(6Q&p!E=0&_ ze8vydXPrJAlK>;2RNM45!6qezoxK>|IKNe67V6YQE*;1etN_M^uTcC}Gw)^Am~I>x#4r5$O86506FFw(?GfC*-kk}(;aWw-?keI4ej*JaZY zB=x`}8NMhUh9ui>NpoGE5}R5qr%!TV{ctynovBk7*0 zdy6FI` zVnR^IakqMA$EGtjpg+)ZaKTS74xLLHxLDZ$h*4|-%69Cj11Af}8$7u5Z&~)ALRDb`!l&p-)y$DgVt{IhxOPQJ1 z3T0oRVJqWW4Vx${!WH5Y8QGcP+PP$B<=UHEd;QM!`TqX%@Hp>tU+0|Hd5!1m`F_0v zv{M!-;sSUx2-(|2f4}jOw4ltxirQwx0;~8Ag18Xdnm@*!M3E+G$hkh5aJtN(8M1bs=>Zni&!7P~hVy*`jn@D_@p~eLs>=+8{GmsfCk? zQ*>C7V-nKFvHlx&ja=nMBLd6lUj(h3*go=|@815EZAo;5X2|ATy0ivgIb*Sd=Ct>h zPag9izpeTnZ~}CTxZEh35{C6Hy|023mMTOH7IrTWRzvYWgDDGa=$_ZnPHe(sXK)Cc z&2u0QNUo0+*Hl99DO+?G=w(; z72?H@>Rg6r*Dw@4HsI*Pc!g+8)|KcSm<}hmc(*6(Zs;7eq(Pst6h-U9#aG%^f@H3I z7}nSHzDm-fEZ+mp21s3v+^8-FT7nl}@%ro{JMtG$_nDV>J-1ok?NpTm>yJFf{%aviS@nsDg|DEEGJY9W+91VGZ`F&%-13lS;c_PUN z!BYvI+no)!&%4|aEs5Hr_fuG_pCH+1FJSvHYi zqd#@wUR{}A~a}^M^Qew z$U?RmBd6^mVq~u#ikk2jc0n_DICPxaELQLB*g$q;G}CBKdgh^>r@vHi5jLgPq@yme z{p8GH?occv&ODIO3lRX(RdX|m#KlfiTYjh0L9{6V}3Pf}vkQ68^u#GRh5H6m} zpRlFg#HtYrikuT;t~QO>^?k=o3x2L9Y1<=0;?`&Q+mJW3Ty`j`hd|-l7O5WBFz6mW zR8R`e!|DYzG|qi{+2@dHF{N_et#u#we%1GRB*!N?g3ohN$kGFW5B*F{XkL`?RekQ` zSYNBoKaH{c>tQ>E*w(No1^y63(eaQt)*nJ%WZv2MpV?)qNAsq6zUL~3QRF$TeHq8| z$8o3bFy*c9hD%Z8SzMeSpR;P}CPLa|g>4L=1wbu+<{ELUoP3#`ELsi^%N($yGp&ES zEKqY0#1i6Gy^r%;tqR#-9#G-oj+YPwN$WC|X$_;UEDCaCV>cRuVU^9@{r755yZ?5< zxn}M%#lt4o?MrbX=?~B6n+wlHKLoB|f~f?}3zF^@f0J8G{t|-F%Vp`6YccUF{87JV zpKSS>*%Es4MxCnvrQOn5TBz4rUSG7$ECy^H zC&TzBhUYE8s`X|*qWW7pF_!&8R2ji~n;=c6Sa~jlnG!$xw&gJ!*lyh9E|<#CQx(U{ zSwe0F;xShtGDWV-si=O{vA}jr9ma)~`y1OXVH-ey+bv&j z5c-)f!?j7?;{gX2g9yx`bNX5|ne9Wjbx^nIoKNf!ST$3)mj-$41w7JMf&0FlEAack z@PgrAFkY>^@FO-kQ(UcD3?0)KxA*A9XJb1$mQT(bgB5=l4AbSxm4aoQ{vD}{mz_2v zZ>}*LE0l@4Ju9tpgu1xn%h$}}{YDprHxEI0^Wuqa2pV{%uTn|~8Sth{I)=UJdB;H1 zw|0(n3~1=aWL^5Uz8&m`p)zFA+8qf2Yv9cV-?CVOvpYW@Nl}Z%jLPQS-`|&+YnH#uYwB^IB%ty>+Q!?Vegx>oqC^;FP4^p{41Ae(UeVy2|0hg^|a5=zD&;FcT zP#f@H4cSej;-+^&-YMyvTwrM|U+Z|tQ7Yt!Yv8#}|3(I=z?)uG+gD^`sx5Pq#A^D_ zEN4~O-$#jM5!R{g>9mowVG~9KwJI{~niV6hvcQKarctjRX2T{%cUR z7Wu~HdX7(i;&eW#fp1dv?)C0mPkckT<$KeIOsfDY>;>qBA8pEohVWA(5<>Q(ijB8} zv#@p8q#%kRPh_4!HxZ%9m+pz=3YRATWF88XlX|Pa% zIA@8F^FAjZllYyBAO12q1FO*X4WyMy0{4lO=_cKCDuLUcd}aDRzIKw#hqcL5_p07t z`TL8HO&)Wr1YgzjAww1@k=Twl+_hrPJ{3WkJyI;8H_h0(nkt*Ziv9xFi1ftV`~ybt z3s!po&O4Rn`3axG3?XLs0tbbpo4}+gZNi7Y9#qy3UC6cZ8J*DqJr*J7Q(yH8QE8Zi z3PZp|$3gRE*g)Ue4haKNPHg!7wjf1eYc1A5y*`33_1iS7dIjQMmJu~b0&x3D0qX?M z`cCC#w?p%1k(x+v%DRjupSyvY6OE4k~rO6Y`)quNYUUFhAi@JJwE$KPu8o%3? z#f8@sgDyAWFi}K)PD!LP8_ylY&rBo6cs)Vp=c$3nPAc=LeKCv{EwjX8m+-X1uEj%l zo@^ondudi~-iLWxw?*RD-_5?yt!6s@4@v}l`||(3Rh&GR-bYa``ze5x= z`T!_4DgZBX+TI8FUYP1akv*p<^y@a7Z$`QT zuG5@^8Ept42Fw@lE-6L}POC2X$dBu1OJSH?WEOW2GnXR=P3jH6gt}?zJxdu95FyUY zmOlt;FDRP%_iUC+e=z5abjz2CAg*{fDa##5n`^!&6eh@a3WX7L(Rpx#s{A# zxjb6@Z=}*iWq?C*qi~r31DZcY!af2F$oob|Vl7=s#ifr>`MFm_!1a1F88W8ILv8pQ zE@(jX@-@h^#AOnVx9|CMTYRzkxNz&3aCWpU|997#O|GgJaMo+it*aojF26p!$A=Do zvb9fd%WdcbU@L!yp0%UO-&}8h*_hs?xaOm+WO2Afm-7nn7;+7g0eHsb?F%4FTn^Fr zK(o3Cdv{KuX?8D&e?^)HaYg?&)3r!Igi8(O<*A{DrLJ!2hK5i(aYgGP29EEY7gFBf zNV%-zoP0TO<<2dpcPHcc+aJCF0*P&Upjv$Kt6!>(c5hApY|d#M(JSAGi%{pXVUrVm z=c{2t#yN>HR*Cl_|1l^Zg1L`mTAdJ7QWqKs{RFZHBS;$R5%S1#nOaCS^eG4o{P~;2 z_-49b!kS-PCV#upsn^k-jC{JLk{NYh30(P~M-1&K?!>9w#a|u5&7pc zrUh=ibmaBM{tFFWrjGVI==VPnGHK&ZV*H0j7?#CqXydLen3xPTqo;LulQ7Zz-s&zMqQ&+3re8GMxa1I<&y*vYPM;?Cu6b#DXABfQI9&$)aM1O`^$GZ1W*BTR} zW6Kes>$c}#S@4q5kG_tkvD!XOgXM5n7rgGa_o|9Iox$Bo*lLqaR?l9K{+;STfLD=a zu33G@2{Q`HpFBIUX@T)6Eq4H$o<>Bziu%U99rk9WZCF4JpL9@}IR8yOpvV4fZ?P8e zM8duT(+2X{BvedF*w?$F1tH0+0e$u1G6Le#^_8`2Kj692lU&%_L zNfWpk`mB@p95Y+iOPiyZ{;1jzH)Hg^;157%zv{$J-Q_S4d3Dc}hPM%glDn8aI%X zFm)@xlJ|VjM>?VK8owVKmiyKR;B#7k$II8HOYRLTK?5zsYxom3DZ&o-?W_`jhyU#j zjF+0^n0WS+x$z5M@oig!97)GdD|UU6p;-H~u`Icju?D*K`#LSdCp;c**GPLcB(@sw zl|~(3cHYT-rZ_1{S`uLva+uGEABIW#(v#?^VoGRa(ZZ=?9U>+>8m=}3cd*MvOxcD^ z>&(f8j^`<{0bE#yn~hwKNjD*G{YAlvr6iMv?aW7Qd~q0_Lg{F$IcDDl?G&j^{}YdI z@=P*IBe|!LhP(i*H^Pt)xLBz>N0(cN5ANw}Te^&@(TYLl=85aIk&M5SMI3+iFolQ# zNd@ls61~H;#hBgNzJgv9tds1;Q3EUA_PQAnSHimltlM0;rb)Zs5*tp~psg3;eM{G6$lHI-FoiOb!pg>iJohJfdN)Okw z0JvID;h(I}IMrYVBJyOnL6ZDppO^b726o4nD>hW1y)1o7{{@IEPAhvx$fQl1}_m z6Xqr#VQvUwwDBvCpyGGPuA<<%>GMSKDo~h#06eK@Au-tFk(!&rmwN)v%`Oi&+O1QE z^pkZbA*#@#b^?FLYk8;cnL9t^;LfQe3vN;Nm@9lV88!lO71i&_*JWcUn}Lz25kT=4>!JF3^Y}S68m54U+vEFcjN;?e7oz`^ZC3ds zZAYr0fnICJhv`;rqW(x2f8Qn&%A3byD6IMYo4~dY(fbVoyap%a zS#Rs-uW`<<4CEkV*!0zjyCSR5PDl|On_qQIkwN24yP zt|MZiRMyT=HjVSsqvRl1G0N(DDTg9e3qR$SdIPJ58ptN+CNGy$AT}FwbF7^7F-4ue zbTtp84||Ct_OIhuger-o&^u#-2~2NQBB9`hB;ci`^yOpJFI~5_M{Auv%>pcdTax zwsrID`#yj1;W$JphA2=Td^3;insea3jO=hHI9ZHNosuvBq*u3uNkOk6*K=v|H1kfy zdsQr(-H^o`_rJH~oP=9%ymdvuLsG2!dVh`rtHt!3gwP`&fNILt`3&<=R*@*io(Hdf zUcmZJ#5_ElNx@y$m2b)Bx0>gT*eQgmkc_9Il}LIoQv}$Wh)hx~Rx&W@f@4vR^Bnm? z+sYtE#___XEsdt%r)lso6qg-kj=a7Y7t(`(LsF1?^-O*f1BN2IGaItW8ecsXju^cJ zH4htV?caX)stCN}fu`2s`t<2*ojTm}Cj2Iv?^e5A;p@8<2W%`dc>z278{O9yKZfXA z%^HzzI}n=)JV#r5L5#H1B%BUw>;9uky51y z^~ZgiuVSZYYor%GB7Cl1Z^@=O0_t_HqNb{bI-@qY&9JyNN2wgkx&`nc-HnKSYc zXY@_g9?3@Z)4|li{&gxU=K#8Am>w4|w5X|!Ah7jq3t=8F3!{@(_UQU+EP zoQlE&RpGVW2CKZ!Y7M6v(||E;Y)JYZrUl zD=XlXGcjfEj^gsxqIW;fOkK{SZD)2P%BE|Cr-PK#T6hcbLVM=i@FyS6h#Nl)KSZ;I zOMTSYh~1yL6ho8*HC0|L_wFX@mQ$K3TKX?bg%stn+SB=e@B)o%QLYZXkE`DujO){h z$LWfX9@Kuajx;|riTkfwP=WZF6E$X4ev*1j9?iJH58i&8*-A&KWHhHh?{vB>dTB2h zxm9=MJ1=@C*xx!e-GBRpNVeW((Wy_w*91o!toTPXrJCLtzP9B}fzO#!GZ{@5l8$4< z9zV?(o#hzA{Ov+4w` zyHI~M$rjH=?HR1~aaLe!^qe%+wGW^c$V0-DGUa|C?&W|yEH5mh?>6xoy}@6*hS4cz z_vDeJjK0mMTinfI?62daWQu9qo zVNx-Ab=3svDg_|B1O%rT=cxh(vYdJbUvWi7TTZ#VxT4!pj&+6&fDI;j#Yg^XErE1D zFjZe(-mzn(&SeC@pj*HL!p-o~wT^14&+oW+K9TDT0t=P_JdY1wLywB(3{Fq<5==Qe zDv@Ly8RL^d+}@$=iS9VN)Jzk%_7>HZ7UHIjd$oQb8v!EsjN_*yORyIN0@u5JG?Z_5 zbNkwwFTn~2N}8|qG+#>MdqxV~>yddjP6p;15oY)L+#E%|*;pA6st~}~uwK}(gG0(K z4m*1FIXe2!y`&Q}em4~=Wpp@mY3t^TI%5M+k>g*u3H3iL`s4IaE?KUKIF%uIVD@8Bl zWI!_4tzUt$5g-|x2jEKJXAWvVFn&HC@o#f-!4_OYV#NN;W&F_pCXY~m;T@(6YLk~> zz-~RGHLrmrwtdy{yK!fPWDi((G^GFE9_S3k&Tm1}K(zv(>N*pyg4cZoOn>x#OQ6`A z${eLzXVMs+GdLfcrI79a=d)(h6H%qhSYQoEmjkkx7(D9h;^LWEc}{JX2{M-l{5;ol z0rW4eK%JA4dx;@E<{`_N)6s@t^uqKGT6UElLuPR{H=-`sO$B#mbu631vEQA_`on`4 z-k>O^Q~f&15xKDt7?Al#43D)FcPpIyIrZ;H*~vv3mDyb)@>~UAFe|zTA^}9f?N~cG zx&jY5!q>6U^5O(u%F43A1|(8?@2VE3%gcUZTZIEqyK#Z#o<4-tQ=C>K(1HY@!dI(7 zXsj=ra@?ps8%KRW)xvI`b5Vl-=1b7L^4VFuBLbq@KRh{_e zbnrwPe|9mIQ-{6+(bW{DTSQ6%vX#p1iXzoGz;()&x{EM-#OV#FIvu>`^bzl%CNy$go}c%cQ1pNB-N!7lZ2$N(TfVU! z5~?lyemeg@_=|~GRp<o}SrC zm>qZ>sRwxmiQ+m*6>5rWlM|;o^4>>wQm{@R$?95g_qo|I#(#JAeikb=9-9|-txul$ zPRaW$i`)^5)RvdK(ZGw~y;Jb* zD?R}eZ^X9pt}a*~4%@Hsgfv2`0MwFKtp}j6{!=K7B=ek=Oux_g&8mlgucah2jI{Zn z^f%;Y7(7n6RTLQ*$!_L>kWQ|BFIJ~@m_(u?p;4i`e`!~znioEOM@i1wJlVdV))iWD zc~(SBc_hRl7tqP-LgMlOD8&VkrbO_kdX&uwwbmo45uGNr9$J>}3^4*n8m8vBA|dM6 z-ve2wz^)X%M!^2DHP#UT$Rb(?*>2^7sw=$1kI!2WkF%)8-fb5umfxnx#>!-Z46=xv z+WbjHpqivD1mL2ti;O_v2Ye6qDmm^7S5N%K%{B@bU|j!2?Arz85Fv9_pzNdsg~!cr z3&)#uoc;OvSlB92WI&WVGdNWUY4C{Y2YzVfm-wQW+JYU*^X?y2P_@YmL-hquo1gsD zB$@6Abl2^IQV+H{r*Ta{QgJiHU0FTAFd~OxDn}g`Y&Lax4AKGl02D2A39MUvqocdr zu4h#=$jahkNH6vM-+s&o7NvGq{bc#}A+xa(f=<~ik!qJ|u-w7CtSUbwaJU%)kI9Z{rA9XwP%@h0keKU>r|2fE2n zc`|kN)}%ehH1@ry5JIaSPP(;bjed6F zSw`{^OwDm_i?5lQLRro3jT)&Z?|-!HXu*2>_PP8Jgakf>a64wWL!P#)y}gK5*ZXn{ z6YOS-(_b9k%?JowEch7MHk)=`sAmljj}@#80|G__Vi+QX_ySuxS1yY>Pz6Ogfmq^S zEILmID*l5^oL<9L`OKPiQv5N$4L5Z$_i3WYkgJ+&pCGuxyzy8ZKvaztJBlc!-#VhD zem7#iRX#QXY0fU-8Kf5qk+`@HLN{BDN32Vl&&wKLx(u0P-4erR&qYdfHHJMbx;~6z zw&k4#8_Ng8QUIwX-(&+SidZrJt zRSv0m_ucEh`rrQd4b_dAA14x~_diYbIW_L2L(G#^AvC;}prQiQN5bcget{AcJm8sk z)>X|GUoMMjHdQxJqzrHl@bj^!5QQ@+V!hEnIgpsL zJiJxTdCfCxf&ZW0NvR*LxTPW8GybCV9fRZ6xF{`cv$`ZPddLUJp?GpSIW`!3?5{=; zONn4X>OVl-X6ndB`x zB6H=2`d=K1lMNIQ5y4t@u~U~uRCqKhV@agoLH!6)o|)~bf-s>rDoho-_^oP27k%Qo{>?AJe)dFWiLnhiCZ^(V*wT z<8b3=e8*V^R~W{wsM z0r@klPEM{{Pv|_34_+$=v2lJH%DF# zZ=HGDZnyURes+_d_E`Wp#ruT4D@7`n-MZfzP4*@&DTb}Fx7V&w4DJ$c_OrJf=$P^y zUgDFd)|C$)=E1=`lJeuf;a7w!=#C$mWcwPa_}FaSa z>sxJ~|7`xN5xlcmBwW<8ceU6#jr?-Tq4`d2h;;~Bj>eC#S<9VVcOpp5HeI1Fm{Y_d z?=fFkQB%V*!xsLwc3_ovr_voabF=TPihVKW>)g#+WtKlxtBe%-%7b~`kyjRRng${J z^9CI0$2Ik%2jGKC@)xvFV)+LVR~O|ub@BC94*6FWN3IWRMr5fDAzqW#8dhhYxAGDb zKSFPL_BVy_4X+%>XHFOO@GU=oHQmpA3K=xV1+^(fTu_X4rM?eXu#XMxg-w#`)(NXo zp_ubu9^4P#QRcWI#Hoi=O$PzBEot;<&<{zUH@aN#i!Xj;v&w~k$9bV#vT~;v_pVS| zm_FFMH@*~JjP2{aB=^jC1LW?A!8ihegEqCHFE!oy7fQG^-yV6WUvWO|Ir)K()XC!buX$Z zZIkyixZkGgBnkq6x89S0oIKC_N3Pds*3(TNXMNy`+xCzD*Q>V7`sH})>Q*6Qs>4(* zcj)@CEYzL%jvRRoK)05|g8S;uQ0j;^Jh090pC~YlE6NpAadE25Y0aEme6br5nrN(^ z3x0ve*~9Bd*NZJ$*0j4fisUNn~&A1AJZxYd4~#Om<}rsK=b$E5q>OgIMl zn#93q3VUbsd~;LfhSD7bvi?Hl*WdVos5EfM59qaSb~UM^NXGzWS(@5H1b#hTgknGP z#^v>=_UaiIcW>G||Dfq$&EpitCWU^`?B&j1QKvF22x7z_f0Eb{RRqf|gPsjZ{ARd$ zl_`DIyz=b}xhI~+h@l>E#O|5rbDijAj(sWdSW8lpZQ)Xg&c9hO-!Pbh?kb(SKe74l z=i*u0gGAQyKQLn;xpzu=u$i~qM)lA{31WSB=jt-JV1uu8%Zbm)ntAWH3YtwEzO5KU zq}uiO;(4^2)ZKyjH^C90is;cKD-pqM-_pG5@I6LxzK&(3gsF@#O zv>>JYF4VJ0^@XIN#`bv#W#tyTh#O-crQ&NU@+B_h@;{uKnDzW>_PQV3Ur3-;Xg*Sw zl)?#f8X}Kq%^()`P8DU}x6ds(GPPU`5AJV7hukPaffrTxqu%tZ4)UkRhVV9mZHZ-e zL)uKzap6L;hbJZ@el8{sfLmMRPHuR5YyTS?5JGLpMH~LMzxVshM&gj&hjdh`r%Sot zMI&V9cP4ShXOEh;B#Q!ajTA#QQ{T;O-@_JM3=ZX1&+DQyLQ-{I77(pTg)cXinszq5 zyk=pMf+)ni^YKg}w~@oh8|z3i{ZD+*JjTkU9=-i1*3Qiayq;ImG4BUyao*y1yz50$ z8dI?KF*>8hp%}1M5X=g38p)6OX{ijzPAvW7u)k<@vQdFHJOQu?X}^^s6v2xR);3~p zZTV8BM|bGcKK!0}Jxwh=cl<~8Lf)L=*VNFSe1lh84j@KIFo%#P{nexgPdGPFAgf1r zxv*2LHX``Oo8Gi-`X0G;1hFkn+Ob5a9=*8c@6VJNVA5AQrC+_@%E(Zs{tZ5IZs$B~ zrEXk|P17~$Qt@vsouj*}=h*g_)gR|N2bfFsif%0S>!e$SlMMj;#1_~52gzvF>_XPe z$^Lao*?oq>l?(81;-puLh3fZ_RK^qOrd#)4_)r^z8<+@Vf9~PKbVS7oI(x3-;~wr^ z{X?Nd#90;(ozkR?>4$F5CS^?N*OMH7o~PxJ>y*R|a{3oPTlQbBrM{b7gr9MKuFlg} z60&4gMjPY90Lurph``(^MnO~G5q&~klj#d5YHgEY_(uUzg2xv=Gk;ErcDfJ$Zq0DZRkG@ zZKf8V4oO)3TRmKAgMIZwKf7p*9HEq_pAxMi%X0!K4gd+qFsEArNJ&f59>(*ijp@TB z7fXedvS8~9G+Vvr*(tiJH@RRinacC4`Rj1g8U_-AlkfNIYo3UTl=3tk+g3=;wI;65 zHMN55_?muM%DYMzbkaPrEt52(8cf-H6)>Jf zDXT?ph`;qdIzSx4os4-yl3NW%k;`ATLtx~I(d7cJ^KXNmBaUdn5?-xeF6ez_@kowh~|Vi`xBiy?YEWmB=f znv8k(lGsPq{Mm8~S2C#Bmtwk{?^hGm zQB%zS^lxaH2n)?7Yd$6|vMOonb+;&rcuxG1C28lA7;7-^P6@S=1Aa(L9J8Mbqga^r zLB&3vHQ$eB;$ZKx#u;~r5U(Y#sjqi}Ohu?UzL`|ijJuJKnJpCVq&ZWmarYOP8)FPi zhOaxh(L_+2ZTfY=o|WNqH%K7ibl=O#uo{mw<0Jkjyt#X6yNqz z@BNh*s$Vv>{8OcZ^04wrzycH+yayJ*p0C*;Wc_q3_<~afw&wa^m$HC@520r@dx35$d zOa4std#8Gbf?-13@rb$B3EJkJf<=W(`Lz)YVs-B+!(de}%CeU#1yWfJ>~#$<9aJNp zs8PVCD@Wn=3th_jkO+=@_LCGa#VxIH&?<>pAzMIG!`W{oG%<`2s(=$|H~3<7wg&Yv zG`8L&EBAqz4gNPGi%_;xe(hnG&y*>2^-6Wmh_!FjN40E$-$Fn4OTpXFxHU#WM_vE? zkg1L#2l=acKF6}da~HA^oW0#%@iU|Sb6n4(LMI3@zmy=nH=$X)Qgz!+nu;4rX_L&Q zPj|Vi{-v0IR+GL}QE<0*)yQT3)6ot^m3%e+XGuv2c!m0EKv+ZR`K3Wj0@w}#gVFS5 z-`OYL4mbBQy~V#d?;ow#S?}(P9DYVQi->UR>n2d9r5B3)JSgS26%6C2hLsDLq)psD zZ~7#!3?t$v9Xa2HdnscvK35p(k%V*mxwfxDnazXQ_0T_Yq&xQQ;yZqW&QZ~_z%+!7>+40%%m@ikb2@{ zt{EgDWMTbljmC8yF2e2oT1Za$9iZBWw<|Frq3M_Pg?xL0B;YR^T-TQNq6)ii7M-x% z_(st9W2`Furc(bNzMI|fSU7EXyV(9LM!@MxltL79SIzw|M7rwio1ADf#1~A#vmt5_ zm=^L13fAUx3-WE1$u~^+U(RF}{K~Ha6M3e;v3Vh=Lw7Ky3hqpEFr#k3ADJ_+2S1aS z9)>#L!lqy84yP?8^abaQZ>zRG5FTc1$?@;WwE2~3x$Kt{JN{2bCFcaKj_~@FZPgY# zB(l;E*C(}~38yLhDSC$@x!vQN;iI;|r3EV^oVW==e5+%xW->!y=-;BR z=rAKLCu*Q8{rgVkV+{`M@zIxSgcL6+x3P=eeSu-tWlb@QA1PgVUG&901l2p0tu)dr z^|R69EE;7^X^V;BG?;*6)}6tolo}e?ps%N|;)?E;?6uk+)PY!S7_UjX@m9|$+>q2t zLxUq1Q@|bh_+bU+A3;I?gEm!j> zT0>4bg10k8*(UwLKmL_L{z%UVReVw84G)OqAMo9yy!CY#a^~#4C#%pXxfd<*@MG9N&Ab}di1*(4;vaDzFFKaXiq9rj0hVMr9O!GzlhyOla*}-xb)z-hC{fOL zZNy5ecR5~}`Z9L3U!XtSTsjjxc|g%S!{X-$Bf>b)hb$$Hb zSb5mGTpkdO9sTYI9Ey)EJrBi#7o2?yJ(xfEGHdKHiW`1eCfz>pG@{4AN9V0pjQj3S zA;X%g&`WQBpG(?Y+g!M8dJ8BXjCb7bgGlIXS2O66Hq6N0K}R>y$M3R`q$Dv&d?NP8 z!lcU2Yii#v*sYIi6#{}Vj{|J5tbq1?47qTWA zP9Cpv|0q>{yygZsL3QpmjZ#UFIuU7rG%n4eEPDTj{RjTMI?9Nr!Eh4Y_;?QtUQ3=E)JN??RBQQ_uRQ-$;k;r;w; zoNw_li1l>b<!zXA$eOk}d%h?;x zq(FG6myGgf&_`4=^jdd;5!(X)PS)Cz`i~WhiL=iT+mZ*SYE31M=d$K{FL--cfSx2b z89$_sW9B8_fc{)?7Y6hvTIaomQq+7U`R3w|^XyA5{hp7`$~sGumief*SOIF^68A6# zAbBDqGVNr1+y^Rshi8E}tM0s$FsW)@z}4?bhpDvzFr-_J|2>&bZZ zwkmDMB9!IjiD?dS7~;$b7V+vZameS-Yc7r7(g~(e-G)DKR^R0C#n7HiDb}N=8kNHS z7tXdV2p9i!6O+xxg29O=^clf_8OER7d(ZN`B#co|#qC-(CcOvrFk67CW~`xW#CK&W z-Q~9k7ff`=aeYEI$n!>WwRP0G{^-ATQng}cw~DgImV8&8?$GtJNB=txhk0F;$D8+p zuzQG_C<^%EQ1STe-%#D6!RH4vp@YRm%df)8PUhTdF5~wGKHz=-Se+CupK~hkfPWBC zFaFrR5g*3rRFJ(8_Ec6p{&(TDp1p8V=#?nuAMCE}?c?iyMU9+O_+n+(8{!ShCQ9Eq zN2J>5Oq7&;exTDy-K%6@9ao`p7i{ox(@)=M_0&9*X#PBR90vO9T__FF#OeK0is1e2 z$xUjo1E*7{<32l|*fGh1J+^leb-!gThZ9&?slOxfD2H39riV8bXrm)i2@cbc!|Bd% zbl&{)MflU~uftwd)ejz=uHG2MIwf^*{EOS?znL!=DjVLD06jD{A)KH)VX)Nh*gG)rw(9~|eh?BJ>AIX%P$1@r5vZ*uLt z;Z_gH!myf!ocpz-%wF|~Q29?|I-gANo6odCEQ3};WG@#nH3pr|`0eo>4&z#}KR+%N ztZm7vV)#OS=1@geJ3kXt4`EJu1lOwuPEF8WxY%~n z!9i5wj;MJ1OG}#ZvhOngm`aPeW$IGMa-T!}A-Sj-RM&Gi z!9(BEGY!>dASHX)Uh&qH#yehHpxWc8QYt8TZnV2Hf)eQfX&VkpFKwwwB{Zg{4_z0M z^D)pxh_719a`~&Ez+A9s!?a^A3;xlrZ<<>qx%iLn0vRNdbU7*WYgGZvL+x; zS>ss(*Ii9Msc-%lLz=~@SP(pesvMl?5ZlUl9-H6Be?AkEv=bd^xWo+QqFILI`arW?RJ%q`KQ z-lbGOjytHLof$Q~!db$Fc^tt3IhHO+R{xP9|GHl%uv;F>b*;f40l{c_LyI13o^LT* zrc;gu{t>hplHiM^I+^`SDI~5TF73A3y|p=Ot+`uZ*U)QpG3 zhe>UZtO;rBG`BeLx5NKcKPW%ib4&jBQT`(ODEp`4^d3!Jw(MOgLzz zBMnmxOtp<ARSE)`m+xdNzWI zrIMr5-XzLJ=}xNssxs4>&MFIe{MzvSc15G1@h&KV3ijwA#P^}Wf&RWu(Cp~4FFfIr zwpS{>Z@65f!8dISHn+3lETW!=1ItLA3Ud5+q$W)8NSYIKN2Wcu#XHV$4|V^lwausi zbN^HX3ug9-RUmVM`-WUDXb@JR!RQe#NovoHmAWpQg68W=m|$Rm7k&Af;Y}?W6l6Cy zJ&Ov!mW3?86WxO<`$py5A??l{6dyDh)Z}4IM`~Jwwmv+n=Z>~%4WSB?K3_fK5=qFoExj4-xsS*?8UJ zZ6;{wr>{tgRlJ}-<5|jc`HSD2^#>4}t_orVp zj|D;Zf-@a7SpFyL>dGMTA~4&~3RY~^^GuJ%h0`z7iy2N8*@sR3_JH9qwwq)?@1oM3 z*TswZo&SG)Y0hqn~zT!+RKH){ZBciHzHH5iA zLo$=|TZl*Ufk9CY4&+60Z+>}s`;SD8?IYwoehY$P)6O5tF9&y`V6Ua6WDDFzolTSg z;!#HVrRx4BT$W_RO*lA@H-FT?r;<~nVDKP^h3&ROt6mE8IR_KahO1&ylFl0}yNoic z>bidzJt@j#X789KBn^5U>reyIcc^lfcgXIwbo%THhoz*jIr^i0T`(rxRrMXJ@$$Ez z>;BXv^Ek{*5#m4Lz0Bs=|PhKl%c)*{Pq!z?I97cYO)FoI+p<&kF>rb75 zGQ(Wpk$wWnja5i#9v%z^CL+2tOC>4$kJE-&`8c^29z?W}k#6KemeyDU-nl;FF?@!?-KS^1Bbq&5!8ac1Sp;5t@^->p|Ymo2qXgsZC`q&YV)wekET z+p~^F@1&GFkPYKZHOwj^Q0HR3vt16mv0(Upk!*j`SEisf7VJG8GOXjQR zc4M^P3t=nlGFt9#_*jh5h}`98{s7vl1`NHRSz=t0B^NmD5=IB5FL^84*|%0)nCH2W zAGXOV;78*m@-s2qs&|Y#|AD>k-~}CJ<#R*nrsPVV88b|cruMq=M;S=k(W2YF$1Jwi zbzL#le_MwM%zirJ$!_IfFT-AflAxBy^;aUK-rr__jBU?w-DW$|ajtgXwfH!1ZZ;%) zA#9=nbduKc0DF{Y{84hC^sGJT{RUh5-Tb|Bx3b1oG^raKGjKWg|wCqIfb zCOOQi_;mVOFg+J;$ANp|d;OB9pMc-mzt2+M9Qt7RCM?#!yVQ{sO^_am6ja$6BuNAw zDa~n(KhScs+7929Uw_2n<{~hspwQ#Q{MP9yuq&{M-_5RqLDLq%{p3xX{xQzt4!DhO z{Z_Ry2%2&)Z_gEPMxB#ie?YwSu@wPCnwSf7cxrvQA5yR>85?t*c+EkvvL((8~M8LctB?qZult7RhGUWppez3rP=l(V0Rpe~5Fx|CEuq^{#{^3OyhwXZiy z6Bo=l6@~9!-SJmRu>G6&?%Knu|+lV1O zK+VmSlKQ5#MJtk3h|{^;!s$S%^M(=6p1fQ{LS254btH~q>08z`Z08U22U2TA>Z@G0 z!+#AmnEL64Ty62?BtzpZ4ee6x_9uZtob^ET(-6YR`$)4SngTNKvEkmF zhAP%mxA9vtLQM&`{~uXz9Ts)dJr0A22qIl7AxNhP5>nC#NDD~HBBgXUh~$ElbhAh) zwR9t}O9+UFN`oxj&C#|*b57-wnrLb%C(&1C6JdC{Z9SPEL&~7$^oSjl zO&h7lAP^t=yZ3`6_{e}e2iA@Uhud~{xEb9oaH^MwQ`6D&PA~Bjqda$tX#N8XrDh=i z$M?G*O|7nc(Q1D(Dqeo4$e}LCg?fHZgz^Te;61yEudDNwr-*?z$K!oI zq`r&)<^81ZWr=qbjg%Nf+=YKL8^m+r^U~M#OhiOU^L4lvp40OCBMDkX_5RuwXs+iv zOs@E@;VG4$igZUVb-A*`Q+Y6oR+@rO@FUrs z)Kzdp@BRlk%BmZ>^pRS47;|!5W`7AbX7hn8~Jbxrc^&^)3>|q@hrVt zf5-4>)-5{SCa=_rGl@F_g6yWifKvkSDpbS4bKUuamn~BWO?6?OM+%YIbUaxNG?=-I z!Q4;kr&z2@=E5-M`M5_dgSc_7>d7EzW&05=E|E%B^!)19aS52 zw0TCkr+^lX`ZYqBN;X#v%;4J~z{i=S+yIDeAcl;j#sqFMymlB2{Dvytm~09p+kvX# z6*H6+A!e?aM+DqghRs={p*m?e^aUA!gV`3>)?|S{{7<|;=UQ)s}cMj zsO@dPp)UIrk-OpoJe$QlOys5WD%Y1JEH@|!C9$Y1L$ayU?4z?+WnQW4^#Rj6aQ2#M)z}W`x-0K5YYY+FYCONqL}Nj^N@mE56Nl`)^-& zZEsg~b_!6=C5>Y?xJzz!*8DhMz9A+nB@+0I+uPM-g?tn@Qm-$6J6(1t1{U1Dx~q#* zv>|Yafv_ci=TGf*wZnJF@xqsMh4JGDKR)@x*F%v;e-PCR9-4IzJVHx2TV0vbDeO;u ztN-<0iAXLN?9b+&u{`E+g|qAjZG5iZ99iowB_fG|I;q%UENOEX)(*de6;Sd36RaJ+ zy9z;(-tSoFs}~)Wy}yZ^doF{l_qDR{%+%WzS>T|u5D+t*lvT4?-c170`OZDb0449% zcLjkFOZ{IA14Ki^vgzovQfY*LVUFInBzAFM!7F4cS%`=SrEwxx!oKX>+7cA3x&Y2x zDRDp2&Si$%8S%Ye!3ZdpF*50(7(YEsng(PSc2LKh{_HBSdz5tJ4D+a_#K-1X$79&@ zsgdx?E+#Tc6|CdK>QBtWUCKH->QPPdc(H}TcXjbizPScy-~H*A9e{OkqUtK{B^XtU zWO{TPc#0ZJ@=Sv;FkGSN{#clVP0u($zvno#xFG>x=6_7^L{M zogyoR@AA^pjc*x6UHbM4qluCJB!};i$UDi3uZMmt=U3M!L`B{oiJZ)O^^=RT4rB=I zH4hc|gd5F<*C9;LP6}qt0(herv}r$P{g$_N6?z{Fnrk^}s-9S1UjJxPOtbal%j5B> z=p~u2)n6Z-Itz6TKqEhVPjTY6Y`Q~jecj1zj8FzBmwI+qKhRL%+10WKkUVeFBPnQi zpS&ksqzO{MnR-iUF6Vt>c?4eQ&#DQR>m-WG4MjpV)q{>O?W8Vb0$p-XKH1M?OJHSJU^)DSp~wE$+hFe`VWj4L2FZ#D|V4 zadW_3CF~qMn4NmRoY>vOR%`QC#LnW11*Hc##1Xo?EJ4f+1DPEIVXdF4!pUxhbq=(( zQP9eW2XQHf9L&My9iybA#4f~{x?F&>4+^Zv+AFCS)ihR#JxT)Z`91Ohs?)QYW&TsW zL&xo+5%hsK-FGE*wo&sVue+Nkb;9`MA;>?yOXW|0hIpFx_t~!8*UO{g$#5LVPs&v_oMBJigugSwd|>7fb%)ZbHM2{Oj2cABGmK6w|GEyIcY= zJdjG({)t*mZDPRj>F060minFvmgcd3FDu}M#9riq zYD{Wp`|LJQqWmFYdDo7+&J{tG3Nk^i_iGOv%vr2%f~pJJ0Z~CJ96XsBL&W+{Y@0yi%8 zpFPBPZPwy%q#R<=###jVBL;~124~O@RHRbF!-#@NE>Ec^r-9H3`&J%a(z`gEb)O); zVeDUhU0tp{EXPDy@AXYDixtONoM_d47o~m{y4l+S(79dUllMTdK%RL9N%pGs!$NE1 z&;Y8+!+Bo++re+%c|q%CE=Sn-womaDsE@c3>1s;wf;J%_$}DAenp>u7R&Dx zcubw#{j{&G&6Hq{U2&C7tea%7O0i$NfhIna1>t)S^ywIxC1D(OR2-N4(lNC7@N$SHbdHRa1r`GZd;l zf*1lMXYM0EvF|}vM(M3px|h0eeaa-9d}SC}7fPl^+_)uB_`4+9s5OH}tOU#cch}nV z!e7T4Ft1X`NFFv&YIZP-1)o{L?t6O2FMo=hm?O~GH!mP=ZktX~g?Hrz$7bWwn%EXo z%1Dljg~Wm$8c`9lJ=%WTe#A)`g!_b@alWqm;#{vbN$!+sJyWI!@M0I@!Ms`8 zw(+Mp!t=wU!nS zfu5tR`^6HBAtok!y+t8|iQ2FdF%om8E*e)eETp{C#i8ytF~4~{_F%qP0ADPVOs4h2 z$6}g|$O?~t=Z=ZRV>&XSUZ&>xSWx36i~d00SL_)qj=Qai93W*rRQo!MR8y40_bDXV z1?fU<;OLdL+KI|50GwT!g{(TOM{ro1t;k1Cz=vx`VM%O`HX_%_hc0Ud<;$3&Eij%a}TCIZQ;Q>7(<6Emg(uXT3etL}~%+Zaf6qks{~A z)}@6ZG!7zQxf|=bckbNad9T^!>}r2;mK!0-;3s#}b`SAGp-@gxLFDTj zS*hYKkZ+_@^d8C>fmfJkrFz6?&Uq5`Y#wQgW4Uo?fa&wY+CE(jet^;(gmbF=KL>3w z&)YPh_1`u$j)`6Lr9dzId&O*v>#$uYx$=$dO7i{;O&7E0Ly#%ea8-F;hBoWP!2>({ zaVHB{GS0&8C}?PR*h)7l2kFOTep$e*&YU-V3Ss~J!Z_bp`*Ke9+FLJ4o9~Xfs91i+ z#09+GY@z>R2=lOQ{*p>hFv{T>pCUk+G7GH?W$60s{V#Hm(u~+C-Ym$x@yB#;0 zlJIvL&nkP3796EBu3yU(Fl#CCY+-2avO(k4Pt&+CE%wP{r>!l;b$DI4EvEu8tTm{k z_%!%bWK=-rE<(0_K-_02 zE2V^m=+fL^A1A`_;8{1_dwm2Sv2RtyjE@w`wM`!7B({g#A zh(Fkz*u^ISSSo*g*c^S6T&M%4hYN*r(H~p$e`gk;ULyX$-4n_*yEGxye&<2w@Z+8U z>i+^ONbQCD4+2vbOdL(9Ub7h1S~l77-l4&ycq6md`HfihFn_DJ3t;K0+bNr_6onEQ zy|r+%w!)t6!o-YoSU8kzf`@a={Bm^ZK%!Q$-6e%o>U}zP57ozbep2?Ml$JU6zvw5g zwaZ_tDzhJ15aPClTpq8=tEvJMjOz|Lc2vw*6DVA(ju+s_{ab9n;3LsO8X^7>(djy7&5^ zgyq&a`BXnS=ob@i)hV0$*ZaCXQRB52A;tu|kGeuZ*mTQbQaX8b+b-jE3#B+ZyzQ+y z3+K3guaT?m(cqlI1kYxh)b^9f_O#rhYsabj086wD!>W&be;xOI3?>aYq+HwC{&%J6>qWHttyq>osJFQ7No=>;`y2=zFkQFgXrDIJ5jv#ajHQS82>18Hx{P8gc7m5N2)^To5^3sG& z_V!89%)5YcUqtFD-@3QnPH7~-ex)~4y?h<^xGHzDdsUE5R74Yj9ya+qk8bbp4f!OKXO*G#F=43gA>A=6@LXuRUFY9@RcmN`+q|ole?lv4PHp>+Zimu7~ zvubA(S*ei_pxG~^;8FN!%5R^jxX>60Lazb1b;hdl@*Xh}LXpagFdCKO`g3oU6G{7D z1QzFKuo)=jg0oHps=LgdDzR6NO@oAq)YGxX_$a!z8*cLk(9^aS{$^93<=kgU>j z9|T8aikQ45Y|&8g=vqTy52gYE$e;FXP@=|wI4BuG;?bNZ_$E=uGF&z`76cf0;1Ol*ej2oPXdvWSEs>AfJdfn8?;_#384iWAV1c)+c0F8tpy7q zt0kTizC5!T1f^*4o|R%DLYl~L3#2DE&M$<>7z%qfoeJ4|HkFY_a`~v8U8RxusCSdp zIGI~2&&1hk!fjU*xg)-|L{g+D>3@mYR#7_8(LWe7_| zt(wym##w@xk10692;Ag1wd8u;-u`6BGz{yb8S%R{vzHWGFlPju6=_@QbQ+`^@W!q_ z(MHv8Ry=9@b?!IAm-5DRdiCO-L=xVPkq22H*vF6qd6J1lC=)SR=@nKg1qi39h9N8o z966@-QQ^8o-bEBBhBWeEr_%)Wb)*RQd{_mh^0TbzDq=Q5mvt-uBdv!Ma<2SZGE2^K8+K`SI_1O_(WS4xG>vP|Mk~ zIW+k)L2Msbu(lBongpn`1sgN|2;c3}sH69eT2qUwG0g?_{RlUQIp#lLkw|;~0y(Da zXA+tlF44ReFta;)3nIHHE;Zto$W)O4J8LMz2jvSkmA$ z@X^m-L~+Mc#ra;xg&A|RTpOaFF-1BN+rB@nBCA^x*(6C7nL+iOz~fi{w75HjNpv}x z61q-L&xA993{ZIf=~~bs?_?HMLwD_$4G#-WdTqAhezUcYJhP|ri3@8Gy}i7g5s%08 z^WRH1D}=e>IEi1p?l> z@Qg;_Byj4o!4d1lvAZ|pp=^|8RPHr2)!HORd$s28axoyLT6j~g$-pRBJ>q{Kq(b*T z#7nnEaf$jAWc#9d4{3V*_^nS`8xa!!U^BhE%p>PPf$a<`NI_Agg@fQP#EP`}0<|!l zpPSlYUq2!luR&_i$4|UU{g!~=FkRyEGRZf0NC5Qc}-g;EweF?p_LO_MpvRl zDT$qzpvt~^Dt9Q~S?f8tx>CcBS*(299BN!RA#K1M5c<>F_@mKwae4NlDSBAXiEu$p z?0H?I+@eaP(jo+ZVY%kOr|1`G&RXH0_|{ zgUe35m{q@V4cCpS7-JJ*)~xyOijRR;pED}bdyN|w+@j5n{^3s9`n*C$8EO7P@6wH~ zFWJ*s>_bYPb-r)JtIEhknYf9(bnSU}MuXn4QT@|_0rY8SjVn7cR;hu1MC>Jvc^zNY zl-gcL5ro5(sc2nUK8`QKUJX70gaiX>bmR5h7A=SBysGujo4Q3lr{*p_w!?@;VRj8u;Yzl6hAy>7b(6=d}i+=yLT_fCEzEHr=6;TB9?cv z@je!PXT$K{@lsyWou5o(=1fm9Vii- z_aaIyKiCCGlP{4S4T?U6j+-(_tx`SwrPi~fg#(tcxykG}`O4cVysX!Ll@*vYhtO|i zLRz~oayU6H#pwq87To-wO~IK1u#EjqebyNceV@u@^GC3g(VwQLkDVGe+C$9gmI^C8 zn3P@5UTDf7=iJB}n^}PZ>{y68 z+I;)?cGI12H1!EjZvxQC`DL{dYq1Ov^sGPr~5{ZB@sT~lR}`69xqY`Lz*R)S!>VazdP&LG zhIDM4%=o=QPKbI$VpHtULsQi19svp{z<@d!O!VzT zqV3hASaftWokHA!jWoP)?Yqu!k+UN?}c*3t47M(=C){L)l_y6S_{%t=LGtWR)wdE`JkM*0pj^0wPITW znRlnl+1P|dmLk1J_-n=Tw=cfXxZqI5wG8Gf^C8{AgL3lnGcXw9)JVA~3O+UV1vIa` z(S&dZ%%bh!GtOVkJ3T|g{ABl5O0&opCu@)CnVQk;FXSF+B4ryg-p2_HI4zq8;%n!} z(Wjuo!jA(!{Q_KmCpe^UI){QPRbjrWIx4J#i-K#xE_a`-9JEZw*h}WXJUF)MJB5tA z4=)r3@4G(+;4N7};EIC^FS32wsb4*|OT9I4525CB{4qoXgy%<3ArJR!J1-Cd>tyR- zTUveqVwM-a+5*g@?o@$ZdQW!<)p^ZYnPB92+uwEXUBmJo$2GK8Pzw%>xY+qIwY(&_ zI)&>4RI4xaeT$2FmIBv$Imi07v!8Z99)@~{k*_9EN&9f*PJ?iPDg`g*mAF-7Ykg-9 zqJsT>01QqGr}k z68L|eqBOVE|HTAcL|XKnAE>>0?sNz=zk@H#dns$<6i@5S`5=_eBcSfcJbgsZI>sB%mf{8Xf}U~_T0VCg zyt7vXyuc1dFFBy9J1-cA@GP^?9M`lVQ3v3*nL4SC2j~d=``lMCk}4lc!Z3J;n{(xI z??Dc6_+nr4^#VlfYN~bq+Z}=f`hm>AdM@B|A-cwSc9&$GIEA0nQnh;73WKv;L;2g6Vq^V=P0Hi%Tmv^(WLsSFeO~OSeX}i~1QDY(f4^6}-sVipd@R4j zB3~jWBiOMyfcQqn^E(xJL>tbmk#@>O;3UT#)k*F{MkKTrEoKAmV&R_k(L0G!ql@}# z^^VYFP_tUd{ekQoPoh1(QWaLW5U<5eDS+kU=lum+dK2E+dNfD*3Ri(bh; znrj;M4T=_=R8hT>NydI@_7vYFDth~Selw$&P(2|OS*X6|YZn$VAe3&DIMeOSnYA%R z4}t*?Sx@!dZ#z7f9r_vj7MqqHvh(ELy$h=~|CSTu>g`S=aff0ws3?-TX2AfWbG05m zpvhx<%DVa3-E-WCCqoIB2VZC6$d*%HUsyfR)8RhLL5bd}P<+ivMbLH#>TyCl1|&cY z(ks4l08-* zhBfDp&wsrkyy!NbJC6uX{U;XHT=~Un0u#HM(kIK0X(W zqZe&z@GDUjob?U1HcZ$sn64Ya;}i`kZJY3Hspgk5LxkxcWjjnV;?}!q&pfCs*4t)0 z)wf^G)OG`V0rP}uI7KQKOq0}Qn%f2RNo<=P0^fw7&X1ul!u@;yhGvBIYJ{d>x80Ue zy72^x$ZuY};JkjzUSI%|lj7~{a)xDUI~^(>otCQ#?#MBY!<`&~udF>=yxzS!U|emt zmj&_`l+bVf-(RA?NSRu52DX=u8#WP- z$#vY}HRR@*VjnANTyzdEZiz24iU((YD;GJaZSfX5P^Op384n@X_B6 zlebt~O9|U1IW%u8_A9VEy-=;@0d+O}$~1h~W_kk8#Wllg!7Zx~sW;yz7!LiW&}TmT zkfa*?^-Gv@u3yR9ISX03_cCqce(dNvW(IAg^_YtlV7w z{q?1;&!>EYa$?@$X!1v+=$(9~xq;MSq-ZZmaOOUEX0d)BUU<3sR;y$6~#T@}o45`2CB2t>8T z8xOVog!^o-YL(NYAvydPf+hrT`Wd|g#8I3}xLXdy6YlG7YJ3sxMLC+9bXRGYCOeaTh7fn z!KomTTQ;O9n${_u!8>^!@B>TVD|J=(HmP9G^P5xQb9TgJ+gxX)%1Q08%2ZbFP8)Z} z;Ii;lRJv{@AtmgII&BqUR*@v`IP`W(TULJad#2CXYt(2qOM%h^k0O%aUMT(PcJklv z(z`~g=I>nRGCtyNSB}AdgR3Kjj;N7ZDh7hKT~> z&3A^CfOXO_p|aOB*C`})_%_vD_tdQQ9}ylRTK)!Cver7@|5G?}sqKE_CDEz~#gwMn zwd?Ut=!=nGF)`nyWJY*wbjTT`O~(F)y`IXJ%rCp9Gtq3a^0E8Nmr736P6CrD-gMXc z4ox&Vp58N1kNCOyG1Q6vN%M&c?(1}@YOu1_ozfGd=ZzYio`V{m&5cZopymx0q*s^T zWEYZm3#a)B=(eh?UlJ#0G-%K4OxgEYTV{_La;u_jW8SW71o!^?qQ`dBUm4?c8PbQsrvXwiSi9?Jpb3hMbGaZt(!Y_m&d9V7g z;A$@Lg!~G9-Ris;O>g$k_KJ#ZgE5-J9xAus;76(`-qk!(VBVq%EbLrMC0XM)zY*>R z?!bw1>sI<);3-X~c#m|QJD_QQaDuxoqP9A9i`#WOWfdg<-s_jBAh*_XqkpKs zV7^teuQ}E`zg=%%ZKW)(Z9F=4M8mzn&#CIxBzmbF8g=L4?`NRc3a&qPZ4(sv=$LL9 zYgOl!J^bBWWQ*M?4PBlOELMXhLUX}0RO1a|4kdaHve$NHosAw;9vDDVxYi2^059$}d z1HACK8Q+gTgNoKnCN&*1Bk>383FQ0`&yu*Lf0=hh-Cxat^t$Vp*sdP@q{>QdEprMm zGAmH2*5aXvBETC|;9h$4v|4MtgMYY?;iF-=>r@V(A3;h2Ka*j3el;%RPcQ|^%6$WH z>qYLwsM=(_gM8@sT{hl9t$gi}aDO3*ViA5O*n4Xh;@R+JmmzPPTf|RIaj4{_s^=Lw zvGz&(TJFGhn6RFN>+X^e1xiwvTHT4lyRf=*4HtkU)D~@jSL9|h5HvxE;SB-S+xfl( zJD?5fPLK4tmu%I`h*eL^Ejc?d*Y#n05g@@R(tK=Goh7?1zW4Z{nxL4n&8_>1AFcnA zK!W-q77l~)vK9JcNqI{o;;lwX6C%UPM7cqHBU1NIZnDcU`-=8gYr#5Z`w9h0<+*0E zvUJO;K1f#+l=>o`3DtdNjQR5M;XIlu#kc*4f3*EuvT@+pU@bkHxA~>$y1U9|_2@z) z`4;gtPxB}xcRNb={&yoG;bNfRbCoIKUlUwEKa>&f+kZnjgv>pB} zQBUWdrjy)nCk>0mC*1sfj&=GACkK+ZNVi$c<$yZDtJrvBkF7X?o2P_~fA3i0>?TT9 zEtg2rH-@wKcp-UN5zQT9m014H6oNf`LXi{OnH8eHu+u{m#RSUNI;y_O*G-r@xV8Bf z|I)UBADdS!@r!>kI)6NOT0hUvS>k-Ic%O&#!yODXeJA`Ib#c*_jrT|2L3ghc0plTw z;n|l9%J=0dRSG&$uLqIV6gdf`kT@vI@T~sA+KxVYn9m>Ak9zCXXS2cgo>~poDVm+} zc-S{YORYrgu?=M%IV&E6Vn#(yJ%%B9+-qyO#NA;@$jg_5656O|3uy}=SjSNu`>iIc zJwc95G4Cm+jPcWMF^A2&-WQTvdsNtZKo;=L6QuYePK@XO&fgTJm(Sx;(9^>HxYp|Z zTa#Y3^+shVuIfx%AZ^$ILC;RTQG|-WgM+9&*Rxnr-s4?Ll>7&cK_UwoqUTj(JgNA-=cQNkIuM4JE;G_`)KS?Rgnjqb(_@KB0-0k@QQg&Y7DSmhE`%_sEo@AbETPfK?Lt`-lqiYkA{!z8NA_K-7bhm?!o5{T26kBZ18ofh&=@ zG87DOGB!*2s8FV5_rRI&LFTxfH8^r_SaqE22&mXZcio6A+oV>b&A zh_OdOyMUy)nGlyNavP-nDlNV~ruc(O&4|GG9vUtk!4S_S0v6%}UVRLg z94Jcxb$zt-Y^^IPF_M90OVSLPV(A_e_OpnCwU&jO6mk4a=@Be_*g1|5?XNnQi_)r{ z`5kiW4SF277`C2r%(?O&#^0r9FXIe7qUT_FjCDnk5F(? zdgrcNas5_d_Yv>bi|nT!11T%t%R+FHq0<(=bF8=(6Q7vYk7sl7cpp4@{Po?xUsw!W zFN^7wDWVMuEPA)^`tfJSJFa+-J2OY}#3a)|p<9p;iF^Ma@FmoXeV1BA?cQ#uKNtPO zvOKi)I{Oi~p9YPF&t%>J29NDi?(H{$IBb-&&>Ym`RjnO;hVg=8d_z+4i))?f5r`$(aB1vg6;rLlWePPX@7kCQkNq1l27gvJ)Ko zrzCLF-Up-#alt=&gzPEb;i{1Zv0>wvYaS|5+&as9gOj=OW5n7m;MC)-<0Kt3A?Bl*ZCeO6dRN3E#V0#OL z3MQVfN1+S8A4CFLgMIsdkr^82zY?{6%ok5uWx#Yg78-)JFG)*v(vT zFq`$lvHI`%GQ4Y-?u<8w>(9j%{bY1|ezLr_dYYFuyb`%E2@}?Fjg~Kkz@S5SVxcVE zGOxCmPj3+{>XvXyNxEpQS0SGJ9npq1N_c6(`OK}p<^X6SCL zSzRDC4bOLL$M@<-M(VSdTC;||Y5N{d9?oEWEuWcUr}txL-6KGDTa!b_I%J0o_$f+% zA00tt_zk_-aQ%ORm(&y|-Cm<}^uObfI%4x>*WNw4wAARbC;e3pu#HuR`3*mlUbs zcTnmo@RKN}bC@(-A0N8J^&H_vGiWX@%)QxVL>mStM!q)5JsR!v|0<SrxIE}vWaPVo8Z|tk2($~b3rhOvc zyCnxE8hGjk*@9_IiGzEagMJ`J=3;roZ zl;E_CVjpWxi(qbyc1%EZr$_t?&5nUi)SRrPek%7$Wb2jhn$tA@in$)`smM1AlU(b( zw{XWl0&pgmIQ^FeVX4j={6x_8NO#amHG^B=XxG1_FZ_pN$~es-e}bk!f>lZ60XG$W zC}XIMFKHH=E@s5k`dWgad}A`%%;?a3Jd}9oV6LH;pkkZtJaBDAyRp_*1Dr^$tdDMFyiNvorQN)Ow1^iq^;d2$rc7qQZBZ0hb#A2%-QoS**{va zLuQ}1kjOH&-Yg~OB1V$PlbUWVH1W*UvOvVkxq>KziQqqs77gcn7Pw^LlNOmo zILeGw1#ubrN!+(L|DFKRO@=-Apfj!4#>ge%PTXuYiZ(UWPyA15Dh^GzWVCat-mURT z6RWV9fSm%c%AS!orok%_iD6(I&KE)0JyNUQ_wL@h%ahP!+1I1pS_%VNtoKYY#Ap8a zi0nVG`QdzK-t{e6p?A2vQy1|?gdCc>B$L48E^&wv#v~?x+9a70TU`eV%qmn!O@mN7 z28&*X1$Hv-ThYr`|V zeesx6Pa-@mpVB;wRdbHy;}0lXhjL_xEI&dkoUdaa@eMOVjQB_lcDC}Z2@c?OWE^^D zo#~#a$8yb29p)k=mN~;3=khfHYKSp6V7cQiBzM+Nzg;j}Y>{y+(4# zo4oe!j*EU{makLOPazCWse(?ab|z0T(;@)u%d2Tgsdp9}7mH&GjTV3Zs7yy}6>no+@65Y&tw2G%l|hJ<{hE z{~^Qu{)oxwgFZhP$@+Tb9K=CiCR_0c(Sj|kCMmw(dh_uEAi+qANlfXj{b(V7e(Ohm z@JDt4+!z7b5LZ|-B9ie87%WmATj<;Fb5YS-Gk$yOONx>&p#h`&t>QmU*gG?1SL4Tv zRvVS^)Ac=bLfCP3Yzl<@|m?ww8xpG@}Z z@j9op9jPJ@o80<+O=#blNz0epp~w#DIDJxO>5`ko^Mh}9N6R!;uRxoF$2 zZN|1H3~;l|kU}?PHkSrQ0jO&US-^}?D&b>^Eb1XLJri`lqr85d^Y{_Y0j*WtUa9X% z6AGHJ$XloCdx25t(Kj0V&bJcO;>Ba?Bom~H&Mp#Y&1iY;=-SJToWsAp}irlm!J zhy3={>=~SHdD|ELNg3xK_}MMZyZ457+|dAYplO4Tm=Mn~ho4r?x7T~xWTLWv7)skE z$m1d=(6ApDkHJ>AB^-#2@B|Xm;NB|#E1-s6nYdEYG4QL-hSjF`I4wdt}pgAS-ykA@I4}Jz+WE zDLkf*JFU)^8AW6viBSLWII_N~eDIgL5eEswI~?mv^cY%Q5R-mx0I^W$TO5RPq@RB00^&lL&4-YVkCydRLucb?Rwy zs?>+kS~&I?i4;eLiqw$82$?vl9hSah&)e)c2(k3z>wYGZIire(x|uX%wqJR!cnd6Y z^=gshOs&%+Bx~Xi6c9e?a-R(G`oG!Y?-vk9U#ZI6o@{-fvw!lS)(U>CEDShD z4-3%`;1L@#I5hNcrFxEvD00&hMo`Z3Kn_w*Nb{_y88xIVEo>t7EpCp?h1{bRSIm$$ z{FCd|)^ZS+)Ca7l8={a|X&NRT1a+=Ivknw8j#DtiAnmmMSoNq8fV#r{Fv3@od zN-$zZ>0jh&0P-}JT=*1&Bh?QUgfsT4H$20Jqg_sv`U+fl3R^OOI7O1!0Q=?uW8a*~ z6C-U|*1Wv(&mCnA%djUc&I2I~He-S|qE{2M|94{c@#g1=+O4bF0n$b$Y=&u>bm%K{ z;1>EuJ1Kr?!pwxgpfc@d+7>QxY?ytA+|8%Py($*y#pjPELYY>pxcT^}A)R)js z$&3dIwJ)Ab@2CHgof)-ET0=HAIGrj1|Lh_Ym0*i;I)` z^LH_~YclS!qQY#yO)IMK?g>M!JM9hfF*en|vt!J!q;S5}$fdW^Tcwm9#&bbbMt0C< z^w5I+_WyGWOB?cg*_Xw%XcvRtuZF@|l!E^8>o;z}tQ zJ$At!ulM2KKqlviv>sBj2wN9F%qJ?{c;&eDfc1YL$f*L}-Mr#QRlw(h8P-tITq@ER z@Xm+6o|{!@(vKL49(sic6D!~Y4-6tqtSA-UmcWc`(>BkuZMd-|? zOB9NLA|@dV+X;CXONr9@aM2WFO<>Hkd7vSF$N(z{SQ9((vcElZ8wy^Q{!899frbYP za%+M5@CaQi6%G6FHE!lj4a4_gl%DxN8>z)l0BQXP#bGaitO#c`0MZrw+MC6fH8>EY zMoATm=gh~?p>=j@r8a7lRYIEUc`RX^y9U{FPj8X@zqG}}^L~Yzl+bVz$$eVvRp))B z_R=DRGJqM-Nx}h6c?~$_U<9mqk{g3`PMr_$Zw2>ji|IN4t06;&F@$9Tgw1uo+D{gg zO2-NhB^w?2T{w*F9*lGVHjo#SaAkI4rbYgDS}-Ctt<0*UZsd~Pz?w?ctR0hldBWbC z>&&v|V)991DE{_OSeick>5mhjRDt^y${ukhse}(OY^QXv8Vq?4*No8RB3d98u zGb#gmmh>1g>jsZT zB)(^YGigTa_mS&IjROakj~aMlXyue22#rm9Hea+7-KEY?<-@ z*fLS5P=m`)e}X;4tWNY!<8#%I=$1Trif=xnXyPzvl)m<6XD&Dh0oAw?261FN2dpeV zPM+2TtbPgp66$WYmY+sHI3Ce+a%v67c2hc{BfU7No6S|S?wyMyDR|wn_1Y98LJqPb zJqHrq+4sQ)7Y^sCFvKWEIdZ->{=Fk9JU*+1Hp2SbR;rG%d z>|zHfB@Xc0toGiR3Szs#kD{2WDyxHEzI4_rku0~iw3y3)u*?$?@5XP68R+ExWq6TV zs)YBpg^f6SfQ#aJurW{s52}NYE9c9AR#LEjw!?CY+8uYa2!C^960&7I`=fC;jdyB) z8~S}w7^-`i544zqC&8XYuQ3pQQQHY9l|?39*EgbppD!&sF5e&~HN02TBhqTc(*rIl zUWaAn&{RRz#^Uew9YFbrJC#1@?Hr`8Jx^Bh?nVqZd>jiN9m^#)2ri;9WoV#3&Dn_S zGlEN0_*+-<{PordRT}L6G4e%pqfzPGC;1t@FREJH&H48jIr@|Hz5AikG>ja`WZOWE z-mxFxj&wnJjt_&-SKjzX5nbN+gjVDx)PR+4J~M0O{%x*H2u2*O!4KyWdsEWv3OYn+ z`v$h})>Zf!=CE>9;g?s*20T`L#P-k3xtN)O7uofDpaS(+IALZG{)V|+C8SWLpv7GA zS1LIdI-c%?Mu8n{L)7;P`@6ok#)Y2aEb`_IgjCGFyh3s|t zmDV|9bcFcwiwmi|$uBP%g8If=8;Urf;Ke4Q^9GvUl82_ji==*@U~U-u0@_BgF7kFG zK|RjBozI1I3;WO7&Aq?)4Jr3f+;T7_bfRLSQ4goDdbCPlVf6UY>=g3cz#wMlM}4lt z>cXqG>!Z52S;Y1WO~MOEo_p-kl>2JLndcFHEupdW(%*WUhGRUx{hH|+w{PR4=IcXq z@LVeYZbWTU&H{riH6fe=_D8LrUfZ9IJyaihL`G3aJ5jr+<;M4KV52W)AKNT= zRDY>*;n$jdapuYP=vlrvSp>d(-&5f6X2B9>sI74V3rm(8uP^i*cOJN9?~1Lu>>E{& zjL4N{mM)X1n|<4gNskVOS==W?lBjEb+gg?iN$TW4baiy86dgi0hy4|CLu7_lzXJlx zhgHv{@^u^{l3f}}xGG_Zho_-WL?;ajr?q=xAvVI=9Dd_l;S-OcLR~~QZPZfVp_b!Y zF_||f^_W79@XA*(8H(>7Q{udr3j`2V%{6@FE9O}I!%*P%f=1!)B7?oMf>8>A17NTa}^ zI}Rlx-AG8cbeD8VcX~Ja#rwVY{SWTl{DkA#vu3S5Yi8E0XJ+rc^ZkQmjl6&Z*-iX1 zeaT&m){VwkL(XG-2S!o%s~e>P1PgqQQ1e?JNx-yuo*X_9;O{uyEQXZFv_92;SUj20 zhp`xwX4M#%hG-%=^LrwsBEI!+Z_qjUtq{TRwPb#3DbhU^pf@469j!dvvl)k_ob_+@pySZke4^5~QxK!Ip_+|OzajNPJ9Vcu9 zEhp9Dp$l&FI750WZ$Aq8lq4GT-krbuB-?M(u@|%J{qQ9umDy0<9BGnz?g*_*g;BDv zDN}OLo*f?_bu>Pf@UClp-^$cFqCVhB~K8;awsm zr}F)0YBYOpWcZZV6{iZv?f zAH-FQwFFRbd2^Zj0-sSJI#UyqI~k|VSYkhSJus^RRwP*=|0#6a9Xk7wcqE4b@j>1v`NHCI1^S-YO zTJmn9<<^;q?{UDFKB|p?jv|h)YA>jvtC6BsOu&>Q& zoK~6<=BYrL_A8GP${Np2AmE2`EC0?|FxGf#Ol21_!E5JO2klN5W$^E)r-g>)G@jRS za6E?QXK2AiLwa7)6PrX4dH6poLhoWFh|BzibVkAJuSmCZ<4@@LQQ~zm9Y(1?MJ`!4 zdvdw4chcnXWC`y{g(^5MjqyWyhE=L*LWjBDiDujTXNq!oIJhEG#jVOidtU$uNX z-<3NFZb zUELgjzB5E1td!kxwrNimy*qw5H^kA4B%mGH|=6I!iX|t)MZGj2!aHu5w z74Oc794lXl0-0*MB?HLFOfocCF6)~>V|_7#k8S}7l+5bagrKk;obMs=CKEs~8Rg!z z$R4shIgd&yz@gJxgFoaF<7AF!qm(-~YlwxGN|(cV_2ljJNfd|oWkzARLZqoYlhvEO zyc;eykEEO}$RwyRg`Sz+4eCA-H~?7^^2wtz(n+ zbrw`$$z=Q7whTo);-o9O$}a`ll<%RSx&)Q+f#ixm?_Y~Wjb0tyF$vgJ-^E;dA-$^} zpH&UL(+!q3X9>aGh$P39l&1qai1^)sYQzaRHuO>%JonMPfy9Ak9w8(t$je7>BI4xb zif1lie==4H*cpaX(o_C087d$O?2+Q30;MIdC-BG+$!4LNa0O<2bB)UCWNz+WlJhZq zdC!gRz%CFog%{1HseigiJ+Wn!5t9V1SB+FYiL2CU@eBXtF(|31z}R}&V1*3BYY?U@ zkQO;b(6_{{=ox4dmdv1d?FL4tC!Tq;*3QYmpi+0)C1e79xiJdAj1?DCAOMQQjveIM zVE#)N6n!%|J#e5T8R`599~a3&nI=;fMw7r*a3sJFs|L(XP{MlhBHFqiSW-GAtosWY zX(jlQgE?>F{ZulOS+LJBaPidJjvZCeHZ*|Yh(pRz)4udPK~fv0&qh&RQV|iVwnL_K02SFH0$%(=fFL&gwd}x-3x-%=p;}x-&quw0fGoT zE?EB6*|~*m;nod5fR{9(0@I}@@a)^ZmJb&6VD7BiOn{h!2unQn_`2$V`Q@?XHmU0S zx;dPmkwM7}Ltu?U@d$j@ED^cy@^8xBT+fQSbN9tb(M;DtA0V{eWiTtvfM+nX$#yFAS`Cv*o@;4~C*ScN8E{{oLp^`j+2NLct) zKZ6pYS#K=-ESj&G23{>Eo$u%4IIWL8<{^@!AQ+GujvzRMijppI348J}hDrud(nyu2 zq6`n$64|14OiH&u{2|2E06XN-_kh)ynCy`Lo4wc88s-^ZGKBjbv7CwLkZ(nFvehJ5 zfDafw2b~0FCSYvXUeX$CL=luG08%QM&I0a7Viq{_Qq;NTj_~c2DDHL~#=#l768n5S z7?C3r77ke6%go=q`@Q=M&q?UFdrhQ2RPL znP7<_?sWqRggP7W7ouVc@p1x{9!2W+uN){ssVfCMkj}xUp2+zhi1H>zj zKe>KbKm##}9VR<8?Kljz^s0Q32ObP13@t%I(SJfyjWEB#>Dzd`eViSQL82Dwx)w>k zRhqTat>@L_x`qtM4(ht*G~EiNa*%qTfROXj;fo$H^A2};ZfmMP%(fn-AqlR{MrN-# zi`wb_Z3a8txq`hY2QLf!U}Z@g4AN%{NumpVY*=(fQ%qEcIm8abr~pgA!U1gG;g9V* z2pH0sQoo9wxPHnUeGd0HJA2)ZTIv=kJKdasuphOAeG&ZIh!?^7WMVnkv&tI5*=Efp zU<}#O7HVBf1>oIBW7wpCXZcNE`zr3gqtQes>!#ihn!{=`y&EZOQReX*Ydp3hv%Lg- zA!Wo7nxxpHUuisk^o^>3EJ!fACJ{E!tW+Yb3NxDb#y@6NI|vBq0VW>A(9Vfl!W#P* z>Q6R=Ief_2n0bOn)wkkJGfC>BjmPtG*m-!Q{Y62pYbmU1>3zRNSpbNVox$Dr&z27qP#1^WO` z9tp4>ZsvqhHJN)W=r7e$nrp#xp0C3XO^wn$vyrW*4tgui-`;XIxn`%l*SnUGg}@#! zLTPhqh6%tH``Aolpd=k^+&y~wTX?4tYC&%?|4n(A{4VvIpmkB-i5uEEdP>S`0K94q zDg<(EfU-O$1UK}Z48(Rw)s@F%XA^yaO9CEzT0oXVo49~&1SGMTA4QhUOC=s;nT_LM zJ0e=8gE3wB3^PYh@cbv!NQ^TO11buN$-iizHZlVc-*B}kX>OsS+e~m=6LtARkO?{S zX#a6Na^?IR02fpQd@3??#d_40$sgL~r@dDBYJ0w+F|fNOX&Y~S`vET!6{_-45&$N{ z=m1R`0Re#vE8s#B88)uWy=a5wUH9br%+4CDHV;?;we(eO7M%@?P>|s#+rWpQYMV$B z7Xw=INWXtYYc6tphOx-$g#49E0K}&0%>Gb$-owM`;VX9hYU!m_ zVq#FxDzHX!seLeScor8&pB;?yRZQ%7``dKaMzxASW^md$F(_g<%dO!VB%n?|YNvBv zC|AKHQ=vLF_{WuCc1;{eiQ(qkM}JP-CiuxU#8CVw=_+qj69zr8^^Lw``7cOGru~ol zi;8WeRtnQb)f<(A{;E=`zQ|$u^wPSAx?W>1sW&aC8Zd9idyuY4!?W=yV7q7)J{Dz^ zG|jte!KaDj$gdQAq5pFr(P|&Oj>+Z39;@0th-vmx@QwWCeq1Frte<+fL5m;X8Ny z32~fZF3=5n`F#^gi)_aNjHZwDKcnf{iFNf=ooYlMh{$e~J-De}bFwZ_t)_I7;cvM* z-lXbF)OPNve)LSb(CW2u?{$s&2BC>N=ex3(mAHh* z^wZnl@7G2&t5NIiUU9-l;}*M*Y8A03TAh(rS-gdx72-{EMgjPN*ncb(YQd<*O~-%M z`2NvZ850z+Rb2sC5L?3P-f788?tY;WSL11$WI6Zk_ZH)fXpdU3Z2+Xbb@R%LuTDhd zpLi#y3?{|jH5-;>%T^ps6>lK_E^jdgsIuD!t8*+nOa@;1A1#uy&?dj_YPH7+G8TF< zja%js@F_fek~4+%*wcSgC-s{MP3O8OgTEM_C{RfB0V{4*etYpjOyDE;L(#nHuI&^R z9fR360Fz%rzrPd67UZMqr178nYW+q_Yp3ZE?Fs|XE(YlKngn7YTrlT$+d_fI$k?v9 zDo&gnorod1W$1%`HP^c#z+a*@GZ~%z>B5i|YD>4oa=`mUxX$s0a(X5dN=K*l9FT|V z2P~Y;0hT%>jfm(xASfBvD3B_ZEuFzp@v4Ng@m-zjPsTN|(ON1xEye! z=!17ORKVHAH%Kec&7HAAB4OL1t0B2ATvWua7z=?@a~F@Ueo=@IM1oI_Q)ovgaE|o( zDf|lbLh@pvkSJ^D@a<>>K#bYdi-WrvIl@2cdHERnt2G<1PN? z42I988>2)y2%^COFpCsoAYy)mQ%xAsx{lTtKmHUe&!aA!}O-yJ4 zdzuoLv9{Adw3ONpkNw(7tGP*@Q!e(|p`kbK+rt9#I5UG)THM~RTeQW8fbE#ii+iA? z3KWaxc-~7)W!M96a%`trRlL9m&tavz0q1)< zT$-^%v)5?Iom*`?I`(Owaeo&OL_$GfS(Z%+0qcwIB6K9I1oL|h$`oWDlg{+HUZQuk zr$kEJWTcNMP%Nz<^b-i;B4blAOtC7_oeeHf=mo7W@}(m#QP-8)-oUR~GGn#sJPs>| zP%fM89hig}^bTCf1Tv!0)qp)-i$K-Z=g8*ce6BM^Ds~=O(V$Evl(vjJ`>1LFTp<-K z(>Xsjcp~Y@63ZIfpk1i(=$x?0=E%IB=w*s&EX;Bfu0B`5UW(<1A`KUiT~j?Np<}ck zxdDey1xQN<3G&2gP;dWiMHZgs7_xySysp7&uujh_x#*#anGV356X^2^{SQ+NE~ z0j`{h-4vNlXc1A_l2h`gmyuYQHNj8?(IGj8li~yc5E(-_D99+rN$egQ-fGm8^&Hb` zQ|<O4tpdAlAHi}s0Br>cy; z39pt|;L&xJ16d|=yPio8{0ehOC&4UE0)@cbQwq_*E#3eub~imZ+(-b*Qj}HF8xoON z=<~urZr^Ha0P;S(g`B``a#rya0yoR6#MmVq1bH!`%5AP$6ao$-z#F|0+Pm6yhABnC zn~`Sw$rJ)9kwNf4T(m_`;ST1i!2HxFp-368k*4}RKY*a;Gg43Bv)?Z0+U5Lh??u?ZyQDLV5dY*ch+c79`4)jxRg=JtEp|R- zuLpQ$X0c6ZC8`!d{B*w@I4{|` zD@qK2aUU(0vTT{4>U&Q=aKWp~A|Uf}`Eu6*!8jrq=Q?=e5!YQ0yWtmw_>MXR%H5MT z;~xc>NsePyTL;HT?1{B1-}+9?9jcIVrBRI9(hnqQL#mMgoK_*qs+AK>+!{OoJ}T5W zES={pRhpk%WOfN2Ky{B#RLd;Ac}U0nfUxU)uzK14o_5p^*meSS&Cc9yLPZ`}gy^a0 ziQ;O9UD1fC`N%=_*hL;sb;F#)$~-o(OW}{TNCpXN5*Y~XSbnsG((3N|)}64{3vq5e z)BBcFjTGlq2cdv?0Sp;JuVQoIHxhRY=+9I=^e^go+HjM5ObMCKTJ(+ys;MxHWv&A9 z?SWLDy~K8gF`)^?>VRI0FPa3`1S(vGE~|^obcRkvT*nEn*nC(iKitQ+WX?BJxOuYk zH?E7Ld4~I*TJZoVPbnYfg4N|-sa?e+$*WMT!u`V8@w3@t5B9F@C+wiWMP&E$m3y_U z@BHyNTwxa%ym_aFcVNWO-xu~0%!DV}Dq~@2kcv=*Bw(U!-)|A*&pDBriYUFDhRu8y z_1%a21l1`fW`v(iBNaWqDc@A=(2NBT7RbwJKg5^gP(@~nb5v@W zGUal$Cna2AbMfRA-KhgRO-rwr7RTYYw_RO?y5A|hM7`=s+t2!%av`x&|0Wk7J$c%` zmORKW%HAm&KgchULrLBMO-DjOl&RveG+7>SXFiFNYOr{TW9$xk8X~bonQ#5NY6HMJ zGu{A_lUpa%QLO&>?Q1^GNz~jIcgUNp-rFydOF0>DRI77@#o6N#CzfJii6 zBD{&mG*w8^qsdEcV0yO+7nQyM*}9R_D5!^B(F0aD>>Y!m#{~Q;Qn&fp6ZC$2TXGs; zUYHu!cA7H4lutowuF@J7NoDA0DNz14woZ|1=6ep6*rHn2Yk_aFUa@y9}6$@>h}U8!cuxaLU_xo}#&70wL|d{ar#PV9mA=Jqk&`{Lli%NhjEA|~F?JeqU)Q#!W{gZlOlX0A zHyQOh_>|b48D<$uubW6pzI$tK8BvubtS1EnBfHR6!BpxcNj9|&+I4O%k8q;7A0+j3Z@ zD+6g}bG*w$G?r&!zZ(_fI4AD2^LmV5B%w!2&RoA&@~T(}pSNAuV*&c`b% zS-yfT2T|%4BrCaL-{8zPNa|5ER`N5cpWY>J%9V$Gr@RiVI_=t%V=}#u*K%aks1q01 z1wDYZ1}Lm?+&$kAtM<8l_^|tJL+oUKFyh5)x@CtdXN8uG0Q^--xfW99kND}8R`f{5 zGfr13<5Tp@7;?Gz<(h8K@i8!w`|1F}4zXXm1MET>t}pf^dxI>BQ5!KRvnE>BcHroPhz-s zU#K&YmY_alHW>SHu6UYQtlZ)|b+SL>&K!|rNlkfuaZ=;e4dU+k&tNDgf?8gY#4_{* zCh?FTSh6&qTnz;#IxZ9F@#4ff_PAvC4#8GXHSp6dD*?sUD)BzvoGY`1(@`M4X7ssa zzPxE*{^luEP3r8T`fBt9;W%DA0`dDUS^D1Gt*03h_B1pA2{;*>%tYNyV#1dyI~W{x z)h_@~n>2HlZPQrzH0d^!CCArmziP_1C!=omER2N@`Kh=UAZ13vYC9#Lq7x%vUXp>5 zc;t;N$BP=$?la{2+;r_> zNdPJQLoB9XH=F-EkEReGaNF2C;tk8JGSap7?Ta14M3SMWK-lQX{vJ^ z8D;xV8qMFgBIrUl=11Roi9~9;GK%w+?1tTJGtAI0McGIQV%on_%$oq53>C|2B+YJT ze+Qd{Ts4?KQ$pktca9jb(RQV_Mn@Q+#>gQ22t_XxU)}GDu1oFf0D*{sJ5GQqA=TCW z{kM&*ywJ%x59klE>hakLR41&0j)&g-9A_X$HzM{-qNrvqgvr336zge+|4Lx74*z)J zR8$g-*RGBlNX-@qEsA<@r|b=tzC6PXXKj}gv)Olv*5xlfHhvhIJ}gZ{_J7?v0$YB= zsSvhHKuoRGIx;Wx%j;*WAZXi~CXTr8p6TQILxPm^FFEQ4|9pskYbLd%r_n;|m`thBFeL9sK3tsUKcB+NHl<1X$L@CO zyA}NTzy9Mn*-%wPy3l_-wezYUL|6C^r_N*&$v!TSQ}l3YP#5{%e|q>yf4*mVzL#{e z#~b^V&HV=7WpjD$WY2N=VNF>0;i1Q|^%u@b-3V*|Oeq&sheF=_u;TpjqKC)+BY?eW z7a9*;9Ko(VUZg6(O%7lWp#q@=ts`|Ja{*H2TN9 zB~t<IMr(Q1hLI=D^4;-r>#v8c=+=kp zz3E?w|G6PfB5EF`oeeF-D1Y5>nqBglJ4N7X@DC=9B>S$}4=+vwt`6MCjaeSn2j1N= z%_G|zF!x4Ou$4!RFTT*RQNdK=Pa*;D7%7NFd5`)NSvrPQYN|t|owlfp1Dk$!IHAd9 zq7xrgQ~HZLkX{>4CxXlM`IbKuqq1igqN@DtxA_vBn!}IiMupuNrI9lP@s^|TfxdY_ zZ4-B;Z>h$Envws$XAcfWMIH_w7v?Dp67YXY7)9ZL7eP;9U;>C?U~qtMz;>*zR_6BR ze}1#GfbH#m>I^&1bK|yP&ita-KG3bDoyFM2t+k`KjjOdHqo)I-8*$Kx2aAb{DQ^lE|>V$51yRK#;xEM#o_m--&?y?@dbODYoAW{ZJ~FAEG_y6IptT zt%j1`!Pt2Gb}zr?+%Ww_dB0xN*#tP<(rHP zOXG0%^uhUi*y495@{xyrRm-1vVu*4>wZQyGp+jCPO%kFn?LN!9vzr#pXxbus zbyiZs(8N%Qs!4yr2Zzw9bYVly(y_O~22B~75aftLEy@nio1fk#2^$#{g=GtV;Ii#Q zO+G52>+4yG`|v`dry1GM0R80c&NDXbIN8x+XA|vbP}fqurkm?|f4{r_>)X04L%~~? z!!xSPy{n{Fzx$JDVWTawZr!IVU6;c~4-a=;n9FzP1DuhU4k6uqJ`$k8jQOjDh8_@M zYOpYwrK2T#1UCB%6|@%SQbep*q+svhPw3Mf7D(h~*{s|@+D;&C=QQ6F1GT+c@aUf& zC9}_&ro4oyGc(c!HrK8WhXQ)L_1;Jsvo}#HHZ*|32b6=-1{k`Z4t>SVhaGZ?2&T@e zQ2)T3>n#(8J#EU-M`+2pB0!^riES$*WoCz~X6-5i>K*8p!cw#E+`M`>ZaTreD;Qg1~_*MRe!SM{oR623cwa`!rg^$AG zpTDd*Pd{xxPv=MNHZCbkounf)Ix2BT8(2?PJQcJlp}Wj&0KcW_cW0deo^>X+zFn?! z@Rj>uQpvsaqQ@3??808gDkFu47#}p`Ms2p@Ml;b{8YDU24y9czQ?L2HN{1_AABU&S z!COAETmLTW@ZOjtm?5)YV50|S?`QwIzIA~ILfkGM%wDJ)LVSvYd?|k@y>wF29EsT+ zMNUGRtQb=Vp`w#&#*J~G=4PWo!9+H_c!eK{9P!XpXm1TJ1a(eC-Qc;^Zo0=pIq)*a z8F*&C1GI~qw<#~O-Rt#tsrW+Ss=|9u2p$aL%^^;s(gNS%N+sxDzYju6=Z&W!eU0I- zj?Se?46UEsr9dKQZs6|i?~%fYb)X{}wg6L|RogOij9oJ`_>|&J753jPwvELyQ!dgU zpDhrLo4aF*h9eB6qWy?}`|6}BJ~s$&63N5`Llu0YKFTwFp^+>;*~Nr=Z8!LrG# zAHE$D)i~cN@|CF@?a3m>3Ydj{<}%ZbHS4qaI@jonEd_DS^)J>0XC*CB; z*UfPssqM|#mb@3M`3_1NY;4h1(1X-}idfiG}1SLIA&u7`rsx zT#D}ihjsf#RWOXvkcm*^F62;CSF}!b=c#BeHuY}$oXD^^-57r>qIPUt1H3Ff zjoT^uBs1TBcl-8OM)fR>&k^-E>uSY^59_N~SQ+%jzB3Nh-p2Wh+ZarKt<#^YbIV`v zgsN5v|CAOj7WCQWvqbsnK>{Ji&`l`Em&6l3qq*1mI(D{qUb_6V_Nm{cb;XBwkgXHd zPT2Nh@cu8<$Hn#Ui`MgLN5&X<7#LP;7#Q?Fi;;)jnzW1&5zW` z1TvScK%b9%RU3d%EBd5RvaDKSwo@%nbTquFYsqUT?rh=~nk1xH;rjvOd0#Bg_c_@{ zgKQy*RjaGxH1h8;M9qvY7K_+ADs4U7jYoEoM-qAQW!)E;brs`vC1S>zCUmce=K0lT z%IKE;I2k9RAYG|g=1<^RyDxN$hDO{X%@m#W^~2Ctj%aZ>dQd`~H}FYh-|BA5^1ApA zM=m^5cX*A|A;61B({r-#zD7f}{6;b=kc>^Z2ovCkea0MVT{leg`)%f31-BEnz)s_h>c_dwJ8T z6om;}uaqL)8=~p!nGtxzQCY&sI8B8h?1|;-!Qs5u7pF=Gj%0^%si6*aY^dL=c+hxa zThjLlY`JLs6qLsfY745yrGfcDjNHRp#3l6bu-7B}FriEI{rSA6nh~l(|AY;1FbuKK zBy8va+F6o%vY@#6NAo9rIu&~e%6f=IN!6$ZvJtbY2M$S_Y9!;}&KK9nL2mE!aW4xY z4+gkvwT!QvYo5(JEpWDqJ3h#mMezE;@*Yyzbr@26Z+e9J6x?1&k^dl5E+okbs`=Jb zjz68n!JXoPl`L4FQWXHpt{2`m;Ih7|GT(Mv4qWv4? zw+v8CXOd=|Ii_%rULF>ACe70Twnts0&c_Y37+t?p-a z`P?Y_@mSsXbIDniaKZ^Xg@j^X$8TRxOWvc178x;RX}}##n#SD?D>{G0WfUW2kwM(w zG5f;7{~e{K;Ec=|fpO9f+?Iwuzi4+#3#qUk<@=VPL2&KP^|ELxs*KR!IGyVrfk)hy znZ_4U-Hl)72miZFRC-j&xYBjQ)gk8RL$G|v13x7VKM@M|@3h@WCwR(y+P!bq+$FLf?i3?F^2q1Hah=K>tD!#&rH_nA$O@rgpL(fcHe zkmL=$H&ZR93C{KTBcJ>;rKTBDhfRC-x>|6uJ(p81a=#r@<03o^W(}*JEv3t+QH!A8 zPKRRp|GJQ;=N@1-k2vuax98_pQ5?AV@un0`ZkcHaVxmmulNyWWn9{iD|7?E$)`KR0 z{@IsLH7$i6VnksTi!H6AmM?Lw`$veDao>ImEUZbu{gR=~ide-3Kcts5 z!vyFLJbRI;7CdfLSzTALuB#Z5zjwU)`4%(F6Y)Yn=<|Fv9XSTx5OQu)xF$|cDQ+Nq z$1vdvH{|X7FIR7vrhr7JO7foCvQ}-*zm=cYeo>(JJGbAh#YF9# zxWQz%UHVRzhnXh;P=rtvk2&S@hYf66qYWXfZywSn_rmZYKErfv)Of2b=uiBtqEJ zvyP{b&%`!0kspFtYjRDii%;-RojwH1j`o|LvF`0EG5UT8w^nDZLlyj`*>A&})joO0 zlxXF5?;<>|BKm?N1Zf$o5NqXi(w!miQ8>)9zMOjN*A7v$0M<5ozi!TfQp_ejb@E2^ zqn@FM4S`$kWCTy$^g(_1rTjz<3gY*K_qlM#KD+zO%$-r?Lk8qEA9t6_-yJ<4%*a~A zm#A}Er5cZ(8l4@lI9(Lx3q)W$mVhY?_*@Sks*Zup5}$2E-AUa?`r!TMw#8e7<3~M4 z$Ow}@#(N@2hs3FN(>Ix0(St&^UK~#lfTKWU5nrj}IFGY$?3~t^32)ud&R9y#fEVp} zgHn-14BkcRQoMf!MbD6QP4+1hcoO-m^Ns>ec=#zNnch%5{RG*_+&;u8dkL2WD@nX; zuNXCS=@KLeO8F$;S5kajZd8thzc7f~nF+EgDm%^^CmE4ScN$>Qwoc*t#r73+#Ohlu z?f50*mzu$<}i?Yo&a+h!wksGBPc%AVQ<~f^L%=^lx;{mJWWe6E*!mM0!!zqRC zdW}k<^}e;ZB{K~?bo(DPH4!Gu!&gq!@2cyY9Y5a`efFS*>B_-Xu=-Iua=hhmdEWAC z&OfN0r-Ucj$&|ohb^v#MlH0)`^EC9TZZ}5peIn|ZON#dVk52*^Wsb4D5r{!iJ{%W) z3^_Pbg?8X$ePKW6*6%YIC*nc=;qOCaq>)oN#E2SxVn+wm_xawit~&msKJpoyVW$Gh zh6M-w`>Q@0zkU1v)$zY3m>%2aFb{xPhX^+qAqOAR9T-AHv{dt$yk^Q_WW8RS^Wihf z2X+$oWtu0EKK+UAJ~!|BjV!nNV}a7Vx+9JM>a?`ls?Yq)m*C_N9tdel&WZ2FTeuw5 z%g=Xr*Kr0wjISB@NOA2D!!rifs4K#RL?qyzQ==xKXi*Jw_+xc8IaDoA1W4g7(clgOs=hpV3w9ZtR#sqbKFYkMW2mDxCii7O zU@7V7H)0Gbv?s&ws?@@kZgx@&7j7F}nG)#9Ln*N>d zGxEb-!Upw!^!W)aGY~!e8^bd>$gna5H2QH%^oZg8@1ZB|=-_JZ;A*7d`PSUU;8Dd@ zn}$VU?fJ)=5{WuERVnc-ta4p4~z&vVgq2HkKX^kM-zZ4 z0$8GlqqD7vqoeIV_4-aIKIs%_(FpVkppSn!o}>J3{HNo;viQeA_(HR0GkM@(w({^` z2>-$b21Xz4KZ9_wGB&yfOZhtx`d>zD5^( Date: Wed, 25 Jun 2025 16:22:58 +0800 Subject: [PATCH 02/32] make the code consistent with the upstream --- ppsci/arch/chem.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ppsci/arch/chem.py b/ppsci/arch/chem.py index 864dc27e9e..4787010108 100644 --- a/ppsci/arch/chem.py +++ b/ppsci/arch/chem.py @@ -96,4 +96,4 @@ def forward(self, x): # 最终预测 output = self.fc_combined(combined_features) output = self.split_to_dict(output, ("u"), axis=-1) - return output \ No newline at end of file + return output From 6356c91908cbce4c0ecb2c945cdea820117bbba8 Mon Sep 17 00:00:00 2001 From: Dubhe-Chang Date: Wed, 25 Jun 2025 16:30:43 +0800 Subject: [PATCH 03/32] make the code consistent with the upstream --- docs/zh/install_setup.md | 13 +++++++++++-- ppsci/externals/__init__.py | 1 + 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/docs/zh/install_setup.md b/docs/zh/install_setup.md index 293c84329a..dac9112aba 100644 --- a/docs/zh/install_setup.md +++ b/docs/zh/install_setup.md @@ -265,7 +265,16 @@ PaddleScience 提供了多种第三方库供用户在开发时使用,这些库 cd PaddleScience git submodule update --init ppsci/externals/paddle_scatter # install from source(recommended) - python -m pip install -e ppsci/externals/paddle_scatter + python -m pip install ppsci/externals/paddle_scatter + ``` + + === "paddle_sparse" + + ``` sh + cd PaddleScience + git submodule update --init ppsci/externals/paddle_sparse + # install from source(recommended) + python -m pip install ppsci/externals/paddle_sparse ``` === "tensorly" @@ -299,7 +308,7 @@ PaddleScience 提供了多种第三方库供用户在开发时使用,这些库 ``` python >>> from ppsci import externals >>> print(externals.__all__) - ['deepali', 'open3d', 'paddle_harmonics', 'paddle_scatter', 'tensorly', 'warp'] + ['deepali', 'open3d', 'paddle_harmonics', 'paddle_scatter', 'paddle_sparse', 'tensorly', 'warp'] >>> tl = externals.tensorly >>> tl.set_backend("paddle") diff --git a/ppsci/externals/__init__.py b/ppsci/externals/__init__.py index 1942d8d6c0..67e62f29f3 100644 --- a/ppsci/externals/__init__.py +++ b/ppsci/externals/__init__.py @@ -8,6 +8,7 @@ "open3d", "paddle_harmonics", "paddle_scatter", + "paddle_sparse", "tensorly", "warp", ] From e8ab68cd5dd03e209e59e84b45db594211bac3d4 Mon Sep 17 00:00:00 2001 From: Dubhe-Chang Date: Wed, 25 Jun 2025 16:37:25 +0800 Subject: [PATCH 04/32] make the code consistent with the upstream --- ppsci/externals/paddle_sparse/1.txt | 0 1 file changed, 0 insertions(+), 0 deletions(-) create mode 100644 ppsci/externals/paddle_sparse/1.txt diff --git a/ppsci/externals/paddle_sparse/1.txt b/ppsci/externals/paddle_sparse/1.txt new file mode 100644 index 0000000000..e69de29bb2 From 6ee370cd73fe7fb73dd57f05cc673bf27f400c04 Mon Sep 17 00:00:00 2001 From: Dubhe-Chang Date: Wed, 25 Jun 2025 16:38:53 +0800 Subject: [PATCH 05/32] delete unnecessary file --- ppsci/externals/paddle_sparse/1.txt | 0 1 file changed, 0 insertions(+), 0 deletions(-) delete mode 100644 ppsci/externals/paddle_sparse/1.txt diff --git a/ppsci/externals/paddle_sparse/1.txt b/ppsci/externals/paddle_sparse/1.txt deleted file mode 100644 index e69de29bb2..0000000000 From 66c095790023251b88e9e0b3a93e187688374a26 Mon Sep 17 00:00:00 2001 From: Dubhe-Chang Date: Wed, 25 Jun 2025 17:02:24 +0800 Subject: [PATCH 06/32] fix: code format via pre-commit hooks --- docs/zh/examples/chem.md | 2 +- examples/chem/Chem.py | 222 +++++++++--------- examples/chem/config/chem.yaml | 11 +- examples/chem/requirements.txt | 2 +- .../XPINNs/XPINN_2D_PoissonsEqn.py | 0 ppsci/arch/__init__.py | 2 +- ppsci/arch/chem.py | 205 ++++++++-------- 7 files changed, 234 insertions(+), 210 deletions(-) mode change 100755 => 100644 jointContribution/XPINNs/XPINN_2D_PoissonsEqn.py diff --git a/docs/zh/examples/chem.md b/docs/zh/examples/chem.md index 7200871784..c7cf7ecf52 100644 --- a/docs/zh/examples/chem.md +++ b/docs/zh/examples/chem.md @@ -139,4 +139,4 @@ examples/chem/Chem.py ## 5. 参考文献 -[1] Perera D, Tucker J W, Brahmbhatt S, et al. A platform for automated nanomole-scale reaction screening and micromole-scale synthesis in flow[J]. Science, 2018, 359(6374): 429-434. \ No newline at end of file +[1] Perera D, Tucker J W, Brahmbhatt S, et al. A platform for automated nanomole-scale reaction screening and micromole-scale synthesis in flow[J]. Science, 2018, 359(6374): 429-434. diff --git a/examples/chem/Chem.py b/examples/chem/Chem.py index 34d4d126bd..51d93caab1 100644 --- a/examples/chem/Chem.py +++ b/examples/chem/Chem.py @@ -1,144 +1,154 @@ +import os + +import hydra import matplotlib.pyplot as plt import numpy as np -import os import paddle -import ppsci +import pandas as pd import rdkit.Chem as Chem +from omegaconf import DictConfig from rdkit.Chem import rdFingerprintGenerator from sklearn.metrics import r2_score from sklearn.model_selection import train_test_split -import hydra -from omegaconf import DictConfig -import pandas as pd -os.environ['HYDRA_FULL_ERROR'] = '1' +import ppsci + +os.environ["HYDRA_FULL_ERROR"] = "1" os.environ["KMP_DUPLICATE_LIB_OK"] = "True" plt.rcParams["axes.unicode_minus"] = False -plt.rcParams['font.sans-serif'] = ['DejaVu Sans'] +plt.rcParams["font.sans-serif"] = ["DejaVu Sans"] x_train = None x_test = None y_train = None y_test = None + def load_data(cfg: DictConfig): - data_dir = cfg.data_dir - dataset = pd.read_excel(data_dir,skiprows=1) - x = dataset.iloc[:, 1:6] - y = dataset.iloc[:, 6] - x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.2, random_state=42) - return x_train, x_test, y_train, y_test + data_dir = cfg.data_dir + dataset = pd.read_excel(data_dir, skiprows=1) + x = dataset.iloc[:, 1:6] + y = dataset.iloc[:, 6] + x_train, x_test, y_train, y_test = train_test_split( + x, y, test_size=0.2, random_state=42 + ) + return x_train, x_test, y_train, y_test + def data_processed(x, y): - x = build_dataset(x) - y = paddle.to_tensor(y.to_numpy(dtype=np.float32)) - y = paddle.unsqueeze(y, axis=1) - return x, y + x = build_dataset(x) + y = paddle.to_tensor(y.to_numpy(dtype=np.float32)) + y = paddle.unsqueeze(y, axis=1) + return x, y + def build_dataset(data): - r1 = paddle.to_tensor(np.array(cal_print(data.iloc[:, 0])), dtype=paddle.float32) - r2 = paddle.to_tensor(np.array(cal_print(data.iloc[:, 1])), dtype=paddle.float32) - ligand = paddle.to_tensor(np.array(cal_print(data.iloc[:, 2])), dtype=paddle.float32) - base = paddle.to_tensor(np.array(cal_print(data.iloc[:, 3])), dtype=paddle.float32) - solvent = paddle.to_tensor(np.array(cal_print(data.iloc[:, 4])), dtype=paddle.float32) - return paddle.concat([r1, r2, ligand, base, solvent], axis=1) + r1 = paddle.to_tensor(np.array(cal_print(data.iloc[:, 0])), dtype=paddle.float32) + r2 = paddle.to_tensor(np.array(cal_print(data.iloc[:, 1])), dtype=paddle.float32) + ligand = paddle.to_tensor( + np.array(cal_print(data.iloc[:, 2])), dtype=paddle.float32 + ) + base = paddle.to_tensor(np.array(cal_print(data.iloc[:, 3])), dtype=paddle.float32) + solvent = paddle.to_tensor( + np.array(cal_print(data.iloc[:, 4])), dtype=paddle.float32 + ) + return paddle.concat([r1, r2, ligand, base, solvent], axis=1) + def cal_print(smiles): - vectors = [] - for smi in smiles: - mol = Chem.MolFromSmiles(smi) - generator = rdFingerprintGenerator.GetMorganGenerator(radius=2, fpSize=2048) - fp = generator.GetFingerprint(mol) - _input = np.array(list(map(float, fp.ToBitString()))) - vectors.append(_input) - return vectors + vectors = [] + for smi in smiles: + mol = Chem.MolFromSmiles(smi) + generator = rdFingerprintGenerator.GetMorganGenerator(radius=2, fpSize=2048) + fp = generator.GetFingerprint(mol) + _input = np.array(list(map(float, fp.ToBitString()))) + vectors.append(_input) + return vectors + def train(cfg: DictConfig): - global x_train, y_train - x_train, y_train = data_processed(x_train, y_train) - - # 构建约束 - bc_sup = ppsci.constraint.SupervisedConstraint( - dataloader_cfg={ - "dataset": { - "input": {"v": x_train}, - "label": {"u": y_train}, - # "weight": {"W": param}, - "name": "IterableNamedArrayDataset", - }, - "batch_size": cfg.TRAIN.batch_size, - }, - loss=ppsci.loss.MSELoss("mean"), - name="bc_sup", - ) - constraint = { - "bc_sup": bc_sup, - } - - model = ppsci.arch.ChemMultimodalMLP( - **cfg.MODEL - ) - - optimizer = ppsci.optimizer.optimizer.Adam(cfg.TRAIN.learning_rate - )(model) - - # 构建Solver - solver = ppsci.solver.Solver( - model, - constraint=constraint, - optimizer=optimizer, - epochs=cfg.TRAIN.epochs, - eval_during_train=False, - iters_per_epoch=cfg.TRAIN.iters_per_epoch, - ) - try: - solver.train() - except Exception as ex: - print(ex) - paddle.save(model.state_dict(), cfg.TRAIN.save_model_path) + global x_train, y_train + x_train, y_train = data_processed(x_train, y_train) + + # 构建约束 + bc_sup = ppsci.constraint.SupervisedConstraint( + dataloader_cfg={ + "dataset": { + "input": {"v": x_train}, + "label": {"u": y_train}, + # "weight": {"W": param}, + "name": "IterableNamedArrayDataset", + }, + "batch_size": cfg.TRAIN.batch_size, + }, + loss=ppsci.loss.MSELoss("mean"), + name="bc_sup", + ) + constraint = { + "bc_sup": bc_sup, + } + + model = ppsci.arch.ChemMultimodalMLP(**cfg.MODEL) + + optimizer = ppsci.optimizer.optimizer.Adam(cfg.TRAIN.learning_rate)(model) + + # 构建Solver + solver = ppsci.solver.Solver( + model, + constraint=constraint, + optimizer=optimizer, + epochs=cfg.TRAIN.epochs, + eval_during_train=False, + iters_per_epoch=cfg.TRAIN.iters_per_epoch, + ) + try: + solver.train() + except Exception as ex: + print(ex) + paddle.save(model.state_dict(), cfg.TRAIN.save_model_path) # 进行测试 def eval(cfg: DictConfig): - global x_test, y_test - x_test, y_test = data_processed(x_test, y_test) - # 重新划分数据集 - x_test = {"v": x_test} - y_test = {"u": y_test} - model = ppsci.arch.ChemMultimodalMLP(**cfg.MODEL) - model.set_state_dict(paddle.load(cfg.EVAL.load_model_path)) - ypred = model(x_test) - - # 计算损失 - loss = ppsci.metric.MAE() - MAE = loss(ypred, y_test).get("u").numpy() - loss = ppsci.metric.RMSE() - RMSE = loss(ypred, y_test).get("u").numpy() - ypred = ypred.get("u").numpy() - ytest = y_test.get("u").numpy() - R2 = r2_score(ytest, ypred) - print("MAE", MAE) - print("RMSE", RMSE) - print("R2", R2) - - # 可视化 - plt.scatter(ytest, ypred, s=15, color='royalblue', marker='s', linewidth=1) - plt.plot([ytest.min(), ytest.max()], [ytest.min(), ytest.max()], 'r-', lw=1) - plt.legend(title="R²={:.3f}\n\nMAE={:.3f}".format(R2, MAE)) - plt.xlabel('Test Yield(%)') - plt.ylabel('Predicted Yield(%)') - save_path = "chem.png" - plt.savefig(save_path) - print(f"图片已保存至:{save_path}") - plt.show() + global x_test, y_test + x_test, y_test = data_processed(x_test, y_test) + # 重新划分数据集 + x_test = {"v": x_test} + y_test = {"u": y_test} + model = ppsci.arch.ChemMultimodalMLP(**cfg.MODEL) + model.set_state_dict(paddle.load(cfg.EVAL.load_model_path)) + ypred = model(x_test) + + # 计算损失 + loss = ppsci.metric.MAE() + MAE = loss(ypred, y_test).get("u").numpy() + loss = ppsci.metric.RMSE() + RMSE = loss(ypred, y_test).get("u").numpy() + ypred = ypred.get("u").numpy() + ytest = y_test.get("u").numpy() + R2 = r2_score(ytest, ypred) + print("MAE", MAE) + print("RMSE", RMSE) + print("R2", R2) + + # 可视化 + plt.scatter(ytest, ypred, s=15, color="royalblue", marker="s", linewidth=1) + plt.plot([ytest.min(), ytest.max()], [ytest.min(), ytest.max()], "r-", lw=1) + plt.legend(title="R²={:.3f}\n\nMAE={:.3f}".format(R2, MAE)) + plt.xlabel("Test Yield(%)") + plt.ylabel("Predicted Yield(%)") + save_path = "chem.png" + plt.savefig(save_path) + print(f"图片已保存至:{save_path}") + plt.show() @hydra.main(version_base=None, config_path="./config", config_name="chem.yaml") def main(cfg: DictConfig): global x_train, x_test, y_train, y_test - + x_train, x_test, y_train, y_test = load_data(cfg) - + if cfg.mode == "train": train(cfg) elif cfg.mode == "eval": diff --git a/examples/chem/config/chem.yaml b/examples/chem/config/chem.yaml index 979183b2e1..69efdace83 100644 --- a/examples/chem/config/chem.yaml +++ b/examples/chem/config/chem.yaml @@ -31,11 +31,11 @@ data_dir: "./data_set.xlsx" # # model settings MODEL: # input_dim : 2048 # Assuming x_train is your DataFrame - output_dim : 1 - hidden_dim : 512 - hidden_dim2 : 1024 - hidden_dim3 : 2048 - hidden_dim4 : 1024 + output_dim : 1 + hidden_dim : 512 + hidden_dim2 : 1024 + hidden_dim3 : 2048 + hidden_dim4 : 1024 # training settings TRAIN: # @@ -58,4 +58,3 @@ EVAL: test_size: 0.1 load_model_path: './chem_model.pdparams' seed: 20 - diff --git a/examples/chem/requirements.txt b/examples/chem/requirements.txt index 0b29998ce9..2b7371b34d 100644 --- a/examples/chem/requirements.txt +++ b/examples/chem/requirements.txt @@ -1,3 +1,3 @@ +openpyxl rdkit scikit-learn -openpyxl \ No newline at end of file diff --git a/jointContribution/XPINNs/XPINN_2D_PoissonsEqn.py b/jointContribution/XPINNs/XPINN_2D_PoissonsEqn.py old mode 100755 new mode 100644 diff --git a/ppsci/arch/__init__.py b/ppsci/arch/__init__.py index 43ba02242a..cf9bc4c68f 100644 --- a/ppsci/arch/__init__.py +++ b/ppsci/arch/__init__.py @@ -21,7 +21,7 @@ from ppsci.arch.amgnet import AMGNet # isort:skip from ppsci.arch.base import Arch # isort:skip from ppsci.arch.cfdgcn import CFDGCN # isort:skip -from ppsci.arch.chem import ChemMultimodalMLP # isort:skip +from ppsci.arch.chem import ChemMultimodalMLP # isort:skip from ppsci.arch.chip_deeponets import ChipDeepONets # isort:skip from ppsci.arch.crystalgraphconvnet import CrystalGraphConvNet # isort:skip from ppsci.arch.cuboid_transformer import CuboidTransformer # isort:skip diff --git a/ppsci/arch/chem.py b/ppsci/arch/chem.py index 4787010108..b2a659aca2 100644 --- a/ppsci/arch/chem.py +++ b/ppsci/arch/chem.py @@ -1,99 +1,114 @@ -from paddle import nn import paddle +from paddle import nn + from ppsci.arch import base + class ChemMultimodalMLP(base.Arch): - def __init__(self, input_dim, hidden_dim, hidden_dim2, hidden_dim3, hidden_dim4, output_dim): - super(ChemMultimodalMLP, self).__init__() - - self.r1_fc = nn.Sequential( - nn.Linear(input_dim, hidden_dim), - nn.ReLU(), - nn.Linear(hidden_dim, hidden_dim2), - nn.ReLU(), - nn.Linear(hidden_dim2, hidden_dim3), - ) - - self.r2_fc = nn.Sequential( - nn.Linear(input_dim, hidden_dim), - nn.ReLU(), - nn.Linear(hidden_dim, hidden_dim2), - nn.ReLU(), - nn.Linear(hidden_dim2, hidden_dim3), - ) - - self.ligand_fc = nn.Sequential(nn.Linear(input_dim, hidden_dim), - nn.ReLU(), - nn.Linear(hidden_dim, hidden_dim2), - nn.ReLU(), - nn.Linear(hidden_dim2, hidden_dim3), - ) - - self.base_fc = nn.Sequential(nn.Linear(input_dim, hidden_dim), - nn.ReLU(), - nn.Linear(hidden_dim, hidden_dim2), - nn.ReLU(), - nn.Linear(hidden_dim2, hidden_dim3), - ) - - self.solvent_fc = nn.Sequential(nn.Linear(input_dim, hidden_dim), - nn.ReLU(), - nn.Linear(hidden_dim, hidden_dim2), - nn.ReLU(), - nn.Linear(hidden_dim2, hidden_dim3), - nn.ReLU(), - ) - - self.weights = paddle.create_parameter( - shape=[5], - dtype='float32', - default_initializer=paddle.nn.initializer.Assign(paddle.to_tensor([0.2, 0.2, 0.2, 0.2, 0.2])) - ) - - self.fc_combined = nn.Sequential( - nn.Linear(hidden_dim3, hidden_dim4), - nn.ReLU(), - nn.Linear(hidden_dim4, output_dim), - ) - - def weighted_average(self, features, weights): - - # 确保权重与特征的维度一致 - weights = weights.clone().detach() - - # 计算加权和 - weighted_sum = sum(f * w for f, w in zip(features, weights)) - - # 计算权重和 - total_weight = weights.sum() - - # 返回加权平均 - return weighted_sum / total_weight - - def forward(self, x): - x = self.concat_to_tensor(x, ("v"), axis=-1) - # 沿列维度(axis=1)均分 - input_splits = paddle.split(x, num_or_sections=5, axis=1) - - # 解包为 5 个变量 - r1_input, r2_input, ligand_input, base_input, solvent_input = input_splits - - r1_features = self.r1_fc(r1_input) - - r2_features = self.r2_fc(r2_input) - - ligand_features = self.ligand_fc(ligand_input) - - base_features = self.base_fc(base_input) - - solvent_features = self.solvent_fc(solvent_input) - - # 结合特征 - features = [r1_features, r2_features, ligand_features, base_features, solvent_features] - - combined_features = self.weighted_average(features, self.weights) - - # 最终预测 - output = self.fc_combined(combined_features) - output = self.split_to_dict(output, ("u"), axis=-1) - return output + def __init__( + self, input_dim, hidden_dim, hidden_dim2, hidden_dim3, hidden_dim4, output_dim + ): + super(ChemMultimodalMLP, self).__init__() + + self.r1_fc = nn.Sequential( + nn.Linear(input_dim, hidden_dim), + nn.ReLU(), + nn.Linear(hidden_dim, hidden_dim2), + nn.ReLU(), + nn.Linear(hidden_dim2, hidden_dim3), + ) + + self.r2_fc = nn.Sequential( + nn.Linear(input_dim, hidden_dim), + nn.ReLU(), + nn.Linear(hidden_dim, hidden_dim2), + nn.ReLU(), + nn.Linear(hidden_dim2, hidden_dim3), + ) + + self.ligand_fc = nn.Sequential( + nn.Linear(input_dim, hidden_dim), + nn.ReLU(), + nn.Linear(hidden_dim, hidden_dim2), + nn.ReLU(), + nn.Linear(hidden_dim2, hidden_dim3), + ) + + self.base_fc = nn.Sequential( + nn.Linear(input_dim, hidden_dim), + nn.ReLU(), + nn.Linear(hidden_dim, hidden_dim2), + nn.ReLU(), + nn.Linear(hidden_dim2, hidden_dim3), + ) + + self.solvent_fc = nn.Sequential( + nn.Linear(input_dim, hidden_dim), + nn.ReLU(), + nn.Linear(hidden_dim, hidden_dim2), + nn.ReLU(), + nn.Linear(hidden_dim2, hidden_dim3), + nn.ReLU(), + ) + + self.weights = paddle.create_parameter( + shape=[5], + dtype="float32", + default_initializer=paddle.nn.initializer.Assign( + paddle.to_tensor([0.2, 0.2, 0.2, 0.2, 0.2]) + ), + ) + + self.fc_combined = nn.Sequential( + nn.Linear(hidden_dim3, hidden_dim4), + nn.ReLU(), + nn.Linear(hidden_dim4, output_dim), + ) + + def weighted_average(self, features, weights): + + # 确保权重与特征的维度一致 + weights = weights.clone().detach() + + # 计算加权和 + weighted_sum = sum(f * w for f, w in zip(features, weights)) + + # 计算权重和 + total_weight = weights.sum() + + # 返回加权平均 + return weighted_sum / total_weight + + def forward(self, x): + x = self.concat_to_tensor(x, ("v"), axis=-1) + # 沿列维度(axis=1)均分 + input_splits = paddle.split(x, num_or_sections=5, axis=1) + + # 解包为 5 个变量 + r1_input, r2_input, ligand_input, base_input, solvent_input = input_splits + + r1_features = self.r1_fc(r1_input) + + r2_features = self.r2_fc(r2_input) + + ligand_features = self.ligand_fc(ligand_input) + + base_features = self.base_fc(base_input) + + solvent_features = self.solvent_fc(solvent_input) + + # 结合特征 + features = [ + r1_features, + r2_features, + ligand_features, + base_features, + solvent_features, + ] + + combined_features = self.weighted_average(features, self.weights) + + # 最终预测 + output = self.fc_combined(combined_features) + output = self.split_to_dict(output, ("u"), axis=-1) + return output From fa853a2cbdbdfb79dc3c9a3c54dec8cca1be9ee6 Mon Sep 17 00:00:00 2001 From: Dubhe-Chang Date: Fri, 27 Jun 2025 17:36:43 +0800 Subject: [PATCH 07/32] reset gitmodule --- .gitmodules | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.gitmodules b/.gitmodules index cfab8e8d79..4f06ec28be 100644 --- a/.gitmodules +++ b/.gitmodules @@ -24,4 +24,4 @@ url = https://github.com/PFCCLab/paddle_scatter [submodule "ppsci/externals/paddle_sparse"] path = ppsci/externals/paddle_sparse - url = https://github.com/PFCCLab/paddle_sparse.git + url = https://github.com/PFCCLab/paddle_sparse.git \ No newline at end of file From 9a36d813fec65a2707a23a4ccb31b2a837d830b5 Mon Sep 17 00:00:00 2001 From: Dubhe-Chang Date: Fri, 27 Jun 2025 17:51:57 +0800 Subject: [PATCH 08/32] fix: restore accidentally deleted paddle_sparse submodule --- .gitmodules | 2 +- ppsci/externals/paddle_sparse | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) create mode 160000 ppsci/externals/paddle_sparse diff --git a/.gitmodules b/.gitmodules index 4f06ec28be..cfab8e8d79 100644 --- a/.gitmodules +++ b/.gitmodules @@ -24,4 +24,4 @@ url = https://github.com/PFCCLab/paddle_scatter [submodule "ppsci/externals/paddle_sparse"] path = ppsci/externals/paddle_sparse - url = https://github.com/PFCCLab/paddle_sparse.git \ No newline at end of file + url = https://github.com/PFCCLab/paddle_sparse.git diff --git a/ppsci/externals/paddle_sparse b/ppsci/externals/paddle_sparse new file mode 160000 index 0000000000..960fe36b4b --- /dev/null +++ b/ppsci/externals/paddle_sparse @@ -0,0 +1 @@ +Subproject commit 960fe36b4b28dfdd29b356f74db5a062896c9feb From 865689ad94680843698aa21ecd6a16c69311ef1d Mon Sep 17 00:00:00 2001 From: Dubhe-Chang Date: Fri, 27 Jun 2025 18:24:47 +0800 Subject: [PATCH 09/32] fix: update XPINN_2D_PoissonsEqn.py file permissions to 755 --- jointContribution/XPINNs/XPINN_2D_PoissonsEqn.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) mode change 100644 => 100755 jointContribution/XPINNs/XPINN_2D_PoissonsEqn.py diff --git a/jointContribution/XPINNs/XPINN_2D_PoissonsEqn.py b/jointContribution/XPINNs/XPINN_2D_PoissonsEqn.py old mode 100644 new mode 100755 From eaf1054064f6a39543f9410aeaf3113316288bf4 Mon Sep 17 00:00:00 2001 From: Dubhe-Chang Date: Sat, 28 Jun 2025 12:52:40 +0800 Subject: [PATCH 10/32] renamed:examples/chem/Chem.py -> examples/chem/chem.py --- examples/chem/{Chem.py => chem.py} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename examples/chem/{Chem.py => chem.py} (100%) diff --git a/examples/chem/Chem.py b/examples/chem/chem.py similarity index 100% rename from examples/chem/Chem.py rename to examples/chem/chem.py From 9b6757790aa723ef73725594f5df38e3101f36f3 Mon Sep 17 00:00:00 2001 From: Dubhe-Chang Date: Sat, 28 Jun 2025 13:26:06 +0800 Subject: [PATCH 11/32] update examples/chem/chem.py --- examples/chem/chem.py | 26 +++++++++++--------------- 1 file changed, 11 insertions(+), 15 deletions(-) diff --git a/examples/chem/chem.py b/examples/chem/chem.py index 51d93caab1..a897f1f838 100644 --- a/examples/chem/chem.py +++ b/examples/chem/chem.py @@ -70,8 +70,8 @@ def train(cfg: DictConfig): global x_train, y_train x_train, y_train = data_processed(x_train, y_train) - # 构建约束 - bc_sup = ppsci.constraint.SupervisedConstraint( + # build supervised constraint + sup = ppsci.constraint.SupervisedConstraint( dataloader_cfg={ "dataset": { "input": {"v": x_train}, @@ -82,17 +82,17 @@ def train(cfg: DictConfig): "batch_size": cfg.TRAIN.batch_size, }, loss=ppsci.loss.MSELoss("mean"), - name="bc_sup", + name="sup", ) constraint = { - "bc_sup": bc_sup, + "sup": sup, } model = ppsci.arch.ChemMultimodalMLP(**cfg.MODEL) optimizer = ppsci.optimizer.optimizer.Adam(cfg.TRAIN.learning_rate)(model) - # 构建Solver + # Build solver solver = ppsci.solver.Solver( model, constraint=constraint, @@ -100,26 +100,22 @@ def train(cfg: DictConfig): epochs=cfg.TRAIN.epochs, eval_during_train=False, iters_per_epoch=cfg.TRAIN.iters_per_epoch, + cfg=cfg, ) - try: - solver.train() - except Exception as ex: - print(ex) - paddle.save(model.state_dict(), cfg.TRAIN.save_model_path) + solver.train() -# 进行测试 def eval(cfg: DictConfig): global x_test, y_test x_test, y_test = data_processed(x_test, y_test) - # 重新划分数据集 + # Reformat data for evaluation x_test = {"v": x_test} y_test = {"u": y_test} model = ppsci.arch.ChemMultimodalMLP(**cfg.MODEL) model.set_state_dict(paddle.load(cfg.EVAL.load_model_path)) ypred = model(x_test) - # 计算损失 + # Calculate evaluation metrics loss = ppsci.metric.MAE() MAE = loss(ypred, y_test).get("u").numpy() loss = ppsci.metric.RMSE() @@ -131,7 +127,7 @@ def eval(cfg: DictConfig): print("RMSE", RMSE) print("R2", R2) - # 可视化 + # Visualization plt.scatter(ytest, ypred, s=15, color="royalblue", marker="s", linewidth=1) plt.plot([ytest.min(), ytest.max()], [ytest.min(), ytest.max()], "r-", lw=1) plt.legend(title="R²={:.3f}\n\nMAE={:.3f}".format(R2, MAE)) @@ -139,7 +135,7 @@ def eval(cfg: DictConfig): plt.ylabel("Predicted Yield(%)") save_path = "chem.png" plt.savefig(save_path) - print(f"图片已保存至:{save_path}") + print(f"Iamge saved to: {save_path}") plt.show() From d33f8bfa6b5ff334f8f92bb6d45c9c77cbc5d431 Mon Sep 17 00:00:00 2001 From: Dubhe Chang Date: Sat, 28 Jun 2025 13:34:42 +0800 Subject: [PATCH 12/32] update chem.py --- examples/chem/chem.py | 26 +++++++++++--------------- 1 file changed, 11 insertions(+), 15 deletions(-) diff --git a/examples/chem/chem.py b/examples/chem/chem.py index 51d93caab1..a897f1f838 100644 --- a/examples/chem/chem.py +++ b/examples/chem/chem.py @@ -70,8 +70,8 @@ def train(cfg: DictConfig): global x_train, y_train x_train, y_train = data_processed(x_train, y_train) - # 构建约束 - bc_sup = ppsci.constraint.SupervisedConstraint( + # build supervised constraint + sup = ppsci.constraint.SupervisedConstraint( dataloader_cfg={ "dataset": { "input": {"v": x_train}, @@ -82,17 +82,17 @@ def train(cfg: DictConfig): "batch_size": cfg.TRAIN.batch_size, }, loss=ppsci.loss.MSELoss("mean"), - name="bc_sup", + name="sup", ) constraint = { - "bc_sup": bc_sup, + "sup": sup, } model = ppsci.arch.ChemMultimodalMLP(**cfg.MODEL) optimizer = ppsci.optimizer.optimizer.Adam(cfg.TRAIN.learning_rate)(model) - # 构建Solver + # Build solver solver = ppsci.solver.Solver( model, constraint=constraint, @@ -100,26 +100,22 @@ def train(cfg: DictConfig): epochs=cfg.TRAIN.epochs, eval_during_train=False, iters_per_epoch=cfg.TRAIN.iters_per_epoch, + cfg=cfg, ) - try: - solver.train() - except Exception as ex: - print(ex) - paddle.save(model.state_dict(), cfg.TRAIN.save_model_path) + solver.train() -# 进行测试 def eval(cfg: DictConfig): global x_test, y_test x_test, y_test = data_processed(x_test, y_test) - # 重新划分数据集 + # Reformat data for evaluation x_test = {"v": x_test} y_test = {"u": y_test} model = ppsci.arch.ChemMultimodalMLP(**cfg.MODEL) model.set_state_dict(paddle.load(cfg.EVAL.load_model_path)) ypred = model(x_test) - # 计算损失 + # Calculate evaluation metrics loss = ppsci.metric.MAE() MAE = loss(ypred, y_test).get("u").numpy() loss = ppsci.metric.RMSE() @@ -131,7 +127,7 @@ def eval(cfg: DictConfig): print("RMSE", RMSE) print("R2", R2) - # 可视化 + # Visualization plt.scatter(ytest, ypred, s=15, color="royalblue", marker="s", linewidth=1) plt.plot([ytest.min(), ytest.max()], [ytest.min(), ytest.max()], "r-", lw=1) plt.legend(title="R²={:.3f}\n\nMAE={:.3f}".format(R2, MAE)) @@ -139,7 +135,7 @@ def eval(cfg: DictConfig): plt.ylabel("Predicted Yield(%)") save_path = "chem.png" plt.savefig(save_path) - print(f"图片已保存至:{save_path}") + print(f"Iamge saved to: {save_path}") plt.show() From a4f882ab5f202ad288487bd0df7bb35be952090e Mon Sep 17 00:00:00 2001 From: Dubhe-Chang Date: Mon, 30 Jun 2025 14:51:30 +0800 Subject: [PATCH 13/32] update chem.yaml --- examples/chem/config/chem.yaml | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/examples/chem/config/chem.yaml b/examples/chem/config/chem.yaml index 69efdace83..18228f8423 100644 --- a/examples/chem/config/chem.yaml +++ b/examples/chem/config/chem.yaml @@ -1,9 +1,12 @@ -defaults: # - - ppsci_default # - - TRAIN: train_default # - - EVAL: eval_default # -# - INFER: infer_default # - - _self_ # +defaults: + - ppsci_default + - TRAIN: train_default + - TRAIN/ema: ema_default + - TRAIN/swa: swa_default + - EVAL: eval_default + - INFER: infer_default + - hydra/job/config/override_dirname/exclude_keys: exclude_keys_default + - _self_ hydra: run: @@ -52,7 +55,6 @@ TRAIN: # # k: 9 # i: 2 - # evaluation settings EVAL: test_size: 0.1 From 41e9f8bbc86381ccfb838afce586bc782e91ae14 Mon Sep 17 00:00:00 2001 From: Dubhe Chang Date: Mon, 30 Jun 2025 15:03:38 +0800 Subject: [PATCH 14/32] update chem.yaml --- examples/chem/config/chem.yaml | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/examples/chem/config/chem.yaml b/examples/chem/config/chem.yaml index 69efdace83..18228f8423 100644 --- a/examples/chem/config/chem.yaml +++ b/examples/chem/config/chem.yaml @@ -1,9 +1,12 @@ -defaults: # - - ppsci_default # - - TRAIN: train_default # - - EVAL: eval_default # -# - INFER: infer_default # - - _self_ # +defaults: + - ppsci_default + - TRAIN: train_default + - TRAIN/ema: ema_default + - TRAIN/swa: swa_default + - EVAL: eval_default + - INFER: infer_default + - hydra/job/config/override_dirname/exclude_keys: exclude_keys_default + - _self_ hydra: run: @@ -52,7 +55,6 @@ TRAIN: # # k: 9 # i: 2 - # evaluation settings EVAL: test_size: 0.1 From aa7d7273f120e5b680f24a16fd1441b92e4437d3 Mon Sep 17 00:00:00 2001 From: Dubhe-Chang Date: Mon, 30 Jun 2025 15:32:09 +0800 Subject: [PATCH 15/32] update ./ppsci/arch/chem.py --- ppsci/arch/chem.py | 11 ++--------- 1 file changed, 2 insertions(+), 9 deletions(-) diff --git a/ppsci/arch/chem.py b/ppsci/arch/chem.py index b2a659aca2..689f3b75d3 100644 --- a/ppsci/arch/chem.py +++ b/ppsci/arch/chem.py @@ -8,7 +8,7 @@ class ChemMultimodalMLP(base.Arch): def __init__( self, input_dim, hidden_dim, hidden_dim2, hidden_dim3, hidden_dim4, output_dim ): - super(ChemMultimodalMLP, self).__init__() + super().__init__() self.r1_fc = nn.Sequential( nn.Linear(input_dim, hidden_dim), @@ -67,24 +67,19 @@ def __init__( def weighted_average(self, features, weights): - # 确保权重与特征的维度一致 weights = weights.clone().detach() - # 计算加权和 weighted_sum = sum(f * w for f, w in zip(features, weights)) - # 计算权重和 total_weight = weights.sum() - # 返回加权平均 return weighted_sum / total_weight def forward(self, x): x = self.concat_to_tensor(x, ("v"), axis=-1) - # 沿列维度(axis=1)均分 + input_splits = paddle.split(x, num_or_sections=5, axis=1) - # 解包为 5 个变量 r1_input, r2_input, ligand_input, base_input, solvent_input = input_splits r1_features = self.r1_fc(r1_input) @@ -97,7 +92,6 @@ def forward(self, x): solvent_features = self.solvent_fc(solvent_input) - # 结合特征 features = [ r1_features, r2_features, @@ -108,7 +102,6 @@ def forward(self, x): combined_features = self.weighted_average(features, self.weights) - # 最终预测 output = self.fc_combined(combined_features) output = self.split_to_dict(output, ("u"), axis=-1) return output From 0f3d15cbabd0a66db8658df77221c5624e75438f Mon Sep 17 00:00:00 2001 From: Dubhe Chang Date: Mon, 30 Jun 2025 15:34:37 +0800 Subject: [PATCH 16/32] update ./ppsci/arch/chem.py --- ppsci/arch/chem.py | 11 ++--------- 1 file changed, 2 insertions(+), 9 deletions(-) diff --git a/ppsci/arch/chem.py b/ppsci/arch/chem.py index b2a659aca2..689f3b75d3 100644 --- a/ppsci/arch/chem.py +++ b/ppsci/arch/chem.py @@ -8,7 +8,7 @@ class ChemMultimodalMLP(base.Arch): def __init__( self, input_dim, hidden_dim, hidden_dim2, hidden_dim3, hidden_dim4, output_dim ): - super(ChemMultimodalMLP, self).__init__() + super().__init__() self.r1_fc = nn.Sequential( nn.Linear(input_dim, hidden_dim), @@ -67,24 +67,19 @@ def __init__( def weighted_average(self, features, weights): - # 确保权重与特征的维度一致 weights = weights.clone().detach() - # 计算加权和 weighted_sum = sum(f * w for f, w in zip(features, weights)) - # 计算权重和 total_weight = weights.sum() - # 返回加权平均 return weighted_sum / total_weight def forward(self, x): x = self.concat_to_tensor(x, ("v"), axis=-1) - # 沿列维度(axis=1)均分 + input_splits = paddle.split(x, num_or_sections=5, axis=1) - # 解包为 5 个变量 r1_input, r2_input, ligand_input, base_input, solvent_input = input_splits r1_features = self.r1_fc(r1_input) @@ -97,7 +92,6 @@ def forward(self, x): solvent_features = self.solvent_fc(solvent_input) - # 结合特征 features = [ r1_features, r2_features, @@ -108,7 +102,6 @@ def forward(self, x): combined_features = self.weighted_average(features, self.weights) - # 最终预测 output = self.fc_combined(combined_features) output = self.split_to_dict(output, ("u"), axis=-1) return output From e53498ffdc486660f57cf3473721fd92984c2bca Mon Sep 17 00:00:00 2001 From: Dubhe Chang Date: Mon, 30 Jun 2025 16:02:26 +0800 Subject: [PATCH 17/32] update chem.md --- docs/zh/examples/chem.md | 42 ++++++++++++++++++++-------------------- 1 file changed, 21 insertions(+), 21 deletions(-) diff --git a/docs/zh/examples/chem.md b/docs/zh/examples/chem.md index c7cf7ecf52..b7d9f86a6c 100644 --- a/docs/zh/examples/chem.md +++ b/docs/zh/examples/chem.md @@ -10,14 +10,14 @@ ``` sh # 训练: - python Chem.py mode=train + python chem.py mode=train ``` === "模型评估命令" ``` sh # 评估: - python Chem.py mode=eval + python chem.py mode=eval ``` ## 1. 背景简介 @@ -38,7 +38,7 @@ $$ chem/ ├──config/ │ └── chem.yaml -├── Chem.py +├── chem.py ├── data_set.xlsx └── requirements.txt ``` @@ -54,17 +54,17 @@ ClC=1C=C2C=CC=NC2=CC1 | CC=1C(=C2C=NN(C2=CC1)C1OCCCC1)B(O)O | C(C)(C)(C)P(C(C)(C 首先从表格文件中将实验材料信息和反应产率进行导入,并划分训练集和测试集, -``` py linenums="24" title="examples/chem/Chem.py" +``` py linenums="27" title="examples/chem/chem.py" --8<-- -examples/chem/Chem.py:24:30 +examples/chem/chem.py:27:35 --8<-- ``` 应用 `rdkit.Chem.rdFingerprintGenerator` 将亲电试剂、亲核试剂、催化配体、碱和溶剂的SMILES描述转换为 Morgan 指纹。Morgan指纹是一种分子结构的向量化描述,通过局部拓扑被编码为 hash 值,映射到2048位指纹位上。用 PaddleScience 代码表示如下 -``` py linenums="32" title="examples/chem/Chem.py" +``` py linenums="38" title="examples/chem/chem.py" --8<-- -examples/chem/Chem.py:32:54 +examples/chem/chem.py:38:66 --8<-- ``` @@ -72,9 +72,9 @@ examples/chem/Chem.py:32:54 本案例采用监督学习,按照 PaddleScience 的API结构说明,采用内置的 `SupervisedConstraint` 构建监督约束。用 PaddleScience 代码表示如下 -``` py linenums="60" title="examples/chem/Chem.py" +``` py linenums="73" title="examples/chem/chem.py" --8<-- -examples/chem/Chem.py:60:76 +examples/chem/chem.py:73:89 --8<-- ``` `SupervisedConstraint` 的第二个参数表示采用均方误差 `MSELoss` 作为损失函数,第三个参数表示约束条件的名字,方便后续对其索引。 @@ -83,25 +83,25 @@ examples/chem/Chem.py:60:76 本案例设计了五条独立的子网络(全连接层+ReLU激活),每条子网络分别提取对应化学物质的特征。随后,这五个特征向量通过可训练的权重参数进行加权平均,实现不同化学成分对反应产率预测影响的自适应学习。最后,将融合后的特征输入到一个全连接层进行进一步映射,输出反应产率预测值。整个网络结构体现了对反应中各组成成分信息的独立提取与有权重的融合,符合反应机理特性。用 PaddleScience 代码表示如下 -``` py linenums="5" title="ppsci/arch/chem.py" +``` py linenums="7" title="ppsci/arch/chem.py" --8<-- -ppsci/arch/chem.py:5:99 +ppsci/arch/chem.py:7:107 --8<-- ``` 模型依据配置文件信息进行实例化 -``` py linenums="78" title="examples/chem/Chem.py" +``` py linenums="91" title="examples/chem/chem.py" --8<-- -examples/chem/Chem.py:78:80 +examples/chem/chem.py:91:91 --8<-- ``` 参数通过配置文件进行设置如下 -``` py linenums="31" title="examples/chem/config/chem.yaml" +``` py linenums="35" title="examples/chem/config/chem.yaml" --8<-- -examples/chem/config/chem.yaml:31:38 +examples/chem/config/chem.yaml:35:41 --8<-- ``` @@ -109,9 +109,9 @@ examples/chem/config/chem.yaml:31:38 训练器采用Adam优化器,学习率设置由配置文件给出。用 PaddleScience 代码表示如下 -``` py linenums="82" title="examples/chem/Chem.py" +``` py linenums="93" title="examples/chem/chem.py" --8<-- -examples/chem/Chem.py:82:83 +examples/chem/chem.py:93:93 --8<-- ``` @@ -119,17 +119,17 @@ examples/chem/Chem.py:82:83 完成上述设置之后,只需要将上述实例化的对象按顺序传递给`ppsci.solver.Solver`,然后启动训练即可。用PaddleScience 代码表示如下 -``` py linenums="85" title="examples/chem/Chem.py" +``` py linenums="95" title="examples/chem/chem.py" --8<-- -examples/chem/Chem.py:85:98 +examples/chem/chem.py:95:105 --8<-- ``` ## 3. 完整代码 -``` py linenums="1" title="examples/chem/Chem.py" +``` py linenums="1" title="examples/chem/chem.py" --8<-- -examples/chem/Chem.py +examples/chem/chem.py --8<-- ``` From 13af56ab9d66c83cf7055c1590db8751bc47d443 Mon Sep 17 00:00:00 2001 From: Dubhe-Chang Date: Fri, 4 Jul 2025 16:56:13 +0800 Subject: [PATCH 18/32] rename files of the example --- examples/chem/chem.png | Bin 31491 -> 0 bytes .../chem.yaml => smc_reac/config/smc_reac.yaml} | 4 ++-- examples/{chem => smc_reac}/data_set.xlsx | Bin examples/{chem => smc_reac}/requirements.txt | 0 examples/{chem/chem.py => smc_reac/smc_reac.py} | 2 +- 5 files changed, 3 insertions(+), 3 deletions(-) delete mode 100644 examples/chem/chem.png rename examples/{chem/config/chem.yaml => smc_reac/config/smc_reac.yaml} (93%) rename examples/{chem => smc_reac}/data_set.xlsx (100%) rename examples/{chem => smc_reac}/requirements.txt (100%) rename examples/{chem/chem.py => smc_reac/smc_reac.py} (98%) diff --git a/examples/chem/chem.png b/examples/chem/chem.png deleted file mode 100644 index 9d243069fc42c220690c2f048e5011365d6d7eac..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 31491 zcmeGE^;^_k_XZ3P2vQ;q3K9~6k`e<-D>a0qbcjj}Fd))3lm$qGl!SDHw6vtsFmwzE z3Q`W;@$T`y^uE8(_5K0R50A&;!7zMg$J%SJIM=x*NaKMbDKR}U1Og$2DamO;Ah>=I z2+kbgMevTu;PiL!1II~A@gAhO@7gl>1$oNPP`Eu3P#ywWwuf@bNO0?#w#9>T$tL&q$Wnm3v{(r{^h45}4yY z@UJJutN%QCBH@%D7c2t;c{$(IgZKND#;gqF_j}Zr2xYME_yyqL5Mp0w2(UnKurG)& z;;~?V%6#^N{QeNv6>>l98xZ^pxa8QMX^;NDLH=JU69Jq8y^`z%A*-R5mn6$|n8v`* zj|8q{$ha| z>Vnv{&Jbd}N4Vobwh_ z|C;KTJ~s}Md5Z@WI)X4{RvrfoYJRz`<%ME)%|>W@U~IQyqo5KW|56}N^-;iZUpatpJY)g!#^d3KnCOQ>=PvJ90n6F z?09nWAnEW4oAF-UiC|ZHAodNU1lvDHBmO@d|l7 zH9NHRqq~>2aBi#b?kbPXR28p?1N!~Z?y92GV$bdJ*v|#qPkzT6(hPv zgUN1axTW?6RFr_pJr|=QRQ0 zpAd%BpZ1+d46K;oy=up#LvZ5#1PG$=h)85;ADH><()et)(i?6I?W@41p!F=3fB*v;Rhipo*X0a}$r z54BTbQw$R>A5s&Ykt1Mq)_4iICXJL3fe*bnJ6V!0h#kmK4A17Ro{mDEAJ4`mWB#l! zd?z8j-`2Z6h3M~ju@h0!`uQC1@v1BgeqwZ+{>Zc5#a+S^)RMcgCj?fL&6(<57Smyr z^(=YTd%L%+-qd|~GryqtmdX@7C;8`AH$T6M*ZJxG2-o!xeQk)T>x4zccAt3M58su% zoZ&a~>GC06LnbV{;JE#m{UMja{wNBVSQ~;o`1X4E;}gY6MD@(@tnuQYBD0z=I57z; z3eu~Ba%w+bO%#+g=5KA~GLJcY4#2mta9$l%3!TRQ%;v{~Zampi5otdjQsK{oPwSa_ zWW~I$uKt$yje<=fw@D4heNdim(9~=F;Zf365u2E?=0V6pQ_dih*+;2Gu2*}@a)jeH}CM8za1{I^(OI^$b?0l z$vs7yLyy*laB0(#pj7qmf=r%2Ka@%2B)ifldvATs0;@PU+3qjclQ{aBVdG>;8!df) z*#A)Kfm7)4EVuwjJX`^>|IufC$k{wylSV*jYTQCoLqHe4BmUg(n3=ky=Z=PU71a&Z zG$V(GhoO)yLZ;S!+%2adv|dR9Zp5d}4u%n@P#)(b5nF{o0y5o0Cf}3IGHV;o_iTcMOWnbdF(`_GI%d0Sal5u7@tm6qM$vgx8c!NzJ|Jt>t5L{rfTCOQkvS&k{y4%F0 zUKrM{LCPiIj)utWR0y_WI2Z2{@Zt_4JohH7x`#Zv;;OUE-3JMz%YTYUExj|c{{Tm2 z$0jemq^sBn{w|Cj8LlvGF1)-} zPkIT{-MMHH!dsPQ-uf!`Y*ACfl&B{}M1MeXqxpi(?k}L*nyx8yO*iLOi|fLq4(2X^ zBMBeQe@vTU$)u-*wKH>y!8~X2r3WU|ED$R4MVTm;O)9@j(wt+(CgO%=PlG(TXuW;V zgr29nV|xAH7+uHBZSeX@s{Q3tUWKQ+&j$^Ean)fZb>SD8nPgK{`Kye_KCS8OW=@${N7(l@@C{ z$8m73F#T%;0=!4=e+5MtWhd)<0piC{`T_4dr`}d=g;{^XzM{+z^RT{8s z&5VXDjYntu!ZvGBQK25^$F=0a|6ZkiME09i^L(+t*7@(pNrzY|5FVyHl>gqb`nZrE zxnDcvx97l(M>JOtw$oNo{5-covT=Z~COuEE}KQaeN(JjvYN zwN6W5vNRu4`MsFuHzBD$dj_>9Tix*SQu9XKjY->*hw1ihZzzP#Tkuv;Ib?{kuQx2_ z=bA5&Wv*Axrl37ndFh2LFS%gQ$zf7QKg*u8w_N+9r{rPmeVKz|N!wWxJc7OOh2Ijn z$A_Orrw!Tq$sB%S;(NNARq>L9P79T-w*0Y^N!zzvV{RMmv6QTpvXtT#MlWnFQ*PON zD+~6&%TKyFU21j0?r6mZ;fPG2!K&DBGUWYTmwaWcgX!K{Z`emWgHiHkF?ARz8!-69 zW-Pa!R_iO$k=)1R-EsVSed&nBo&^8Pe@-}ph{`=mCI(HO-hCU@Knownl=r{slZ>-| z=`~1*qyHLr16Wr*c)3+SU$@)$NW~#dMaC>7v#>S!?$5!DH8B*p`LzpS;yvFfFA!YT zu}E56lKFG0Pn9yu zYvp+$+OoObt80XicSC#RBoqa3rnm=GjXx$4S#Mmgni^_4+btqtx!XYC?=x^B@M!z$xr@tiCO-XlihX5kz^wcqxd z=NAY-@@WfrOlm7)`VK$Y&(yjUkoy7OFPQ9mb~p^Yc(4^sCB*yNOS+%HBi|}#M!wlM zm)ITFhL`t3GvdsU2C=^TGj0V^^mm?TCRE~)JnM4;N}RW6=Gaa%T;urUXrq-(#lwzv zxfD{rl3&m@`SI!3s3AwH->XfOqj;%+H8Am0bYXpz`JcJ^dZ)P@=rLhF~kq1FPOvdfT*k4tu)gcUVk4p1~8 zyYoGwtYXG>sz`VUUas`o9LIgTS2XVI-EUu-|B`)3^*Sv+tm0(?t&QRPNn-Kk#8KR5 zz4WXp*C$<#E-%V89rjKxoe8o4OZd{5ScdjKkMsD!MYJGUF~+JC~alM%i){&5G~%)!MZhm>WgfFb1hj0dq!Yg75gZG4t+!(a8=eJlGp zsYt`no{#ZEy$pz z3^MGneDTg4XZ*u`0&M)dUfO!uMlMnj9{A6ugg$={tPo)9IlHqWZei&e;fUtCcRQwf z_Qs^!Y-Ni>s?VE_FcGIrWqsgXfc28#_$t~b%p1QWWXZ$g{xnE_Boh|h)ioEIX0Y)w z4L*`jvx}LPMunzTLf(ryq~tHWy($#5DdcG8d+PEMG7VzepC!Jh1xLv*2QD($r5(v; z;Lmxk7U^AI5tbN{u;_}>FsiWDyeTcc?S8V^S@4`e#70gnQTW4aCeK{wX|j}GK0AYo z-9z?Gr>-`$jSWXLi&xbaI~W{v))#I*M*fPxKYdAOs_z>1pQ~V)$pOiVY_~~3ulBn= zjaEEx&kTP#3!NgD-0qo-Y<`OkhOfurmg@eFF066B?5vK3pY7yG=RNFr&5W#3-X@kw z*|Wfby8W6SI;RX>BY=*gnY>D)6 z2VN-3d?M42Il7JwV{}3)@X~bbdG%p!-TGib(d zQ1zfKWy6w9CqGXkS_Yge7R371p=NbUKPpfUGR)4l;w<2%?$qJ4sI@Gs{uG)JnHJjR zgU>B7@|csu?cV#plZ`Y(rF%t8_>W3ZIfLhGUW~l9V+o^j`V-u0wsn;L{|z75W#?)9 zF2#N0RV0xy{npcuUXGOFzi{yq7w;vasf#yqFH-_L7;Z?^(?7441Px_!pTpUGo94i+Xgf{!jsIQi zV`3R9UWkWqufu}XD?Nej-NVEB($KV4h0&Hrn;-z>%pZi&BI`d1h0A&BF_$at@adYn zvsoiuBVKZod`SyU3#47HKmJu`vw8*1t>GZ!(6G~g-sBPPC+?1{uN)z-=}mnG*(W(mB;9`cE$6f7c<96cF`<$ zOGpvQpS@4Gr!#9O-?35rKjRiZPdWNK!r|LG{0^-x#}`(F;j@ig(2W$2tl8s~11Zwe zW!O&Z-O%OLMaE$v-;=dekIP9#&jKMmmw#v9f8vhaSL=kS2;|qZ&9#}nV z`wP~N(eElU9Z724#zGz@CiL4@wWv*wCx+Ky?c#s(_dg?ebX@yde)W@(e5|nNjRa@O z$}+d6Rk-PPn9Z=mPxWI|iMn?QYieqO3u%C&K(A9nN2kfZf-t{&O1JNQ)DL0!kJ1i^ z^(TxX`P+I1yI3-z@rW=Ltl9cc_V7FuddAP z=qL$gLGZA6ZMQ&Civ?1C;eTUWnIe{;D|*@Z$+KUQQxAhUwF?uuy-KF`eRIc_wYj^t z65ZNA-sH<*-Ay)QnE&@44P_~Bj|iruiwwGimY0onqzB7rM7v7F1^ef%Cs(Fd^Uq?}SkC#M-H}giX{sS{9fdE>E zs!Gp3OjhsZtS8|fS*%S^G}}Pk=L*^1Ksj(YOx0yt!dmpF9Hb~yr+e@Tn z|DIWOF1Qb88O9ql(OA^lpI*2o((CZybxOX-wkIiW zsVl(D00vcLSgA}D<bz-FYiV z-8~xysiV4J;S$*HOGu#!Nak4X46IF*tFT5W=DIr*FQTykRidB;ac}(XK#puMl5-NZ5`Cfdj8F!A71lp%C-MzAbimTs_mvW!^|I9ZtW^2c=0$2y~)nEUuH3gx{ z*o_D8mAG|>7kv{Uev|hoMPxH{qr?^l1?U1aC=<{NNVWj9+n$2SL;>##Q=lDsk8?e+ zBJw}bnqSd{P3O`(Huz4q05^ozF+w2TLEMn6ZPgHNzl!FCrC<7Q;3(W~BYT77w{CmN z++HR|8cidtRhxAwSR@e({%yTx{~qw-t#b-o_Q#9sp@KHsVHdY%(XkB+;_H@`Wn-1? zaEHqOVuk;G{KWfO>VgPdXJC)MHmtn!LFT81aE%rR*=k<4qBR$U?5S=D3-UhX2JU45 zLLc6zG?2ZgaJPT7MJ{b!B)eWZgf`jJzA&}{o`AF;&|OpRzne|;7oLoe&V-M$;5_E; zV<(}Ef9o8I&UujTipcTf)Ad+8gq54fcJPN;3BsY1;qxkR6G@R!N^D=c7 z66rCO9Bx&M7KBRggBw4EXSNyeb}U5+RR8#ez)}L@4PaR zkqT`mcn`Gl9eY_aYscda%p4LbOUuS=C2*}N^}e)A_S5`DS%XzY0^CK4ufEWsV_0Z7 zRdu((xcP)>wK#0vFd9<1z+5ufWs1^I^HjvSjT>(Uv=2gn3!zT1!E;9@lHgg}A`l66Vo+5_JDFCY(ej?}SCg(}V$i_!QpmQndc|s`J;+WbbtM zKte^Kmhjk=wF*+44abRp`9=lzkD~aovJu5(vZlSme>AOm$H@ZYb?z=7yCh0ayHB6} z@JaWTs?2?9HS(tG5Gow31R15jwe&}mGB+p@k(D@lQfo%&K#oR2Ls5ZYn_p5AvimnF zAHM$|Cs@+tv7Fpc%V&v<%bv2y&Ovzh8sh)eG_FpCp3*O6*Nd&q9ifkR(t~$07is*f za7!J|H+R0fa+dg5Ih611Z2HbBL}w$-Zroa^LiN{}B3nj&h}Z$kRQ1mt)KJF@4Nk}~ zGan}=O%g9(kxYm?&ba>~XNqT~Z%)*L5v- zYiDI7rX{s+5onrrB?(D9fF)4Y|sy^g0z3?04R&?rsP`|`|EH(|Y#VErHYuW1Q=PEb* z0cNUFWK^9Ez}SHoL{z!zsZteRZzK(rSpX3cW9>>VoiV;s-uv7i@F8*$&CvEFy_%0) z^A)}`f{Ez$ve!$=4Q#jicFToR`pd5-JXs)3-JQQC^?J3L#I1gL@}!NKZH8(}bHV1f z1U4e>U;YsN;BFpIky#VcdF^9o{;bDNxZ+uH)7jC`=Kp{Z+wW;E0n@3$uf*D>TMGLj z$E*It$evc05{k;EVDxc@TVQ;doy{`c_Ew}KUg*%%R-OlLNLl@?G^F*tH1EbpBEy6; zZJiNnJ(LMs!E7#cY9H^d!@opaU%}*Q=j11ypB-c3lB@U8UMK)~=)KD}X!QQ{oE9Nu z)yGp0YBVbsV2>tRuRJiNM-M{~%I@@ZzXlh|^d zqWbJW>WEmhwjA_>6h6VJ@!Ykk{T`)~D;s!m7+lib?L^z+n~VT(85`+MmGp9*@K9v( zU}tQvD?1i=HYei}2A6jQ=j_8mRLB(&`VwEAGtt7GH~%1Fg|1!y8GfkoGTdbQ8nmIz z?P0;F?d-&HToXC30e5CkJv6(-D%%QPwSPUg$*#Xdgarc5d`tM(D+c>qBy!_3!4631 zwWq2O%by9DbvhRaW!^$?AuBYnwl<<2KKTo&R--rbC&D}v#^z%i#(ufo2ZBAJ>fiFQ zDw*hGxxQRPR#Gg23RIPjw#w$bU7RCF->8ZBJxSR?9RN4FhKd`HY$hLv@D5`fv;EXS z;mgmnMZK+!#;*9byur0{3USqC!*hd)V}mK)ZrwcvzRt+yWAkjzkEK=Cz7cX(`d2H7 z|8rjJbpXQY8geJbX7YMXtMV+O)6ecNi;3gX`h(ik`t{Iz6Y)hF-@J1eH&y4lOsAe3 z4^P+{gh3SYyfm&>{yz0)st02Jwa&|0lNT)kh*t6WBEx_m&IM%;cZPTC!y!9q?Lo{XwWlL`_Xx}u}N8~FS7 z4yMaHW5M=n`x#{hSu^C}kBZ0b;bHlr0lx-51kxWagr@E0n=@cbe-V60)R1F4b?Opt zj_MbmZMBImw1Y!^!A>8#*@=~0USlGDXD!XXU|3Ifo>$o?uq)T%_hq%j!G;S%Y+b@n zjg#&yh^K!J3^I%>S#wdNmbG{?l%!dq5ql({G@y{jkM~IGPk6Dm$3LVT-Tu{yESC2(Rwa+d60J5?Ynm(JJ|XB!!_ zHUBYu>#yVVV{}wVPk%#hE2O)}clZ0V^qbl}D&m~V$A@n473ddzye)J<_INa2YS8e$ zXYd_m03-eJ6_7`-*|`pCf}9`dY#U9iy3JISUJZ(V@mQ;z#!gGpt~>AHUq1m6XTQC~ z*_P>KxBl}<=K#sf@d1d&>+H!*aXQ%$o-46AkMD%~uE57x{5ikn_sY^%H#kmd7yR{S zkXY+)b3O`cp4q+X@$5ABeU8I(IIG1Ikoq~deP#+JO`7^EX)s{KV&Od4MJbii5TD)Y}YZKnswnn@@~6QuqTJ`9dg|~geqmE z9icbhmx9=8!GrYDiFKb7r=u&L+tqm!4X)U{@7h2e4oCCJKvVf9L5 zx4&_gj1Lw2)14|-#|sv;8T=#EezHm}89lhVdx=K1KfN+?Y`ME&Uk?loq;2Dd9d-`2 zq_bSkZHd|(2%F`a{beX%^e?soiuChFT(>1}uiDrmas8&?#OA#9&ddNwJm(`K2SOoA zc6|tZ^!9LYmdAO+X@Yax;ArZ^$uHjq|9^JCMXyC38K99;N~2Lp&y(dvW4Y*4aWh1* zwWlIYF!1P|-d@;beF;X~%XY+g)|dAvU8*T}R?#l;*M6GRzR(yVU;d9&%dQe|p zt;@T1rr1ydy(b>F>>0l-OFNqFiSthorXifkk{+0p@T)RyPR~{ev0HVA2)-Yu(>SKz zb{IsR)EI%6R&6Jlyj)!-`ZKxVz1Hd6lmvzx&i4Eh$zz9n zOl7@AX`Ja=V-xPoIJFh1kd*#_!&5DUIxoyF3bK=EI#E}DpUx1KOVKO+O_lz1x+2)y zw0yHD)AB*a^I2F%w(P!Qe!;BA&>cr!o7ux4l;?j)70+1y8T|y;jF+{t^`Teh>Ua!) zD70Nf(#DGIm)+LJw)ch09=R1|ZiClfQ>mpCPCi1BX)&RkPYuRl$9o?ScPpLMnw;Xr zmS9V5B1w%XpNDvxoquj%@GELLJKRvSxdM4R`G-S(%jr#-f@K0Q$WhVSZk3u%AT6gQ z5TVx-uM{&gm8FX_zDaxWO1}ppu~BRY_^f;YgEN zC^N0pIw&&Zwce#;q7R~`XZMUJRuM@%{V$PH>MQ(BWr$|+_y=f=)Nt=SATv+st;rxf(cei_l;2m$T zF6N1lUL*8V-lBYKJSwC#?eQ|wps{FTxK($qMv#`3zI;8JlQ+{N?s7VFVl^$mBv_T= zI6|!_^r>;a+#AhKJ&owPT$Z+)a0)-tB6sXZu}exjfPJ$-+{VobCSAr%miML{asWM} zMM4!+nQ}n2BjJJcsq^HcK+>N3=gAa%)I+8!eqDG3B-X zTBpkGI4O5ugxk)$1}$^aCj&VxeC<)_H2wj@XM1duj1*g!O1oWX+_TbKob51mo~mnxoTWO08s{@hzD zbh@L-9#7GljEqlykzMjqF#+b>ci*XxBb(MfSg~`A!&U%0uhThD2 zWZ2+tBv_?98_wu%$)9{@v1R zXQrfeH^L^DR*5m&C$D_vKn*l2wjqxBVE+s7zg4X#oqCiHfklzprk`x-+q{320_FXQ z9HETDeE&eI>B!i$t@IFYjz-T|7O*BnOaPGX{28B41C;ZZ-#hi)E+3I7X3TxVd=C;u zDO$>TSdX1w?khj9A7Or6uSh_j2vL3@klsp%+MByB>37D^njg!?Mt4qwZsoT0RQWi5 zt5u#K7ye3Si}TfO=+~Id++E~Sq$=K=IHq2EZyy}g1C_WlF~-{fhsIVb4N@fLCjEmC zuH&%~MY$7_G{X8~r3y?@kB3#s1mQc}KfRi>AIL5x#_#Te-0!%f7ygaIm(>uIzlQ2x z6-gd zJ}Grd_o2-+=)`c)KBgSb+!|mt%rHBssE3yCo^<;#Ldq~mr+wvLI;nAXyO@l|+!C z4dD=6lX=bP8c*r((OG!`W+ny47lK4-6;WVkiK@`cFYis_JL!sB`EFHv4fE-w&lGOq zpwXm^*e0Sa$u}??cB=TLfgg>|W^9TTpL=Zfg`j!)~Ew zHckps=|?Sg8c;-%?_zS}*P9@)4`VnXLAcT%zH%Th>Sqo9clzm;^5@cJ6{xY>_uHR%f z8dLCRjv4mxlZ?1b6IpKe)3-U?Id4J)QrJrxi zEivAFc7$(4r^IIJNu524SEG~4etI*xpx1bUVXw~=6;SNpwQeZT1g$BQ#(RVaJV=g& z*a6m@T7sbbX&=cwP~BQ6TL5f^jBdWVLLOeSW`93YEsaeQIgboOM8zljQCwq6H|Tf; zH{r4|xPf!0_B;*cZUQb(NPYNB&>84$&`(Va9=ex)@2 zu$6221gCxIRJ*K{ZH9*B&$__G2$?q6W@lN*-cv?s3%LeO#wWDLL(v2SmLksw8aAk^ z>@1PX^s>7bD{~mFI>j6Q=;tIzL)m%AUxqSiU#U~U1D)DhZ;h9W5=6ApuVI3 zY#BMqy7iRROUqJiH&K92Hv4<)1)o7J4xSMTO{Y?qi0lYsQ~?Gk zV=2|Q(XL4rofNTu{o3uvcxv$ivXTWg0G7$U2?zOQ^TnH2CV~SDJ6*rC6Kf?veYk&^ z-+uItxf(rGc~C>Db;o?MXn*)NRgvD_X#T`r0$twud<>j5Xs4#Yk0X`9{ zYEXRnjjmt|=Puiz+G^s8j#OmITUqV|oYvTK>!dqQmfsi&$0K8{m6L=6o@&f0aVxtB zT)JyECZF#Vzid0+a@Z@F9$)M<{2b(SW5PCa&U+JbPD{avCASMn@DTx@ez5dRtNjXo zr4XGHg`b(Lin}kxhVL;312*49Tn!=1o-22k-aAgI8I=CE1vsCCR@9Y#+qw`ZH zoL~zbp?P3s7H~lA;ahJ!{;R5Bm ztHjNP%cDgMSYo9G9NcXJmtWHcr%YDFSJ?d0eu7!FJj`i#4UtM?`=)G4N-(b>Qs&d- zE-Rt-yb>!m%PxSyOlMV;jTyZ%(1VuvMNKk~>0g?vN>Cf?Roxewk)oQ~4(y8P+ZXQ? zsnmE4=)3xzxEr9RUaXJT(vP>IDceB!9iU%aj8=-xtmKe@AgZ*M7jJY-;}1=`X?+lI zZZaQI23gvtp5o+5Ri0DSI*K5-vZ|OUd~CrVZllbtQ%M*z*>2&;m>H>ss(^gon&^SX zXZy6M?w6+>+j}83t_~d{muj+US2U=9$7BFV4XOO}9Wb)86+-C;(_{mOjg*)Gr)Bs3 zb!p$+Z-__kcRjyzxrhu5x7GE&SI{O zjqgxByD;x|5jVo@c(HP^HwnuK9Zn~Qo+YF7-9IgYqPuRA^mZ;9us02=h~oztWZixa z$o}{ins$ibr~=vS;2V1OtdpTH5u#@CVnG3+?v9nW;zDO1m>a#aPZAB)+e23#TSzyS zyTyf2maK_jw69fINB{cJ_>Nt;*OkCPV+L^~T6k&s2 z0X`2;)+Po^s_j--+`1=uhrH3*k5n@-lRAcDGp(2ZQD^&N;^Jx%?hKdsr1%l+5S zi}m*U;DdMeumZ3%qs(j^M19FlXl!|wSP^j}_m`Gr{h}M<;%s31hm?6YV?i3fB>URD z$*A_HAndkr7H}92xd8JgGD9}2#IW4Nw@4ZkgEplx>xP#;8KLxhyJN?+-2jKa0h z<32Qt(9{w~A?8GBNc6}w&-L6&T@D7hICB4p3FY8TbeEYr#4~6X_C*I$T#fcznUJ7D zK4f6-X4stv(Sy5wDTL(`!)zIazvaNY0H1gmZeZ_@M;@bjM}Z%ro5*ltb+TsXBJijK&G3~wkKp*>&ncI8hQW|&@6I%3^NE??Lm z_IgBOlPbJrMNH?AswBl=Ipi{NLdz%DP?hR8vUI~>VwE@AH%^Y)3`LuVK$m6K8}sLt zWYv~0H;RH>ZBGb7i5*{q21lPCd&*|tHE{{wC`4XJ`tBy1jh~fZZ_A^s|9(@6MV-KQ zq3JaVCZh5aG>)cX?$J<1R*>GFXHMHEZ`jl26WM#yMv8d|O}sR|S=If-B1pwe&RjLf zfhPZ^I8h!+GYNnFz~q)fgP2C7yi#Q(Zj$KSD1+8iV2=WMTO~yRWO%+sl}u>&(4hR) z&)M^A8l&^!uU)2mLduevBb`S=e7gmD*0A~B>=O3P{O-t=A69&g)pl}1UpRlurxr2| zyho8;)9Id&dXdMbQ`*^a5n z@-CD_rbX4P6Bt)am+9nlmGnrqu`s;5D}h`~ z%BSgydlg(ELpidv-?wkPS6BaT{&dWi%mfwXrV8t(W-lp^^oDf|_JK zZPrY-SXA*7Dc?@Zcjkf4jB&J!_v!Pbpca0U_?zNSa1MG!X9qyVrliXM+qzlygXH|7 z<9&=eor!Z~Lc^l}1st{MM-=<@o}Vq*9-&%}$5L#Yk+z7=B z;KJJqdG$uWx{!e4pPWBYBrRjJFO5+6ymH#8a=K{c`Lbws>E11VpzW#Z4v7HbPxI zy6tq+EW!2%F=yMW3q8;4JJ@nyla!W|`#n1>B-;q8rmB_IPaF6~e|I|npG=4)r15Js zXB=`2cLQ|Md;r-U{U=-IvCHg5&L|whVJ1S_&4{Dkj-EzRs;hPsMqjKMF zi^&Wqo8`@_rf%NPFG1AFBK|C+X=^>0yI7;C9P%EDpf56QNDPzMQo4Jx(b~Hv5^hTn z)fj&$ma9h2Gk<+IDRctQ>BiSeeQU04bXIsX;OX2olc zc!=H6hf*#aGHKiZ0a4{7xxG{azfe4b& z@ntrh{NT444p+njT+ru3J~_|4-TWm2!*uqtWl_p$;|$^ zRK^VugcTT8euI;w7ZkG^&M!J~RT_*u*RhFE5!v z;(zD#PB{ANVDOtVvJhOl@r&~{Z8y#Ca@qpGeQeEF8vyPB3;isKLY;-K?5cwtK4=#y zY=~jn+Or^H)@N|~oF!7Akv_``hqwfvYX~DeZmdIY#bTEFqcBOR zPm`IxiRt>|U1)*wO@=9{&BQ&J(2Ol~EG*>g{rFp*$+>)9Ewjl{0%`{Ap#uYG)f2Dj zIl^$xWa+_1N^q-I491`*%hjLBD3i4nSlr60?w~y{@_?*m&(M6)#|pEAfM#@))yLT* zzcbL^_Kdt`MwwtKe)2|J-Ne4KiOLvhX0&}(JM)FlE&iWGJq)?#_9V*HWg!*PeJw5WqwX9NF`2QnLgq#*xY+}!hb z4o$1|i0qn+?5jW09?ksGl&1eX!%P5N^n2?#;{|&ByL@w=kjpXBnKFfeC8AesqfHwg z7By4p%$6Rt-Qm1+lk51ZnR;teQfJ_nXV7yK!1UNaAQev^T+xf{_g;7~UF-4j=KWc! zl;?X7{8#b095HH=^OL@^n_$fv5W(%mUhNEj#AL+>r+!In#?&I%jTOUn{Kv~4B?9FX zxW$op8oD@|D|*t~#Wrz$WUEHnmvndx5Bj^%H>D@;?j!?zC2Ga8{H|H8`j9l;WLa%$ zWw1wy2~kr;7k(wlyj#a_2E`7bN=MQIsh9Jpn zPSNm3z~seE@%uOjpucVKf%v){AphTP*cnn;{AP?x(8W9y*RZ~Ai_{IBv*W&I7Z5)z zYHl$;+A7RjW%p7vyy}pMtDFzX?d8Efz|-F`&Utp{*2y}RgC6Vey(8aC5J2`|$merXO-yx%zkH8=(iuq5wm~g7E;o>r(BLEN0~T zN_gp6AoZn|9^ zJKcQ#`&#-+)M*e2%~(#sM>8$a4tZu7`s0$hZkfopQ=;%#5}V)G@+m81m1gYn3;y&K zHlnLddQq?z*52bABV(g>*-crl@u>GupXfF|sU2pJ!XqY^x_Xe$JxsKXX^HZfdy&-~ zMbBjs3&5pO0JWds}Vnrf=52e_Tsf+ ziRlWlf$8man7Drvp`DBxhhnO~Ggn!@vt`|N8$Q+x$>ey39f$A#j6)5PamkH6?)a=s z$tjz?`csfx`Sh_ZpmBKF0Mdoyafj^@1q4{G#nX3NQ;7z}LC%cJ-jLuISU6OR&Tslt z)jLsRaSe*Iz!F89n;Xh5L}kJ&SE~6YTXgrfTEcQ25J8F3UoxN1+p#lRJsk`QJs>m6 zyCgIC3=6HV;FbawN(E?`1Za*@tG{trv!7iAuJ^nB2l*-yaK{g){q4ur3&hRLdV8)n zi`IO$BEM$UKAUltboSxcZ0$HSmWaOMaw}n#D@ft=jwq^E0@?OyN%I#QWLN1QtEU*w zWKu@J?5(j*#cPDt^Mw0iOXG8T#q5R$*Nm>48%e_k+|`=k zY%5)+p=vWOX7S6W<%~hwdmj-3#>v{*awAK-fWLgiG1n zK(U2O+{4n%6O6Y(;%aP%=Wl!R;(EAh=#>Vgpetz6{=(O(9_@i<35FYT6gQYjeF7IT zRjHInlB!jh!}+QFM=8%#X#|->LmF%Nn6p=}vrxE>gSN!P_v%lxP>+*{VvLm~Lv#3x zL*ZOzkl@q$fqnrdwG@d7)of7F%6kI;T#~-P-~*^2aYOHTjTEesU1%#2Q|~~J#{(=- ze#dDsCYcj(;T$6aT?dg_SCyio8$rO4PF*kS@+^u??E4`D97I7=V5* zwBKk_V`}osR0|<%S{NeCOie&QY$LttSxb??zUo8vehvomH6A<`GRWHeDC`bU9yqp@ z8ZdB*RUs*k9ao1dZS_1s=Q13mzZHPIj$-m!({wEc&F5|Hy5u$Yi0w30$9l6#kpVB) zyMvtGAkV$hZ>U*Y)P*@OpvT$_Cv54}Ig-Y8jnGq8{K~(Z=a#R3t9uctVzIGX&l)(@+QK~F z%Jh3)1$tPAPc9$1-$Qc0lE$N3j9zOWu`TaMpwIi^WGk7p>XVMQx63nd5xc%6Kkj1ljvTdH}(Y@vctMp5%`h?Of0o zEbOwP6N>~lW=NKR@u_8aFX+i12}J-d*4G0EbWEbJFDNQLR>OIMM_$9;&F>5=R9d7f z43O(}2=_`pE9o@1eK|{OL>_uYw!$*>3b{$rM7Zqk?qb}QyV{C&d4CjeBj-J<{y&xW zXZK8R^I_t;J^d?XAeKD8tH7XD^Yt)>*RX<%s2NcE!tXl&eDpo6AN0lj1XAHC2boR+ zsF#TC)|uk91n3Ak&ddF{G-8SuC9DLQdtE`XV^Q6or?hRTQMz@s8bf3@Se;`kQ;;P5 z;o8u*>9N~w{$UE25^{%DlK zNgqQ3u?*8`l9W&F(F`}_kTaYxJIIYV9t#MzVX*yQQ_UMO!ba4_;^d8UAukJL7C3TM zVG!lTqG*rMuhEU#pd+B#{{C?Nox#-e9+i#mm3Adro-a(yyUCPBcJ*7MRrOr5n$<8$ z7t?569L`VSvhhA8yrH)KPiuD6<1G9ufBS1~oZpoLUV|&3yH4|SQPpJEmvvfg$VzD1 zn3=DS&1hkQVdTk(o>^=%+qX&ojqCiEp^v3)O5?iT+SW%jW9W}lZzmT|oG+HeHk3=f zH_+rX!DO=PB-u9=^sB@mQ6#+1jOl1Ljcdk-oC{?4lm$`dZ z&HwOwC~IMQeB-O(*Hx)>dGEz6a9IwU4zFImzomI#;rzP}mj^K}wi{ncV3@5 z%!y<^_mI<|EJ9|NT#scQHl@nDmL!OxCsHjijvy`m#hm)&C_U6H;L$%*4i|zi=a4!( zctS7kTIgw{XbN(7`J;%sqsN6Jc2gWPwqNTT9~tRJoR7c90AArfaRE*34V8C1nw`H>D_l(tNI%2wkp*p* z2@n`h@W7UMr^@&tKvdFv?R{?*>E{sn)R~HJCO2rVv}gD_rjYEIw(356ZLZy4Dv0PO z<8CF(n@&3>TA`wA1Zmel^SF!}yk6`}w!MJ@x=^&Y)C~&FT)}={gm&r2(<@ZExDDQy z1ssj^V=jxiP|i$jH^DWdbt@_~g=&4{La_9>yf0RU?>#?yq^ieczYLtDpvy}c+RfZn zD?h!a@}q<19R#wEO}JShp8`EFlE_8S*(;^(YnvqP+h~YC4ce-=D-tKV(=4-iE1<41zrU)|ufe+mYd0FUIcB z>_2}UA36nrhun1Sf&3!QGa0g(;#w@Kb2g!M%{~y?F7*KD^Gvlm+@d6Pd61L$TCDrd z+j*GR_k_h{OVa#Q=-7TJ(H2D>wY| zbOTfFpqC@swZ$IAeI<>yY4FCVN_U7?Wzu*7=lk6gd zLY6EcYuQTHA`}TBdyAzZRLDMtR9dVNQH|_NmTZ-5$rcHNEEQ#q$i5B3_dMzKs?X>5 zdH=rO@9p;WA7kb@&+EFL=Q`Ipk8?lfk_g}KnyLGg-@q%#?JIYxHhK$~4Cx0>Pi6Or zh%~&}?r64`Qh%z*-h488oJrvl@vE??n(zz3fr5@YyQYxb%BLp9jyo}SQd}>aF9*a= zG86`c(c$?QD@r9rudZpaxTl4Gwtgqii>PfKeb{hNftWhho}Tc2e`r5V6zVy5R2ls2 zI9hcgnUUas&Dn>dLHA{`PJ1k~@|v{4OBW3WE<(`Yse45QvCU#gp-UwODlMEraH#tt zo+7jV##;RpFNJ(%Oa7!>3U*113;7nlDO6y?S-rWQ8Y!R^|6cY&4Gn;4?tx1RXOPco z5xX3pDPua+_qP{gX3nL3-kmpLRkEQhr3P>fmvd2MJe75^V@$y=8$W|o()v=R!e(?z)<5Ws@XthjmiG_s%l|@aSUvKAp%p!Smd z*eC^ZdP{S`hFfF&mK>&d!kM}qioos4sQD_X!or6k`f|oKP|Z=l`_n#_sC*gCL|Mnw zxz7}<&+-24{LQ)wp;}{U6)KM^j!f(}Z%-u4>3&Y_PpW#Lv0>&dwO}Tvw|X?nfxFyp zn0Cx9>Sc>E;$5|6c+{Gq6uG0ZAFY#iW_Vx#F3>-JBmT zS~2N5*!ZHR{b96b0wvJUjHMJ~}ax{a1 zQMzoP5pPD#b-8Akmy!GE1BJ7YzA}|RoImiSoLE}2aQ~uGbJnVELeZgpxrZJcKREc1 z9ue6*kLp+cz5w0{^J92X6*OtKNOF z%yGbnSy@skuRf{4l@X6>?a0cr2yO zH1^fMRLz-qQZD0=<6P)nVfN73w>bYnP12%rCb!?7GyGq7sNyBa)$x-_i_{H_=(mO# zoxajcP1F;R=b+leh2u;T#8Qu+X@q4jbUNh&m8UdF)t+pJQY%88ijdMDSGy&pZrqt(>tpi#gqN~fhuI3|Lofg^4&NnYPyo+gCa$r6aa=Ps<|EyWrojJ8od{+`l#deY~ERZvy@+Z&I z`JImVT!c9$aQwdX17l7!yZKG6B(VQvhm?y&plA$N*ZcyS(D-~I?Q^`R-PCRF0QOe} z3Jo7rI=Vl-_szTfYbS7RDt%VlPeSI0W1co+-vU$RgQ!eosQGCCT7iVhQQ=@>aQz*F zLPn(eWeO|q0Y6%;!@_r~q8menh#imtI$uo!-I7I(qU@<6+6UIysA((I*$|m<65Z zs_O={SzcR<wo>4{06IfQ}kPD~v z#`j`{h1|$qT^l|DAy8I8cUTo&YU>?KGAp!4Su@Sg<@A*goqmkl_vGT7qD)>7-Mymb zyqpq>&L{mnM>4vai;6Cqnt#o|X6G7l6`!v|uITdQ#<hs5Ryc^sRTTI%wC1Y!j`_RMHI#{Bp^_E8CZ~hlD>Y z_oym2|J)T(?!3e`Id9R^%*oM84hny&z_-1AKVZ z@{q?W{)%uo@j3ph7hF)fohF#?#CXhIdR5kTCrC*PX`pEJyT3a0 z>-(J8N7{orNF8@hM~ZhiA7YoXuEk3pA-U;h!dfq1N6!PB(T3~N047to)};bIwXS8p z5bd%VKjllMc*X{-vkWLIikA3@wHSQN@E!j3?v%Fb6ms?1VnOzLeXEU~hCi2ABD3N(5e38DaSFg6vq zXdVo$TI&TZLT6p8#pM^2%}0x zn)4~m_?(+L6U~hueC@<|l0s@6i6`aqtn`w;@FF;i-v`1yz$8NV7DvuQ&bjCT7>Wx! zFFb%lZ>BQM0=sOyUDLxhGlwUsCFz*Hvcn62FaZ*h(L>?wU5vI|-Od>|6ZaPji@hOM z+V_bi_=+t+fsCd+W)M@zDER`j13_N>vm4lK5W6}G>`sn%WbWj? zQbp$8eg1aLqmw>Mzw~w}O(Q6VMHbY}ip$?_WLms*DFURA;nV-dGdv}6q6dgCjg7RB z8oMU9pO4abD%!fEz&Uc9ujkwCB>I;rV3ntSczc70*IJazqxf1ecP+Tse0 zyiv?38V`KTp(}CZ^ck{<;}N0ew;khO7kCbR>H*9H{})!0`haD;NmsJ3y${PtBjlh# zFWkIN=J4_pHaIFf9<*y@^(aIWx(f>7#F>72c5#uW`s84RukJ=>=1iDF5!q@Yp;>Vc zkK;TgjuH++V4?kNk8exIh&*LqQ!r-kCjy0=U|< z*DU*lVlAV_K1unJjxin%hWdsU!(KEB4X_*0ruvD0LVRW~{8p|Fi}%`pd_2|xMvGmJ z#}Y;YckC}aOR`HUyDnMqLRCE%$meLQ zvU5xC0)C2n;R~-+-B}Q$W(I7m+n02Fb07ggN;#yG7lZWlus!6pfnEeQ!-J)zSJ3Pn%aKo_UAyM!9xJ0~fRy z_xoC?i?BNFy|m^8R(9JO71#}RI|DV7z|v^bE$dM5U6k<= z=-k`jsE6cQ^}Nw@6!UF|%J2qX=6HGqUFH$PBAo@DaXeIovTd3y4J3Y*M7eo-ki6kx znZn0wTg;E3OLD<~Za;wl@6C|-AeC}3DMUWjjto(g+pLk;N1l&{mzZ!f63sE?ghDq= z91l8l;<{cscEL}a1`5OxBxB#7%v0mRDk?qd`mb;F?-}XTl*(K8tadY|Z#dXZQDp*d z0efKeWaW;y)sw2GvGWU8d^RdF{5eK;N%}Y1nE~lo54C9I0MY1Tj5fvI>Q+^j-VqUv%q>xtEen~B*c7$U`F)4~S|A+4t zMgCm0LifAfJ*XZmWILuC^^?L(&SbHdj!WmPTmK`BHGL`GQ1w=)zM{fieVs}ABfF5_ zd&$e@m4ACfGk!j`8&_|zA!I5A{v=gK-)@|u%edt7U* z*Zzr+>b+~>WC&k4dOicuHW_T>AIa|@daEc#)=uJvpDU6@)rey4Fl)y zz7ni9D%~)cI)ETU41b-I0()8pD1IL5YO71uGSIj!7{4Zd4^03qU5&=xBWRfk6bppu z!(2Z_j}bHg=!Bh-HI5{HDnCr1QCFs4jai;ERVy#WINvOHBJ@J3mBBCl$E#7SpAC*? za(8ANK@i)`qx{w_4(P_lzwI@l@S3ZA?GoYLFO|Ji-;Pmlc{&zTZd3e4x-h03TerGQ zSuPM8K-Jmb@?n*QDmG_{ZJu)p@7-Uf@%nb(wL_KON-kgj``ydn`l=*DCp4LN>L|#S z6gxuATOr?ffrV^5WsKbgqmq9A!ow|2eMIe1q7{t$_@)o=`>ZG+LwM>;SSd*gj%+;P zD9pQJX>7VkeAI)FKjeS@bC1pQH@4~d>7af==#)+OVzFKi)~@HVaLjkQ`y%AX>|Yaq zE1c0r@NKPSP1sKJ}*Q+Px*~0 zRr)?*kJ6`D5QER6|FbeRImn`eazn?;M?bjdJ$id&qO=%Jb$1^*FYr$m3Ulw+dEQTa z45*qxG7QyzK@!hxfLh4;*YnUQgEMUPu(O|De3J`hx8dbehTtafEIMoUnU_d^AL??+ zQf4~BXtCf!S3)99ez!g!-O9> zoo^jk@ct!1mp z%uL8&Qgz!G<8V>vyR;)~TiB5-N}iKD#*a;H_s!#Xl?zR;7kb8OyTFeK;c+w+2yQd+ zvSlwqG!r}_)ntytB(U4rmxfZ(3dByztSZn&*`QQ#NN1QbmZ@C8%u9@aBazG7?s29H}7p# z4Ewrv$tU@4AM+8{H`PRm4@rHzT$h9I0Iy}-SYI1oDUGK_hW=XPI%Fa)d^=dEMvnN_ zGx+Kc$Cwm3($j^!OqHe-FUPz~?Fq}EH0xL$qA`rH*dEss_6{ktb25fBYmvV^<)*}tDv3Y) zaTKWY`3h*4eQZUbOxk{>gq@YCo#~Q6n&Ja^vuqN(EtAEH#%wlU0DSu$f1XesvIf@O zgc4z)z13pAEi1}7j(f@q6v8MeHM+=KzofU6;0NMgS`n;o=bEksjqhDCj;P!Q$f>_J zJf!BH*x86HDmrT~^q!VM;d5-fTykc*B^tj$KPSPTT4>N=J%}Pk_F(#|NG{Vt$;nb- zb$?`*`j;G1`}WTS?NFK=={XFZM67i{DzE$+&=H%Ta$7J_K4LddgA_9%LFRH=MJkTR z7Btmo^dopv0>szWImnQqzs%R`v&uigQF`Uco20KQ5i@l(&dyHqxJ3&3E=|E;b<(-H za7u)4zfFg43SH!sgmFt!3O&i2Ni z`yQaFnvE$>h|ZaiyLm^@-YfVt{a(z>4cn386;~PkxJRR=J268=567W-h1#b6@VAY~ zi0@)l9^2;Qep^NfM0=$d?D7uB*HVOr^o$hyhegzik9LZALDJXU+)AvI8*TS6d-Km# zb4y*VZ@!YA#I^5hpY*j<@0WAdJIR%!wPSt4uRo@I&EXZC1<<|6IzhqOgR&!8rlq{t<$2~QSVn>U6`3I1?$d>*P zE?PS-`i{uQX!w;h4%RO5`lf^VjPe_ZB~GCP&(*0)JQ~JO1afR9dDwF3G&Fa zvr{GYj-S6D^qcHB6w|)3^zFtD zG^TPCV%YJ3)A{(A7hg;dad{np->l&sD4ccm^Ww+SRzjn6M#ua~W$eb6b-JcyDxJ=G zNW%DKppS1#D`4*zK3Sk)2S-&k9%es~hI|kYghBTp#ku$5+|CnG zhvs&M8bs!PcaY45^B=PL8Pnv_v~bqyP1lc)U)ZTFK3#218o4d_aV>Wz+ZZQyWS|tC zuRl-Vt#Iozdwt5TFXgYWVK6j;m?3HHL&wR`2UV^3FPBcR%u6%p(TLoxYTXBmg&u|$ zr0wYc;OPI_9q;g~{5l%d!9vKMXt@@g}QwPh^CujsZ7nj<|S9~>d!46Qp( ztNvxS=*i-0<)Y4o(U~~8v*1#yV#zYB{(N}!%|lHCoIp&wZ@+CzT>0B)cw)pV||#ef}`;nf}&e2H%i9`a_{uFoyQP0MCDa zwZQv|(zu%@lR(c66$n55`*lF@m_(;U%Vj|OC1~@7%_K1#k|K6dUz6A$xShBM4wtcI zHB++x(3JkiPAyjwuSrkRA_(7(-x!d2fwRl7>iDGo{auORa91U7<|_O(k^MLF@CO_@ zVZjVh{)WatBMEd7T5rCr!O0lib(3}T@tr-2W?;6aci4L`B?VuG&@3^W72)IhYrQxy z$2iQ;jngmlBWiznP2TLd?Z>UuBqEktm=QkzKW_-GqEx#!eV824{<9`&0_kx>-`psW zh(AvP6ZtOW{o7F*gzqn>PN|Gh?Hc;^Hv_FOI%=GZ!8chE!EGDbWQZVIj*e1Oukay` zJJ$p#5t@JSE9kHU{6kBDUxs$nOewPAAV9Z`VYf7{4|qPIs{%wEK?*xVPaIXQ~!~TuXH}; zyPDT^m_tN6R(8O%FXJTHk9v@f=eR6F!FBHXIuwF;^h25PDbW^?uxmjWC}X$i_Uhto z=fSgw=#2{M%8A}b_CS{s7EhvmKNN3ur0FFTC_#If6m;v!KmV?gAcqp}N1dSLH_eE8 zAX3F>-a9lL?dikE!Gxn^#-a7}D92oRDxKpu!CsvAUxK}D?Z1(|&835z4@92tBs*l= zz~tCh=KbErDCcypSQ#ViliV(Q5NuYQ#^%uA_u?=mOl(JbL)sc#P zpaaAcjJ@hVj3f93Rsz~$iExEamkb@(u}!VrmN`C)_BlefZfh$Pl8=Jkeox+(vO5St zJ9ZH^gZ*3Y+I$4nh<`mK=wpg#^Zs$dp>_uTGWM6!{5d(*-yZ{(CVvMX{ADW@Lp(R!udXesd}QA9h`<;v6dZ>;iCw-Vl^$06n9`Lb;E8 z`asTq1%TeIz|YWyWMd-aW()ubh#9E4W9R*!^xVT_ZM%U~&|Cc-o(ZYcXZWxG0wai5 z(Jmn1p+7`hXAou}Jlq;%YOMnm7-#Ci(baq)tS>3HEw~T`s|?%~Lt9{(_`gCkR>r1Q zsYH_-*syvae;x&Fz6cma9e=IuyeefAukzINh)34R{O%yhoS%162J1 z!H;GK4HM*@FF^^Fy~#NpfXFTgkO^4#Ti5~Q^rso?htWDc<&Cwud2r=p zW-Xx%XTM0;?gK&lbmP?giSy6&@0XGFC%U64pH9`tP44Ip z;$a?u$cS43MN#&0x3mq}rS?2%;>G6-#jx|AEL_{-4qVLeUHat?xvSWiZPj1}4}fZR zqOe7h)I23R@#S?PryiM09AuV=gzp&fRW{PDuWIbw=QpR<@~DbxT?*P^ zDWn@DA8!M+h!G3!hl7IUwf90rUDm#j-V1rn2Wy=|;iHC}OUf-GNz2{m8 zr}1pHz$c~I#zg$rb98D6SQ5HHc=s7h{g>t_f#^(xu!)(Rh6X2M>s*TK%Y&T(YE=wi zf(>kVc(KdRhzVNB%Tt@W^Ff1Zv`k%`p&`H0`pR3%anC$NL-o!1?3@xql5Y*OR)!WQ zG%TVng=?IsPlrNR;tKFeq>i{wl`o}1$F4ka28wqAI3_!$K}JRhDL_0(`|?j={nBo;Ee4^mck1T)^z%^X{D%;oDMP#~{hOgtv1m7Dd;6u9=QyVi)c%ywt|O zV42$GwfY344W~$9*z4Q!++}wKPRj)PoH?zk5dbQf=rFV!%h|_F%si7aS`WPF;&vT6 z@Xuuz_?h~2)ZU4Ge?cS233Ev0y6s)q$M)?LL+FW&s88*iT08ndrsa$U-F+VA_zL29 zb`+%5i=uC*-}1H{Mg`_J`}gQJW8|5*r&y(LU{5i+zIfM?G9$ZN>e4{w0)U-j*oCFb zLa72=@9#NmsAn_V7ds-)$@dr_s_LY4sLQnA6_u%%EAmpp#9G%b+FcVwJ5yt0!+N?K{GOaT&I&2 zM63a_Vl1tUmS|{GyqDfwUizw}BJ<1dHpE?C4?nca#&MKKy4+^uKEu87)T7D`($svD&3b{j#C~W2B#gF< zvOkC*A!xAm;5)yC2NW-ckgr0eA5AV#=>whWz<>H#BNFf!RJf_T{xqIZV49 zG*;z9Ta-3zXE*I#Cw5h4d3yP@N0l|TvzlSQTx*i(C~>xcI^*l2lTfx}G_{8WsGp|U zUKTB6eUNt}QKF9kfR7bDel3xT#$JupeXAwPuOozC4|R)%{WiN$c@Dt zT8r>cs)|Dq$2$7)SfVJ}%~AH+koLmMe%5I)`zVpfa77;2zluv=+%WH>gi;67kjF~i z2T9O2(bVG$umw9Bo{AzXMDZ`OfUPk0_OgxOkUob1$HHuU36$(lu#1KGwk+f96D#Fq zvyEaMDLOmNOWiZQ!MP+41pn_0r4ho2wl2Jf4zU!3xb#RvZKY@t8#n9*KQo=Wave(_ zq#relQ-;A9?*?Kq)=l%u&9aG(--lzpQ07@R%<;3(=exV>EH z*~IuQ+}?>8O=c;=H)aKArO@6Af8UtVub%zJ0zqZFgKGKiNk|z`@U?@nl*{^2%Zd|k z@bc-jNC(4{sRZpz?=%jA@Yr&2U&3WWlYLU(ray(^Rfj73!R;yV+gEle754w}*V;Uq zwt@9VwUfQY$+SE7MGp_(lma`?DM=&l#1S7WrOU0wTIPM(~> z%yy$whc%xxiFwKyP$sTS@LZ%*!}3JdrkM%Lm!D9`LWbr5TX}l#IlaZNo;F~-oJZ!o zvtqwf@&-U}%f{K+xjSgp+F#chVU-TQFcY0Fu>msd7-=G`4MAI^l^BUb>+|o z97im+29y<7Mji)U9<$k>uA^I9N1>0;>dy7|Ei|)@Rw^G`MrZDxv-WyB-GMS1GuZ?4 zFGb}uO zYqtbA=%$6uWLc@ywS1r1Z_V=voPwnjqbd+89)EP0Cy>6Es_tGWVte-Ti^PlHYr95_ zUjvM2Kes~D`Kyhe>#lHlwcYBc`3Ao0^WapT&biso{p%$cp4E51xY+3~tGaCkPWG=S zx>5cg^vIMv2}e9ySpG-``IQKtjPcX diff --git a/examples/chem/config/chem.yaml b/examples/smc_reac/config/smc_reac.yaml similarity index 93% rename from examples/chem/config/chem.yaml rename to examples/smc_reac/config/smc_reac.yaml index 18228f8423..68cd4300d9 100644 --- a/examples/chem/config/chem.yaml +++ b/examples/smc_reac/config/smc_reac.yaml @@ -48,7 +48,7 @@ TRAIN: # # eval_during_train: False # batch_size: 8 # learning_rate: 0.0001 - save_model_path: './chem_model.pdparams' + save_model_path: './smc_reac_model.pdparams' # weight_decay: 1e-5 # pretrained_model_path: null # # checkpoint_path: null # @@ -58,5 +58,5 @@ TRAIN: # # evaluation settings EVAL: test_size: 0.1 - load_model_path: './chem_model.pdparams' + load_model_path: './smc_reac_model.pdparams' seed: 20 diff --git a/examples/chem/data_set.xlsx b/examples/smc_reac/data_set.xlsx similarity index 100% rename from examples/chem/data_set.xlsx rename to examples/smc_reac/data_set.xlsx diff --git a/examples/chem/requirements.txt b/examples/smc_reac/requirements.txt similarity index 100% rename from examples/chem/requirements.txt rename to examples/smc_reac/requirements.txt diff --git a/examples/chem/chem.py b/examples/smc_reac/smc_reac.py similarity index 98% rename from examples/chem/chem.py rename to examples/smc_reac/smc_reac.py index a897f1f838..90287fb6a2 100644 --- a/examples/chem/chem.py +++ b/examples/smc_reac/smc_reac.py @@ -88,7 +88,7 @@ def train(cfg: DictConfig): "sup": sup, } - model = ppsci.arch.ChemMultimodalMLP(**cfg.MODEL) + model = ppsci.arch.SuzukiMiyauraModel(**cfg.MODEL) optimizer = ppsci.optimizer.optimizer.Adam(cfg.TRAIN.learning_rate)(model) From 38b1a8b3aa067fd8b7ff63bfbcd4961668808921 Mon Sep 17 00:00:00 2001 From: Dubhe-Chang Date: Fri, 4 Jul 2025 16:58:34 +0800 Subject: [PATCH 19/32] rename files in ppsci --- ppsci/arch/__init__.py | 4 ++-- ppsci/arch/{chem.py => smc_reac.py} | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) rename ppsci/arch/{chem.py => smc_reac.py} (95%) diff --git a/ppsci/arch/__init__.py b/ppsci/arch/__init__.py index cf9bc4c68f..16387994c7 100644 --- a/ppsci/arch/__init__.py +++ b/ppsci/arch/__init__.py @@ -21,7 +21,7 @@ from ppsci.arch.amgnet import AMGNet # isort:skip from ppsci.arch.base import Arch # isort:skip from ppsci.arch.cfdgcn import CFDGCN # isort:skip -from ppsci.arch.chem import ChemMultimodalMLP # isort:skip +from ppsci.arch.smc_reac import SuzukiMiyauraModel # isort:skip from ppsci.arch.chip_deeponets import ChipDeepONets # isort:skip from ppsci.arch.crystalgraphconvnet import CrystalGraphConvNet # isort:skip from ppsci.arch.cuboid_transformer import CuboidTransformer # isort:skip @@ -73,7 +73,7 @@ "AutoEncoder", "build_model", "CFDGCN", - "ChemMultimodalMLP", + "SuzukiMiyauraModel", "ChipDeepONets", "CrystalGraphConvNet", "CuboidTransformer", diff --git a/ppsci/arch/chem.py b/ppsci/arch/smc_reac.py similarity index 95% rename from ppsci/arch/chem.py rename to ppsci/arch/smc_reac.py index 689f3b75d3..8664877ec3 100644 --- a/ppsci/arch/chem.py +++ b/ppsci/arch/smc_reac.py @@ -4,7 +4,7 @@ from ppsci.arch import base -class ChemMultimodalMLP(base.Arch): +class SuzukiMiyauraModel(base.Arch): def __init__( self, input_dim, hidden_dim, hidden_dim2, hidden_dim3, hidden_dim4, output_dim ): From 74da604124fd098f555fe17d8af9020e52fd8c95 Mon Sep 17 00:00:00 2001 From: Dubhe-Chang Date: Fri, 4 Jul 2025 17:19:02 +0800 Subject: [PATCH 20/32] chem.md -> smc_reac.md --- docs/zh/examples/{chem.md => smc_reac.md} | 59 ++++++++++++----------- 1 file changed, 31 insertions(+), 28 deletions(-) rename docs/zh/examples/{chem.md => smc_reac.md} (77%) diff --git a/docs/zh/examples/chem.md b/docs/zh/examples/smc_reac.md similarity index 77% rename from docs/zh/examples/chem.md rename to docs/zh/examples/smc_reac.md index b7d9f86a6c..bdb0eab119 100644 --- a/docs/zh/examples/chem.md +++ b/docs/zh/examples/smc_reac.md @@ -3,21 +3,19 @@ !!! note 1. 开始训练、评估前,数据文件data_set.xlsx的存在,并对应修改 yaml 配置文件中的 `data_dir` 为数据文件路径。 - 2. 如果需要使用预训练模型进行评估,请先下载预训练模型[chem_model.pdparams](https://paddle-org.bj.bcebos.com/paddlescience/models/TADF/Est/Est_pretrained.pdparams), 并对应修改 yaml 配置文件中的 `load_model_path` 为模型参数路径。 + 2. 如果需要使用预训练模型进行评估,请先下载预训练模型[smc_reac_model.pdparams](https://paddle-org.bj.bcebos.com/paddlescience/models/TADF/Est/Est_pretrained.pdparams), 并对应修改 yaml 配置文件中的 `load_model_path` 为模型参数路径。 3. 开始训练、评估前,请安装 `rdkit` 等,相关依赖请执行`pip install -r requirements.txt`安装。 === "模型训练命令" ``` sh - # 训练: - python chem.py mode=train + python smc_reac.py ``` === "模型评估命令" ``` sh - # 评估: - python chem.py mode=eval + python smc_reac.py mode=eval ``` ## 1. 背景简介 @@ -35,12 +33,12 @@ $$ 本节将讲解如何基于PaddleScience代码,实现对于 Suzuki-Miyaura 交叉偶联反应产率预测模型的构建、训练、测试和评估。案例的目录结构如下。 ``` log -chem/ +smc_reac/ ├──config/ -│ └── chem.yaml -├── chem.py -├── data_set.xlsx -└── requirements.txt +│ └── smc_reac.yaml +├── data_set.xlsx +├── requirements.txt +└── smc_reac.py ``` ### 2.1 数据集构建和载入 @@ -54,17 +52,17 @@ ClC=1C=C2C=CC=NC2=CC1 | CC=1C(=C2C=NN(C2=CC1)C1OCCCC1)B(O)O | C(C)(C)(C)P(C(C)(C 首先从表格文件中将实验材料信息和反应产率进行导入,并划分训练集和测试集, -``` py linenums="27" title="examples/chem/chem.py" +``` py linenums="27" title="examples/smc_reac/smc_reac.py" --8<-- -examples/chem/chem.py:27:35 +examples/smc_reac/smc_reac.py:27:35 --8<-- ``` 应用 `rdkit.Chem.rdFingerprintGenerator` 将亲电试剂、亲核试剂、催化配体、碱和溶剂的SMILES描述转换为 Morgan 指纹。Morgan指纹是一种分子结构的向量化描述,通过局部拓扑被编码为 hash 值,映射到2048位指纹位上。用 PaddleScience 代码表示如下 -``` py linenums="38" title="examples/chem/chem.py" +``` py linenums="38" title="examples/smc_reac/smc_reac.py" --8<-- -examples/chem/chem.py:38:66 +examples/smc_reac/smc_reac.py:38:66 --8<-- ``` @@ -72,9 +70,9 @@ examples/chem/chem.py:38:66 本案例采用监督学习,按照 PaddleScience 的API结构说明,采用内置的 `SupervisedConstraint` 构建监督约束。用 PaddleScience 代码表示如下 -``` py linenums="73" title="examples/chem/chem.py" +``` py linenums="73" title="examples/smc_reac/smc_reac.py" --8<-- -examples/chem/chem.py:73:89 +examples/smc_reac/smc_reac.py:73:89 --8<-- ``` `SupervisedConstraint` 的第二个参数表示采用均方误差 `MSELoss` 作为损失函数,第三个参数表示约束条件的名字,方便后续对其索引。 @@ -83,25 +81,25 @@ examples/chem/chem.py:73:89 本案例设计了五条独立的子网络(全连接层+ReLU激活),每条子网络分别提取对应化学物质的特征。随后,这五个特征向量通过可训练的权重参数进行加权平均,实现不同化学成分对反应产率预测影响的自适应学习。最后,将融合后的特征输入到一个全连接层进行进一步映射,输出反应产率预测值。整个网络结构体现了对反应中各组成成分信息的独立提取与有权重的融合,符合反应机理特性。用 PaddleScience 代码表示如下 -``` py linenums="7" title="ppsci/arch/chem.py" +``` py linenums="7" title="ppsci/arch/smc_reac.py" --8<-- -ppsci/arch/chem.py:7:107 +ppsci/arch/smc_reac.py:7:107 --8<-- ``` 模型依据配置文件信息进行实例化 -``` py linenums="91" title="examples/chem/chem.py" +``` py linenums="91" title="examples/smc_reac/smc_reac.py" --8<-- -examples/chem/chem.py:91:91 +examples/smc_reac/smc_reac.py:91:91 --8<-- ``` 参数通过配置文件进行设置如下 -``` py linenums="35" title="examples/chem/config/chem.yaml" +``` py linenums="35" title="examples/smc_reac/config/smc_reac.yaml" --8<-- -examples/chem/config/chem.yaml:35:41 +examples/smc_reac/config/smc_reac.yaml:35:41 --8<-- ``` @@ -109,9 +107,9 @@ examples/chem/config/chem.yaml:35:41 训练器采用Adam优化器,学习率设置由配置文件给出。用 PaddleScience 代码表示如下 -``` py linenums="93" title="examples/chem/chem.py" +``` py linenums="93" title="examples/smc_reac/smc_reac.py" --8<-- -examples/chem/chem.py:93:93 +examples/smc_reac/smc_reac.py:93:93 --8<-- ``` @@ -119,17 +117,17 @@ examples/chem/chem.py:93:93 完成上述设置之后,只需要将上述实例化的对象按顺序传递给`ppsci.solver.Solver`,然后启动训练即可。用PaddleScience 代码表示如下 -``` py linenums="95" title="examples/chem/chem.py" +``` py linenums="95" title="examples/smc_reac/smc_reac.py" --8<-- -examples/chem/chem.py:95:105 +examples/smc_reac/smc_reac.py:95:105 --8<-- ``` ## 3. 完整代码 -``` py linenums="1" title="examples/chem/chem.py" +``` py linenums="1" title="examples/smc_reac/smc_reac.py" --8<-- -examples/chem/chem.py +examples/smc_reac/smc_reac.py --8<-- ``` @@ -137,6 +135,11 @@ examples/chem/chem.py 下图展示对 Suzuki-Miyaura 交叉偶联反应产率的模型预测结果。 +
+ ![chem.png](https://paddle-org.bj.bcebos.com/paddlescience/docs/SMCReac/chem.png){ loading=lazy } +
Suzuki-Miyaura 交叉偶联反应产率的模型预测结果
+
+ ## 5. 参考文献 [1] Perera D, Tucker J W, Brahmbhatt S, et al. A platform for automated nanomole-scale reaction screening and micromole-scale synthesis in flow[J]. Science, 2018, 359(6374): 429-434. From 1e2555e859d376bc58162eb411694e80bb1477d8 Mon Sep 17 00:00:00 2001 From: Dubhe-Chang Date: Tue, 8 Jul 2025 21:00:54 +0800 Subject: [PATCH 21/32] update mkdocs.yml --- mkdocs.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mkdocs.yml b/mkdocs.yml index 7f0c496405..4fc334cfc4 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -111,7 +111,7 @@ nav: - FuXi: zh/examples/fuxi.md - WGAN_GP: zh/examples/wgan_gp.md - 化学科学(AI for Chemistry): - - Chem: zh/examples/chem.md + - SMC Reac: zh/examples/smc_reac.md - Moflow: zh/examples/moflow.md - IFM: zh/examples/ifm.md From cf45ac5d588ad2214abc3a7c3c6ad37ff880dad6 Mon Sep 17 00:00:00 2001 From: Dubhe-Chang Date: Tue, 8 Jul 2025 21:12:50 +0800 Subject: [PATCH 22/32] rename eval -> evaluate in smc_reac.py --- examples/smc_reac/smc_reac.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/smc_reac/smc_reac.py b/examples/smc_reac/smc_reac.py index 90287fb6a2..b87e5a0c32 100644 --- a/examples/smc_reac/smc_reac.py +++ b/examples/smc_reac/smc_reac.py @@ -105,7 +105,7 @@ def train(cfg: DictConfig): solver.train() -def eval(cfg: DictConfig): +def evaluate(cfg: DictConfig): global x_test, y_test x_test, y_test = data_processed(x_test, y_test) # Reformat data for evaluation @@ -148,7 +148,7 @@ def main(cfg: DictConfig): if cfg.mode == "train": train(cfg) elif cfg.mode == "eval": - eval(cfg) + evaluate(cfg) else: raise ValueError(f"cfg.mode should in ['train', 'eval'], but got '{cfg.mode}'") From d30b1d8dcfba3a0e951631d2d79bb3e66c086e00 Mon Sep 17 00:00:00 2001 From: Dubhe-Chang Date: Wed, 9 Jul 2025 12:40:06 +0800 Subject: [PATCH 23/32] add dataset download links in smc_reac.md --- docs/zh/examples/smc_reac.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/zh/examples/smc_reac.md b/docs/zh/examples/smc_reac.md index bdb0eab119..7a9ead9599 100644 --- a/docs/zh/examples/smc_reac.md +++ b/docs/zh/examples/smc_reac.md @@ -2,7 +2,7 @@ !!! note - 1. 开始训练、评估前,数据文件data_set.xlsx的存在,并对应修改 yaml 配置文件中的 `data_dir` 为数据文件路径。 + 1. 开始训练、评估前,请先下载数据文件[data_set.xlsx](https://paddle-org.bj.bcebos.com/paddlescience/datasets/SMCReac/data_set.xlsx),并对应修改 yaml 配置文件中的 `data_dir` 为数据文件路径。 2. 如果需要使用预训练模型进行评估,请先下载预训练模型[smc_reac_model.pdparams](https://paddle-org.bj.bcebos.com/paddlescience/models/TADF/Est/Est_pretrained.pdparams), 并对应修改 yaml 配置文件中的 `load_model_path` 为模型参数路径。 3. 开始训练、评估前,请安装 `rdkit` 等,相关依赖请执行`pip install -r requirements.txt`安装。 From 06b4c44ac8eec2a40b12eff69cde26ea6c7dad40 Mon Sep 17 00:00:00 2001 From: Dubhe-Chang Date: Wed, 9 Jul 2025 13:10:30 +0800 Subject: [PATCH 24/32] update smc_reac --- examples/smc_reac/smc_reac.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/examples/smc_reac/smc_reac.py b/examples/smc_reac/smc_reac.py index b87e5a0c32..cf81115181 100644 --- a/examples/smc_reac/smc_reac.py +++ b/examples/smc_reac/smc_reac.py @@ -111,7 +111,7 @@ def evaluate(cfg: DictConfig): # Reformat data for evaluation x_test = {"v": x_test} y_test = {"u": y_test} - model = ppsci.arch.ChemMultimodalMLP(**cfg.MODEL) + model = ppsci.arch.SuzukiMiyauraModel(**cfg.MODEL) model.set_state_dict(paddle.load(cfg.EVAL.load_model_path)) ypred = model(x_test) @@ -133,13 +133,13 @@ def evaluate(cfg: DictConfig): plt.legend(title="R²={:.3f}\n\nMAE={:.3f}".format(R2, MAE)) plt.xlabel("Test Yield(%)") plt.ylabel("Predicted Yield(%)") - save_path = "chem.png" + save_path = "smc_reac.png" plt.savefig(save_path) print(f"Iamge saved to: {save_path}") plt.show() -@hydra.main(version_base=None, config_path="./config", config_name="chem.yaml") +@hydra.main(version_base=None, config_path="./config", config_name="smc_reac.yaml") def main(cfg: DictConfig): global x_train, x_test, y_train, y_test From 0482f7053e0cf90988e82b304e8c1cdbd24156df Mon Sep 17 00:00:00 2001 From: Dubhe-Chang Date: Wed, 9 Jul 2025 13:20:08 +0800 Subject: [PATCH 25/32] delete data_set.xlsx --- examples/smc_reac/data_set.xlsx | Bin 219784 -> 0 bytes 1 file changed, 0 insertions(+), 0 deletions(-) delete mode 100644 examples/smc_reac/data_set.xlsx diff --git a/examples/smc_reac/data_set.xlsx b/examples/smc_reac/data_set.xlsx deleted file mode 100644 index d556bd35e6256955b2bda30ccf4b48aabe30a7e3..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 219784 zcmeEt^;eW{*Dob7(nB{xmxOfJ&|T6EN=Qjd=M0@nNlPOoCEY3Ap_I}H(h_oR_&(2j zo_C$I&Oh)Tei+w+eeY}6=d<^b;8HE4=6#*Rq0Re(w`76uT6cGU-0R;g89|0Z7 zNXFU8-Nwn?OvlH?#?6Go+wnC`J_-_ZE&>uD|Nnpf2S=bHWkjuq3s?41;a29yGo_VU zNp!x$kO45Org&#>a(|hXQMR4keRkvxu5=!uy+93ke8u-fv0B7I80Ye8dY0?`*gQHvma!GDdGqc#2nFwjFjXhet&~+)L3Nl+lBJrTjtfhe$ zc&q(iFv|=}ga((p*3@2n#}gbboyOgq#GkRZWAQaT-F&8q5>4y7LgZ;$r~%Q}X0g=z z^5!~Kx!sMqiD%Z_3h6snL_!<(x3N{!mUaOuRbH~%(x|gC!~O&ZwhljIzqcLCanWZQ z*UkKdwr0V(ky!w%O5 zcSU>;CVm=){`&oXH}LsiR0M?kdt?NS{|3u?T`syapw?6X)L{UyG<#$7+KrRr;q(8) z@qaJ||HJi)BvthuF3gA{gA{Bmss`UJuPPB|rL_+~MN~Hm>$jtn?dv4LScqrmJl$WLIsO{YIG9$ISY5=P5L^1OLwN znc9azEm~gi)$?)r0fyZ3K*KuT-y>N*y#k<~^1yc^Mb0URa9Dyqw?AlwQZIU#<9#*NlG`uvuT>Df~hJ>Mj<4la~8@(=s{N zdEU(nPT>e2RDe8DoI=8-fltO(1^uq;(Q#eItL%8To$w3Slb#9V{c zAY$uJ%+t|H6CW_P2x6EK{~D`w*&CCeR}aD=8KiAj+4waEOY0$ggxV`hy z0OQkx$bR)euP+UOlLy-Yi$RBKUOkI1Bv|^ev$ozxt2p8Ruq;1<^{h8!wf;;WlCq!V z3eszSMNPW;qmx8F(b+!_X@p%0E#@UfBQ}Z8ICH%(tA>NGe4Y4}{iOkBX5-;9 zedsT%mWbv-7o*<0SENRvmCyI75q$WUq@HC;rxcF}HD;#;H@u+XcVJVB@+8H;(flDg zk|riQ(PVV@ZRiKe028%Z_xC}G(`-oSYizwI%x?=Rmg}u=y3jA#mFLow9hay?yCAZH zXH3X}1kC|U(JL;3Pgg7!FjnjYF{W%2#->gO?R&G8f8Q@;89DUdAdlC6jeI^IH6c>H zPTU|4^INF?#1rH^-%!zZr?M8J-W}qXHA)q4yVEd;j{U_x9-JlSG zc@co@|LmU%Emh@SF5EVZzc4PZOix0*MR!ix1KoWb^r1SIWkv{3z}YG}YqN>|oH8e} zbA->)=z!-XANB?^{&^Q$Nfa36E&l5xktfchr=!T|mWQ9HLgb&|gAexh%MUTIGu-gM zh$juQXA=l+T`@C|!U~FLp=V#FB6#h`xICWQyG>9I0|Y%%_ao>!O#9$dNL^ziO$N)$ zpZ8y6mM>6qLVTzvPf0MWE1K7#(l2!NmMh#w^FGN z;jYh0E^EKRMorR5%y-B0j=@=#*d!f4x-643vZ$C*jkCkGq$jJySgd$+dwc47f=IWo zuU%~@Ks$hD@%wJ$uH}zfX8ory7j*R4QP-cwnRu2y3rh|mmrfQtwJFA;w(Xkvu429L458!I7q#_ts=tj~qSDWcvRpnHO0wHL z`_?7HttIMf$hTB&gXMZ;*Dbh2?oD^Kq1hjA#>_Gl+tBQen4oxcYIDW&oGyl1G4=Oq zSRRPmlC1?LTVrzBIDeXsO}J}ApnHn0LWCpP(K>NcV*Oa9bE*=gJ}P)g%()dFhUujV zp&OxWDEdU>|4+v3z^z(Ga#qcWZbOad7(<37tz|li9DT8d=kj-J=+{wQF9SL@q z45Gwyap@s7y`5ltr$J>iyiNE8kf6sAlrzH+o|DJK)U!5=41ire)r>H;yzb)s$KR+)=Mgm(_R$g3P2Cl9C z-8PFb`+J)M{C>P6v+vTJ+y7?Tio1yD{wtBT{cpX7H0@~W*EJuBBo-}Ip9NqLo$jYP z#rNePE{9V7<#IkKLojLR2rzQ8t(CaCJjP6um^dKieeG&Jff9ZAbMtGiMKt$~$Tf>f z?dr7=*W!`Hz{z5x4aJz9==<}P{?q+eOoE6?0&1y_y(~BT_f*-lE|$?0ObsNT=Wo#& zotZ(*7D>p2DzqA#ydRc%zwUpD-+b=(Y#F7b5V>T8AL-p;$kDn17IVX3uw^v&>g%qL zPgA&tf}b)Xmh66}I^}j=7N%SlWZ#~rRSW8-9&O9#SVq^m;JyOR2?EYx{?|FWxnYXf zI$V)#_6A$7)D0fri#_yVReFq)>sd~@kB_L{{@)(~Nyg^53B~)9Mpzk-@9qmQaIODv z>v3O{OKZoDxFLCX${aX|@!tof0S9Swfw)GgEAo+ZPVNN9-O^#4nCnq zyZuIC%Pw_G8Tb?On+DDtAF@;*lj{C4k7u$NR_eI{q zg)QTPzTNg-*;B;%EDSglI`+5{jrG(pS$D)z_nzp_04F%W4!3^BqB<&J4;Qo7-LtI3 z)b)N(ytyTv_ZQ*!`!(PyUNrSZ@FMJ2<8Ners^q0l#CiI(R$`iFKNdJR^->*0b4i!*p`VV7Imc~blQoSX23*V%Jg zwv+hZTl0IgQfRjSiS)BIy|9pVuZib*4(>Y&Q2!v&j(jIt za_gjC4zIWJUk}$$d{la~#YEGPl><_&^hw`76pY9tXjc>a{nzpKBl79-Z1x;3K4Xed zTfjus>UUE(ReC0?Wp{)n#^*@08rQTr55pq24`YsO{+GqsLHGOtxCJV1B5ML@&VuON z7t#bnD}dQ3w_5{)0^K!k3zfO|O2w9;MBr1C^gmwq;u51B>Ty3>=|sjEp@|2o9=)QZ z*IrQpD5atiq?3K(>wy>s-K9(`9kDrt?_oC=b+Q$VC_}IJyoJE$TTI!>17Ygo z&%>vmQyY~YPvak|AF14>3{+m!Md^tw32-{Im)!OQ-=^ICyd!YB!oL`skI%VV@C>`= z(JH6Yf``|w&)^@d_Kzct$msbRK8d08xekjom#9Ks%uISB2Nw$QKgHk0vzN?G-1Jmo zdP5YBvp{(GR_2Z+tRFIybJ97R{o?nAVD_;kJ}cT{=AQQx31@Osi!R2b0)N3v3#xHD z@&PxF$)QrFnMWUrjSuYzYdlLhTXBV+Ka{Y>>a9M_g+XU@k(D$BKlQdgw1R0N=X=-R zlvk#YqSgN;YQB2wXVVvPs6TFzCU7hBn{5vNK>f%(uxO$_YhD0tR#L_K;?ywlODWwJ z`o-wDOgv*R{&q9O&9!3iWdCSk<-t#NMG+j6Lp#SM1J!& zLn~X4)LK&cNz$y@>W8|`a;EP&GMJPb^Fr}vk8x{tHvMk<14S`S#pg7)dVM_Litl$W z1*+fq_NL%Ap~tWm7rz;ImB1R1mUw-t_DAdonR&0v31M#m(C>6UibAo=nbb#&s0)oO zK&~J#JVUuXZLxbb)&mrZYz5^p1}8M}E8}8HT|Wh4$Dhc|H{ZMnmM;8-C+sz^wl1MJ zry@&cCksG-!J16|rc;pD5>0A>oRW; z`Snhw>g`T}l{u7|pTcWda<`!4toOR{cPNR+ZG)A+n}(4^ci6&XVtTK)1}=WwN~Myt zU*6PHUv>+;-iyU<7YBGiRIJlTV&56mHY-+IN?{;SBnq zRiEk$5H!m}1%Pv$0M6;9NNoSZ{a4IS!@_9WQJTV2?I+P~wtk?`z_oIAX+yDY?@|qr z!WAHea3M?oQ!6un_xY^Tn6H1^ac&IJd){X&0QB)fpKAA!%KojkwY&T~|Ltrb>>58q zOM_dQ!HYV=I zS`?L9ohN3;?;4J^ANht{+>&`O-?1YBZ1jjS0}FSDPY)+7J;i5)0@875bPEC!iB#*| z%)>Njm86t?roHt{3qTWWJ(&QCPfuTDf!GvIk*-R~`Cti>sfd#NX1j`lW=<2CK`Tk^ z1yd(gaCnrN zN+tzn)V|QXd&s+2Q%7LLonXSUnN-iDr5Wi!Z!JtW3t?fnU zmTz(npe-~ea>rV;MArwmTNbu1EaCrJ?Z>esXm98N5QEhZZOSnde=Y!2H(++;mpgo~ z`uCficMVV4su;lS=P0+2u>J#yQ0a3d3P_K6V1J6m5OwR* z1Wz@B)R7)+>%Br+C1=z|ryW$b>|CP!I zS*ygpsJcEE6Gu}Q>ef02U(370Z+)M)Yud#=b?Np65LLGpA>bNCO-vkm=3;B(m4`K zeB=^&fJ+*B!xX7@G4NTbK;+V2Uza78FZMo*Y*`hlol$@X2N2WVwJ}XSpv%@zUt~%E z?SX~iq#*MCk;Dfa`kM)CaVI$S3t}W3Yx%r6dZg=`2fBV>Rlvy&x$X!&H-Lb*EmoOH zELt(!uLcZR8x9{$SbI0{>2Y$IA~yAw%i7r;|9cUTk#imAce&4*^W;4KuDL)Y*Tapo zvN~HKFp%luQWW5805gdX%&M?({L~E!Q2UtTjP!Bez%w+8lr&zAVfdqRiur}aN9};w z00ziG1KJwRuO6tb`-bt*TAHTx|D!`xDoDi+@C`L>B;< zw=Q)7CjU@a#5Ij;1(21v@OG53&w2{a&DqtrR%7CuZ$3cQv=@e-`G@N%4*^DH@7H|$ z-RPBRk_;F1*Jt0)?0Y%_kDXsyDfW_*W^XD*2CG@7XhsH71H}3OcNo#Xa(MvclJa59 zS%=~BkSo&$B+N;S*(^*r1k2TEjdxD|3kbFrfc@NbG{9)Slh{@oM{;A)8xl*IvWaHF zw#7;}6xt^Aj{oPQG#>ZLUsc#bVB>)9;L>lv4NPDhP550-QfqFNkVPUdmghk|AOF%A zn?-sK|2H*6I*!Nmbgi3858?p3?Z(ig|6D4=Xe-r_AIKxYJ&D}PB~BeQim?=~47s1g zKUwWJD`$`AIB+Gb_52D9f6Bm|Xno6SYSYzEBO;YNQ4_DMz;rT8!c(wWQ^2&<5!*pC zf{f{F99}_rhQTR1ZG>|6$A{&X%epJx|8CMDn20LIyA8E=ed`6JQ}Aub-{mWe+E#3* z`0|AMf`IjKZvggsq#aQZ41S^>HEI=-P6RbuT*<|J+*4{X?t91EC;HiHbpyJcOHv)j z(!7ADL4c>#ULxa&tnKUAKZSCS>fO~ZcOw5dZ*sPtxaW&iZp>m|1LUI!j1SRfeco-_ zw-b1dPuKAPIMU)XcP&&SM<*8oYat z4Y2<%kLN|$UE^P00T>(FJ9x^Jl5par6W8{TtV>Hnh&ml5svp$s>H#nBvbkF}cgdNB7@F7PlK#p`D zaunCJdMkPWMr40ZLXs4H&t|$H_Wx4+*?QWYwpJ1)+H>qp_~k+54G_P<6Pme)G8wo7 z%7g(Zlm4oW@q3Y9t7yhvM%X)7G|Vy>29c~;Vy+K~F`(*S+GMppYyw(q8PKyIo%2~> zkP9|vWoo}05JRaxezqTZJMQIx^J~(^j{31k{$OyjhxL~FiE~pXcC@icdIoepO9gp3 zs9BSAr9b_I@=E=vG6z+)>llq|{|73WA>C1V0Tl6u*a>n4~ z)UDqX7>{2FQ6oT~V2DCDI63>x56vgXR#VH)&q4X(mr?(gG%IJ4nMNRo+I`)`3$5zje4UWi2zZl_Lt)H6vFzsy+3SIwqLWblBDo=#>zqKrTHNb`* z6p1|Ujv=_Kr{cEyfs)-pciH={JR!7JI>0vBUbBIt=J}+lFSm&=6p@Sau__D_NaM*=*_ym zXH$8k^czq!6;W`sHpraoheWfl*ws>aKx3;r^rmvsWQro}s+PNU{f zngfE*`&Yu-0kTblFU^ql(?0xdRnV9WWw&Dl zc)sod;aHr)MF?MeQf8dRy(y4sb}(pPd~2DuZ6{>X5N)F$nG`m_+;{3SYV`axv}yR= zn(^E)eyvLWx84F54o|_46!RTXYg)}a6Et^-#F{(0 zBjzda5nfs-U8e@0YJg4&3`khdCu*|>@Gh|`xd$0`noFRp_afS&2Y3iX{Fs!FvdC$9 zry7yaC*CF!Xwu*mG}nGwS1s&4K(q`+<>YCBJcj-Pm%9o=S~!c4dBm6u16XS9$*97- zBbG6Le=+uXH#oRW*7eS;IZP>G1_{g?ms@F)tPuc+n*idP@8MO^?62jVD8J#5aG~5j zZK+o#Cb{c$ekvD&@!RKJn{Zw5JLp?_&ptBF#@YEsLu zHhZKt&!+a?{O)jw4$G3S?!v??wFNFjZTl|F55@qZd_B8oVJ^bh(doLx=t zi#eb_e{0JFkLUo8Jjq(aThu#{cL#W-hwk8&5#)7Sze8mcBvP%&O6Eq7<@xRgUStE; zNl~4Jkp1@NtY(tNb@vxC0@jgB)7O&Q>v)m9T>AGhQ~MDCyM&*`qXh1A%NFV_l9Ny( z1yG!0t&?%_ww!@OzPxtwGe8!`iC3tLXk;r{=KrYDPg?xCh7hwxnb;i^ug@aH-W%pK zE94z^+(V$Eh<(}!7zCcQ7|a}NLWT3%cU_ya#&N`g*X}*LMmGM>VH< z7N>6UC};Juj(fWl^TkCiXY*pocw#Ce{EfYsYvwurtaDFNQzwJ?tux zxYK3$&iE5;rM=?*mC9pRYm|-FQg%VZ3xzTi6*>9$N}m|9SvGbG+V@bN3_2-W&bodJ z*%baAjpBZ^JT`sN#an`J6<+0PWyGTH#dgyygFlVfFk#KWx~@UrufzFNb2kB(|7&*& zGk2Acx@@koic0UqBhHrUA0FP%3!<9q9Co zS6R{wMbQHOBhSCd8{g?*iWPg#j9NGfM){_1QvObcsRX5l%h%vd>vMN1)r;B9Ex);7Xa#fARi;J5)kN^uB~_ zQzpSJ5LW5xZAIsppAyR>d`Tx<7PueAByNH4x@jc~!IaVDJPf5|H4W#RFLV@T+~`l! zR_GLZn@Mrq|66P*gATihAkJfj0D}EJ&cCbatwK2w>Jq~D%EJs^fy6lI@!X&Q zNXjZg4Hsd$(*0P9NV?wb*_EUT`K?hFZ)8o5b$EYKnHha9 zMis_HEVH(ZqjTpIt{Dt&g{qYI`$>|h_=>p&*3~jkgsw|Cg(ruhDax%X0FJS?c3B!t zg@k@g=5?X{)iCuiVSqz1Nwc)vGlacEaYeRtnUrv$U&81RvtR=B8G{37QY#Tgm4XEC zYbuQ48@ls$8ONwendsjQi%SboAic?7E;UF%JEE-+JbAf;=xDR~ESVLsTy49;7pF+` zOP>sM`T}rvl`stYJnr(srwp1@ddv_*B5afL>|}LA{VDlHXL!I_(6jxRi@z2<6k$ki zJ76WbUo?N_G_Tln;!2Ny;wr}-DW*MmtRT6a)<|IXYT^*SK2Z}X&4>DDLrF1m$IFhf zz?VNb&QO`6c4$)mJWZo}2-+N$2v8jb_|Q6mAo`S<+ zP2nmbd64ujleP|142lcY_XH(oEB*C4wZS~P%AVhbAB`Jhs_&IkQ)oQx*L#G% z8!Ne|2cPQDfE8_4{z|KXOy=}G^xw)V+UR({MM{S1kTPH|RenYAZmGqp`eLsA%ok+w z(@&J(g!%T-)-G$?)(XTSw~3errR1s*LeNf`z2$;K3}v zfB5ZHz0E}6Xyy2RCml{G5XO0DG~+z1YHmU{qQJ&9;i|Kdizf|DPyLXK*V`F3RjcoI zpeRWkQxS-*jIAWv^gcj3(2jAa}8Y1M4KvErsr9aL2`YG zX*52H=S?BtNv>4%V^vR;7^&>+ELKO7ocbz!cGQPaxOrO@>=dG;M+SprUw*IdP7eC@ zZKy6Z*rI5QFx7lon*O8C#U9v3F8sKDuk+ZpM;IQaVFJAJVBx_>#-I9%B>Yzvsrj-E z{f8j8KK?VmotGgrQg}_3MtG5er!9EAl|u4N$+%`O*LpvpTbfJf86XT~qqrO)q4f8u zPs}V~uYgSxE7?5d+#`7lq|+*dhFyf+GYaCbA&~As*0-S2mxz?U;E6V6u=yQYw{T2A zS6e`@$CJ6h5yF;I2iiHzF65;OY%)l1n$KwD1s!lgjRC?#qkVaLvIxPA=2qp2=vXUr z0X0>U<3+ya2{T3)lqXM4uz9(5Rf~M2>k6S=Mv~B+_^4B+@8V3x9Y}lAZR(yw%-bBB{+?U@CZ z+bz~B50K)sB*#<8cVmOsC*AlWHD+hVrsUFX||;6-Mio_)babGIL3Bg2NB{Fz|b zi8*+KL@C>it-8Ha@K&5z_tl1(bZz=WgD&U5^ORGmrH#jKQ~DU;>ffU8bvWhOFO@7Q zFC+|UR-y0CJEB-kt`G(eC~4p8T}E+NB*uSflouBY#fPK3gb}1esbPX;HDD~ zyiX(*9%$7oJX0a$;<1+;09rZ&JyMT=i=ZUSkt7a@;KF?Lw6jU-cb@mINR+xalu6^0 z^8P(3U329^iblcjcBaDQKJf*|m4wLQ6p!rxP_M!(`q5Wa` zC@fJylpzz7@|2~=+4hqo-`^p~7Hs$7fVX!uMa(zW!|>LcW}-ma^L&3d@+NeH6s+^O zH4?F(pn174%WPgFQc|e2gbDJFFA`(}Lk?1Jy@+N(zY65opP($7lELJ9f@J>#yYSN> zO7P^QyvHOI!6@1_YlC?5FAa8EfTle)d z#R>|~=4DU}!-^Lrcr%G4td7_%DyF#-rI}bFaUe)OF=D2K>ct{wb;LBGiEuCx18`({ z99@8uC_ybCjpN0o{(}3~d;-GHn!maa26;N!&3bq|XPBC_-5Bfez&_m3TsSVnUkSEN zO^7D~0wY696T66p2F^%BiLmGUfgP2nkQ&1K(wd<>4#^xA;eMuR=8=ygWy1-fdzO3Z zf_7^C4Ns5iu8H22E0*+pZ$s>igTbzj;vZ|00Ba^dza;BMOkd!r^>Ts-px#;MuLE*B zn~$0d;eD{%Ew;m~q&p7TUGBt`>=Ou)p1k&-utQcE`i&|4k=KGVg^Abp$^I5SpJ3AFhCn;U-YWH)_0XjSXYH~$bM}?F z)$xGHO+7X^kS96v_yCLRc+!g(#@x?~#?u_aMevll1?Cy}#Tf0UcN{*R&UP_`(eoh_ z-|85o;eOcK-QBR+khv@Z+b*^0G)vLH*cluB$zRPeYNzm{eH?x(;1@q${<~Vn)9{Z4 zQBg{<-9nK!xlF39`sNF^r9ZBR7;UhtwRg?eEgy>qpuTS{)6rp9Kmf&z-`(=MaurI9 zcKX1f`I7#szgi+>X<>0q*gor8jBUyxWW!&FG74-LAkvinGkN&X5G6AuSVrC{ZX`7F z?M-ec&Xh!u+<^u&3pkR4#ok>Ph&Jy5Zoy-;X{8D|3*|>A9#kt}t@FBr%!o|{zUK50 z{ynMPw|}VwEro0(QiHK|_cRwKAbIYkkzp(r_Wff$^kBO^*V_;JZ`e>ex%}Fnri#C@rwbJTSL8-`q;bIJpjOm1g19fL79H-4g-%ALoeGG^3>Yb;xyo>478w-hXFE&Zl!tK?{U5hj8@OWOd;*~rA#Uq&l-E{L zddgqS-Q~QbQgLAnI02!JMD7&Br)BuPkZt4NHds3py^yz0Ck}ILsJYtrGI%6Ln}+Yt zH|U9D&`c&62+_f_z$`gY-;$a_UFQ3!|2wW&pDrwFJyM9gi30{ztm%($IN{il4x)@e z-*k=`gY!G1J6Eh4Xz}xIa*+W$tmBBCA?!#*26Tl&teX&|F*dc)xfMda8=UztfrKK| z<{6giQBZ2_`XQ^gj2hrHstY!7hQS?TBTp5mJB=?s^G!lJr$*&2k>7<(>KVr~pJVdY z`UZzg0cHvMbHZ4%2_bKpGlG4bAPgRyiRGvh;$JAlA!cWX?w^_;PLMmE3#s(Sm;RGi zrpIe9xxI!pnM?`ZM$_Kw)br{TFkh3)6S8ccExs$~SjYNUiq~47yQz>ZP_c$3^fQXn z;qig#)50DVNDN}=$zd~Ap-Yo;DcNw6bGK|18L;w@6HF<4M@?%yS%Uq+n;g~?sdav0 zFo)bG*wh&)`S!-Y^qClab4nT0w130&J8Q_ssj&N3u}DB*=BBO>$XfPx^6+)`vC^&u#7l8n1;ud?&5K%-w(M5Um)N} zExwF0SHZP7h@sNfrv0-?Qq`{wzd|@5;Ex!1t$gC{(%~7zSu|juZ*1)kZBT(;X?6rH z^<5>%eOYLYNG;D?&|_BzZJ)`kR*U^$qBKt2KApivCJQ!Mg7^U@w=mw2@GCSEXWqpJ z+A|-PcapK6j~A)N*N$3eLdn~mr!B^kwb#1$kOd!GXufu=MI!o7RO_2baTjCwGE{wpT=U7gfz+3JLkB z&3DUZvw1q7Zn@C3m1Ze2yaCSA@yA-}bO~4)ZM`>3Z}();RcFL@l%iy@@7+Gv-KGIc z|IQFf(<*zpBMR*nUF7Db)fSLZJ{5lEa$Y`o+tOR+s{FIfLQ#`Wyh?f-xBNvw$tf4v ztU!e~Ia+U|{(0o`&UTc8puY3OymXote6P8;O!Jpo?#PYOLz1ys^XQ9C$9k@6QQhO5 zigs8knza_om9xl3Z$Mf;Y*0T_`C-ja-0o>;b4rK@5yc%)1r;@yPxf5p@j=YUb0l%P zB`1Vi6w*7@Fd_Pl8Nr%U<}rvG<@W^zR168@E3h2zm~NY@B=L@{zSRiYv@^`X{d_h# zU{saUal8@?%~y-v%m8XR>#9s9>n)hy_)4pLx^+)<1)^DO%Ee~cayG$bgM7Li<;Fh| zyIQ}O1n!mFHHQg%cw0Hly7rFO0e8|0py?~jlzJ)ThrV9dZ|W}$ZFx`PkSfgg_A*(W z@<}JQ$KFoCuzol-hac$SQYu2fH%Uu%{wWQM6v9cP;cC(a9%m2e-h%U9Aqk;Pb6RaeSWK)q)J z%q(sSb)}pP!~{mJGT6_+J1U%&%^Z&QXG85B?!W^%dZEI3M5+@kj`mS|DF?h+Ro!rz zWN7AF8DA7(BhEj$gr6AtFlL)ZCO+?h9qrvIw^xPpP*ln1@>D4UberF%yR$6XMF2P1 zuiEfEHOvgy2h?A6mW(WwQUsRG95(itbET&aGGRAE{==hkjt8in@`jUn(OR*LTBXGWcxWH__hQ)WGnUB zxjC$a`P~VHGIR9qZ{MK#$`Yq81#m_pO#KBV z3=$U=iktp&X935H>?y<^F>4JA6YzoaTSzNK(Mf@RGHJPf(_sVz%oza>aKg2*5fZ12!p04H`ifgV5Qj;;yZdQWvh zn)pAT7reEl6?L^uLf51t$NWb0bwXm2&n2>eFmHbYE0TcE@cXGQQ&`S!n1_PqrsdO5 zDwS>^7JE~V$2)9WSXabyt!3N<-0!%LDr&;RgOfpP#+mAs5Eud(oe)yO?ZyF2k_Alq z8qEKiR#uFb5|b1;h9;}4J-RKzM0zHfg@54m9Xt$hc#Q%oeYI{X19%XG-=>VY(t;G6 zev9l>V><8jQstUs`{zGJa&597=P;%j?4qDKaM>zh+1r%@_7L9@5iWD%Iwcr`+YOKg z<9+0r1HI%TgsrNXOwDu&U-(-vll$*b=tvLhP5)+n-E_ckR&M9ya@&UmtPEXE+7NwA zxJt;F*jWX`y_=LX$*|1^_KkShIfSuSd^&_n!VOt$$HhZ}R7>mf(M_@zhIg0-g13Ixrpp7rz}Xb+)h{cqwc8ZAi%*-h-veslxh!lz0RfBJQOH z@q3=S%`=N50Z@65oiL`j@X^6`w`H_JsVCrDLWoJ^xmT%{dL--lZ4~%0$fZO}p4wv# zC6eZ^4W|PFQxp!tM0zn3y8(B+yKJ1BtntZ<0A+_6ld>5S;%thC6Lf$Rgxx8wdF~(H z(VZs1q3_OuJ8y#E^CR zC_wX-#7t~nlxSL2W4QB}g)hJE^Wfu7_x?MWrdfp^()YZi%Amg zwW>%$e02?&cpELJUU19K7b;=BIvULwC|cb-wO2ENE1qUx9F^8`o+i_zs;BEoD15+i ze#%Xm^JK%9`LpE?qM>UPVKs-Q{Z~HKgaQS9Tgg)=4-J4<7gIF~ixM+_6VET|l_|w8 zCd?0yYAP)iS97$K>FKKsB&bcC;{dMCrRl4frg>UXMhm+5gC3 z1nWHj5-AktAv{RbbpUOlS^aTG!rzN6c#A(!p@`Ahwd8vOMb^|yubp+`ast7DXe_lTzN{s zEpGz-^vpeN?OF&Gy#N&zZ*ar&5Ez0O&p)K1Rc}$Ks$#3!|AhTXo%jg&Y8&QL%Snlf zZic+&s<4Jc=S2Ko%e}Ggvxsae;OAY)+=l|9b&W!~Kpsl4@jbX~253=%@&vglart5U z2xE3VLtz9ew5jq;+(Z$VR=g^w5}OKwQe66%;k9jd>$Mcpm*eQyLsH$iu`u$N#TG>x zSYIlOilFbwkhAP2g)%ZIYG;}V`yiC(G#Q%Grl*k^!w>7a8Qx0VqKFOVqC_YBTGgL8 zCi6CMXgb*!?T|6eXcuPf&8iQ1bf`rsrKFEecs6aIj`6Wj#>(IDTbe5Vj@h9FM?OCn zb>%&v0b|?jDJ}q$+3*bJ!xTC^!@Rk+G0N~4E<&&G9*$#@%E9Nv#y~AR{11ffC5wT( zWH5_~J(Ak47JjUPN2#WpEVCDj|N5^f2iK3Vhk9crcV1Ia1ctM0Tx?w+2O_xdSb^Ii z1PD)LSXF}_=9}jCl-5|v6&=d&-YV&ozM7hXkMBEI8gHv5R`D5i$at#QbC+>E_HQ# z{8>(P64N2~9~nv*mC=&)Co?2<%G;r#l7c@=p?O7Wt49j*ZAET(3iUn^a+LPQ%#3eIZz)KSK}{8f7OyzMX~ZM)2X;7Kius0x z#i&Q5cF#AF4_1m;5LQ(WbeN!r-#EGHFEO2shpN4(~KQc=X9N9`VS&aZ@`$Nd^|;L8X7zzak`KiP-C7J{D*zFHDb zVTG3GpN0?2JQU^e0Y5*MGQoK3YXLVORIAHt)2d=^M14`At8~uwL&w^+jQaJ$=HZ6S zTf-YBlR+jhwX_S)6INXW7pkagEB1|~1u_iJ(DgrGSL)^Fv_(aJ!bUeQOAk!7XERg9 zV^Z`KOEYQF8zq4(uJ*v&L;A1NGV~K)d=554+9}|_3DqMxT=!V;in9&Kr4^*6S=Moj zffe+IMcIRVh+3pJ{yQM>HG3CUEtU4PRC0&)lJw2JrQFp&{N*ZAh&|!j!1%=@4alJZ zwSf5n;DC*prkyC?S&WE6oTmFHzstZ!N8sXq*DT40{q2Xx9qTVyjXUK?ZQ$)EZ6aJ- zhT+Q4)s`%RO3`37S2=?$;|%0NfF<(rp&`43$Z6c?B^A9~6Xm5%Eb7mYxt~XR>S`Rt zQP{Sm3sp0 zj|?^GiF|N1u(q4->Vkw(Q$4SOmIAPVajAP&o1XsBJQK!7c~>v8YXT6b?FUKuSO#nZ z;40LvCv@`aQL!)%3^>zOf+U=LP?L^vvF2)+TTu@}n^>4_pwR)8mSCCHy(TjN~w|$JA4JL0LFMS^Gdkv)tq6|EN}f2B_)BT33Ef zVpH|RV?rseQu#fZiKg?8K4CXR`C-im3hwfRb%D=o{C0UJ=Iwwu|Lc`jK*=8sQqUR&1?}6~!5?HD!JR+AiXz|eX8FSsL;X6Zf6gsp-F;)G$&CY$l$#?Mg4c!#ELuU9w z8s{<`rqO@Go|+FUAZ+IXrkr2I=ODVdP0|I&92eT^D_|7xFF!GnrdKeb^LCJ$X9N%9 znA~g0Aqk@;!n3ofKq^Gp3D6$3wwJ}kU{#WpssE6Duo)i80^Dsw6`<;zfw^$S6f9RR8~R)~+{m3Fes--rbDYh0y7E8I z*{hnfeeni)J8(c#^FSYhL0d@jZgJ6)r1S2?K<>sd+aoccW$#SV1>Si<}A{dwz6Jku{^BUssyQWg|09SEy+k3ztJ|zH@rrjoC|>oG!(r)P3TXs^w7~ z@Zwuy@!rO%04tBoK$qAvk>s)f5li6zu~j_rTz5R`2=lWZb^h&zc7Sb+)KfUTsb6!z zLAs*lK2@kI#$-46>!aB-Ak^l z#25HpH67RQgUjXonHOt)E*<>7caBK`Loy!^H;=tJ@%K4>udcz1gKI>@&`ON{Kag3W z_q^L%9Bq5%clEDMUzkjF4S+9V>fe z@4d43k$EV}mX(l~eT-v|Y{yEZ>`m4=4x&OF5>ZM09!Kxb_xGRM&E@g99^<+m*LA<{ z*VCzRqU+mKbzk-)o~x_M$X34>zWc^?fD%%h8h}uFX00Ybt}Zy88=<$roeZ*mIQ9Qb#+&S z+3Gd@v&O5jahh|FOY1WQzzl=03PQqFW?2s(p4)7t)tRT}q~@0XN;T66(&x#TpWegJ zqMr>HtPg(YwY5vYb@9|T{iNR$5MeY1_)sjZKSfLRENB4>oZl3!e_(m!GW3quCK`%3 z%Tw%su`yJy^Y&Ufl311^Ij+7na#C3^_68{1?pWgwuWw&Z@N1J?D%IPk#n%Z1y=}Wh zguCmB)R9i-zCw+yh^twQOELm+8BG;6IV?Snv`Y^K;vPO)U>Dgc*xpjxByY`25}dMm zvy4%rbc|ceTmDI zvBRy8X3eCl~N@NSL!$aPR&CNzAzeVG2;;JN_sA zzV3+CDJM|Xlfd6RR8t-4Chz=9%=&MWjr&SMkBENe>kzxLewU)n)6{01N}RRB7AJ%D zq-<{hQT1gSIEQiiuvLRaDLd(0=@nL79j7+$Uf0JuT7JH;tXZ7)-8?!%YWdwCNGU}} z6}QwcVdl)@Tsp^fi$^W?&0dnT{2hgsY|N>A(}gYEGc8ZHN^o^o>E>STvqdj_aCePf z_?h#|<9*+xNsT7)%gZiB)AZd1tRj&6A@7oc3q|yHWoBCKNk>_($zmEL$+gCMirhp9m`vtFD=^tsXA!%K3XDPQdqH*o6nmDQ8LIqX3t zS&ft{C(ur#6~py6=^_J*$|#A%U8j-~+=LKp1Q zZpXWc9OpCM7TLJ2y<4Ab!Q-^3GdUTB=s@)yB4RAa?yWQly#GI)ZYO6E*pet&$W=I% zco`|&>~_%SUVXF*iK*Xnb3nQQS5EnsoD zgl|nwyYr^MXTLu4j8(~Tec}u99`8Z@L?@F@?wM|lD1W6><^X{X%M?M1fq;RRmvMhq zc^@q@yXQ8@m5WR)$ygW1-uUdI@)g{!4-B72PHu@HZL`M7+_*CseHFg( zjmUYc${e8etBvDK9=?x%9G(AqsXx32(-1nLF=yRRwR`Nm^u*=YBus+!|XS|j7K+}WhQW4u5 zgzB-ijFU=*S_2EPJTU|O@OijKhaudxCL8^5c%uC#)EXSM=dE`&GGdGx;lmP#>xJs& zk0zEL0<(M6qcluDKS0~ltjXsr(#;af7AwcNJ9vd31YBE7+#OZVDs?`|14 z?dw5iB_7Po{ghYi1a%IJ>OyoMcH@F%r(b?vZUK9f*-))nk$ItrW@ZD;F}9?eeUrY# zlt#!#mpI7tzB7eh^S-p$NgBn|AHm@>XVZ=<74QCuziZf%6Z`Ww6X9LICmraz^^Izs z^L~az6e_~Bw1)J3{Losm;5Xa2lM3coTr?rA>8f-YbN-=a2mE8$>}ZV2=A?`qNiW|L zqYr;-ld4{FvJ>gZ>EuiUucrCsxGNz|2$r%pNlLwg@D&B40zK8F+2o%56p4Y}iagOJ zeOu?Zy$c6G1SQFf-dK|%K%qr$2oVvSjqr}rFC0^R$c!!gbLvYGc}_{-V_V$f?L)9M z2q(94HutGv++r2a2$%->#h zM=dL!^*-PkAtDV6`@2r}BKq4y6SJn{v4<66OTPvCo9P#qIO`{}^xg@Ls|k&;k)5sd zwVZ$S)-rdmX-4{3>05|7-Bwo4`0i^|p|s?8GY0qW4wS&#xm&HSKfRIn!AyGNJeV=%MYCjaEGfW|wt}XGj-yGxWqORoF@bTK3 zk}ym(sl-EXrsuRE+R?=Ist=VWY9}aVM>=BfIPKjZv!pnddAioQ_B7bntP~({swZd)b_K z(MR5=B82bY^qwi>giRN#j9e1CJIbQd2~8~eUD@JX>9_QoCAC(cj05+xFM!XKXO(G|?tv#&*D=%1)wv;*XG?*XRUtISLM~CbeZ{pSXoP)mrq^1r|bNSJ; z>J&CR6R)RJDkF?_Np2D`5u@-= zq~)Mj^uYlKgGo~o!C*Ws_pS(it!7%2;HlY=WdV znf@8fIUp-fXe27`}1eHQXR+JtKXz$o5=^ZwC>eBo0 zR|~{Hez_|`IP5Fh#pptv_7XLgkVG*2C(Yo{l97?^nEThwAZ)jw5;SPcLC*gf)e?G1 zgzF|7myMH4DqeUvNokmmKhlZ9ccoK?&piI=ii~?D^YZ&nVoaj95A*UMJP{%wnGQek zOY>>7(=aF2JQ?KGj&eSVxXE=TkM#KBKBW0ba8GU3McG{3`dm%k+)|fXoJ2FXVRO;6 zp@0-=l=hgM@ue<*oS-(8U9`W8-pX-%0SxUB_tqtuCospS4M713G$*z*LhYRqir0$a z{J%f?gN5+ri`2e>2YC8?1=8epj#rlgk~b{Tv+HxZF1k!EB=v_xZOPOOK~gt*dKS%3 zbTah2LkKge$}~_u74`jnwe#vTiMKV}_-3h#;wxc3?Y;F2X0IX~QvM4>g@!y`8P9RQ z8#nrSj)f7i!pv%lvI?U7g+(5KA)9$n_cv%dtpz6g(%-z z!|QIiC38x~yoZz4zLbMrZTWz_%6Nu47WlYfAMTTst;cx3v7lohJ`VhnaNJjvB17e4 zPVTU!$PGavv3zK9IpvoYB??PIQ#kKHHUmO|pGb`L`PnOHs%`v<3bKjY{Cbq5JP2pe zqD@$S@avI-lmgCHMOhVRANIs?<{^eet0bI8wF@{gO^T=Ej52P(CR};cmE06%oOcxL zn=#5H#*!Fpi%ABzQp^3?jsemZ#EttuXScK_I}KxEuajTD-q^8fYP9M-mW9fF=@)~vZz55j^421c5 zo$=uZ5+xLG7Zyy+`DZiLCDW}!a(`F$O;lWx2}qF$(M&$B`_SJLNF_Ef!S@Az*L^J@ z*ga>-4a=*-5T7UhAynaM=1V*t|nh(+S+$ne6TjQ_5hL^dmBwunOB!T|YQU;)tC^aOS@r9s3f#p5FewC88#@`UUBt zrcQ*I`0sfVMPeFD5UL3=>EhE)Q?pCmMM*uHEXzUBrbBJzKBiyNvLnN>fLK|PF>EVS z=HZ`VUfq3@d=o#EMtEVN$#|iL~-lxtiT2oBP*M!Bf|){2yV~T1C$bZGWFe{Ml9fd!YFDJmSx3?=#%E z?gs0<%#&m>nt{7(NGn-n5bts}H{yPXvbn&{TZiU}uKzu|Z)c&Xx%S-P5K6NZ@*}@X z!J>jwRx913`7|BX2(7-&19R?FtLisw4u0!(6?}Xx@4gEqlN!l+N_BV4 zH)TO>PraG$h&*mjdiN_8CPUIUgebArQgyTp8O@gbvc8rV)_e=%JXR{((gVT_tUkkpGp=p`7E8xiZ#enPU=NnH*{OdwFyF!d2+h4hV9-V zI7SyX+g)XNI#;=@FCV(xe0HEs}7u+(7*5NOr4ougL zKmGy9n4>pctU6!@(|@br0jt9_Al~G+K56^uxAjm0?}RG7Z6N+0mbk(7oAw7`nbpT zS*5lmk^*Oy1fBEkmZLi0V0|peSy3-)Z3@5ZY%Qh_{z52~EXA!`t;#9l(=eQ`&rjdd z2R{dyS!?c>t{aPi>V3XyZc*#O$qzdadCgcOGe3mTT6b*C(tVj`gqtve#44XZXCE^8 zW2k3bdT>Z{fOH(;2jw+1RPbO~k{7CIL?4&QkTrku`_DStjdG(li*@tl_ph8gIG3^LI%bs9#&+bwqpJ8^_u(5YOREDXTbZ@oPzE2(Gx# z8-VP6}G@Df{V2#Xg5&c?nXPq|2a@vImk(y?h{W$3Z5wDT^uk%%-uRrZkRMsSI zMRrV|#wzs3miy-ICll2oW-`{pvMB#eg3^?DJD+7fz0aT9vcX&(YOFeVN2ycw;e()_ z-G**KDhX&INVnZbxDe$xd)m|xu?82QK}!vR&VvzSkp4@>z+Du&UwyA^*868-X+ zy1XVg&(XE7o1n~-%jKOPoK0rlx@Ol<{6zB!QR=HSDGdenzjZlI-n~RNVKQ4%zvm0@ zV?^F8_|mgyXZl#FPGl{d?lbLxs13dQa^>2Qc)4lV<5%6hEyscR{zgRBVYKE03YcMa zt>QK;KKZaYTUni(aPbMzjzkYX;Q#dExSQb}{2)+jU}4aJ(#G6&lhUSJFF??p-3Itv zXK#F@gq$?R$}OZ-CiKky$Ql(%9X zqs{tsE$+t2Wb^JDq2clTUufK9NU$OCanls^H=0*1uH->2XP^trE;!dxozwflbrA*6 z-lO>Rev9Yhk*i(#bPQ2-RAb6;9}QDxFI2tU4I#BOndX0qqE&wI>1XTeuEvSMJ%r#g?+u}G&Nc)s1NShS)xU7L?tH`hUSQk?M6siD*Z9%$q6z9sQ&QlV z7v4nN($M<+v&^|S@LPlAPHVd^+QqWDN@!C|+^`l|u#U^bhwalt3Yf29*;g6j>t05W zr?*`M*2CzKU$?RbA?3d5H%c+-*@J&2QCD1GvhfAIAuP=&e@2xHhIDxZlVQl+z=eq* z>(75txnI-?!gEQ~NWP;FW`!GXx>XLN{W4Mu4^{7OQd)t7+F@Bt6d}oFx$*1(J7#K1 z+<4N}pWPXFG@?h*_uRK=U=d;wcKtkTO{6aAGf zTbeb=M|A9#T_yQ)>F)==noXo`RNKyWpe?+=2Ha1SP6EZ{1xDr1eQ0p1Rb&pZUy3Dv zcY@8j4zoNAM8Eq|)J=@g`-Ifn`{0Gklw=xFWD=JbF49l}?@eT1W~aYl=4TGmCnp(Wjni3_u?nb%^hqsq|Fzq0hDz-Q&uTE`4ZN#;)_i zFhx38?Eek!u@|)??Y1MVW03jh-`IMB9d-W(w+qoQPv8l3fL){#jZiAaZPmXYy4W2E zdW_EYv6*#X>Gh@q@LY(%r3rJ)4&IzB!Hr1=8us`dKCDzFm?W1XN(82O8WdDjA`G*0cSzrZ4%~Z zAHvcI&;R%>Y^x7VcjH1EzRN<9Z6miuX4EMj3hXXxKTQHyO56fzv!889bS*fdT@fQ| z$Esx@4jTTFwNI1Bo39L@pI@gVa!B*1RTA+~oRm~wctX+)U7T&+qZ`2QjH`Q6+f=X% z4CzN-jL)i>K~ z>l*S^XVgZ}r_i2W^J~<>*ltl%S>c<&QNq)X6+FJkG3V6XbERqiU@uEkV|{ z4}KdG3vbi`Cpe8J5bEF32Oh^7T;v}2XLV@;{~MjVxUEy;;p_nOLCylMoap$u6m3_j zv9ws=5z90>()(jrcA`1652_=C`lWEjcnRLTtbn#fxUEC3m@TURa-s{L+tANvkK4uKZS+Rqx~!7=H8wobJ#|cJZ`@>rVrt-C?zm#kVHd}YK;xE zS=dlz&p@JTdyqp4WDwswe)a5Y|+8PkYCyY}{pvF9s`sN4Wwk&~z#&>kk{HKOJt`@fY zznd}a$9aTp3HJ*=c)hqMvj9o073PwV99|MB6(q9ApFE8f*nT|qclxBhi15z9yhd?s zF#i+(co3U(D!4tXNV}T*FdyV7Jhti)FnCW6MR5ti#tJXURMSi&kaKLI&dmp#tH@W2M2L-$D_VE?n!kMyg`;B z0c*aI|A$X%eI`XfPZV+d_s<5S>drgnmq zJ@hK=PqgPGlC>;rOic(}hND2ySL$V!%BnBCR0qrD8|psU)wN>dJ~FOU5;psW z}8pSFTejO%&IoPN1gP>krh*&-!L|rKtbygUkmeb zH*OBX;$pn8Th`V}*y-lyv*j5&RjXURLWV{^=gh`mmnYnMJI`k;#n=KaPxQE7+cTW* zo*Jy3=GWX}OrC+ne#f6%Mcsu}Kj_Lq`c5oFeEjB=i?&MB4}>pyEgpEaKJ!L9l^DM*r28;Gt_;Tmj>TZO=Z98yJrPr8Lj< znM!Hci4@9z9nJK9FF<=!VGmp94PcExV=xr|XT$Ju%oo14o3<;2E*vr#npn5AKT5V?AR$^!;|C; z{p*c{ReG;1%Bhf|IB_N4IPD;j%TA4T*eiJ2Fz2VzM z{elfjL%h>rVnUCWM#} zFrwW_VttHW8U8O|KG{a}!4>G9?W0TKdxMC%ewOa<(McxKIEy~dIZ0E6U>0FVfu*i5 zy0b5~8)aww4OquqTh;nJHw*mKe}A0)Xy1(cC0yLgx$`@f%1h-3$B*>f8_@LAkA7%6 z=Gqui)L2zwRTlFZG09>f6)e=ik2#R5+j;2iN84V2okU zSCH5I`f)UtKFiJOd9uPQ6ET=CZGYDbo=8evu%j@!8+`iL-}PeDE@52XsYQCO-LI4D z3Y}&wzTLTtXfP3g3H^xv^p0z?xD~^*m?h~7;ePuCLz*>~kQh1pdQH1)X%ka{3rv~6 z*RY?Qre8!$SoSz5#Os?HX_?-=FNgA#&MrE&XR;m+ohg}QoG@$b(X*HvS`E)|PiB&w zD9X-jOusz3SLgAZQWJl?UrY%-G{?o7CuGuR=z=TOpd2gpM(TJE zwDP4aJt?T@{0!tX$0~d_j-7wlkCel5{ModxXy=PUp8x!i{EdrRgOy8@&|^v8!6^x; zhiQIWxxYkKohxFZ?AW)kaW!B3+Eqzr80_;--|C}2j!X2zVj{cSrC&L1p4xQ#I_*I; z@)d3av{c=PYOa`VeJr*VhYXE+HbMCqSc4_nXp@7g6Ut5OL(dJ4{bn6;5iauNV(hPXP zF_XoDYDp{!z%rG;R^w6wHXh%638`APs#$(Ns^w>n9CGb) zqUo+mUWEZqUr;4E|5Xu3z{ok+>L!9qCL9gYUlvV}eb8N6%&8=T8dKw^y9z_2wiorw za-sl`zSqM4tHr-t|Gu2ZtMe}}MB`<9zoa&iqt8ZbBv-ouF#25hPa;dHxOmJo z6g5VcBx}L|x!=)0NJ3z?d|o(pKPYH5U}&AVz7%9#0(1m00iNdjW=SlOsHTFST95Si zV)7uR?uSG>zFb=VR$ceIs?{nf>gvzTN+2xZPzdEiZmqggs%O7uWR0)@$f%s*Nk5z_A5@W6NJ05%#ZU z=w@1f+zW+M6jRyVS<**rq!efpv|^I`*}FBxFHdKkG<}9WBDqR{UJqT2ldU{cad|pc zx!QFJmUO52diUd?-3e4-e#Bpp4%f>`ZNJ)IVNQ2~8FnbhU?%^R9Tmc2{;1JgnA;v2 ztd|6B(g$he0>BH>e*|!bR$({*x-|?7oSk6EC>&iha>A4*DLU$N(JGM~XFhM)Yeiny%d+1ncf!E{N=K}o zzjIn9PykI*6)+%vv?)r^SFUd*LPqiC;a&@saHYs=>q)YtK-;<=m_uhQPN(xx}JHSoL8WLq1t5e|>MMsmzZ)k=Yw~Q5_e#w%k1C8v~x@c9ghfU!F_H z=63!aM}>BV5TY{9HKm{|Vu2M-yWlf53#YmXddd)A70-xtS@=+YB9l72<58dC1GP}g zAVdYfEx$o7(CC8^&`4@edU-^T$NvKVQwv8S z(3GOwKz#B`8z|w_YJ|CMpU~hJyay&$k90SArbiLBQ>r=385?Qe%q(nF;lo56*Le42 zStAUH!|AnL>Zq;(HzYuG_sqZaTE3fHP)NBT=ZAs97zvUoKG%Z|caqODg)}F1T|c_q zpA}2r8wA)9-ns~;_-NKIjUOq0BbOpvc?9nbk^H_?Mnr-^JOFqkxLX;>*8;W*iOhjt zePrAxbOV0Ae{%Z3r7oyQK+j{AZ%?9^H#FF-ZCA zOCq#Lzm>hfPRlS102S3~JC=8MxAf`del`1_s?BZL7jpAptL0fCaSr$!v85KBkK;4G z18^yj)(8+uoxL_z^S?RGPeCZ6**uobf_U9=<>Q{2G_Xij3@;?E#M|3rJiX<6ST?ya zalna2E;qSAW8pC0$GnQL7zy>s3O%c(maV!y0D!yGC-tYds6Uel0H_ogTjuOLpJLWY zzsp@Nk1I@Vo5=9o8Bm#p#QELbrqAdQ(@$%`s$F9X{D?1Z8mR?9r3P;gsJzMetJb}- zkDG(Y`vYtKDOs4 z;Q)ZI5DS8tIj0r?WW*y$X_BuJf!5Dx-ckIC+|qwAq@S2Hn(doKZ{+B=I4G%~Kru9- zaJtW_{;~TqX{8UJKG?-pc|xIBWQV)}bvV41Q0aIwR3IvU0=r<9=Gw6g3JRUeu<>%v zK()QO`$yMLB^X?SD;~kx|4f(VbZz`s$Zyk2=9=5@c2FA8A`q(%O3M!{-!3o;dk=bq zj#Fk-G@bA#awf0%bUD*zs*lf##6v0Ul%WF?gA)oIkbdO0M7#ytt&sgyV7#}o$&97FP{ZPq|H_tk&VK%R*mvWNWW@){ z*LJ7K$4Pm56#I;+8@JC6tmgh0A8S(J;07R`l-$IWPZH7GVl9htS28o0y8lDLr=4XxukulAljeeg#5D~zi0WGfrbnoUsY!ZyhF#h0*Yd*OmVg}e&fP&l9oFZU z*LVVKJt(J>Z|T#vXpvA<(9@JF0DxqI^A6$*ftfRDE=VaM13R*W7+ee-^LVMHEcEqA zs1Q!>%rQyO1zJr4LQ0rL5pbK+pcqW`d|o*w*!^}K+V^@Zrrx$qu{xc^_c9g1USJ4T zk5_@j1-ce}L;|#%7)WIk0Cj+htMYDX%WhZ}y%Lb<&IjMKMA_0kxVgwra%I}{Rf{fH z%jrz7lmOA;1N{3NEmAu9@xzb5pqikmTD#Reb}tkZ!@Eg95l%5ETf zdrmP=R5J_!rKQO#iZ)Q<4Z-@Ky%W@7pVC&1+n4WyNaX9pImmrCO1To@OEJyYi9+a? zC5>~ofWjllHE=P13K&{%KJsbB3fB8iDd!q9Fx;|hK*xP}qi)eJsZRjhL(LGS7{+q) zH&5G2Q-M>-A@vBqZOe;SE&h{mM0d3qxGsUUy!`4KX3QL@LWpd-ki^?@C`wXM3l^gL zgy{3n;v#uh!rR4_2Jq8Ngp=C~$ww$yRC(vhu=>bti&T?8|3{K6InGGe>Of1>mvu&LY{-6n=^#o)~tO_&3M<+bV z(Z0)&RsXZ)-5>9?BmxN;ikEAr1%N3q8ImTvyz=1mNOaRVHTUs8>_J32dcR8OYf2}s zj4$%!+B({Ecvj}->-C*uVQmMu&A)i8+zncVYjXD)tm2k*fbD$RS*`FwO&*7h&UQX8 zkGPCbJGYEj)~ik;nBdz46eQcwLqp14>uiJ(1Kssn`nWVaho**X+9yL^aDp%6*;8Y& zDy(s)sEBOuKK!8jr@#h_d-3wIqCuS;d+8_VU=Slp)Nre`bpDJL)wt6~K^12%A+HW4 znW*Sgm%jN4tq|!^i5) z<;^cs-O4s%Ak00te`5YlQ9gO8NqrJ-uq zGDn0CkUGNPJ}O4<6GaAcSZ=TjVnnIkU;8dO=Q#1Q5bRS=Fen-XsVkCy8hfEHHxQNG z!=WAHHGnKNNiNp>+U)6GJQQpCoW%hJphd&?xn2qGgxYWNHd8Ejo)XcNZx)6d;Rzmd2EhM-?ad%%?m?DiJsa4Bj?^9bu$=pK>JIMf_CXG{0;@AnX7=SvQ z#tDD{OGJaR*<0sGmn83RAI6%pYL*La!2a6{As_V;bGqB(gFBu`K+t&8`NE##VuExq zfmb(%X6uqfiZx{4K;(4@AiQw4JQr_=o4fFU^PR|+2K6YwN~oLMGcqxl$Bx{(?> zu!T%=r^uZkS%0H0vq}nZ<+Np3hR)WcZW*CO%(yu|v{KMka%1`?`Q{--@dE4LsqTDP zsQQFGc-fY8YnCTNzXq5yr?wR^q~;0l7{ufGHM9&-dFLC~*7LRB$mk>-uD+ZtM7oJa z<@I1cHp{C{fULRUqRPhNh2MMA_qiwXI_3V$HLc;lA}6CAf4k=nCP#{f$v&PW^*uRc z?0PBp!n^E(&mnczI`@Z9Zs&}IJbG_@(krGu805p~{imTDr&UX!19Gbt9@)j8c&KW0pU{oxs2zD;dBWciJOUFM&7 z!Fh=hD|=|-NHm4R+zHLHppg;z9BgXznSXpiPmVIkH9&>*^YHRT-R0(^B5ZH9GB($h~~(XOuB77#R={3 zv-km5@u<6IV|qjU^5lYgwqOO>`rZ(`%;Ccs`FI-khXUv7Qsw+avSB7=_ec>FCI2lo z%4(=OQzJk$swPMzd-gLhY!$lC)4BU`wWa8@Ox&bBkt}P&J!l;{`F8xga+l+&c=j9d zdy`qDuc$x2tbAu=kS4AXCYOQGUgIWe!R~E{LS-KlUR_#VzdR7cBA@Qy{(=rcmMRm9 zkG%Q0@1DKM8Yn5>mcoNu zscxB~h9Msf4QZ#_u$hHqU-xj~R)#?bKVuHr5$@~AxILEFsPFjLk-GDe0s&IIM ze9bI2pgnc+$Mukj6>x^S1MR=Z2FOz$SqJ+N6;70^t0ezVA)>pi>^kU?Bl7-)tY#ZU z76*!=9^!E5L$j>|JJKTHNrQ~o6l_{^-nu8_Rtz#)xYt`o+HVwS&TOc>g`dDq>vH_JDqcS+`fV#`~1aN+a7$kiM-!m2DIm|z-8S# z^Fkpnp2gkCVG8n8naF?CVWl%I zZ}kj@5lfz6pG$i)Ix4wD9!JZRC_uXGSOi2mm$*wFI)3AXt@Ch=ooy@s_Y9}AGTj1A zr}?)bbnvjK0l`=(T4YIoKbh($;1;I;cG+GEh}$4T&(6%9bKk-KN%o>1-K>1SpGmjr z)VT^`UNHoGpGZ0X8kdZVWUc{uG_q6ots? z)l5Cd?IpJvP}rbOy5I$|x%C(11uCK;ge5)l^Wu{(3cOd* zwzFTl=#@WT{Z}U^gAublBw(xm+9KGsThaFRZGH#Y(!F*!j-jQYLXhxWpE9rt@tuB% z?G(>=xkO4L9)?1kYi%{yv6}Y64~j5bsU|}aGcI;}^8CX9j`{>Cz*Zh$6)?r4OH{^~ zYDn5_6Yzp?a-$iy8kEhGj>QUHR&LR3LH5(wh7Xya;D2A)U%ETRA^`H#vFaBMS&|tW zD?6_*z#u&Nst}^)Ty$(+MpxcC{v|flG8K>V`_Zjqb}^A!^5`UDB?}Qi$&PLPFzYiZ_#eH1Ok3Q1pMgN^e?(ao{5S z?8R|`e7sPa%Sy4|C+a@Lt%g^hbwCjwf@tnaAmQ$LIb`4Uw~qd%ImKWnc^b_)OT;@(*Va76o9g_&}0P^8$=UABi8 zY}-a0KVWUe_;M-c7ZOaV%ygv_&X1K__DX;GtwjrX9=>&rQ*d;XS66I}SS!rTg|r+` z(uO&7KQ`=%L~7b=i9a<5;Y58$lo_Kd(WT( zIiHTVeh*P+G}x(m&VAbNfn1}J(V=P8Imkq)c)pdC^IgUj+{6?QpeB$#Gz+o5= zXk!&T)>>lSsx=<^UkR@4I{&j{-(gVFDi}?18sLO?u#*0DS2 zQZ;pWZQtlbQ9fa3ZW*$b83{11eY%?dl2-L}mP}L=_76ZBPn5rZJ|daTQ4k;~;ke2a z77}8@iEAUO+(k4^%cGB?KGEg%_o42nxha7f#~G-E$$3$HVf5_!iB#mgF4vDaZ@`*X zaYwdl)n~AK+lFGB05yZqvz9}W6BJ&~0T9ma%%DE@#FEtV88qr$d)TRg2`7xLu+drF zl=;th6IgbFnRFUTK;jZ0=$br4)_|d|JfMBTXoQ=BKc!ktcfKDRWNkv6cET5%G#>@Z zMdW&p%*z2$8`5b02lVaI0QPF2hIh!E7*b1Pjk1 zfm_OGlk9c}Q~C1avynwsXk@G6L6LFzY4P*xAt4a5UI3tf)% zoPa9e=Xi^-Sr#~Pz>8Nx7co_hGLeKHtMaxqUZJq8R033{HQS_ifJnhWG!Ag#vd5E_ zKD}-Z-6>_D_IC}nCUB+rIHrH#!CceimM=xQo=M)OGM5jq6ST58u@AOINba^09?A4A z@@#Z}p5YYs*YNSF$a^UIlI_P5-ICVD$k~e0Oh9_i_*WYZJF!(+zh>uX=nJs{gY-%;}#&ANd|a0M~S2WI&clpxYH?A-z>DK z@EhaVguvwp5$T7O)SI%Ig!PCkFmfgaU>v41#0R(r+?V$&+`2GtG8bDMzG5bS#P=D%=wDSqSa4)5d2H~ z$ITHdYC>+8^jY;ZOs2NZ)l<@mqMj5ULiF-2ZNv~i22(3uR_Me6K5oqjWq{2iB&sMJ zw>F>qMAfCNC_EuOJxN7px&Ar$pP!^{AiK60q^^94(%A+bT=Bz{`XIk6y~(YArd18K zYw38iImw>j5;CTLv10P{t(N`|AUpzg@YXY#Pf{9RfUE}v$k?g2Mn;d_D8ijWz3e1B z=fafM4}I^O*V&3|kUCLiPo6kX=HK8Tc~sM2EtR4=*HNd)Ty~e@b)cx!`w0VQ#mY*7 z`W51qR?tFSDV#UHj<2?Mfjmf+VVlC2P&+&FyVhl#Z}b^NpY=qk5iAv`9tJJBd9l?0 z8@(I1EpKm)G~BeWxzt5&2#Vy^a~!ws^)mcKyQ(k*XOxaEnpM(Qb4!W=LVQXM&^}Y0 z$fewaahsh`;+h$}7j9C+z25r#J1Nc^z#ilUP~7d;5oLbK$W;UdDOhk#RYUFl?f_#$ z5t|H*PNI^O+KIYxpyA0?%iclbjc zq)e^a6KTnrOi*$)6z`Ijz$8^T;P29!UB~joXOEuzJw0EHs865G2km z^}JxgPPbNb-`o=ydNvaCD6nW&hav~**#IVF90VOk>;-qqe9{?GD}-f|fmvL%bPM+^f5|#S;?=qb#}e zmPUKUdtpgcJ)}q|^)&dn!q@o#bDjh=PJL*iA!D4j%@v!g{AG;<3D>GBemr9_>eTgz z}r!8=}U5Bz>L^yR$PdK|`i^+hX5)zH^%_Ktbop0lF(s zCh-L7{v)eaK!qP(?=$i}2#@o4byvN@m3$ZVA*uW zbQ$Eo?nxt7%Yl41806ZVQu0`&Z`sCJhQ6G@e|NDy`~%6~e>x2D;r%s9?hUVtGrrnp zQP3TIucR>A3ua9m6tZVZYRYRMh0)4F0g0H5;YbO3Pul-qdc^wUApIZgv(@I?yt}f4 zCTg&n?C1P(k$_qErVDV2Rj~_X|C8R(ea|_wWYn3G!sXizZxzneoo1q^?gy# zw`c>{R^|DVB%8hU`7T#_fADJX-I6WD!`fw4pp(?|=sz2-(pNvP=S2s#`y6&0?e6FN zK2`{E{d*-kq+ed^Yb$640HiNq=;@HC%4asZEHo1-Qr`xLH+Kd@D?-r;8=G7Ezj@!n zbC3BYHQ6CA{shUA{PBHoMI3+KjMWO^)zJ8xF=0fTF}9hCz55@Ly5))_g-}kTpR2AL z5c0y4FQI6<>w$OwCl0k6TIkm*4_nZXv67100(N__%qd8G)C81%azKY)z^M3zm4lkH zXW-9w`sb1U_tF`zJrq@ds=w!Z&3aM?dW-K1qai4#g!RkXJx@#Y}9o@2waYt}HbCuNlyw)t1rsJ?Eh3cIT3dLtV`9d{hmbmtQAvSZc z-`js(5|+DMgE?SEo_rET-Ogu~|7F2!i^2g*v}*Z@@D#)>|5oc}vx#|}vt{kh#n^XT zE?OKxC{`;~Z+f5*SXYCDq%N1S)opF6h4Y~@`MVFV&2zFIb$AGd&+QOBUsQQr z|KNe1!6ue=TL){UNYei{*uHUiU4K6r)1{||!y*@6F#*WWjr=2D39f@oJvf@U|8J9m zVuKgs9Y4gihvTnw)W9)5H=uU-l=m5RY4}`;h{Km`{T7i*phk3PUpcfw3#`CjHsb~f3BkAOZZe`6(n;d;%1*i6wli5SFaAr4?kq|qC7|9ky1e6Vxd0`% z)!8^-<(zvI78q<+;18Wcj}<luKn{clFpB1vJh`;Q;V*XD^GD5l+Kj7nv>Tmbr2 z$pyYq$)qUnrVt;fiJ&OWQ`>e_kHsjf$oI9HD59bx`UHQBRZ=NK(SGs z|Mooi1t@xM^I7*j>PKf`9nt9n{x1aI9?2?vauUD&=Bg_1VNZ|1fdT#7buf~^1WImZ zBL10xn^LmYncFx`)kRjJhN3H^;S`i(8Z4%DS3l;zCs#mc%Izk2W)4>JPRarPVNvq{ zXZ;3J(LKhPW*6!?A|@^|A{UFfioji{)*F$YNC6dZb-Jts|3=A8K0$ehPe*uBN$3`; z|9Jad^lo}Gqs=o9x9BjXk^jc-rC#m-eUXpsF94_#7@Pqp5%aFgsTxe_E6s2+qbdlc zcwFUL)4eLk{{O2(&^s?mwDCQqrgWxo6@TNoSq=sxi?|y9cf7OJYhWysXmbI}5%AKAN42Xa>rUK9+2XDj7Sl4*@;p*|qeoeOdpB3fch6^gJ z#OqXn}00ZX+7SHkcA z4IeUA{I{I`4V68UABle31|2c$-Y=)Mf5Sg(by9DhfQFa}W&kDI<(dJ0Kp&X7>@B|K~G*9se=Hj;ZZy5K#Kc`2>;f&>!N`8}ATACVnFKY^_tDS6 z4C~eL$zL{+CdMHbrlA81G6JHTJdta&fL#S`WZ+Bx?^qd6%R_i$V@K=d3*u!Pn=t#xqEd542A5w@3~=#+`sp-Uy}c{iQQr<<-pHU_wkBYinlO+{ay*tZk>Vwt0F zblxi;)vpDA^0!y8i=P=ZkoX*tS{1J$&EAX0zAOa#jmwKau9=YFME)M5wq`C!WY{wR zQ+BgIV?K4wh~u~0cg`%r&$;PXK!_BUv2H`PX0avHpov=g*|9lu*-4BygP405_30ff zY++rCg|6nCG1VH=R@kOZ3<*TkVH4dBoVMGkmuM=(@-&CP$shS4F6yGh&{J>GV7u1@ z(&YcCrtQc;N`<1JaiPLz_ZWKjm4Enw#Z$=qwPx>-vIg60FEDaz#vNlF1}HXZiTbM8d8ne7x_aI0U` zi@9!zID}3zRM?p_miuNzslP&$1jZ2 z1D|)WawS0|(I6MXO3y{->$o1tc;+8QEj&)nJ*Yo&Ue|437(yxWHLL}fT@xzEUY+fg z3VoaPjL?h*@zwdbsPdVCvF#e|NabGK(F92xpJW@Nr3_D&*7nJLED!J;FU(ycu<2l` zqjqof(VQtrr*cNjSuxeG9*}hX3>lhj9K_EH%&U@j2r5MxO)z7QOXCWSQn z;BLXUPH;6|%$8t=*W)K@?!?gzQPfVJ7(}(qgm4lxu$6!C|3ui1UM^Gs#vNbI&ikWt z?9%HZBI9+nBNf`AljXnT}AMDH+*=d;UeyZ>}>hm8BNVya~3=OY}_zu7n? zP_JMLWu4k2|Fwyb3MxPgGiZkS)wdrTYiCkOjdW@JTkgXy1KKz@h3_{t?^(N>3N0O{ zI&&gVI1`XfFCTq3c$_|@j+G`rO5#V?{H|4^YvhCb0b4YW?DNM{V4u(=X`Jbwu0|4H z%+C?(0Pa*Nq8xs4%9)NB&qm; zR18@((Hh|@H8;xaK_7`?X4K%o;TH9$o0%Jbf*eNM96Tqzh72>C3Ix9opby=W{0e{l zyq~q1f&BpP9l$ny(PSY1I!jOd5VU|Yz5!;Hd7wQ1Thth9bk{Vq;x|9b9m~z<%!QuV z1(?X}Xf>%csYRO1hN*jk=&x#d3P`hzDz`dXapSq?yJCe-m2S0o_9)`=-ApkS- zd|A5pV<4PJ%dO*TIKdYt;T<&?q9x_Pq?{E~!FPT5WbC^YCW%7K!c7C5MB&N+I`g|6 zJRE^s`)uSjH`uVJ#1PInK&_~U-feL4SWp)!-qR>uoOcIL^M7%xmArwWx~6aL6%<@w zb$LbI*!L{YZ@n=fG4uk5RSFL4s*urQ0pTw>IDKobKu9B$fpMgB0gHu1KvQ(M_nStgZ*d)t5LHx-k>OnK75 z^Jk!pmsCIJ1|DFNhbbMiBKu-_NbhO{%A?@WG1>KbX4l|FTDvO%hR56;jFl zNQU*_YJCYlc>H15^9g>blDa)m`by!zo`@vH;OH;}G^Yt{@ zwrb;XBeqswm zp>T+J&3_(N?FnOx&O$o2m5GA#A!J)hKYd7pkr&1&%hWl)m7&Y%W}eNo*=Dv)-1|Zf zlfI|RIu#$e8l3&;oP-so+4vp-T4N=rcG^mEl4mn?%|+3OOM!pn3qzk;D00w7ES57) zFh>Iq>1AwfrY)LY=A*z1m2H`mXF z_`<-MCJ}$2f<(lkK)bA3eUj;{+9d85-lPxN3%(1=bE0=51jN={O(D-zBND1|G^F2t zvPLt~D`RdA(M@!QUnonou2jbxT2aK=KzYErOLR_W7GWh;uBmer3^N}X= z0H;ub%`JFI8(V~cSy4~aT%M4|?-XO}ChTyYaQ%8J>3%OtVL}R+MJat%&bBi4VyI7> zp&!5OR0N#DkA;+#GTIa3%ElT?mul*mB-7`ZbT9hF6tm!|?|n+0{@V9dFI4m=UKlE9YRAF( zM$?=ZHa7UX5%NLG)eJw0%PkkK`kbrVYk6}QFJa+f)82^G(3_vX_XF91Rz9VmPS{tJ zA_@IE-zD^I3fR_MoFub|LWk4rDdF8V2X(ktibY|dB$jBdCgv|i4Dd;jcPNTcW(cjw zd;XtZzt9`Nd%1kB!Ucu2ag~ep`>}Tt=c?Cz4kyd`%9Fxs`4#KupFM6c@8I*|uksct zF*Wd6ASkw`S_9LQSB`r17xBaPySz5vTaxHGHgkMxO`{9P`lJ@U9u=6I%Z3!fX-LK^ z@73D=c-La(XBF__(os+I41N+h{fx6(%N)aR@O0YjHLv_%uS#iFmdsDHxu+*8 zlSqLMQ9EpkoFvbBEf@HgV8%t;xBqThX*?@uEn$WE)%R2vJQI6Xm!3raS3%_WTVLf1p zLSf3bDN*r|&?u-5k3D8Zj)VHRZ07UwLR89bQogky!$K!$UwbrwC#!dsef4U$t+4dn z&1-sJ#9wYa?i;n}Y2rFS-_zwV`+f93eonK?bciG^Sz-NN%x(ppA|-4Jq<0>Relb@w z+iGaj=?AOjrm)stnP>e+eVN7``q+&CnS6Y{Af^=Bc5CY9kDpa6us|eBw~5?Q8b{$u zpCoEA3Y2TlI0yL&kWIRbPju_#>dfvYR%sl-Ri_=iNJtaaQTjn*Kk_YqxSf$1N9SMT z;~^7F+4%8**80p)g2CNNaV`?fAJrhb>6&C#J))&9Irz}w@|EJ5K;NWXZrE)OlbZau z6)DPS10nT>^P)E_$*2IM) z8dKQkWpA&62gZ!8-S%{yOF0m!DrU?Qzilx(fBpEgjD9uDhoZb;3A_-rP64AI*ewV~ zmq>Z*6|JtYN$*NWV3yugXYLl?!*Xb>_haGrj>eMA%&}gAetc39a~B?LUi9A$ITfro6Qyms z)t;gdQi*6QJKcXVc~eP5W_1sy)-m8DQYWWKx>vC058C@K!HYy>(f!Sz&FF6t^8J<8 zd96fPe{m4z&ytR$3;|Z60~K3DAgrA$0K$gXS@&VRp6`&Iv{obSr>mWr<^DKW$OIxW zek=!=8PV0aCTA2i#M$1>Qdv0dpv|lR^IWJwXC&FQkKZLWZsBTFg|42h2!~4b1B@~Z zP5Y=3+hDim^2-;L%KFJy5GxE|O8OQ*xfODXJGzmq$l8zkb!u?Gb%s!Z)_O)f@Mhfd z9xPTuabKj6olz3ZGCrT2xn_i;jYL-C@Oj{BndAnb(M;~Lb@dnc2NdfI{@WJV zPa2fo9Iol&pU!sTmVtIH`}v3~$4ZVYx-wn+n7f^dKcX$mI#H`6FGvFjuPa6!Pyrkn z8v@&)H3#4Ppbh-jF9 zt9U!cblUYG=vEwaAW^}|a+hnTgGmShX`PVDrTg(B3B~RYK{{7h+C0Uk=o;DP>6k~> zf9UzpKB}cIB)05g-NFk$cKPc=L(fM$dy(Dmqw8qxRH;A$&i=9Zvi9mv%Bekp9Nx*h zG}D~b5l1=AN~7OFNfl+psyy&!&n=fe{2YT@fPFXAreoHA)&&Y*5@ zgx%)Y$|0m8xHIQ%ahro=~ME8kA39v4GTTyVus$D=G`4N;#ZYTgWUR zg6^wACII_jtvK8ZKH=}CH?Ggsm*N`?wKMB)>~Y-dJn6jk^H&q-8gp4~c@*C7+;U4@ z){bnb=9%4ivdrK5PY<%v)yaYjuWxA#VNn`-GAZ6&Ds@Y)RdK0rnnRlDl?Y~4zlMeE zJ&Xu8uQ*Y1>T7YW5p($79$WvO7KyDPmeD_=ugX3s)RjFf`!S#-9;yEJJJ`yVrQdaVvFB8)bq7Q1ar^zu&#>SIrmtK4s6XBubEzd7B`Zeq5 z8G;&@W{27rIGOy0H(JMNvLf=ADoBnAzVPzp6>^uT!B(Ur9ayrW8;!ppiH^Hpv=n=y zx6vBgevcQK+E3z)fuf(FxHA@(kaJS9B^zV|>6N74!c)#feci_rt7Eot;Wo05xH~>F z81rJrkb<-eVQP}W_Zy>_^fRKJ&0D=}X4yjl8bp^1lbH`e1$KMO8D{N~XXb+=IR}r6{A4^3ET|IZC=vnes9C;$l z{l4>CI)M*r0h5x;OI{i;kk4*4xMcSZ+%=}vHyYHCu z{5T(n&hG5YyRB-@)AKgCM(TuD1YjsYZFa5P^LB_#5{^)Kwc9>_vFN|{3ml#g7T)|6kx;N-3LkR*}a zeeZB;AoW(L_O#`_3wUnG!$3|CyscKpo=bwAIBEO?v!&iKR#0|+=MKSaxEcp68FnPkw5r38l}-(eR^Eu%_J@4~a8rKGWzF~GZsbx`i=%Ol8l4l= zTU29@v@AWQy1Xe})Vb!0mL98L6>%hZ2EV=y6RFdf1`lDE%IDAOD%fZ&%A`#SL>-FWe z2`9Z+o<%p_V|Mt&K8c?11A7P0oZf!QX3siT%2i6~k34S94sdTu6Tp4K~RSzn)A z-XU|8yc*9WBEzVTJ7E+UpKPGa=Ie}fb$2K;b$560FqC^E(+b|+d8M{RnXmI<*y>RL zLh))_;Zt&4F#3z<0-E-2_Qbllkqc?}Ce|C~(BuzcJhZWAsUo75N$_aVJ*A2_pAKQm zyI;$d@}4++^Kj4lZYLsLK1^+-mz|zO7W~;CTfl9=Ds12vLG{YZ-=3?Qx|-kZAQCbM zs}Dq$_uIqjc_W@0?c7@4ZCBc$ztxCZ37WQ%9NxdpynASx5FTq(?|w5WXW?=>EKxK^ zOp1BG+Y{pcCw3P)7df=3a24}QqfNJbsn$qt?6}{|=Ud38k8li+?sHVj8zQFf(Q|oq zV$X;ciin1;1P+6hXf%lxB0u=f*giY zp|Y(g(2LJNb!+*dtA)MLB_|EZ@5BB*f3^~eWfGfSf0n5l6ElD~C6eLSa!HZ=fo+&#MyfSV`Oza8&G`(6Q^i-1L zM|MA}%AYop&EHvqOkC`7t?Agj`M_ z=B^PWmYK`YFDGGtr^xhQ4%{DagcL3Dn+*!_FmPd7T^(Jy8IkQCb_Gt~TKjC4wPCL@ zL7k1}*gDfJbMP&NG|bSFj1pl{BzwytJ(|^a=*PI(0p{oSgSajd+e|k6Deu=%Q-Y;_ z@B21rhTUSGstw79QgW;r20eR_~yYqhjzAEfzk6HB2gOrmEdGtHz}pdNAc|SvO$b2#M9L% z5qtCJn(3t)rg^`PJ5NE0@$%DjIm8fG4u%ad+mg$Us;%0V5seO6R{Z3@+$G@3;160!SKYo_U=G4OZC$ZTeI9Dw}axl*|Zw+Wn5TNe*&{hq7xZdj-%dUSN zGyZqKTIzcWCJ(DtEHocl3|`l;8LhhCEXoBReaM`|^Ym3VZh}9Wv%K2|1d0{$9F@^- z>^Pjm+sKm!;xvaG`2yTgp>f&v-uR9_7i8}x%^RR42(?2OaX_Q^A373-Y z>B5om!x-;cNijUIif-??dkl6!D;&D&sLf3pk#9Rv=COr5DRn*j;+FI4Yh!8W@%yAF zPcK*422BX=g<3tQa*A$$uq+VfjoDz0Tiw%+JzKDU3^y6Mnp zS1xZ`b35AB`DElB++mPdg^*(l?U~oe403{$y*loEJNAYg1XzR=bFZr`^#fUN`=`Sk zF4vu|hag6W?yKJF#_ikO69hGJCqx&x18Za++W##N{@M}H$#?8p_^tBZvZKcYNpPm` z+&g7KY1N_HcfP$EbEQIWe!EZLikprH z*uMW9p*JDH%3fSItL>dq;uo_tC4D)k@Ra6z_Z-E^h@McfH@X-|=~I9U1Jc_)5F_w5 zQuC;vf(19-~Sye8*2A#F4w>xw&m zpJ2}W7`%u!-4+sZ{2Z>~O72$AoU2y_qx2D1H@7uX<4OBzNaHnipvorAEe7|-i92QW zoU6F;DD~nU4BURjutMtO`kyQt6?KYO^57&!w3VKn$c^*~DjcxZvEd2w1Ou~>0HHpx z-1ga|bfGJxeDEIFR|f9#Vwi=tZDJA$$%S+NKpAof#zr9coawph#wK1A;KnltO2$MQ z7tssiE<*C)-myb#W=MB(iPSa*Nm&!B#Fad_XDkuBMRGi-8Op)AW+qnGDkn0A_kOvV zn}Msm_$!0tNC`uY5)|iWjNyhAPVd;9_KFwL6D#?ZvEkpH8n5zNM-!XlKvY)nw%kj7 zg*N$zWk1#J?8|3l8q|(AR5A;%{IW1y1UVL}H6-(dA-6XK016UgTXxV!-ny^yjPg28 zX&5yzoQvt=)_b8>39Nv z9?V802>&qU4Ome_voa3Zkh&bji~b*}_*g53<9w`)y$LcO|AzGEF6osGaye<4JQvb< zWRVEJY~g(jf3)ckBjxAUPnQ&`R#VEi%I7Azd1v5Mt~Rzt<{^ywFL#0$rp8mp^A{$F z`%j9vn&T55(CTGb&0CL;esA_b2vau+@qBN#Q3~*9LPpL~ehEHk^28ZXwlH_kS?RQT z2e{gd&O5~s)kC0Ex{5ph*x1*lCD7bXmgqZm3FNF89>2Xm|~&z#hZx1>I>Lq!%1vcq9Y`78${b!a$T-G)^)GdPyLR@Hw7w0&aZKcRE)Z? zIl*-@-NH-vRd#NHKhm;_AE_E2xfk2DZ5UK(Z+3_6;6l^~l2B5*VppuwWLVcTXDd<< zT?vEtE3;6~=Hz9J~aY86y_&<^k)rJ5xp`|oiPro^p!q3KGc@@P`Y$a%&rR@qjsx5%zf77eSFv3cRGd&b z;m+D+No5BKO~@R73cZ#vE{P(~AtWytrOpCzE+P5cboHqCuaD|cFS}!B%>ogrlOm{g zTJC0SQUp}qF#A&s;~*r%2?+H0KV`$r_T82WC#)YK7CH-c88pzn+J245Ca)#$jo%O@ z9s*%#u{!<)*YN0i3kbSh2oF~@A^beiA4^&KF!MgNc?}xzdaQcW`F?BC&Pq=rzb`J> z;vJV*oHKQVUhhqKv7%IHU>KgP6nISTyn3~$V!?x%o0*0cmk3#k!OF{we&ms34Z;2a z^leM_mb=D{3Y|6dcDFx?S)tF9*a?z6U%igAh1D?PX)@8x>)+_6{NU5yX#Kr4tl2uv z_FzNBjz1hzKk!4wUw68&S!HEEyQVp8L^f*e9}eRWA5Ubh1IvN4uud!IZ^JmtXM704 z0el_R`GAsp@a5gmlkqc@!?!N)%|=`+`ILX44ZZgir^9%gv@>Vv-IP1G69-M)53qiG z&xCoa8{OMAg$QQ+IUxH<<>Wh^Dqx4k6_U9?21&kpT0iDx4YB0q9c#kdu%r+DI~hWj zeXLJHj%YfQ$ZW|+{NrC7&CoyoccH~WZ9$JRxaCs@1i!j5b=eoOA!J=#0s=SD-Yv#` zxs2{71i*CTitl4AZQN*}xFrtiL8Dg(eFKdw+|7@(T@3;8GePkGxJ~lepCV+S&!lf7 zn0f(7CDAY%Yx?n)KfxOP7WYOL>;O%gE+Y{@PK)tZ+yJ4o0tnrIls{V)X+@I=rg|p` zrVnNB(l8sT2=nZ4hY{C8sh+F$_`|*B^s`cRePLbR_g$tq=gdalDf+Dcq-o~kPOx|Q zRJ{)g|GBFDH%=prWSUHWfbG6g6NHg84h>LIm*{miV9Y4b- zXXZL8!?Ek*C)`v3(wqR0#-w#Zx9+t58ya*8`nkjY&_Fxl?ACLmK_9PSU)EnaFmFCz z7m2klr=19ZKk(oU>ouTLazhZ}C6Qp1$53RM-5!*e@>W z+iV{1@II7Wr-PC2`Lrj_tn#TKGE;BF|;G)#VoxVyCbS0neWa+ z8+ftm`ITG)lOd$_XJJZ$1HCUsyLpDYyPak>r8G#{f(c*1^}B%Alu{~zTqwuypi(WY zpWJE?{5>3lx#`!Wl94`W*evT~AJ-*>7|WvDY>(}_m*>ce*l!Pcz~oiqEWD;+rIF|O zA^m7^n9KFS@2blRg&^6Vs&`Rb#{Wo z-8r<|Qmw-;A0M3o3Figyw>YuneEh6?7IgBKta@i;7bB%AfesAiSInb{s(0s*0C@B$ zJ(www7*BKok+h7aAhC6t;C(#-ILhp2MID=_AOv@_Mogy+q(4T>Smn!<)1y$ne`3U# zd!U}MSBiM-hurMTHVi~4-;aI2yE88tUlp}85l$ACyZO0Fw{Ulq>-oflm+FjfaM13} zN&st-tPXE=POrmXIa{1LWqP)QcDGaoJ>Es{`@!pK@J|b(0|JG-)MEM)6feV%*pRnl zSA0_NCX*FY4RJOV%f`Q+nBtfvg*AnGIca?Ma*lm>^h>Knr-kq0)MShJ z7DysW0td6Mr5ugTXS2w>yQspDee`Y89KPZ@6vv(ppF@#p-jZR{mj&bKu`|Nd-HWaWAF zYR-ECyxv^#%;&f9sB6LsVWW|P$IYmmcg>(dj){HlH{Yx!nVar!trKv|vB@rP@IxX_ zdYVpma0H;oIF~Z{Mf5yAJnd}8XDpa)ueh@%8o*y-65+wvcT~k0(5{DKc<$AB#F!)Q z7@uAK!8q>0BDF-+A-j}3fCv9_MH~!olHDM|-q%Jj4Hr^>(c`zeyXcVBB!=`se_?gu zeUx-v>=h)5{A{>tkFzkQpLE`z%-UB|Gv^VtAKv4U%kgM=qD8zu;%Ai+tkqhgo+o

uF}i;Q?G2Vhg%=JlPClCEHGc|&TyhlO8q96oHSi%wpV5RTO$b`M@hEWf`-fw zF`Xw0n%Ot=a{6b=@99X)Uj3+8ds`=efHew`zq<=yBiBXW#zmrItkwlm_+i*pyt!f^^v%&C|FB)YNmev2V&P_y>;foNQ!a`}s| zuN(ET6v}eo-(cl4qt9WEG zwWL9a8uKfaJP|HysJ7Z;L4&iYa;31~t+B062vh!gE&SRg%-+LL%23F<$-1r;_*np! zPdd_*9(+-B<1*mTXYiPTYfL)I6+~vU^d>fbe=a9b=@6w!U9VSjCl^Kp6Jl@7V<{wL)TNuj_=to zLtS~yu8OrjF{Iu@-W6p5|BcrJn+I;NzP#qqf0=JR6$0Z~;bx=@?QHt*3eA!{cpmX^49e}75 z@dN%jd21j`M(6m}f_TgPKan?E^ARnWaNU)aN2Bt~_mDfOv;zVRym)qc3005Km;gxc zUm43-3tAldnWYPkS>lps!1>SF>;;J-n+w_@=8?$xhjLOTn$j=7{wUKz(n7eGRXfQj zf?;DE4+B_FqHeiY4c1q8WfXV@(l=0HCffrWEw9$LoLmjDy1R_dAt~mjSKZ{|uYQ}K z@I{x@BR#PCNgAo7T{d6E?o%w)9*xoqH{dUG?T#9^&3VToxh`VfBX5|@^!h{jSI(rP zjGM|`aF(y579B@(ymERxL9B?`Lor=07m5Ul-o02`o6n#6thUB7o(}RR9_Sm1vq9Ek zhKC%_Ffku2jb`J305J+Lq6eH3WYH1rK}yq3BgFbkD;e9%cna0)moFgU$h-c;@kn;Q zS~^#|jPx$gCYEq7@bJn1{z6&jN$KG_`)6Kh!*^fKx-%7wWGn^win5ksTX{tt{nsfy z=%z<`+u{)XyCA1sj$cxj>EGfSeuH5dcKT@ZQwE%NXQNT%@jxIbMBQ3CY@YM}h^=m8 znE?kr&_kU|3*qqWP`?v`hsCS^cQkELq{@z=F7uJqk~cf`C{(!=cVIt_b^_2p`A5I6 zUNcQh{(?kH5!NU&mrmd_?9Jz2qmUGB=DD{(pVSHPWtab>J(oQYq$b`F#`f+oy!Xdl zZ;gu>oS{{Th5~NtIuOXKfjW!=q-P&Ua*L1!I^uH2}7c+cEwG_J*2i z1y{zH{af6>nJH#yQMX8XcC0H@~hVT_(mhB8K*1mWTX0py=< zw=TyePs^n9@ye-sD1&Ei&Z*PY*ur>4+9_=~{(#ZPWXkuWyyq$@{C#_t2Zra90b7_$ zdCkA|lD>NY>7T2ZLpMDt_V|bCk6RUYF?H{vvb<>84-DCmB%iN{z2U#qhdUt#_^l|S z;qIeW)rcGLc#mVjtT&|v;XC-GojUQIkAK)shA)fjHQF5+4MFDlM8puJB1x^z?T$$n z@e3fAWR23c?aKZk8Qr! zen6VRupA#E*GUvEJ>rr*@aF}90;8v!1_KJv;(ZFfcm(QG@8SRIQv_{vEMGHBkNl<)gPG+EnSW%p zK1sLB?~8wUD&TxOw1Su}N4Kb2LmFlazSwPzt-sfMpyTpbX8>1tV8O?zmEjIxrWbZJ zIlSG*;JKHtEDwyZ2ar2R@uumC+9V_8*>g+cpxtcEMywG(flFk9^Mje4@d`)= zzvz`rf0TWjOB4i;MVWNc_GlQQ$Wy%)n~ZF67RT&E00x!#AVzOn7IP86GuixZKvE#U z)l_Wb=g521UvXS)@lp3oz#@9_I;!`;3}L?fj?&xY%WCVx42YpzbEEpRlNtW`^@TP1q^{$r+% zRvKK_(eQ6RduXI6Zl6y;J=xC|u3XNvXMLIkLLWF;TE@|3=30%peiiY-0oyKY7Gn)$k}EFnl11&4546 zgxiD%x*9_vi!3ZNDSXMNQKJWd5dpVCdW_qr^JW!9G7V5O^r;@x^R~69VKjztbG#7S zNF9&sO1H}XtB>X4LT)@buwL_fk{Rvm0j!9#zy>s?xirl-7Q6he%vdr-0*s_qTnwR1 z*ybG`A((Ruy;12*T2B$_k9^>ST!60af_)l2fwkID-})i%Ka5?fVTE4b5+6UF%whg= zKP}seY$*R3ucNvp{vwx#%^3Y=8xjQ*qzS#Gz*H^8#;FTr;=KblYb-dHjPMnW9ME24 zas`sv8KtiI4?%eCv3hyC_$vL8`BN^WC7d0`W-4f*?Xo%apcbvSL@hvtQ~#6tF;L;a zi?@?5M*(iOP@VaXfLwauZ)&Kx>-K?L=Y8WO)=F0y^wdmtoY6+3@m{l|CYJ?@=M947t5NEy5=Ptly~ylOy#zfQ^NN#LU)sdj^h0t^ zg?U70+yUFxF_15rjc?_iO1%yCp340ocxRtkJ9a!7lN5V@-Yip%9DlJ2{7t>s*9ePQ z*;}5F$#g0;e6HTR{b!=BxiB_k!I@75CtH{BshU1-%OpSBu}7)s^8u04Qz`r|M{x?7 zL6Q!Q`wHp|!m;tK8JM0k(FEBnyWJQ-1EWfCvG>!8ptsTL)e>`S-}2FM6JbutJ2`sD z*anIwiu_Qc>p&IDbLNBdjj@n1xI7?hjnPx7^3G7vdKNL-kE;w9+_~t~sfIOgzT=H< zYdb3W*vfI!sy(%u>b-uxxZ4&6YS@|||JCczj?bq597=_qChTM(vst&y{w#cd$B;Sa zZB)|ZTA(APvHbtHm@o;>$Wt2K3VHd_r4!e1zvEcpJiBwa!W{I2y^6KBUU^i;mK4#;m*`+5;5}%pSTnRh~tOBJj=SmFm>m991+nB_c&4+_di~HZ zj?UaDjR9fdalwZwFElIipq2w1N92B5h+G}I+{lYro4Iwz2!`^&sXGD6TM_iPw)|v3 z00)mLf?~SS#Zbyv$h@f(^50_m+O|ZQvVuF>(tNCf`PpjWK2rshwW}2)XX<2LCK3A_ z+z<-SIbHKb*$*Oh6}{nC<(nP3VZK-89ozfqM|n!NjD{Gb3miK;few}dng#*R8yEZm zMtrUtZ|vV_(9iuTpLzAy%jlWw24C(s=Lr7J*de^pottolNkVFhT$_{Xk%uaxit+w^ zgHdKvCd#h_#y<0_C>1^P)LT}OX_!(AA}_n9EA(w3wkpC_je^D@E()CfQTN|3iEl%2 z!uRQ+0XBhi*S~Jgos6_Hr;JT`AH+`_CHrz%3VFVToEv4@{HY%K>)5-!65sfl;=-9z zs*)|q`KsJmWg>K=iVM+;w9+EgVaK82#ZBiYsK!3){}qd0Hu8{L^ySm5=p@{M7JNH} zu$4~REh+yzxAjI_^6sMN8GLs_5(X3XC_ZnW-)DN8d;avLWNnty-|}yg>M;fu;$!TB zMfk3Q&s+{u3Vh3(5|7<-Vt zYb#44o=bDvXe_6(D#@7={Ix~@5kJ0LVUGqu)Tc^~VATv0>t0LwvGJuMmUu3uXx&o- z%s+vRJ7~e9!c(+&`a9T8rJ2vZj%kn_@^jJaKg0<&H{uf-zfaKojLGA}Re3P+*zGvt z+jxJ*=5Aa$7`6suopl$fxF%FHk~s*m@iq6|TZ|syOK{$&ZGL9)G*>btlXXgTQkZX$ z)OKI8>}!%AFX>bOd{b2R4gv)$xbd`MQ2OBCXU$@{6-^gPHKwBz=UE^=DSVW{O4rfS z{)CoZXo6NyWRCYwpPOkhf4vuZ`5QhprxRViieGU*&c7oV4X$x-{{E$EQJYIbpNKC* zr=vvJ;^-N>F!8RMZwJkvFyV&RWh-M*Axouq2O@#hc zMw4Dz$j@^b69l3xXlx`pp;F-|dX4V_&hq789(_MJt+t~1EJ$*Lqeu2{t|Jq&D7J7m zX@|G};y7J5CjQz7&m^|IUV^m|KbsbycJlhDC6+N`zcuXr=>EZy_G8sLH&cab8Ta!| zT;~8D&CSB=*3cgp_uX8Xx_g$kFKtvbZRZC1Ax>G=whiO2+TpWBO!^ZH`7iV`56bDo-NTaAl(HTbp<-BrkkB)KJ0abNWY9 z+9Q2WCaFFYPh9=u?koLqwI?Fa$KyEyR1T0OcU7J?jxGr2^%;FF+RhyrcILo$!|3@H ze?p?yju(vo5p}D7`D5@IX=t`}>Gw;Or!lcR6=>ykk9XTpujceBw>|jYOV-++=+^KHgy)}%jrF{$bqHzu_r(N{u9zSqQuS7C#LVt{n_rg;*lU|3VgHB%i}dv@>P_c zCLX3+VZ}&CSM2^{2fKCCPm}YZajm@Njq1;ypT?@^8VW%<2~PChS^4jFB}wU-rv;um zpGd7JJ@z+?ujT&mD*N|D?DDGcVX3yz@MI3fxOxep3+a6Nn8qKMfDy7^f(`Wd68!sq z6n;?Wh|q9k%*{^}VJVqaXwcfDIp>#X881O?9=$vBt-i=a%qo~Ct=-vHe}Ipo!LwPh zmnqcNJf@^+0jg1@srMT*@>Nhfj^Yv>u?WGqK?6S+?UAkN*z!O4w5${BxzT(x@{!-B z^D?>n>Tt`}(HU=J zT&vesbbyAu=J*}8tm0R_-^>3Z_?U{jKE3lTR4QdHfT56Dl5s5+ugI^-ZRN91Rva6{ z1zyMswUyb(e(wu*9+4!v&UzA!^FUM_*M<{Ru}bE(@H5g~DbhPTGxkDT{o%I5j2C!S zD=oamT9}{z99`t=M2>qNbOq@KTt<51=*CZ&*~&VzT}6MRnj#pABi1sa_-2*3H6F!Q zMKrjL*U^?#YxS1aE9$FVQ&09+Zdber4X{?J6GFI5RS*;YS^UOR&?TFREm~yGI!=5lzjQs zs***_p5!_{m)MD?A=Sq_Iwt;9N|x9!-rX@PgNQxYW8C62{lz`Q=#4k{9%pr;cSLwr zE)od|7Kb}c322$$A?8zF9SU!>JO{R!dZ$Y?&=soOhEsR^-97bn>w(Ij>CTh26nU|) zxE|qOMVYil+2K+6e$}$s@qU2*2lin98`FCCiwxd}HxpqDiM-zZ|sbUyx^?hl_6%vCuS?Y``d&Ixt%py>Azh|+Qo`ej|t!3&`q+Ds1 z_K}4rwfkmoirRg~Uo1@7syks_-NrJ43}crfF5G2Ax<52ChOGLW$CZOJ^77Naox=W9 z5f*IbQ?ZnUDbN;3%s=_)^n@gKQpQI@#`Tj*FdMB8;rI1ISL;n{n;kqN|0Xy8D+?uP zh{{V@L2jj?lrGs%e`C)RmBf-a)>N_F{uc7HXtKX(4_}HbT0a0>@U6JHWhk@DB>T?I z(vEuOdCd6)CjK|OMJl;{7h+faltZUcB z5f_~xg$JIwAliQHuxi~SRQJJKDr8AV9q)BVx|uSCOha|cd3rS;WUQBE##z7B)AUau~l{Y5I=tI%Ajv{ zmnt^RFXcOv*srj)em4CW+F=WN9_Uw!5n9bKARMZ;hQ!OVK4$;8=0@THj^xdKYL8>z z_6NB!T=Cj;0muA|0H33YUNnW$*pVw4QF2KWu@-F~zOd$$`~$^Fjc z$6Y^-^W4jAMU5A~r**vXh<`TCW&@9PzvS}f%ckcuO7WLkf+|Wwsa;w|;u3f)#(B(V zn)Ojr{F7qiXdgM3ut{{5oiSX&c(^MxHAp(=hyMkwg&Y_)D&d~O!#p=n?g?ihZd`Ik zDr5$qLkQeiN{-p@k3Lt?)L`yGyx0TGd?QlZBO(7fjCJ^s++@YZ=w2F+xoF@Ofyv&$ zq|F-me^a&Hoos8Y7|U~NXSfib7B*SZ-Wqg%8$S~Hjzf$b1T5+LW%kGHZ)hUr=HO9f zThWr`R4gB1Tw4D4?SPx|+qgeGI>&h;)_Bh>^J|kov>p_nl0O%(A(nnj@l#rCoUPAv z5{MEZ^LyU^xBW2JT<4r~#&H}oySug`zaxFm04MhB?7Uj@&d%&g=5C>hj=P%k z=us{DnV3_g|2`Ffetvy}_J`{^p9h;p-%@A?LVk}r&*In|z!q9Q6bVuMdOeNelzhaC z3|y!aKvL+^(KcZVY?3U$uJ6yG-=fI?pT1fv@vrxv+I=9E{!ZQLdUBuYxVrP0vnD0C zED|6O5A+J&ocsqf&7Vup)}OQI;3|JBFH(QGu0n4fDe)HaQM#fQ+v0I6KHz+ab`2eg z`XZzXjT$XXl*DrYQ9%&*#$jny#YZ)Q_UP&j*#Vs~_t}(Zh~FLx2cLnxn7@F>!XSkIF^Whm<|pdpP5i zq@9opCSQzXQ)1x)0D+5W4}b1u{_Vc}$@t^Zz&B6pd(=1X-n=$VKF8};L28@(+2Y>J zne*MPfLlW0p0!H8HY#*Kt@RS=+8j|21p=oyqapI#O~TzXb~J-6W!E`7A|82QHY9J@tSq6|-h+ z{)*M~ZHskqAt4px#x)^q!HCa*$Zj=g^XfaFcC}#Z&#(vtkpOb?1 z5!3H!(UQMYQ=5-)j=5R{F>e5@pFKaTMVEMl;xvv|t0mml&WZ2Mcfoo;JTR9uET$w| z`3bT3a?SFToqEadqq)i|e@}GV*o|Gs4MuX&!!z?BlXQUv&>6U|67&+&^Ez@cGfQWe zZNNLGDJdFQpS8Hp9^K8Z!)|}RpqceO1D!oOZGb15$ou{dN&Qu@thrtNd0ZA|aByyp zycr!K`(#A!H{fqgkSLi(khJ?$-6XXCJV3fN5$rofDaa#_^dxb;c<78Ii;jSY!k8Wa zx+DB}Sfj|5E zoIR;5{IleD6bM}q1NYNbfAAfmW|Vo8vj4BD%j?4wUCI8{$BL`MyRt)XV+aMn%j zKI%mqV*K;NcZd_bm8Qp!25G!vZYMN%e)UJ22>ebWeg@t!62{huPLym(@Z`O3%Q)BU z@U!98bUOM#>Cd$VQQvNnkVFukE>``AY}b-5d(k}d`7yEGlvsZr#~fg*nEU)~Rl8X- zf!VaVzwdkNi<^iFTt;+LvK6*(? zm7CVUc_{9$O9}1#Bllxc?uL|;{wD*0oXAz{y`2?!E3w8`5Lj|ZZ+{^7mB0R z3*v+G?OxC}Zs|Csj}*VRV9~YSVmB7s#t%N|cD;^(XSbj7-MDtYX@JJ*F6|FF6-56| z^npFhrch^o`IZzx_{n!aAxX$%MyuOHMdY*(O9yUL@n;2WFy6KB!oBlcKj;3Rqe|~m z7FK<(&%0sixNwO9zQvUMl3rFH{`#~4iwGw&V6fKEQ1GoAz&*coIZjDPx+s_=eD(zarmVfEAu3Qyx z!MloPf|gVanpiQ#wIj9j*?Hb@Fs?jF7vUp)OnQq#NBcdPUc%5zVI~P_b2h^+Mt0p0 z5!ox55MaOpRjW>;ahN)ty4Ae)t7%)#+EuA0&A+aD*GXYOOE)(a8jowD42-$Tc{m5^cG5Z%*jTlLQgS9FVd1x<8eou;%g2+_#8g&&YEwp|EYV5(xy<>l z-Lz}{JE*+s(nnn~kvj5ks_vZ4`TQ?yn;N^4 z$a1<1(XhB~H^S?+As#RopP7Aqa}!D)r^ zO<*XAIrCSv38tgR5hY`g37w2*2*)qXNIZEFRG?R}U3@isF#~J6SIA>;&x9a@3}}aL z{@R)zT{uL-WWmqIt*kFnRafvgx%Gkw} z{|>d$Dd>}~yxQTg^QUjVHX0hqnU{F2{{t$|_$yc^x20a}7lYwZti_fePagKP=*^r< zUr#10ZP$+ELw)KZ+&n5gWn1~>7Qrpg{f=sZtLNn(tQn;Pz6@dw!ijHvgpwZg;!V!B ztc(Aw#_9aagcPmS!vA%4$_nk^_&eAc9juG2 z)tO&0L1jfv-=UCcXXoRmKS=4!DwU*lq@Zo#YQ5)^mPhT%WYh05jM{FA68EM5Yqv^o zD>AL#1wxzN_ca>l+BV3&r!Loz7DIj+zRpYX*)iCpA`1lmGemTNy{88TXujT$*&>ke zz79teB=A}Ilp!jUT2TW%`|pKP?A(*J%ci?8rtYGbl`0?e#v#$Zo_on#iR_ihE%=5< zUs`3wX#yvyP078=XVjQs?j89jxK#I;7vnP7p_mUtkw5(|m~szl+IxPNO5YJ8;cZ2A&*7srW2w)iZUu?5C8oiKNl+$^ z^SbJf^ZsDrFuIfRoZqH?{ZrlQ-eN;hZV~o;`pd{qo#{^r+H!Z-FiPD0R7KYQA;Lu_ z+)k~nKXh)^Ua#_4toWJ>)vS8%nc%cU_*B5%3lBc;^Bh&sO8(UO^?sy}wEL^fov_F& zhydqoiCk-SudIl&EDzPnc0G7VA%4Gi(@|_%y#WGJGrPMs?fIO*nVy`T1)}zTW%Jun zgTdl5a`>yYiDzUPNYB&L9BA|7mjE2;!4ZX`65fj*S~uD_*k0VtyGy5BqD>-Io?4%F zCh>WAjbw10dwyXKgWg?}i!1Wyzp8n7?09i3a&^P>PlXwGS9@+6MP|US^iGOa(MrKw zYyX|s`}Ec35JIxoO-Ed>_s=XgREXVECqeOxH7-fFVjAciem>}M^y$%)lNpb+!LN;8 z5+~j|MSrM&Bywc0V6bK6r6h1}ul$xwIeG676dR_Yfi9>e;goGwoSU(_`_Jg+qCd0R z5O{9|Yx1{_hGz~}u_foi>n8yc8==Obc{fVWcY$;sqtS+bnA zZDl((JfxW8$w~WhY&LdZnrB(pr5qd;zVuF6|7Oa$KhfwPA<#(N;m@aBUP563~mSKW)SI3L&5?AtlYB4h^$Z_JWDI`Aec`=d9+2qC0N85Axu;Q-FZphto=aOOic*A_)H88YX`Y-5C>pS=YNPkJR6rRB z6;&Hb%el&0-#(`OW{2tNgC4P7#F+NXvJA2B$4M4r1nLy&%WdwfG-Ue5prc4s^F;V` zSEj_{4rcdN>9bI=B{5C!#kXFh7F1F1zUV_sWDSC$O~3AkVX=d=T`|U#)zPhkgBXlh+&GPkesGHRXm-IJ%{Zwqs zt?mYBC>TBllkMXbO((CQA&px`QdF=HS|S4l;}iaP zmK=jjWTGU*o=+N-8+Lnj=az^y2CrWa|9A>a z@u#q*u-;LFLCusWZp|f&*?K?D#uHE&lZ`f&A8Y2C>Q)+)bdiKO{9}uQdQW0<{eQybq4}$m-DR>Rct%suZaZxWj4deHGNRSGFg@v}Il-e`LMe35l9*vNLacn_{No~A(3NNSP) zbf99>?|ld@XXtOA*&ye4CC3@am-~FT!;AE$5my9UbW&z4o;Yjaxf1nL#O2iiq|(x= zB$GgzO_QSj<5vXAfQ5UpQ@YOdRfLtO^?R$@ zrF+g0LTE{9eU${1>{H(OT$`VQlqdd-$pJFHLh&7Is|du_+fR?Vfeg5l09r!wvllsx z{7~WFPkkiK9n6uuLB+D(|Lwz(tyvWd=tSIo2X)^+%+nIXjT$MxYn)31+IicrWh(EJ}B|97lXBJFqi&#|(jab*`J zdGXFVx1vz9Zg&V`x?QAB-B~*<{vL3+2NwfH0c-SgY0iT(*0BCf-Wn4F#5GV zS*w=EM;3=(+q=?~De8s@yN~kvDPNF!+iyJNq{>^xiB% zE{j8v2|*9L$tRS7Qjw_b$(_w-v2;RY$D=pcNQAB>$9ku3d6|2YOjMC5|9l3La2HTR5Sv^I{sV>mz6Q452rE*g_nj^p^Y7z zlq=RS($wO!B{c{3E#}Q<#0FVoz6s!(Rz+@4zd!@W(?K5cHun3g_vW;lv9QYA`#QS4 z5TCR$Y5mWVaoopl$`{=<01twy;hTmh&?W*gKiOD1fF`_zh^CWP=ra-4ut#2D;^A@eRSM~E$ndq8KwE7;Np55%|5EDv(@8i;QOVtA6)#( z5o0Qw^f`ILoCmLi=R($P<)W=Y%;*HK|0u>pU}FgT%qjQPSx5KX@g2hP8l}WAIjWv2 zv?C~rSNg5>(H=M7tD;(^VI!v16eZU{rCN!AwksUD&;P8vPvXx?b)YpGtlvE+2K7PL>lISt82!UE2glg8` zc>8`pDZi|U!~8@P1N(OqYIZwZY=;S!SmuIBrEU*P+G}l4FdgqIQJMxoP}2F?`oWgI zK24f6`!Zi-&bb0+5{vtc2#qs8e^4xbw26?So~1Y6s@b>GF`X&T3pOt!H%BvZ$Lx*b zJvBsM_OoJ%5@ExulnWd`6^Kciy2)0PDJramVoQ7(}tCxp$9zjr*&>->=yu{#u~>Qi>(~ zDg4>m9?rcl_C451eh4)`N%(WgD9yK;_am7M^X})Y2SSr@-p?ejr4*la6O9Y?3AT=5 z1CITwz#~MOY;crr<@^0B~-G|64j@djBG(dY7r~uc9Tf~}2)Ts=(k;ZC} zewO{8C0gQ^jjEHn0?=Dk@a zU3f#8X0TAMo&QeFzBuW1V&xia*Yk#9yB>(mH+g<~**(DE$+|i*3NhXG(Ls{?KEoj6 zbDO(17Y_@cl7l8gPI?^>6?JT6V$s9shqWN1b*Sk34Lg=HfL+<`Z&4gDwRx^%9sl7C z?UT@FzSy|OHkj2{!wSoR_oqwCA^~Rsd9H}HKJbR~BckV|Usk%*aX?rK41-ES6%K%9 z0}1ILlE%XD2Y+ul)ab~OE)h0^;e<#~a6@4>zrOQp5{&yZG+$U6EuYJHiYd}D%_y62 zKezgEb+KLRyL=Sk^$Zlg^<#s@?Dd#4*KtFek=X3@+H^|@VKcd5`nw8%>f{Z&;#DbS zdB|r_zZOn^2PBJHSCJbtnfVWcMo*cHL2)DmjPzjf+@pIu2H0n?A#o%~YL@*8|7GEQ z6SKt9{cEk%?EqpCL}-+7J{g}5Mve5K10Zb(*|C3?ZHKoNQs%o#mP_-NEK{|?t%`i< zMKfp}BKjtXkHpHbvAdAXzk;0qH~mASPkA1S_8)%q=_6PDOt=AFj~LH-M=ymTL9$1c?*f8`w8Q)2mDMHb{h{rB>dKgJs5bbS9 z>pMJ?)tBq1n+;VPP19drP#u*^c7UOgcb41t~9K%}LR{ zoB)7@o?C6*EFU|d87MlqzV~kH%6l<$?>4++E84lW%WdWcx&aHersDw53YaOnvpd-E zQ)xkqJaQw4XXp*L2KINizqEoiADj$~wJ`Fz`067o!bp3_s0Mwdgk2*`!0#tjBJ=A7$;Bigc6?0SWK`^S z(`_=^RO9cVMKr1VbJne5;45sUr_ZAd1K#I@Y{{3C(P@yqYWnRTFt3x6FF_@mt&;%v zb}Ay#@rcrb)~$A#Kx>c8!d$U5)cO@XO5HHT4n_y77%rF)f?v@*&;O=d^K{;s||qN%x0^?9J+ye9AN zeG~qtZ@mgv1t+?3qg`oVmy)COJ;KC7V;H3*NJ9sO@(Wmn3ZCf+t430o)Bv&?lJ92b4~@woBDAo*^XeTnt$^g++e^~BmJ1Y9 zUp>Rb!ebtVqd3|-hRqMug+I6Fhp(G3Lnnu?&IYLnxmp<4pZr}2@Je$f-4O7Pc2e`d z4UL)vQWbSO(uHdEVJ%QWu;P;-zEj??w3ZQ3Ur-ZqOTY)4q@DmEYv~jLS&NB!n+r%_ ztCtcO(4OeRbp~TlEWR>HMH05r9Zog}$=CXqbk$gw4G>hNrT)VG#F`2v$wkIiOnU7~ z2kzNVlGj#!&6&hCJ_mZNI zAv9dukm@r#y&+=k!6Zw>SE9E-h_eKgIjC4xucoNc{V*`E?l%`@P0Nu2Q_qjF!O=w- zVzB_nm^TJ7<@$wizanD@6PMx1W4)AoUkjf|UAqn!A_|RR=L`vUyT*jjQ)iK|eCf(| zY3}=l9QuJDw>ulO+%HK(agJy!haBx1N<3Ft*wwg4Vz2Ngk{; z2qPE|qc^+2w&{G(rUjuyPa19+nL3zUy#;TX6bG=l_^f*FV5B>cCYX>dDp9-=9SM)Z zy8n`EE}70KEpp3C#1G@W-v=#V8Ohnms0e>Gs367tEi0;(@Tq3=KJ)isDWF3G&Nk9O z&?(>(yv=v}4XGlXo<#8+0NUDqQqsguW3?^!2dI<2AS!r6`xbWzTL=GvJK1?m-4mUsistOP?pZ?4ptTYrh0HUzh2P{+ z&d*`A$SO7&U76J)DY|t_Qvu~EsQ|DGBDfUx>;%@37*=e~$7tR+l(-zS&Zt=r1iCcD zC3`QHbVUVVX@j}9Y|82woc=k1jPbG=9Vk`Jpj5@|L_#&^1zMG`xRAE~nXl4_q@y(; z&yowNxXKdERi3iDkF`X2Mdc@~oRpAQq&QHIA`F{WWqe+lck46YoQ0xKPViK#IqC-_FjwaAuA8^ieFLvmuO(^q zCUy!YIXiYu-am2hN?NuSSu60bFgHNjg%GrI%fPoeEi7M2wkyp|}Oh_O*_hL6O|&=v@8Y zEeN)HuRaWMKs$$D{W=8w>=q8JNuT8p#piuO%%pzktfN)Pm;){!Ux|jX&OJ5PqKhT} zi|LDl#tRzo>GZ(=9FOQ#R6)YB_xG1K3w&x@io)beQm=cF(h}B?Lm!3Y)i=L90ENwX zot6bXWZC+Op1EXO{t#|(%#W?28z%v9pC^D6c;7DL3%Pf6lp}(^Rx6g59}YDEHsJru3k&L z%Jw#Dl@VZDTS?C*xK1%prI9SR?rBgma0^HwiZxT$2mTIA`a(Pizke(_Qx`vt@M zEE{K_8RfrK6p-xXa_}ypzW$b!1G)QTd1bZnN9rpxz5Cjun@T+}cE`q2YA$snu%Y;NOI$JG$tEt>9Yk(M>@ z9qIJ$cYhie+MFoD(7Hn1y#wHgYW$hkC7JgiM}Ca9$|dq9DK%#`ren3jN8Tg~uQC2_ zL={a`lpDe#H2p(39xSR{4~NXM_z)ah>n#<0*5D?i*|9Plq+zXjYY8XxPNonhp;_8p z_p)+T79cDJT+D@2lZN>z-@}tGF1Zr^%Da26zjxMJS za?BLr%Zn%_|BlY!jk%HpLX+Ixneadz|M#A%`g9mh$-%1PfRoxx+AL)wHLUL`npbw_ zoL7T#2tqEsgUU2O^Qz6>)m>B9{^Q-Z-gPCManIeA{GhDLi?W6Ctz*oZ%|}s2O`Y2Z zU~b(&bTPSp*}oV68`o zgsjY3ZjAQNf6n8~%(nHCYj?9mWG`KZsW|S>2fMc95;IdCs`h2}jE#I(X#@mbQ;+I61 zN}trKui2!#B3=mN>Z4w18Y0R$@|0QMEJCrKOc^-asrW?l&9G0xzN4o_aZGjG7YC;P z^lnTI{)|^Swek=QG4lNskW`Mis)!FYxf45kep*;JPHkD1R2#XOn$zkI2`#yRQ3KZZ`J6#@Dg(<#v{FsHiaepchN%Z9=+a0 zkEqiHmZ=5zUvS3dMf`=}XS}P8zGij8*Bga#{k1T;Qprh?1qlyeZ7k(Y&OU9`njrXA zLJ9nIe}xM9F+XCUHu^~FYsW7*q9&mz3(r6E=kAvIfUt&NTLEse(J3x6tw#KxP3|z& zpn<$%B*-AG)3>_}p#`*UyibNzOSdeA8x+m3zAV;sf+3ocLT9U1J`QyJD8(2JGl#U* z;$yr6z(y~jwd?%-c<~b@yno+_ZhawFD=w2~nvZ>F-d=;3x%ANBF}# zSRK3U%lbxk*AyG4i=8YClyTuSPWr94$sU|- zWec&t?p=0_##*CQKh_W2qg5Zhh1RIo=K#(NsT7Ftou6^p|W_L+wzOGl($Uy6=HhE-utUEg+s)^k=IYS>3*7iXH~Lt9p==s z)-c(BuixM4E;eAx(m5 z#r@pi&3njZC2o0kpXn*rXZ~HoHrGbf2<^-slqL6hpP6oz9z^R$R-)GGzZJ*6Bx}Gu zz`mk|+z!0rR*ta1@P`n&4E)GBT5Ijs>C*kfNSO1TP!65D)hLt!BCuS31H9lm67p*{ zc++fHIaxqiWGJjbDvdU5oi2Ik1Vn0i9lWj}coSp9T|zyhtmB`g>r#RbEl=;gd6Z3< zuz!@Ty1ZqO*Grd6C|qzrKq+B|z!8af_k|_~T}o~z;Bu$`$htw0QaK;vWB6{zlH`uQ z(iao^EUhpRWugUk{@oPF$~b9q+`cQ~E4~@a9w?56VX~=-!S~47YA*N3p9zusH!D5mwr(4Eo;|?SS{5%XU|s)qoKX15Apxmx=>kvl zI=QW;((WvrEE_`NEnw!z>38&kBkpm8;i#77v97j*lblilWdANBGI`9K0Y(ufMS`Jy z35_|b%0%hPYZ@Rr)W0feP=XqjgvuNVPP3h>t}qf<)QErw=ZQH^#I_Iz0kiK}EHlE)EcdjAmpWQzIl4x7r`R=i%V5hB32%Ht;h!+TQS+CM zD#Q3Na8&RvE8fu<%@+Hv-74IK-7g0+B@v*OnecWfNJp zQp&uc!8?ne&O<^uvW^ttn@i|-HJY+JvnJdS<6u6S^UL06+%+7& zAo9P?7-bPAauXT(<+iLwoH@MTq#)|wq0aQUr?W~G44Go_LF9wpK^phD zwJToc5dN!eNA>)OV=upN&FIF|X85UD2{v&8 z?yC+?%U;Lj^f7G>zx}-znY`xBAg_&(@^Fo;LXG*C0AQfFoX6f4ll4t``+kuR_oG zwNdpL4F7JN_mU(Uo&!mj2Vsy8i*ZD(_4omKD+n+es z#IuqtD@Tot9M)*g$+uc;gV(!;Mev5OsEcfGo{idva~=0j?) zMq~r-iI5rCuqk*$SCHhBFM>E&oYhz*|L#>ogIh7$W6 zWD#>=MFLYdX}Czbs9Goaia-s3oDc6DHk(i1jv&MkI{#^0y#+%?OA)It2rNn5C3kLA zBa5Frb&~U!BEC-vZG{;TP=Z(VNi8U68Zh%t#)pUvaSOf(BEusl^eM3MAQ_(3Xi~o2%=x02Px+xMO@}uB5RV3GKdJQ0&ex;1gime z(-^v2m7ddR_ukwy_KReat@0s(#k-nk+eEr?4Ov~--~A(R7IV&N27KfGry=rOpm_e+ zYmngW^QLA83DXE|2Ht}&xxCZ1GZoa(A(+u6(xg2>V`m=%QXWo~35S7?ATvC(=TA27 zo;#Og;NQJk%qeo{@k%ny0H zgz4qimCgmfur3^P6gA_i7ww*BxqLko6|8))mwY;k#~^BcntyEAIlOheIrn{&bPD`6 z!0TIU*4gS$_jLk54Rk zditrg-$@C(y)v9PtnbrdD!@ge>ml!@PT-;N-H#d*1#7R_w7)%)|I%eRwQg(wO2R4W zFNIzIk};Z`gir(Y+OmIJ$Q;<~T2yL8XGpVQo1fsXH+ZS;Dw!ghhuAiP&=U2`v|}Vv z#_nFsMC*hZptP6vL&y8oId(m`x0O04I?R6|<(ntgXKj^Zo5AIf`QFxL|3(+6S>P$Q zqdg9w1aFw+)lVwGr~I^}fI_@jGqZ@Cd;S@;3bY79}b3MN}rlFuj6@{Y>0q&I1 zIeDF15$tusD5&3pH>OJGhXYG@LBK^lKP0LD0)U4@A@JUE`ca1lWdkl)h=glVB%>?= z6?bCtF}+CUK!cQfZNk)P33cQ7g8ZaXj3l`V|q zcgcGwEtS}$h}R+F=DIQljqKduzv&j-?3<$Z$pv+i}_iMRLOnB?;8=o+V$9bqn zKi!c(F z#2--xoRr!PxJO^H)7wNGjcNKw@Ey=h;)b)8W_rD+o#u8e!>}$ymvi83XjNcNSt*SS zFcnf-%;VEMfKh_>5`whZ9uOUU$WGTe`J6d=BTkjIp@CXe`&?Rpl6mlxp%gCBg;x zYIJhIhUhije1?;XF}&I_#z z7oZ=!$-)&#tZBI%S+*#~N0*H63dHcaZBQX01Mr5PL-|Ga$l2DR+C#$^BI$!XmBS(Z zZ^Y-tpJYtgR@@E}jdv9AJ+|O%do|8YzL=Zeszv!HF5y@R3qezVUcjjh5ML8esfhWE zAL-j#$ItW#48P5s)e-clbHjR#&_11$lxso&o&Oh31Drk@dT!#lX3=^#K~h=u^sk0} zCb$Kb_Ep^Cdg-5GYuY$bz=#t7e8PS@5?rj@zCRw!UKG2)WU@Wyk!Hb@qdtA$8vD(h z-j}nDOb;$a^ou+(-B$oKd84l3ts706?UuE_wA9v@UP5|w!xJB#1(-)xbc1UYkRAZc z%KmTWu-g<)q~ew3kwowH=7g{m0E|V4R{+d^%9;W*AXvQAE&yGK{F*T-8{U3}dnNi) z>762sDocYOR}%U(03KNp(ov7^601AW+pNSdg@xG-^05|rp>Vr0*-!!$o)Np5$~pb| zvMLKUnO}g0v5skA;e98ZXCY8HQu@i_zMi) zn{bbR@0hrS+?K3C9{f#}AhB{bfLF9=7PpTH!CvV@P_36DI*u~Awp<+o5=RzNiUV*U zoIwV_VzVz53_%`1nlTID=j~6{<+fwD>gN=ePuPT|N~nEv7U-bXs`Zj`B}3!z20v{n zb4vHS`~JD=S3b#YncS<;*!vOiIJRBZCxzTroFD&n4$TT%tPeWN{uRg88t*1Ai>9px}k7H8OY!?vv@r1 z^(p-^^-Xm)bAM>Iv#WdT<3k^uP2XK1OORNV4o$Kxw|G2>^f9yV&Ga1-aPpk|@Nx zBnzcm_cR!VB^VhTK8SHleq@D4oZMKnO^i@U_b4MzKMIAQp^dSt;Wr}!?fEF1oS)d#S^Tq665)5ZRFih_1I0yYs zy-IeAYxeuCh|6*)IZh6Fxgk4ay1q>A{9if(M&mhq{Uw%p^5bn_Ek*qyvRcIXkboF8 zOdL_2Y!tYaTmf8o3S4N;ISOwemLu;gqPkH(;CFnqCtizaQO{UTRHk&w466?Ya!>B8 z<3ETZ+^n$O0FnN&Fiyz1`W~u1HaG*R$Gc;P=D+U5s*2~|_`#+E zshAw4D^3xD$^Z!d8Q9chOn%f^9kdh}L zH(&vY__(vfLQVTemBdl>53!b|(f_;xeTC%KOO*DdPI+tupStwsW8fb}M-59d@?}z6 ztrSV5i?AVfZ*39VBBUS5L$NtdFng(KfLhM2RmW|UUKk546>QgkySj!QOL;k#au21! z9RuVQ4WLkHM4IX3&L)=U>&TVqR#kk*gb0sp8Kr=_`gq!DS-VkX(@r2_O(&9G*3BTX zNJtX?XtR($T;@-}i~Yafb;0F1>iIgpI)@^+V}ZVn;Zutm&fL88Pwk9MZEwc81xYXA z+|a*@z-XxB@kTm(IjOzfP~`VAD~jm!k-oGF`9IXzYD#x%j9UJj=SU4@#ixZViy zQ5*z+_+?e*xC9x0>r}S{T=x=*n+9j2g`5iuF!J;3D&p$NqCm1LOGux8PNv7QbQ%9( zi|u)|Cgq#hL%xb59tcb^`%byd01f0GPuEctPK7(jhxh(_920%~ZX1~LeDFwMaJ3&r@yz@EHIdknDlm`pDZN^F?)K(7f`p~H5Hgsngf zJuMAs_zh_I4{D0+-e0pDtC8y2e$A)^)aU%!V*0P89}magA4Xc##)O7S-7>2ZIRmm4 zgRK{Al=|}gL%#WxT zY0{V+SqVmdg{wyKT{;P-X!2~43`6cLAa-miMFfpGpwI&{5U95Mu`OR&j;xief5;7L zOEwq_KgDL5GufC3+x_$cFa%qg{~hd6`I4)`Cb^KBjY){dkCxjJ1T;sDJ<)Xf|r58KAM*j zSLZJ{bXLI`w`X8TeOk&fp1KzlPDJOy+aESC?C}~nT>>X0=dnNaPhlTZOUXTV&ucv> zr{ie|zR>he=r9R=vsFQh{MUMLwF}5{m7rIgb3%b2G6|~fd5oh>%EqLuW2nemt5OV=A;j}@VUW3frcZs6xoB!^u z+dtgKLpleR^;OV}8yaH105Arc00D9(f(-VxbUpGpJBHN>`Fq^v%xYq^lmOc^5S-8c zl#if~uoSoYJC5aGVlPPif14O~0nU}BwF_goO;fB!zy0|XoJi(Uu+{*+`j0gr@&B!& znm(EtuX1?Qw|OE%nQr-7r7;XDg&xKI-|#_Za(M}CkpXT5)SsFz%MOT{EEWMVo(6FyH6?$RbF{LAi+p2-bo5vkZnG1o?!HR?v0 zSUqL~9X){q@N1SDij)g9*N~GPQ>`|VkQy1D2rYq^^p4AWV=Q~vG((_yYFkw1WUvZU z8zlVH9q55rd1$mGT&r!hZJu7%@qK^eFL=ZdSVQYVtZvyh{p~qIFB@BM>eTu_4*+?R z)6nv-0A#{pwliq$CzA;)Z@OH+cKm}_`8ATAa)V6X4t?%ZECr@o8Sq`_aGcFu581Ej zyY*J7i5BxoOMq~#67`GG?H8$=XT+LGz6MaB$iHFvZzfR03NRvP^N-a8o#P1F*FSkz zQRn7$zTO*xyqLzGj1szQHn3d%cWyc)pyP4s77LAbqvk$}F9X+wF;E<=J7?AwQO$tyakAY2cDG6K3F&Z2w8neBZnL0Z0;!p;{Tf z9TKNtqM#c={4?<(m8xRHbYHZFbnPP*P0F$TlGUP<<0lufnv}2KNkPeJP7H6bxUN%& zF8;;cGHkP$`NNVF)HfKXxO_AYuE6PS^K<|Oe_U!t>Y|0Gm1&2`KmOa_6sIG8kGLOi zbgq`ssj0s8q@1^B*#QM2Y4^KqGqCK38ua1|498W}P`OA_Tz~Ibv|znK|fz zNd(~f_y2%%b>ut!ukBba@lOhX16R;5h+Oh|j`~Kns|CFCR2S}MVGG;`Konw`2V-@- z!Zq~Zii-cyo3JM+71U0#%zWd|4)LzD&5^S!*)MsNlp70qL?r{%Q*NPAz!7z0?K~hJXF5Kj6Fa#%mE3S;NYA z_Kz&#Jw$skTN)Ma6p86oS-x!$!ryBou0UL~I&n7R2vx{D$6yb1O8CYY&T9A-*Lt#O zYw4&)&Dm~ic2t*x;wi*Nq*5+f`&6m*UpWpaGd`nHH!wM|;jALf_vSAYj(k>Pb2aWW z9QoCX-*YB=y(@51>A!D>X&Z;xj^m{^Za`%;&kAYXlYh!&77=g2%6E=tVerTqdRGD1 z@wFCO*JK#DehIE$Es-_;NUwGwy^)ABuMhvt#3W7;A3_(gE=Ja^c9*3s*^45ih!K&) z>@Q7w&GBnH6D6yxG)E~<=3ff2h{wjbp;&CNZ*|yUwj5;2;Ur6&e_c#&yAsiSNgad$ zbr?2vSR{7kLXOYtAXXQ#(~2D`4|#vMV0TvAli+@(E= z4MiA^06BNF5Lqbbz`KAX0j2LBB3pu}6vrn|)|)++%Oo)$s@CWqjN=FZ-&)27aKei- zFATEl62*a)mb#X7&#^T#YfIyC4YIV^IY(6_L;+C{_U);yvthS-MHVbVnbrUDCIrx} zK&jdRKOH2kE_N}gS$=zxH<{n(%mi#QIQ`-#eZ4#Dp{eJr2z59q5v__SuyB;yWoz~ASfD~gaCWoDL}4Y%3T7AZGQ)_K2po)% zYc-GRddveSQ4!2s0u@X?z*%a~SN5B1Z zbkxb_=JwkedTW83M|0wl{S|#^{r^W_Diz&o%j@l3e+ws)5fO`L_MeiUVb?=1ac>q; zEHD;$WRxZT{||#lmE@(UC*#;_o^1{XiD0%SCuFhB+Zr(|GH<60PUD|hh^H+O{YGU=-$&13JPENQojNUoxUfRp# z{UTLY(iLDl&;N0>_k_6?2`t@}Sq2R=G?@JuVb$-=>F{XCh^`X8W4W^R=GuFWThk2c zeU7)AolZ8e|qN?G@XihQR0Ud4C^U3NN^7;0WWid^zI>`=jPn>t4cev zL0;Q;C#8xMdkUIxD`%b#WLp+1838_LiNxnvBQ5B-8W0-!)&n16=*;J*k&EAejcKmL zPOJ>fJbejb)hoPq@;|g#ESO)#bOcS=+&N5nw39(!A~{JhYzMX|dS;WwS!M3|@NKo7 z4h@AJKt=)H0;D_eBE7`=Z)>zgBvHN-D2%)Kj;H6V(vP1G1hMh6Fw3`sWz%?LP+u9s z^r7uPf?U_5WA+Py|H;Bfb)^^)77gEb_V$9w!GO;cXh&klBY>VBuX^d&DyiKHVr_&< zD3{7jCkEc&3DFm?HQ-_)pz*G?4^r5?-6tGiGvCmxzWCF6Zx3w`5;z@08`U7M6o1e4 z%hwX{ZhqNmnBk6e@p;VkM<5{SRu8G7%$-6*kQ zlt8MUiJ+Umuj6GSac;bYx&ebcVYhBR{GjEm$VKK456nIM1FcZ~FsaA*Ea}U~u0-_* zQY27J(o7xe@e}TvH{0(@6GTVgS8aKxZE}2Oq^xnnbINOFtO)(Zjj-%5nh9g_#g7x) zD&U>>KZ)Kcq&O_5Pl<0uxNnn1)uFWknN>Fi*QjsY@Bv;QW;7C(^<&Y)UsNFXy z3LRehO-b&+7M@vY>EZ6JnTDMKr&42jV!d62fchP-si0!1-gNs98W+B}#g8g+dU`2` zkFX^Y1zHkr2b4$>e_ar>5Tkmn&;MjpWUBP~HPiv~lrg?N~Q zt3DUtJg(1bG$`{Vk6LS*zVeV7O*+-8`uM|?t1aIJON$zwu41~6tqVNKaBjTuxNRBT z8Hcb-X-k(B2s?Rzx7YjuGuFEVOdC9UL@V)zM$ zNzG}6$Kf?X$`sb+=#6_i)*)@Iof|`h%9x}FdJr!ZUn>3nf-|46Hc!qK?4osIS|Qkw z(0SOitAZ5weKHr#^Pg>6uz1uIVe~*_5-e9*WK_-t^`O3upfZ2c&I+CyGl@Np?^knJ z^!>3~o$JXyd@V{xHqIgrggZxxc9&yK!|3L@8!2Oj6qqPm_$0Rl2keY3lIL;|8bt)- zz6-A|+`Gz!3LwtqL$=gzgxfho)}2@W_aM$^*frQ`O1s6=`#cgMF5|IZmOrH}8F~KT z7Iu_Shc(LU3kTiCl)^27*FV}n$=JlbyIXF&E*hiFaQfkK#iS9B!CDirSkuu*J=Tfs zf@(%0QEO?&KWM{}ty19U3Ib20BpHo@~%!I(SfR=!fl zhDe8ozDf%76W%U%3v|Hx>1vrRPgZeuCa$D)Es)!LBS$&#a3@Cmeea?y5y6&(kDFjG zXiRDA0E?555Rl9XhiKzVWy<+Zfz}0}hF)|Y4ZMt5rNrwR{C(@-{rkP*35d*5^hjD*t%?QY%xON3sAu2Ir4ijp;vxTF2kxX;X8EdK9fO}7`tGy@m>~ zz`9lh6AHa~aty&ns-WWsfBD@JM9{a4*WJ^#((QT~+r2j81_df(M0%i16ZDWS{svlRQujN$!^_G6gmE}#Z~v0E5p>$*fL?y$vZ zln(yroCw*}1HFCe74g*SNY%sR5=ckDFJ7qeFB-{S%VH zCST#Y+VXZDeEFC%b|fPaEqHTaZ%RL#)dDur;oaLM#YqGb`+s%Mh!}@z6(+TprKAT}2#hil5 zuTl3`Z=AxXC$a+@$OtGo4+5LuQg(aN@O6j&+Op$XUz>YUMvM!ifBV3)FEke6zbs6} zPIF+o_+Fzdt#g5|^3({ml=BqO;wRMKai1N^dIA2yp!Bz4f}T0#ePwV0W2U_eCCb)X zQ_NowIRiP9P8O;`$yg@gYylO~uhBv)`B%s{((=pw;GJd}fRNnPmTsplSo~cFgqs{y ztlvqLvUuEQzM8^FpQ_Z@cLq>}WYW8YK}jeNjJ(3b?tn-OOpF}-6(y944IkinFCHEu z=k(-T_Vyb?vPs}g*RX_oWX&4LIsaW;d;kAx^v9I zyuH8>h`YU{cxkh;fsBM9P)a^h=4MWjfh~EK~}(X@E{a-zJh3g6H|D}TDdf)m<0DfR-bR= zX5`LOku@Rw<3@czL6eqfL%E0pAu?WW_cE{p(Jbrb>;xmFnJR-^8Lkfb%}Fc4ctaU~ zKK31L2j+?kQ^K+LRgj-Pm283!+j;6Ee^wej)$E9j?@V|<_cQ3-c?xxZx~onscYL2o z@Wgj#e2I?Fyql=%3H zI>op3l1!P28UH?n(^n8lwAosi?j65UuhwlP`#mSZkrIDt13;nHbA^;2cx>#j%d{Uy zZ7wMw6-_Q%NYcxaUzFkDznv*c^Wu0y*RsIeBzWX9&is-8&+0+K`svL3Mvqvtm=kr~ z9svOwSvh{vD=c69Ii9EYBNhD=e{^Ff2*f#Rj?0`dA+d8*8dyePtQb-xj;%L5 zbUJQA_=c@Ln;JIYgK*v5BFc~u+z!Pf;PulW=E8GAv3_BD!;wcat7rGdiIeh$&z;z+ zB}35art!k&mY!FXkvi;?ehgdOH?~0-C9+%K0>$JKp0Ji*bcr@dn-_MsVd=>DLX$u` zvdqMf6;~{Zr`=1T0ED28klpVw-qDoYmOGVP{NTFqoGbCj*XYvwa#ASB*-@|9!^Dpc)5BJg}a0U-)C$c_~=xQ4hjBU?j z7Ea2U^4v9+(RHQAcbUI6oA?g0^Br|qgAW60d!xna6-98AJ$&hJFD&=8 zr#0j=LCFZZzY%kU1|H@VrN6uADE!YAYr7ObGncA#uOP?9$8}RO1+|}z4xqCHl?foa zU4Kv3MN4#3+_M5khc_8lccnakJ+#+fZLiv6(%XHri0=fbnZXrW;mAszjAwOtBfPto z{?vFk9F-BWy0G(=J*uL0;1n_SD^*9fl5I&{=@7%@!K$#X7>Ro`W5$~ zZJ|uZ9U5^yZ_}AxJS?|*wZ*JZC))+=wTcOLyCQ@fH~`|bYm^W( z*2SFenxbg-^yZump-+kC3IY*YI%gcfL2lIql@(sI#Ks&9@{m|%9cD# z3)oUcCDVMNpuPv+%5gl*1zEH1D)szq%fnKk^-_vZl|k)Ri)%E8IOxB8R7h#BP?ZU= z2rbTS{@BAvuJe^{tX&fTBjM{&ben@sPQ<3E)Z-%}R$~`6G--cu)1lRkaW?Jkp@Yvg z2}VhCj;ih5jnt2rNZ<(65Soj=o6#({I=E-&2TU4;9%oCUb}zCOm2cDsy9yOF{a6){ zJ$l5ckYZI70<$1(n{|g@+K9|2qN}wK$}R!VDs*pbgq45kiMKTOl7*eQDl8?5GP!6t zdQ77D_a_M50Z-G1rNOQsC$m?qgpXi5^u!Y^Em`b_T}b=4RD-qJt2Y=JyP;ZY?VJrq|LxfAzO)-g8ovFTQl+b(0dard2A{SFGIvTM`b|OC?kAah!=P#a$1# zE8eHl()+PIg;N?|FNf=p-3*r*L3HuGz_zrF@1^KT(+|^Ib_nuks4iSYwTh)LwcNW0R+Yy!#rMah35&Jfn}D!u)`nZC6Vz z|eLI#v@ih+@u%)k=FFplON|g>GExp9*`Zi;`ADA z&KCgyT+nrME765VuaWB}-$;qQQ4mC}!DsotBv}0CgbrU(* zHg`3{g|Ny)d){bo2^h!M(B#yvHmkjHhlA@>h)cHxqn&?vL6G09em-Yro`C?y!PCr1Z7uTl1sRib^}Mh;McUboh6%=17jxwys0RHf zJgj0lUj7O)F(*Rn@Dyj>rhM5n0I|VU%|mH^Pc=>b&wR%D%{{}G zogi@Q?yVbv&kMZ3A3R7lTuZLXEWD~-f<84)9vOW>aUZ-a^hNo`Q$f;*wt+o?wX8wP zVTTd}3i4mFDCj6*CX-9s-%)jT&2s=UO|o04__6PbDr@E5+BV9}2A0UrjJO`fWj;Pv zQD*9yKs4<_;a0vC!?R7yl`BRbCIVV20X&E#ag0PvY4_}Mg{VFK5_ccft3MEvZHQh4 z+?RF&&&d18B*Pb$hy<7?C~0S2f&*$5-*tQ?MRFK{&)xIj$k91m`+7|VLSS#+AfUNxvJoeuaC52O$Ok}?&EkM8mcR+YLaO_%)`00_K-uw*GOBw zfBV1oevN+9SgTG*+EeFybFJZ>p2oBktosnEWM$hpTw_?d3Se9!0wOGET+AXYg}sk(stT4ckv=SbDk z0O$8E7n4YkddH3$s#2{T23s((Z4IE3p9E??u{2*AOcMwv|7kA!upJ_bDE9+4Asa2tY({xqpSm!NBaN^LPX}- zpoTJHWZv52nm2E}q5qNQOJ$=)({K+d>-+AgqJZOne5F%+XOsO{k6|_cma192_(MuI z@%>2G+c`Q3HM(&i27OCW*LD3-nwOo(NB>{z55g9mc568;z18{J=1V|jctw6QcwFrS z?_c|(3hVQRfC?XLr;>x77m(t{T7Ei*?pq^!%_7QZ-;?d8`@S~4j$8T`5n+F2`eN2+ zG^PfmokML9MOM?Apm{N~k?LYYnAD3(6$|w9t(ZJ_E!>h6W-;|!*ryNf1S9-if+52Z zo`S=embMgZNpA#~J3cp*KQr`2_t}^uwIE7|buZY3wJq?pcsp}=xOkP7itwDDr3qptU`FY5n(?rh!a5~N@h%ALHO)sf{as0$8n~I+ zXpqCCneCqK`qVbeJWEH92r8jGEOj#ytOqv*)$8BjP?V~yXv{;e_~c1tsz#c!0rVzK zX3@u7!+W4g2E;YRL)K7LZ#M40{UE!vE;d1+vb)q*5n4%MLPU?>1uL6`QKn!esQ=#! z>j`pM@CTp8p^meQj-{~Wr=2*;QM2t6QLUOD&}c$btfa01lwBg$k9jd{r5n!h^@QaF zVw%Hsvln*bu;40{C$Z(4j)z`Mnjvr~q$J|zi;*W1f1S^Uch6LZ#L@4}xMw_yCl3)@ z$<%$YRuo^2Mt&d};^?#~=d#kPV;91#GF=k5|6 zu0u^JpyZd^{m%Q|@L&H{px&e~xBa+{yjv`$;;lHJ4@*`Y|-V*Q0Qv=e=ie~H0&Gwm|$dw%kh=3%wM{+I5Ytg*>N*sk&dy$Cm`SWCw zdb}hQ1D9vw@NzlAY-{})oh8n-|JTXAV0ZY_qkY)|g94QI!z%djkc-oEgjzPwgB1Ng zLYWjMXpIo0<{+lDHEt?Ey$dB03b)XjcO znw;y?xpgT?ZF&yx1;~N^xRPn(y)rlNm8jaf{6ea3Dl5&I>8l@c`mZ4CnnyRAh7SSj z!h?9Qk)U$KEw;eeZ)n9l-6BgY{GNy8|7pE(;?aEA3RO?UD|IM16 z+Tp!tTiPu2Wlzy!$XB#S|0kYg3kI1 zivbyLXVy5Q{A&ZwzIzwAzrDY=EE_nszcbLdzw=nOx_Nmft}2;Tn1XHH#+Aj*g#U2{ z8$wopzbEbY6e|aXXtVXU87?rbeS2}zd2d7pm65kyp(-)L1eLN+Zua&6$PMW=W`Nyj z;p}Pb1RhV}3V0_@cH_5YG{@SL(wj1sDl497FGxV$Br;e15KEXS1K@^wN!(Cym1DzN z|n;*OJ8)@^4ucz;3gs+X;C*cII z8@=Yd&gc}Th5*!8$H$)d&H>(ach>M_om75$-phq#cerM+!PDF=*_W;%(h|N^XMZV) z=NT`NWD>$(&rSrlPb+ofu`vYvo-Gkni|0@Aq*?d*$e!(dJ_f^R$B!nU!>dvi{BQOA zv<-X{QH)yzS1em=#ukm0Tz1gXS8&bH{`GF7o14rt<>~GjGdyv*4mAjZIFfkz{`w+v zY7X+;%AFm?hm~*(5Il5f@BQ`!x4rDEIC;YF z_}0A+df;s%>=p`(@-e5sZK^pO#^0nEp2l7IV3H*WYB0@n&scN;q2PjSK+nD0#sQ}= zS;04U;c-;TXRT)x=3btPJf0#&Z!$u237me@H`MvnsTsei_BW)~BWMR@m$zOpR{o|% zsTDl8+%{|-s3#DBR$=uyy&IIlwI{f`DW48e3K73r{=MW$dx8(+UD;UETLDAp`e1(c zjOy}x6m`s)O%S=gR-$Si zpj&ByT|gx>_zhER@F7Y5$LJhG$L-ah#(u(;OD5(I?%dBj@y5XLz0}4~fQ$%$noVz? z0t;<|xkp9sN_-=Zno_S+EAVw);2>yR#)85)lj;Ka_lN#7vP-)<>ROVk zQW&;`MoA_FYLnrK1k@1*r9t?`R)AvXx=Y0slqNTD3JQbS`_d-h%s&&Os$xU;yxB!u z^(O6IL2L0mk!A+MSm6y!0F@%_)QI%wgY1;7^}oN8<$b|B?e4Icc=s?KKB?*-P)4es zFy(1K(|Wk}JF4ycBwz!hoG?}~&1jq3T@sQXgcF?K82u zflHT+v4jLg*Ye(o1`A9CcW>AORel%QBn$AiA&2d2zU`%Jukw=SZ&?g)oQ4{GToafp{HLZH1Lj*0=OyOm5c(vJojX+QfmahAT1-Sv zy^n~{C{teir@IZ&VH>Cr@IL(J6V^{;yYW?9D2?p;FfROacL1)KX0k%DWxp7SG|GX_ zfEs`{B+X^kYATx6m8F3zxaAihsAds>tyw_MdP%M#XTWhY3l_M8%!0E&Y^mAg>LeE%_u8(6EKTa; zCn?O8(;v`9#ScRL?0MCK0*)dv)rFKH>V&+nj8r$eCb-MkrvQY_1)dAQ+_u#E?&082 z(W|UGpPbx&d{0p<2K(x;2KypVF+Ij*Amn#kwgIk$dFI{S+5B+xkb}zMYzg%li=XG- z@-fI3XPC?`h00eFO2rMJ+Af(ZLG{oC0MsiJULJ~O=1;H`gN;k~EGDT}g%ZGesJ6ph zq%`(Px=S!(fB|H-Nvaqi82TE&KBiDV9J5Z#Qcb-gEnnOWAW%i$gm>Jcu+7S^VdH-j zFjUJ9rqLS(t_K~7d1STjT70Snbdmte9@=y=+8Agx{-Bv)h>j5(C%FNH2#mF|Mj{p^tgnb`NQ%1|n9-P#jUAA+?3A%r z%jy2L9un9<-wA8T(5F0c5v4MR7BbTWi-BdfbLgYsmwYP^R)2+D?BL#hLp=ibBDXn4 zf)&?MKS%V00T5o9SWlFqp?|5v8ip;!h?fp=ShIg3snQJ#1S-oL}m=XogF(+nc&F^?&-2vd55_yYr^nexvH}VxqgzYx)%y z%*xSE>^s#e{T7F%ta0^&YZf+~qIfuNt}PCUD@v)vYMbIKFcGZ=i9 zej_us`^%QC?n37T1muFNl>R|xFoc^x@5N>;+`TxIVLPdSPZqS?O^ARR$T=xlSNIa3 z4?uc>A`c=Oz_IyOwfaffpfnle8K%D)|0&(VlCQYTe0|n+Q0E`8)SgjAyQZcr9tDI} zYy+}>Vp!xQxHn{KFN&2uQo2-!`V-l?Pzg+ z#jG>BA4wPmB*fOiod49E5WWO#>x`0vGQ-^lVMZ{^R-J`~O#m4O(XkMGaZKh`+O0=s~6<80S84P-@#1djTW2>!1o0eu63c-1V^n)xqbiN^w zQB0(joWBi;eK;v*Xz#-&WJSoPU!{7O>b--?{BZSHH~A64Ai&O?CRqTaE`iSI6x{;- z{bhoGAp7om6*8)Veo_~@z=XJjSr}4Y9x?5K;Q7bkq#gD@_#7r6GwsH9?uPqe*Ap>7 zzJk{CJ7`$QGuantgUtB{YLQwk>bGxd1>9ip3dCA8fhiNs2nazk$`y2?BTt~oOpne01iTjICRNts}G7p@(L8yC} z_GHK%kM2zveV>vWx_mhq0+RpmoKVjsqqPe4qA52_&NO#g)&tNMa7ezD09sLBk7W9| zfNt`l7g8VE$^+b@H4e-sqFQ>EBHVr1I#(oK3I5$T^CzZq|#$amAM%=NZy z^tx#w=_#)zcwa`@Lnf+cyODlyu7zK~EA{eT+^X+cJhvTsGMMbIA~= zeRU5{LYN2eRM%BNP4&uF+so~^lI@DJdN9`T(q&S|g`uVSM96~W7r8G0t}IKakrj)a zj>p^a0($HZs|^ulygfz;>j#uUL#(h1FD`2PKD9XgX);36a7)_{@zFXR`R zn+EL3u}H5z69%^)q3PGrki|$OoLL$e>lLw!7T}4?pE2mb*3`SSNg=??b7wB*6h~l36rc>07TUz zyWe#^^pB%HP10nd{=neXfpdZ!7MMv^m+-SK^Xv2ReDH%>guTn;)CxduPX9GY&aDq` zXjiO5oH5feDb@)KgAk3!TG6!0CyFMIQuUz8Q9lz^l`-N{7cTzvvdJ zDq)-=`QEoI^**=TJ6f}**O5SG_Yx3Ne6m$(6Oium)rA1o1p4n<5?g5V@17E0dpKXD z9fE$KZ9ESi|Ysu6M`80bM_+q=O{VuI%7>M*esA7efuzCoR4fqA@xz+e&%gREb3+P_{X`x;_FidC>38Yj=4^6sL=aSO(p%qD0G9F58|x;V|l;Q`MW z945GDqqRW{1(I1%PH2=a+;{^~xwZf&kR+`F(c31|{M_@C?^c#)r}pwq{Q^|f%>q=A zYZ>S9zNK`jdP)#WWeW<8F1``Ll0jDdy^!?hTk>Y7u01{nQ?GA_EdKGdB(AvI{2^oxWR z!SVANX0^>I(=LUi&@N|(dEcKx|2@;_MyFgQ>;EswD~<*sIQ(x82=tK=7m zbNpPMo4cmrs}jHz0kk9)s#AqHyN9JBQ$*y%(uV1kGu4Knuq4b-b`OMhIV;NWf{YkV5f&nOyA<}cFMtW$MpTZHCr&YCLYx&8 z0)W!CHT;UV>`(4b?{=7sUR^+NI*iw}%DM12CC-n8S^IDNw8=%}^nBjn<-yebZ1bXg&(|9P z2`JnqPJZ#r_!8TA&irzlNpnbWfsAWLrjxVaWd?e+)IgVn{vURdmC!HrJ_nBgilH-l zB0QJN(F3Q2(B3T)S|WfESr3cnE0F-g_V=pKLUC-IY`sRZQoUEaz{jYZAQ2@{x4qW? z8PSknM182W{Jq`0NJy@4EvO6)`&RqDv8VEcl|tDwxaNvQ)FEe*?D7nzpnorYnO*1Y zn80KH+mD8GpGPsdJ3xI42<43s80JehB{?5-@nGn_q*=>JYtysf)kHn2jj03-w+z~K zTY}`R=M*3U5NRCBKqUI1wknO7A9}&TtDgKRL|n3d4c9|U3q`BX)la>>Gf9Wm;cIa) zkp@g&!fNT5BEib^{dB@#@^9{9vcmj{-LnM=CShb& zIDHvOv2{s~;E$ZQHJXtRF>0WIDN>wxcZu)%7QIJ56hA*!4dGZQJ@t|zzX2S1qh;d5 zEGudh7XUR005W+^cN~qpk$hauUT#Kt&Z0+8=ueTEf)D*#Tb{^E>kfkMCwi4X5nWbj zB|pRXre(|HmC4+hpc%O%#<}hxmwz~)w`F{Z=sp=gh6t#z_NwK&)KN%56B%8l6 zTs>eUV9oMQ!sVF|i3uPjsbc{;Zx|w$iBD8Esr@PW$%o$J=bpm9^)U!3q0&RyTue}UZ7?|@IZozU~xy>B45m8Oh&Y6EBU$!bZOC_>{z&v zoGu0WCNDF>xJ#B`Mi)(IzzplV!TAP`N(~%6-5k%BkSiA2T}@{^>CQh8#96 zz!dIckC*8aVuycjG|hi~&{dAUvLmJbf2&At1%k(9X&&p@kq!`@pDMh`VkH#sP2%3UXEsE?`EOov3Lm3e#N( z`ky#8!AtB}nfTTdc$(tuDZ%JHh{n8lr*UCg)hv9v{u|-n`a7=-(zA~g(Q5~;(DqcW zO<2E7+b;mOb|}74tAwO)p~?A-&M|jf_Plv$_1*F23%j=h0oFFzabq;lQa(4&Q}8^a-g=syhhD%UmqD)MXp{luhp-c+NQ1KauYb2LldaAVx`0xa(SoA?0lZ zB@3!Dt$|TICAgUJRr9aEQ#G^7KC|0UXI?k{d%goB;8=5-X!;U6J`1##J%U*I9?)9| z1!|ZdLo0%kgzXD^Z)+_IYp}@5ZL}4l#8}HT(VcQm(c}5#bnpJa z)9M1qah-ZO+o}b?(`F>*)5jFKZ`p?9HXiR99dBN>1^=f}@*)*I-@*v0Gh2Nha`aH3 z(8#aW+)O!12bifpL||tyCGqT6yri`)fJ`J<)>%eJKIoI#QcW zdZE0B;vkI=NwFI0(>qZT>y1*Y-FBakC)qv_jhMau<{P0jePX5CD5)IKOi<6@fH;!c z^37m>beo!$FUqHV+S^PIG6f)_=xa8*sPG4&{j2DG87O2kEqC}@mPmH3tVbsB#Plkf z2jGLL@|98b-awBDDFqHr3@8Og2^4Nyk#!0kF?o!;?>9)Bc@73ruSnbFIFr;{{|XRx zA0pN0@Q{hYb`THDl*ozKl=0g`5SRiH)K`1o-Vkx0(390)z6Od}1aaHAzWhDl{%6Lmu4SU>a(hN4G4OzZccsXYYycy`t3-NI(DW7D=_#&1k z9z6{yHv)r%)yprIao;2Mz$khwpqXf~*Mb91ARF8mGwxR&80kdXUw{-cL=pIdEyU+x zc3AF!#NsQ80@eK~bjBRO+hyxp`05C+={;fobob>38e z2BnJKpuhZPuKn`AJ$StIt8?X9e}t_q`s5^#^0ph}E>McHh(>|%YKvY0b5)CAAhMWl zlLnKOkwo*()o`CJb+q1Sk_G{Q>1T9bGg~f5_h>LR^dkV;F&gwF6{%gh%$unAg~l#x zvd=6`S3kjmE(Of8&I!Qdh?~VgO?#I#XbJAyYwXvXAtn07b@m+TkLzEC6GkRL+n~;4}7%cJ zLUdqD8#K)D7-e)vaD;)*aB^emf-*nfrLuZ=s1=g_h0Bz3->YKXh_~k~yE6!>;igR& z_rrw%Ag_eZ^l@2#4EQ$`F;bHeTfXaF2lONZAjuSOXhNd7$Kc(3A-x-Fs3&qUNNv8% zu7SgC=YHV!(VBu=z>Go7of}eFgx7uC;4yomNTNj27Mkd!6PSN2YcSTB%W4!@_j z6a~D`#(>W#eYt3wYktylK76vi)>IqPOMxq(kunI|bnUU!P5Ir^u{w1%B&ga<)8a_t8{akg5{ZTR=cOO|5YCS0jA))0Zc^a62Y`b)|JH9YAsXV zWiSw>$4JK)>*&08k| zqXZ72=q_4{-H#(8;f256Ka0ln8!NbC82Wfz{6 z=Yqp?okb*uy=;QzsB@+bph>L%CG4D{imr|=9g%UKlsbQ-BWyYZ<6ef^g0xYGE44~N ze6;%~KH>zI$}`CgeB!sb27QX2#kfP`G8O?UgkwX1-XPS8#}I+emJ~{%rJv$O*xy*0 z(Vg^Bq=8$ z@Ec5-IwWD#OP_uB_(eH>&;^6I$yYowqx4*inx{3mYNQcsrGa`P&;_;dOZ1tx z-~v~Om3q+9Y6pmEYea%&OOR#M?KlHIyBRv7*2{ev6X20@pm`u9J>dqv*K6-1$?uM0 zoTV|@A1^S%XqCV92YgO|YoqDCahn`C1Hbv51#W{C)8fRw*-E(&q1EtyUyx^@P;t-16$#4$JTqH9v<09Q z3cs9qgYI0AF!Io?jiHuNt$hEz1YY;mf(F+rRULB1m5@z|(x zm~qQJwXM2vGB6dZMQvLO&q!oLfDy<*(g-9_aq$}0^%@^?ucn!fb$c08Zm6)j!hFM8 zu9i;Gzsm**W(G}DiV_d>R?&#(MYMP+pJ|aVDe^a^@3FSfRk~8T+fNAtMUk!>G$@8@ za_$%|F=i&5%-flJeCCP#`nx>=cX`%3?I8`1dVSS|R!Sm+=6mU^%qQg>MG;!9fVw&k z<`Vz;4LhLeG7Gh+5Y?a}PYU#bn{3H<={5ur>%Bxn9{85RohJX8m4uUKCE-ggooqg~ zIq)_H@BDem1&r0XNqw}s$qTOUT|M!urLsT!9jB)GwGU)#U|e#>bFC>LAfP%`Jqb>8 zqh4A}YZJ(?5=I}V!0$@MZT(Ix4e$yD;;WCNlMWAX+h&WpvHcP3?_(&~hz-?*^3QX3 zf6OKUb%mhT8KC-*rp?w4yqS>-x)#v*0j8og!E0sCEz>+oz~1NF(GQK9`=rK&Olwrl zP6b`&dCJxGFiLs{TJN&O)&1C1qt6^tZdvQSr@wDX#GKvL_RkvJj~IOLM(+4~St`nY z!E0v&t|rFAhrT}sU)Gv;2nN5ffprqJ@58A@#)at9F8TW@8dPtP!F&g4`tN3aP6So6 z$y@qJS(8pp5N)^ZHP9Lp8JCKdig=u!`TFQZR90ke`?Z6QIekt`8zNYdIUu#qTnJpd z1Won)YIj)7bLKY_#00!p|E7L+BolnC0xwjE2OnL!#1ym+J{f=${C>JM?UfA_J^g(I z)tqVtP$y5D$==W0D9#)R#r97CZe`Y8mo^qE-~VjniO_pye=EfL0?LW3SD?poPp35d z1FS}^8@hG15;OgcH82f;My!zvm2#{nqN1eT+jox$1!>3=^a(u#Tf&+*h|yj zw8Q1a^ZFF4Km+jm>FKmrHQDzD+4l>!l3||3anJDVW7=E-BYnDigzrza-lbx^KO+)e z)F?3|bpDVw?eM6I6uR(0zb*~J*=zlpu>Eb{Jaet{1~l#rwDCOYrV0d5h#LEm4X=0p zfT4Ozo*9T49zI;we+;#S9wJoOC$#IBWtuK?QWioimLsWmPBRJd2NTc^HWsmR3L*u%01HX z;LxBwINX{q2HyxZdMwFhMdz6XKin2tcnJ<5{fkb50Y#1cqKwUQW%o^1j<@twa?k;5 z(Qk`j>;)qyYccaH?{lYp4Z8$Aqs!fkUm}2Yx?ho{*I}S{_{w^R>Yw%YQ$$5FZ`6m0 zZEYIQojSt_BK>DH_=?Q=i1Xvx5J39Ry_byZ?%bQzd2SrFeq+BobWg8%gN&l`ZNso& zhUC=I`|D+CI?ZH$hj&Z+DZr;#iiIpUs#Np8-#3^VRgv#1+O3OsE3WWUMCP29TX+Ep zXMU5dbnA=4z*j@ImjXWPJfT6HcV1uDlBjGO3vqCLT5V|3kO60{k%v|z>9@CT?7!v7 zmzge>WD%tjaMqu@I#qCTG2gH;Kh=QKGPB``;Ft1(Z1BaL7Hjt_jiA|)y;n6t-y>_{ zsgCSWXc;Brac*p17?Ol7B$wqF)_)lQDQ~*=CIDm7;~|<9Oo_uN8TK? z#mIY$@&%c`fu)-G5T`iZvs$JFXG#&V@=X`O>NwDA+w&B6|uLBenuuH{7$}y{tb~P~z{IquUER+k+_)avw3G%h4 zS#{QBj|3id1DuXD>EzFCCkr*u;Dp|Oc8}t@yIh>U^9#(`QvW6k@6s3jqHWsTP`44Q z%=_uC#};kW^Pe02<+F{zkV?T~4wfWiKiBzOMFyk+T?%_BE$Ex<{8G&h{b7(cpi&f^)XVVxPwc}s8K|1kbVgwh?lIK3<& z5Quot<9~hT-H9LF1$SU+51ugWg30{o^JrJF#;GYLM!4A_g#W&Q^^($qIi`=l$ zj<6WB_)sKv@VEtj#IZ}NU(R88UxT;4_>dM%9PzF7d##j)PxO=d>zwGZ^+Or#u7fZR z_+#h3%D8IW$5WrtD`eqk{UA;iZK13+V@)eyhBEGB^ciUX`5ThJAHiJp+-y!~RQh#C zFemr)#*riB`p|5}mw{lS4(%=R8>e)fbu(~EtskPg&U+T$@~h3XJatIv*GZtXIBX$z z=pBX>LhPd=?i;iP)1A~U-8c&ypy8zp2ebg6I4e~#H%uoC_M7ZjWZ}&_B5@B03y$EuuS*eUNSUuSlcVru_d1Vtfn?ua({Z-|ey`q;k6 z-WwmB+T7oFl5M{rqZYL`3O=xsDi^~4z0X6de4PHmuM)oXrnvqK#~vM`0$fMXe64+y za3?AzMnOMF_*g+z>RdZpWs9i)r2uW-xNkQyh#Pu^_8)-{P$z8PzHX;TO@0J?Mg&1u zv!F>aCcPSRJ%e2T0nw)4+%C>o-^S>OAfM;HmYFFqZ_Qs|xRSq#;!Nb5#-o_L&$OgK zHhlA{Uh=Fx8^4_Ewpl}Z5q(@&9=VpH%e}r(#U8W|rykJ$e;KShUR}-*eQDs!_ zJ)fXDI=ZfxL(en9@R^_onON=TIN!c1E&AK`U}*Hezro^bFcLi#+vR+xhZeiHM_yg9 zw|xX-)-}O$$vNv%{6RtoQLY`7Zrkgts-59~I@UB+Tgh-@bn z9>TX?X5TC=;0V6zqH{e){WYIFuTv-HLV}FK}O|Ebd#BG*9M8G zn_K$KSeZ;{0i(jCAXTO0N_W*VCmz782STlT;5mqE z_^-y&GztHYt+$SgvWwbA1*Abb1OY*55K$39=?0M)x)~${0coT|Y6c~wrMtURfdQpK zI;Fci2F@Oy_dVY^=Xd_&dEECs``)qkwXSuowZ{}5>mWi%bt4;>4mr$*;T}s$R zIi+vp&oc1RTWkVg61dGOzV=FvRYQp+Y)rpIFbF;Ta713nz| zdRHeH19){Z@(eR!F^});48o%R^xhmr!JR8~n66ff;eg8GiiG@WDJj2*j5zwT!2bX- zFw@(syi}VHQ+#6r6K9_QV0Mk010#d7ZO%!sqB`|1--25h_Tu@*+X??bu4{kffiQA zPbf7`zBarM8B&@zKdbS*3Rt;?lZKb%K}{m++n#8Z$=Q`DRsT-6X#zCcSs-y)vbbk{t?vfv!z`&TcN2V_A2ati zWJVoHA4Vls7%0&%Dt}Vkyf)E44}^OBVkvSTFJ@O2*E!{_z}7hhLlZTw`RMm4lM`cz zCIxN5XQMc5t-tj*TAIyiuB|x_+}L%*HA=lRcv78HosIbj!a`B>TG*-%fQQ< zq!-IoT@+w3;X|H1o4!M+xne#!4x3(25-XI%!2PmNOIP024lJTMB;D=n!6eYsSAL1w zQKFz^CJ+7M4;%>6OUufobLBt5u1VP+I6pCrzAj{+u~L<6$yoFtP>l+EEtuul;$76u z+xd?;Y6nfK)X%3LQV(u6bwS9?g&EM!F0d@VMJ5a(JM#~=^+K!oLI1&+fbs942mOcT z2|m4?99=+WP5v0$WRxp8_=3!s_Eg+7;3G%spqe{@<{LkMQoTPB-2|vG3$kS;t~Q#v z&39FcLHI>Oc9LFHfrNs%CDS1It%g15a1j)enWd zW3`yP=2}uK#F-B6Kw&vPJVx%1Z5Wohf5MdCXn5vHi9#2M)J^mBmo+36WJpdDYMVOj zP8E2o_akai3os?T*7H4h0(8h#p4H~2-axH3m)~26MEFI@AvTQJ!%zz@r1ZI+`bYaV z8iHI2&G6@m>@>ad+x@6_)7NQ|26<{6`)Ml+bKdQIR@)qu-XO77Sm!(H-B1qDk*TxW z*1XZ6^>q|w!sIErKBx4RCGpHBnUbo4qlFX8OG1j^+YZFxg-oS1=l;)6LDt*18S<+l ziV+ITO!;ir%+;zR(^Srm{c7~YCOL8P^3&H63nvk~x8DKfNUCj`CCy4WQbiM2nHcsC zv?a1yIN^A;ZW~K&@;%rCFsF8xNGm!>h(}oa_9K`Pri&x1jbWdZ(@OrLOy+r>voGVE zE?Lcu)>TsyoVSC;4xDZXHQ8z54)}kaLygYx8_BOmk!*$$xaaB%r7Si)nln0g|Im=u zV3R}M0-s^-(SMH#+PN1=Z7&Zo?}cCYGnTn?RBL|+rQlVP@9v_{172YvE zZ^u|B6_{i#yfeRa#i26LVJ<}TOLCEomk@s-Cqp?TKsL@t+V|KR_ zLbP^`aK0QiIoZ=wTWdOStBB`En8}HyxE$FGMpp$^1#2*BDruB`a;VvWS);?=Q`pAN zbgeF&GWoF_DKYFO zxXr8p_w-w{XL}pCjL)Nu8cGNex+>+G*cHXG1>_E$b-uRI*-D-8ALnQYVy9h7&}iE9 z6fg(_BF|ug1n;F{yCdp7Y^g?Oatp-aDim)QMP-sKtY`sFEg^Qay|+p$D53vb^RZ8Z zdJVb%Jq&FU*wn7&GQW3z`uCUXm)SOXcFfh0Pp#NavW)(FfXyryACz|gR0EF1v?mQb z*n<7t+fJUr*K$E6uzN&2u~~P!YcwlbTEqvMbbDR``Br#I+}7>a=63INvo-UKadf4c zw~y%}-7`8ab`uhiW?ru-8Q1!hV`r3t`0maOYI1E)d-f{#fd3KeaSXs)eBp=PRn$eC~>0o=82eDS2WU z{a5JT-~J18>tnT|G(aHaDpd$1W&qkF^)u1k@Zqh+ZFF#qmPssUyw*A7-7rGW%5UG0 zSzQ3V0v{0@sGK^OX~uy2Q2Uq8m%fYetts!nAAy&lkYiS&a$o9$mMsa4l3}Bh4r|8B z@?;3IN#8NcyT-zYg1Y*$y}%e`nh^m|s^16x)%U8wN;!~tc(0Ea#+=@dJ)_}1$Ix+a z(dbR`yY{AqH+2~h%abL1oQHW6U}tO$m`+*l@?`a}i&>{}%xnF>#b7;41?>Xk=PXp%NDcvp`xKlBZ+R&$ z?L(2O<=0Bu7th?YE%^v+m-w=W;A3+cYhAID3&6-j822xfr{xSE1EZEjRZ@cNPRL-v z<&onH#eT#QQ0j~a)=tX9hGv08+qJiU@y^kk&3~OQZcBm0a(PyR5@-C7BH1ME{;9e2 z!Kd#ztQL+rmev;%*z_xO-`Xrixe9>?u$%8lh1Ui(WRjI=iX{zvmp@cLEysHw^Z|m_ zR7vT+f8jn{Jn`V&7*xJ7*#luFTn-#jxfp}SYv(~{;y_CWjb__vjr^JrKJrV|Ql-!` zxklk+_hMU z$8Gjn5SaKz;RB&k@cy@1sLssjll>%LyUM#1Uk6vgHL)1*Of1kNs7?I;c_a>(U&c53 zN4R>1QlXqoaHusJK#);e?$P+txPCPUS{btDij4;R>S53xqHtFsP&w;%+<$ZA;%~Kh zkUF`5v2YuEbt>~mBl_j4-KNjvyRkk5Ao8{odgrg4tS(Wt-_if=PqVQ=yL*51Ea#8>AluvPMj;2O~?I zcg(6(s7QV?Evv)YhG$M6DEVA#sGq)p?~H7QJ&&Lbj`zCUuE@c}*K|C>yi2^anIVti z83FvCwId(Dow3vF(dBg0T>Hqur+n%=h1P4|>K`z$)57u|*8?`lg4_}RHimE=uuItm&Lb}sccQ_AH4h{3FlV(pYGf~s(u zaS$evQu3=kaX4HP4*`0czsyrwy{V~ew*y1@Q_G&H7ofd6)wWW5YACM61LAZ z#5bnucD1Ocr%%#__P?&HZ{$+s|ad4ggScbycMGBTaV6r~6 z3oxYF>=0&`FQN)3>MrU5<$CN3Z28kKOqpiTo*DKT!`;|^^~>!iAy^cxVTfP~m12Z| z{^#*WmSZ39c*kaUqoOLOCBf1WNBya4y^=pp7-c!ZJjmLv~!+rF9erLt(Gqx*1 zRn!bra&2JAaIB-#oRp@SykzbRo$+>!h4^H-K7KOSDRqcHTv*_GsVzwpXbZOzKaDKh z^h2;Ul_bdvDSF`Fo7yj2p>@x+~ z=`y^{+>z-XdNB4ZfVetaNODgrOK@R(Qnh-5S(p9#Ul)m(35I#kIWklLKdQKXJ;ga= zC23!)_kV`$mQY0v@V8pMXt>u0`EM+Zxev{?XsfdwZmRBG(XU9ZC(J<0Q^bDIpFCfm z4svcm!k>dkSJvM;?|d_jJOL-Irp$&s-((Z=u}Pgm3}9jz`Ps+YQ0bKa%h#eB-{iIy z6LeCA!s?x#U0>3^A&KQnamfUf^M(eBPw4CVOvwz8y>I4SmlXlMPL&nZthK zQJ*;%2esmebg`BDeuWEU{2!(}Wd}7t=QFeZ)_TiItY@2h1WeNOxHAxSo@=O}*s~$h zXsee&kRkN?Bgcw6gE%yVmw-IUmQ~2H^X8UUb8A&WI{(pOj~8*%DjeFyXLH^@3(M7i z#9a;QzC=)}mV^zX2)u%QZv2w^mzvM_! zDIn3%>q!!Z@qx0tTz%qMDN+yVE4WL3I0nPdej*Q*C-%d}Cip@wZF%f&Uh%WwwMVrJ z=^#uL)oNgZg14-Ypvl6rd;T6wD^r7TQgLNy9@FvdtB*QGfD(1{}PN4%X;P+5cLD z@Q*!1oxMiM{^XMD{=>~jH-7PDL z{XFrqmdG;Iw01+1A-b4owUDzbikX6Gk@Wgr{w(`2z<53yxpRQcf%1i`$74B}W@B3l zCR3zR)HRt7o;`j7X4;UmYu9R@e$^^dZO28V0w}t$macEVfQ0IDQrY9paAQHxM2=aD z2;Xb-H(Bo$!BE~69}<$hJvo@hvn=vyvvt9uoiH9*?_#GyyMitXcLg zm+{*cskDG#>_R!iPfn`RTs;kTI95i-*PdZ+cFvS*u$oDoSCg8gKg_Z(5q()j$|yy; z&X7Az7MSM<0$$7pFKo4!@i&;L21gWy8g;>BwKR1#RlXEQJpf)`D5;zsc0?4KUWnmI z9IU`m%Zm6jt3$dcET($N7)@owtc|Y#(d7T9mF7atNNoQqS-@)<$B{|pmc3P^#`eV6n{`c%GfcivRdAo7`0z%fJST z-6|mq(Le>h^-$H?q@h}w0&2F%4$s5`ahg*fz#2KyobJYUOniuGi9F}vWY0=zYshLp zz-K$jTqbGTj_N?h7T%G1$j!Se^I!Y%*u0c{#>@phmgQU#iLdwD(4tYxr{E{3DF)hi zA$QE3YwmmqF-^&|2R;_q$M%Nr{|>jfN52oJF|^^Rr7cE~cB8_GmEablB-J|2kk1>= zz9*Ey6lwAV|81CSyC#c;LZ-yi1~u+x@4M4hm06l~-jaDgY~`2fi;QG2MuLaK)czY$ zwH@`+GqswOWX7)UI2zHwh`xlbCP#|P*MdVax4kLPkXsgx6j7}MRAq@)*?8Wx;SD+W z2FYr1@6{E!b$RshPIU7+VlLumMAHxTh}b#HVdhL)VhjBm`Ya59Z2;t+P9wXg+h6yERcG!h=ZvX3 z%Vt0ZR>YWx)`(RnKYgnn?8_5(1}(TB2b)q-J4%{m`{M3H zNW6rG4^!#?qxGpngso;_(KFtWn6<$ByK7ZyaH$2befMADjQ#!gE=%VpRVEVf;^1qy zx}}(v(sJFD6!&7>8ZIaUr%UpiQIlIGxlTKH3~Y26QHHX-$GwucobekVg2bi|{- zL{g!hpWGzb2Sszdr%(dyRiFrQO<206;Q@sJyyzvl?P<3ngmIRqX>j()yQ73y{SXzr zs`l9<{5{$Rs7xQGPhWpuZCRC5&cnm4%c50=w-^nL$}Xy&mVDW z!6Gd0oMTj>tf1~zfy$al?+}lE&PKvtbs>+gJkXr&q#cz(Jy4c|650Ynm?0McD3Ie~ zD828lvY1DXOK;huEOG97_L0IDMesf;Icp~rJoy66BnhuXHx#}3I6?=vbA54^PS=Ed zf(;+SXfQO?loGVm`AfWRT%QH|3D`k05PYr^>r7U@>xS`~CRicU`+9)@2#d;iu;D52 z$0_Pg+;sMeSHq28=j&>v^$PE|U<*F2M+x1rfMY@a{|Z8fyqPN zXdgqZ5Fu-{30mbcRsu*RpoPwFL&?@gZl*h~_%X_qMW| z^x`|@b|8FpjaoEPdRyI>+7b|jdpe#P8>|duYaDk{ItSeLt@%6Ops~-$5d7gt4gnDb zqNtsDFD%1P72{p}aKSHM<`kY&1}qylB~H4a$mT$y4`_dAl}m$W-fH+ThR4JmRh0Y@ z{K7*T*hyOU1M6esh`WvrIbZcvu6p4C01hXan2&8mf1%D~up}pMsdD-1@v&-Bg^@yZ zTJH9S4$k}m8`syGXee>UC_0oo)mVuw8`dtmF03kY_ONG%x!?@XxWHMUKw3rUdl&Il zlVe$^7+2X9dNAlq5hebh5|Cd03JmOs(K^Hhu%7$_h!5+?BdebcaGY~F3EX(ipKIwA zkW>s@c~v!<)Do7TRDx3VU=)urI6;aHxRor!m8Xi>n?Wd8?xlu_e4qqJ0n1gZ&X!#F z{D=#KDsDMyhHjaDBjA#P6T+IUIUs$bjr<9R)Sk*cx6lQe4`p9Gbm9A7Y2R;ysNn}R z!>1c$sL1GtGM*MYE5R?r!(G3$T)Ruf3prYO3fq$_ebDa(PXc0k2!wvld?|pUp8I6E zL5i$&m@`j#HM<&;_l+QhqJRkmHNa2!-Jgo$lXw$(xPDN7M zeiBZyY$_XzA2&ZvNXb7XAe6`_HasLU)Kk4ID3sq9I3)fpe2bJ%N+KrrQO9#5UtKfg zjBs=2p1kr|sC$;9-J={D_4o~EdWq=u0C>Gx2%llI?rdy{YK~=PeLjx1{40ODB5&&~ z=D8JY^8PA;UT7fHuQs(L{M@9VLK;o?m(H7P-LBmJp@M8{q$Y=h7f?y>P$FZJYwuzv zn(<2Tis=n#rh0$GWMdiP?j5V9+|57dLJxu2XVUB5S?<(ET5Q@Eyel(a zjN702eZ^PsX!r}02Z{0Gn;Y;#kaozW}#r0Y*#&!hBO@@@T*vTu-(WG!3Ndut0{RD_LPU4>{R)I`c zPeko14s@bM{plU#H3J1dtCR4u)%=ay;oE?seHM|*HD%ngAg&^MF`1H-4F7lglphna z7)1Q9({khYXxfy`rlJLYXNcx(bgbG^b@%X-_fwEuJV#>>slnGq&=G2g-m&Hb*9uXk zy7tt)k<8#DS4N7}MDkG!!`F$Lr1iNT0_Bd#;N<*!HD+Bj*gy=!qSQ_HeVAy>9mFH2 zx=m*lp1RR_MY^a2mX~beqYUO>>$OBdoQJ;CI>-T~OEi3)oc!N!G~E4`{s0B~CiB~EO~3L>bOZY%Kad%ITCbGr)&7Dmao^&h(( z4HTu6cIb0SgMO_ZZh!g)c#AWu7Tkx_n9=Ch+IJQr>Mi_Ay5z{0-^hddApv}idZW2; z<}DXue^Qs;)!i~!=*Jt&H^~y@V9AU zylYrUH#|qxVPG+U!U>pYC$Ompy`BuYvXO-|duNbTF1nO)khJ(fW1T)HYj${7g3OzM zLL+W;CjlkT4ZA&r=icjn`13&jtV zh(VN63|tFyGO}adJF3Z0i8Z-y`jOgJ)9LK$S!J9@+tZGzhx$Z6;tLm^-R1lT%&C#A zVJH>6NT+pKNn$=l%1?8-BFrvRYh)fo<2ew?t4?XXAM4KlSl!b`gGGF}b6QbyA9W3i zp;wUWI&`opopgsu@Y!UKI*k`vYw_Ey`<1qL8I}EIuYTh_yxV1;d@|Jh%Kxd9_EiJH z{IM!a<)&&xk3Sw|C&|0$KyXlYSLe#Uc_Z#)BV@fwZi;3YjyjkSHh-T7uD9 z@xx6~ET!K(8nl><(MH@rk)u}h=7+l=(jS=j9@t>&d}~#tm|GS3=3ftF+Dfppo(b1z zCdfnV;&y(O<(X#y)ryuYgF=D)TZ2Wk9s-KEfXv$}RN*r9DYt^$VLRa!F5 zzSMIyVqwW3eq9Qpxt`SUjr#*{vB5|OW<4&ajei*eCD;6RLK z+Ol<6Oi@-)v#PMHFz~Cm+M*O&Ul1*`5eo{hb-rsT^rp{J+;m7WgS(*SE!Vr+zxlsE z*~(MjG-qS!YX3sCq{Fm*J6PsrCDj>oc;$!${7rYEj-@)hwpEqNZ>ZzMllExMSmrw@ z4gCu*E~SWtZHQ{|TXl3>;$*zaYOt9p1cDTz$s*2o11>&Y$g2OzvzyX5%;Pu zPwyC*3akB+Ib{z8O%E{- zP9JcKAF=8Sd+QY)e+!~&s2%<5g-}WV;E+S<0szhdLE*uaXH??*)!ErGYDa7A>0LcM zskJ0NNN3y&ejT_diP5b_$FK1=R1S7#{PoO>X3a<2&HM;B2giLV3OYTSdm1W-K2V(# zUCaArXX^{{L6Oh(IUsE2;h!X2dY02YQb@fyKB;2q1VDlB9|w0@0vmY+%Nk%rgx(gY zTl;p=UuEF)0G6Do&d22fDQJjW^vFnhpa&m55{D89Gy^tigZD$g1mg=z0e$AJT5kW1 zlLktu01&U?7<|Pl*SefM{e#I3cp#E{G{Tlm{C5%nH~k+aFzf_4g>V>TtMA;ROj=lf z(6fOAJnv|>xT9;YFVpJO!rrK*mqx81FB(Sw@`W}aJPK2{{(`7I*wCHnMIK@WGzJ0A z2^~dtJn~3zl{g1)i;uCAuHV_$i(*J+K;a#URBqOS8(>~g)X=qlb8<;+*oDsHf|J3e zwEbf}(|M|=0Ko_~fATLEDr!@+#Xs|_1Y>&Z;#SQRsw#!d<&ql+d0*^)VlJi#)|P7C zUZJ)Djnlx5bb!gcyq&Ph@mQCNK%;0$xU_mcVf~RVp^6?SaRm%y>D3T zIKmw6Zw4w3fQ|SzT$(H&8S*)toLmqvFoN5fTD^AI6_i`RT{~TXr7=FjYzJ+rv)O>B z{dRa*JDNS_U^s`=+kc(d^r}O(ql@Ffny%vo8r#%gn|CSzHCEobCQ*py(6qa+MPSyj z_c`I@F>N!A(%OYgz`u(^P=2@gQE0?0PUXo3m_cCR`k;oiu*hZ=24Cr@HC&lc)2Ab% zoOWN!w8EN=sZFd0s&b;^pj;=%KTdpAs0B7fR$myJ8h$xvd2gy&-ELkowFeqlTpOOwjg(_-Q^Zd-X zp*&1o9Y$|k%uu$<08XP4Ez}he*mMIz0k-f7=XWToyi^#aWZb6!;@MO97b)gewRT%! zQa9?3SL#Gw#@r{N%xrTO+}YIl+P~6pLYYV1CsF3W1o%4w{cw8DZrug-{#UjkiuZWd z`1+vdXVsJ@m5QyWyy{YbtXsPO^YZpXE5L^ilb02Y21McR@Zj!UgCevQTEbv&Jiq(c z1F#uGT>tQa^#jZ66|e^9Gd%hC8aBU;DvLT@e2q68@i!bF^pxp-UQUVM(+?7Tt-vvK zfU>M=$%tkeZ;kyEiQdQr$ie79BBOUPY!4VROsEnLNNPZcACWCU3+ASYLR+T|th?0` z5E`_>wpmx#1_{dSq(zp&tFN@xDdrC!>WDU5uTk4gLCb;MiXLu+xz(>RqxzG4?*XZCrw&5w=6E~VMB zRk`0|vWe$2Yk1sFm{S?>IQtcM1sKv;DP7g0$6pz%tD{cqJcWT~l&TW!S|;UKt^E%u zsp>GI`HGYg{Y*L2VeHQO=G#6a*C4VDm?Z<33{dU zx2!&M7pqXgQCnlc8a}Zdi4vvbq@v~&Sa^Z8Cx4GcXb>7|q!dnJvW_EPP!6vhZWq5! zNS=lV;8jpGC-^i_-4JT{^*S#{>b15e9M3@wM?}(&dl8BYsK`flWr&Z=*@8gTf2+Mi zDWmOsxW!6`;5V#d7R#wY8;)-IT0fY~a8ZJ#Wh{0#*L3EBT1xBVAJ@K|anlM6q3s9_ zVe;eOMd)R*9Z}N^G_POV{K+lfJ14ZJIHzlWIGx@SO>W1?d50vpTYD7^1bM(=;H`;Y zhwtq6@5zB!`g~h>Ls#o96#yEGK=L1@03OfcW6061JQdY>$$~n0r;4#3(jj>@%X2($ zj>R;!QJya-3+PpDjxRcq{L+zQ6?I;s)UQBvSCw`)(8EZ`_BnU^2>4beKcy*9$cy!f z*tIlD%=jB`lnGU%@F&%6b08yS93uTOyj1XGmLf%@98;YdPy)aWL(-!D9ZWpB-T&#% z>|hOC@Ic+Akm*ePj`!9YY5#5KQ@tTyzCXAA{{Pc@BO$YeGeU%dE1$=k{E3D!=zGmx z%gkVzD0b!cy54OxL_G!NCkMkltFY+^1I9cEopdOoPZ4mI>ymliDBfayg*E^EB3FFaUS6@xX|nZfe@YTnZ~p7e+%q39bzb zfZF2XtbaZ^9uP?h%l<~h?g*|Y+%+oz>+6v<8lRhJMBzfbY}ZNI#`i*||FqB5(YfMh z{i|tT6i6<}QcY%6!TG4XMfN1f_|Ux&$MWb4Qz;H`vn~{`uYP5atdR}$3nBw#mwy3+ z01nDEzWwZcsQ^KmSV51V+yxW?Rk^I}J&MZI8z>Bl5HP~5L(y&jd5`z(D=D2QErqkl z@yI{lh}soZA;3kcfx=w-fl!V)sL1yhFSNv#bHM6v-9;zo=fs{)={`(rDn#nxDcG3c z>4UC;Tm#*EUbHY!ZjA2&I9j^U5!GhC_~i4{*m37yscQpLw$pdWUP+%cP+*8bZ4p6@NV`)ThgJGltr7@=ECxN{Q%( zi0UvMqin}s!yQC#ei~bvWY(1$NtzuG&Gc(MxYY!u)$bYzN$Z1$*u3VZBF_Q%V?&}& z+S~SB{FN-N-rbMS8f+-Djykha)G5L7;CxQ0*lp$g+L}B2fLkLYB@Jxi|MMGN!@KMe zB2_6HW-9FEcj$tqt=hpL8NfTAH3vfTm3zF**rWk>1x^T*{|iL4$!>uIj)X$No;yU* z(xHoqxA)1SKvze$$UsKWQ~(e~i%*{9jFgz!|KAA;WK>nYF^sevb$=g}F|=pl{RYb$ z6zsYRq$7l(HS=1LdpsG=Dp5|G;B(cjN96$i8`GS7#`76XiGh9>d98uVm4rOumiBXM9=L0NXen z3q|bN2n~poq7=%2llTt+KOrh2!={7QxUWN3MU%FRObzTfg4YX7^5nd4O?5@5lAnUse{S`vAZBErEUKy@Uq+r3 z5mtUK^33~o^M9+AC16Phn?KcB>2~+uo*-sJay22Hua40Ac4S=`h1alDHIuZB2vrf)Lg16h;zCHeb z)xyh9SMz%2G~>~oRts{B8cD}2017|#p>?zstxEPwRBJ|U|NrGN3EK)1()iA_3d(Ti zqMT=cNqnKV)3hJ>@%j6fM-qd=O@`(xi^sS`Bo-7F_}WqjHM&S?k``(xVy27yqQDP_ zuwG7WFU%tsh4vGJXQfsn-}?DS z*5cGxn;IFad0&MsgJh56z~kO-ILE7&QUdotQVN~oR3jzg zY-*S%GW)X;IR)LG`LF#klBlU=cKuMM^lWX#n$0t&V$}n1dfa8ktF&mhPshM0GNMBl z3whCRxHW1#!1NQ4iBnWbWUe1-bKLTx$~8^FReZhoAk5p7>eFOV{s%dR75qVmWF%|o zsW^9;hu(0G_XCHW96~n4w_p++A)|@19BgQXEEx#xKi`TDSWlLt&sm}~6tPO{{oVtiTv*WIRc@dE05aZH#0NXG74hx^8inVLwtg~z&C8Pr2J z1hg)4|M?VhuNqnV{hv)PHjoukXH&zyn;5KLnB`{vNKrlb`_H24j8Jod#vD8%!K*f} zCx)O$w<-5S5+pK{y8@XYNIT-X)~AIT5)Lafmv-5{AX2UuYwN4D0->1kej)3+AdW<^ zmNf1bx!3CcL{T$fU2pt4F>HA*qPW2wc@N4(#H`R6HlcnO1-@z|qbYF{-Y+}=0sGo> z05P`mG4_`>ylu}a{#gnP#WX~oxJdC&{M)eNbni2Mc|xJf7iIB<;r-T}l@pHqNY&XL1&9oE;-i=uUazF}7h0OJ1v}_gFg12hZx$gpv9cppFd6!UG3WIs?0Tq&Y z>rjG2@Z7nMDo<;jF@Hu)u7$zH5fyEN3)zK4^5^Tw6(LZ^l-MLn?E?J8=C^voz9lY+ zB{Y|rVQjpm#v$gQISP!c97xwqk^>R=Ye5|}439gyl~KD%(~B@hZ$WQ{w-vW=$4_6&m^-X_G_0jQ9f5u^XL`1^}==L-36L5KQB(CQ;e zt3lKxZ@dEy)~kgQ;Xn16Fm4%@9jGDqhQN!xi|f;O+>J zMVm|MQH9}&XU*;zxky{T@e{<=nxmOHW-ZE9R%}4J8Se%9^(}bmK{WT=4721({?avh z-Q{o^Ci!1i%%7<2OMgfPmt?H~hwJ9nfYpJ_;(mE#)2bv``*P>>Oohz&AeBVc?v{E5e{{XfQQBZ4C5mEp*fq($)&l|vT zlQiS-A4SyZjR<)`s?r62A9Fe1I6W!mrBdkJt?A}$cKp63G(7KsciNP+ks~5y5#OyS0x`5t*h>9r3fh|b2d|h?wywQPX-|q8m1kZh!;1;y ziD+&(!g+B4ZA=1)i+P0>M*Jzq0$*+fDS|WI^&zB^xexUwS-Dq4H{pz?V@25;M4@e)#1KB&kRGy*&s(|0CGIsUUbIe*}K)f<- zTqSNJ?EnIg!)P`~2j*YyuSg&s0D^9!JWzz@RyTFCx);~4hlFo5PL$Z!fB;(*tkzXd zm9-UlKa^;v1Z2W^zNc8i*>m|dGQCg1Tv*KWm8YeEv?%i?&=`8VxNWUN#Q81nz1v*OHoMS^pGF{HY+Cm5ik`)5;RloKM9Q`7jZ*r1KoYIZ z$P8gitNb;-`4?fBf}M=4L2~V`!n+vg6NWSM5n(Y0LTl+(yU;s{9OWH zA<HCE3e|MF0 zXoCA?h~?6R{ctqG_bW#3N|wVOPFiB9WOVaw$U_nJPmcXEuwn<1dQ$(WM7QpDpH@D_ z3y!#dCm4ibihT|zEU`eA-yi2{o|5*l<$aBOB|J+v-+}kB#sYF9K8AP7a=27hrN^|P zul-5XJp3q7T%zc-*C;`IZ&z8h)CFxx_U*eWKHajif15rdp51HIHI!+6=BZKmsU%dowAY%?nSa3Id2DrXJcd@FRW! zh4Wp0PKIF3th~%6R$M-z1q+{R>UO?XXxUjc!6KaN_jKDcKxcl!O#I)MztITP7l{?2 z+cVAIQWWbugFcx*Zxx(QKgTlm_=fq_-RG0Tb9N>)P0{m@eEy#<7YJH=H|8E`shE-M zLQKM84+F!9)XzQ9+LH?^Q3YUBg z7&@p9X*A1J0Z@r3@aYgLY$w+QF9+$a7$Dsh42>l#XTs9?6T1^+8dLDTpk4+?ObGuM znhW>jO75nH9>KRNpnzqVBhbxV{JbtO!*PBqb+KKO6S&AGmYe33fzx2wyvkqByg=*2 z+lrNUUW+9cCOkx5>UZ!?m#z%$?k_x?ilWI!3zx@#4r`{D;mjmgz9H zGuNWD0*vmGywpM)0|VMlu2{z7pNpH!pA-c9&^(vkCGdu@VBr!ym`-ufP<6cPp$rix z$7~+ucB*~z{OO)d-oA#`4)pGtKaarxf!wLB zf}(a%67tA70~RzzB=`45X`+t#b^oIYrzd=O0wgzR=68bZHAj55_Rg9fq6@y=A!M3Y zXnDcp)O>};#D7pOd$&)3H&er)t8#*%<0bZDdz31#miwdAjr<5Qhl2rQ+&{Ac;hh@xKd1z^V)r*T{_jT)M;gX4~bNVRBgz!4- zqB;^TBa%48>$h~{EkNge_x2AtcI8;a6Q_~iA?)#_rKjEYX*WSXvyb$-md(~XCLF|> z{dU*4HCfHRz%mAYtVt{0FVeLBnpitQL$}hTnkIeO{Z~I(VZPTXtLwq>HD%r@K_LYd zO?~4|8~IT)a{M}_$h2%lC6XS7QI7fexZuHuC!U$%DgdjkLXOf7c${ zaAn?4pBLNIP@ikx|8$q|w6)M4B-|C#OoqQVgQA)IiKQi^Vr>r;SP^3K()VIo_Mh} z>YZJqJB53o%dTGB5c|leEyIxDv35y-kwl8au5^6H%<~U|JFWgMT`GM`5YN1;%Uo%pJ>dH!+TsV?<*n=Wh<)@I-e%C4s&` z){@t}0d=+_)!zSjE&GjjzLNEFkeOlk2w_xzzPNZZO5|qCST=OS)@QfFINyF)7&Z~8 zXxkV^%86xE;@8jhY!9oVOTex~eU_^J(`Rqf={cw6o71qv2g}J;!WYUI*8bo5h2MUl zFzLb(mZ{Tp_SarctmV00Ac+u^Y=+Gpksp4kXMW@m<=%R-CP`RD7(e@pH%sHek^yFl z|1>RY_OQIq^mEx5O>BdAzq)4PdtJ-PR?9gk&~F;FUJSoHb6adz(*Pr1mH5rhHZg99 z)7Q@11c#%^J>Zon#+-Jf<2QTn0`jvIvSYESz0I#=5I4{`y>B9vZl9>lf*BXp_y&+f@CVQyWh=9JmKe$SLiu#k3`n0Aqz%ocFep1|+H)FPe zXwo6no#v#Ce{k<3K@;eBv<`1}=$1pUJFp8-+~7-A3L|om0^0$9nN`+Al_>C2=&&yX zzv`qTbrYhPP%5uu7NNU&Y>3+(lcYd0O*ZKt(Zc z$*NlVUC*!&pY9K`2#>7&i1^8)JZⅆ}f$_$}WqPb)LB} zX-vZkijy=uW9fnyx(x&jvCoE$VKP33N5vHK@+*nd4xD&K69o5?;s!bTq9`4!)gQ5% zRmxXX2}I}y91X+kzYXzAMg66X7P~g>3*oJ?kLBO){WlPDzi8oY-ao4X@Cg>tW}BMQ z*lWoXWiNJkql!fkFx)7IccBv$0plOBfl zbd`cm>)&>|tEt*@sWnvNL$#&Nve!J{59hqdP2KeopqqnO1{lK`>s7ue>hcELJhMc^ ztNt3(_?j4}!r42@Gj;X(p}Iw0PGHCbh244IjqE%g==;g0IkLlTy) z3b5W4w!8`^!5F&!s*@J=y>3TS9xI}ECh07oel^VD{q$ew;&(l{6c{D^D(Tv$k5eS3 zt)eV6?*E}8KE?kE88D2eR@m%!*rxisKPgm0mZ3G_X=Ro2)9h0%83dX|y4qjiv5hVB z8f-`O#bzTKHrt$jbQ$|KVPyAtcW}x$Or>;sS?)(J2V){6>e6gO??B*p6R7>dUPC5f zg0x4gRb|Q7zpuOpRPipa`|WmO@uF>cCcE@MaQ5-USgtgx|BMGa(e8F90uxBpO*hP# zfs6C0%Y1(c)BKx*5~FBs5=Fz11_5ecA9tQ)hpsuqp|+MmUUkCSqx`gwUubD41wX_X zFp?3TIn1@^YSJHMH{E~W1@?Bh$+zlVBQd5B>p80@`)@ zl{}rNu;@kMsD+4XXeklIGV5CF)BOA>l4>ztJ&3KQqsZQ{y;p;mNm4O$M7VcU24Ms-&J6dcLs1(Y0)GDmjC$DQqKb`se zq<6S(@_Dx7Z1MAD;h?D&8@RRbmN?Fj(5AY-N;-RiGhemP&#n-9><5b5?4|+hl_78e zM3qZ+=;ZV2#M?vi0^HD_B&yz-}oc7Z(7KPM$QJs;Bo1S`lOstDh0_3e0d z8`f-6PPC8S4beyBgY#DGI$4}?tEDFjh~}Gn6~gAUMokkz>}k-q@pn_Sm-WCY9AT-{ zfG|&^iU`X3_2L+|XzZ zz4|u8SKssdW~-iISDBaUUolxGTL0UxCxg7bJTYPH(`YjPc%OAt-Jb_)G}#h&`19{r zl99Wq_-d7RVhi!<)w|HQmt&B{j}ir7%0j@Ynt9g4vD&V7UrThdg3Nfim=jX?8(&#@ z>}%y}_(@bxXHqUZ8pi*-ag-LW%|Gg`Cfq8F0WUNG4@@t`q2w7^rur7LRlHvudrdQac=>3XkvN|o92<00 zzk)~~3T%LfIm$>~Pl((!Mg#g8weSrpKYU-?X%}cxo!IXUu_d%4=%3wcSA!yaHja@U zT&gS>#mawSJkS7ubwJxnQ{^c~0iU9n^wa?_eAF;vg){bz@Sk*Gg= z_hXtI(kB7@G}o}{$@ItkgU|F-*n*GswcV2nRqfrDBnzg*w`K^m-LGULQS6&9*xXko zs{l|bn*4Uy^}yw9-ld~PIo4H(9hx99;qQHqDiIHk9Y(^cl9g;jgD;YkzbDec+$LlZ8jjm zZnQx}!W1TWOiwFt5^Nt*NtkA^yF%XKbvYF%Anae}@=ncn>14h-pMHf1gS5NgB$|k+ z$)kzb?jE$i6j4ib%s2XP9Uu=Vz(cV9r&~$*ipd>YZJre4K{eQd13ro*uT_WYXurC@ z>2ZK&J)52CZG=Q_Q}ldV1;|MsOoNICmN9A?)=uVy>&hGPd%}aT)fntch)lIsh)<>t zTFhh-!<7_ZAT;l6{SrEGLJlvi6d96sd0B>+~XTTr+x2GAwA$Dg?Vnlkg4pa_Y=u>j3 zbJ)C()K$+ewryYQ^&~77REu}aZnZ_J;htR_cU40A_T3D-b)l!*D+*RA;{js1Pvc}w ze!MkfMid+NaAceScMl&T&Rn{=`_|Y^G&fB=Uo>ZJantkO>i zD+^+rF3-m<(sbK*H(7j(tORh_7}SPUEnX*XJ-Z~Bhqb>CULLffCnlvvUl<}-UlpQolrMWo@1oOmQlfLFpnGD>F z&cR@w%1Qd?ny6mCYB(7RjU{1R4n}t1)7borXuiL3NTY^D{=D#yg-U^_Mr*YVq=nWf z8HBYPc$LY)I>~KDig8-vt&S;e6_*F2BJ;~d{;`msWZeL-`viH9AjNpL@Gm|VQGkDn zf1G3-b$b_&f7p&hPmm zW(n08Wh>&s33(o4pKfDvMz8M#hCg@Oc@vXh1%b^ zn5VW3{T6?hA2yP?+ZRki@9ue70DW)jV;l1X$Me=UGQ!N6fT8y1MlEtIO^l-a9tm zL?>6i22TIN94eM}PVsyIffmtT|6bshUH9!iz&CHO(-Y5!HRn^L2*FSLbFJCYA?GtX z?qWCPtmT!C`EFo9XE!0PAqT`5t=PTfmI)r3>`d{>$~VSZ#E2b7U%xWyA4|);gB7sr zi?zRunU+DC>$BamA;cV$P07*;>ihqUx7Ky3$-hCBZL#Nz?1>5 zlbcEzEH!ot*(B(FO}z5F_b^<7>^)`e^+ln*HgV(J<=s9}_xZEdGd)?&)KBTUBXvo+ z*^+*PzjuLC+dferj>8c6`e0?@Y_BBj4)^XuOsNptb(Ou|-`&V{3rx0B92i3HQgr|K z0!5+3P2X-Md%+lGVd}B7#paEvbBULb$FTCAXym%aIDAdUY+H%{wa;(PFU-CcS2EKb zqeIE`+@WlblPDDa!N>H*Ig#BNBC|a&HKv3)?h2<_Z0CTlOEJ>v8wdo2)0`yF$U?~A zU0y4uK7L{JylXcN^AeF|5Y7rjPElQRHK|oXY2nK=5-XmpsoMaDrb6g=gF474f^Sa* z;;x4rR*0IDbLo5ngVdm?brY-kDS$uU8xP#T4aJ0MPG)luxQlJLy)<`r3`=U8 zWR5;u;1B4OZ!3&wnXD5rBKhFjV8-uX%=uQbJNjp_wl>zH%()B(!!J34*=Pp4E9ZbZ z8eaMy|rJFz-111MF-ehF2jeX{CE4))7ZARSKd~t`%y0RPhoW5EdPFuRV z#P$8W6o;*dohynq2n7)g&$T^_=eE5y_+ruIVAif`13j5bfVZ*M*_zAKcXE%G>)~+I zgU;2ziU)q&#kLSeux7R*?^4<( zfordcH%~^cQ+(>Nw0tPAbJvMUbA={JK>NjAyG@w0 zh%STuIn9}PknW;n(&jZ$tCn#`Jc)hyg;vpA@9!_7hV@HPZ!?o_Z4FWMP-mDeh26#n zKJF#NeTp@ajN?tom@2>wWm;|jSlT2DIRol)&jSU18J?*m1s1VVk zjGp?B>7ohgD&nzS=hz(|Ic)^Q%CyP5A1r>q!;b*N`7j__lDlI^2`!Fv5GZfIH}G1$ z!v2_r%(`KFkl)B1!3V)QPRK3&;b*Ye7m{cK+pDX58oTxR0?-;Q7@seg>3Ov0=%ZL! zMFa89vq#1%LbH?c!%zRob-hZUa z{(rp>%o*#x_gE5&LAqJ91<)#7f3U=EK9C;i?6?JOCsfUb$R*nKX=}+bfUQMtr9C1~ z)h{N3*;9P`H4)nW@qKeVa%XZ?Q6GEk>QuoRcp=zYY6wIwje#2d)iiBZTTQ9yMF~tv z6w}P%29NG0f0M|7$um~XYn3GvbKmL1sp>y4;?f-Sw%{!a_9R<&wK}b5Ol_{aZ|`59 zi5pDItEP2wM{2P~IXfhVcFA-rTkJ=&0JgFPL*>G-WTrvKAOF?;zmX3^1$pY5DgaL| zUf~!ekt9ti?QzDrV$PF)Pgl=^;hD|1qft{I37Jr;T#-_{{){*@tV~@e4?BEKEjP@E zxnkbyAQhzp$SN+@Q8N$|sBh*Dx34todakT*%gyrM>#lr*c&z2c=FBr6Xfy>}o|!Jh z9J{*K_AD5U zIr`03z@SyRx5hQYNm5)Qc11t@lR~*m#LUo;HG^_a!q}*mBm_0&@q_oNc6y>kP?Aa; z*K7Lxfk}Ub^>0er3%w7U?JBRcr{ba9SDYD(u{hxk=(io(@6Nbn{4Y{Mv)9@}`37~k zr%@`ZO>?(R66^6z5jsxiFztB**g>0Zr~RYC(+GJZ zeTln#)#CE5ps}uKJ{UBTKs6Eqe||G9k?z&o)p-|xClrIhT}S)$wU?lZYD;`E%_IMz zXtd}9?BnBG!?=@)wDg+7d_~F2-c2`co!5(8J?qW~A8e%Zg@5zG90?!d-_-9kV09&P zf4)Ke!By77)P-{12GBf(lad`VQYS0Gp)>ECs^gnVd35~6KY^W%P{UR>CRV2 zdF7;CzVgpv$+b;8D#Z6BRlC$C?(6n4_>SVR49ybe9)tr3(Lo`PPBI-vGK6nkI?dZ8 zSA*CM$CA^gcyQyjc|U_6*0oSDHHGmk8pn{(w8b_j9I_ytEL{ z6%F9{l7q_V;`=4{@((oe+8k1IO7~;LLcvarH8H(TxLuHR|LN3nT!);co@(duv9rHvLlMfwQ}6Kg)O|DxKpR;{ zT#2qC=~y~*R@$P7wbT3iLS6+pgnjJmorU%(qi9)|FlnnqlQJ&nSQtoSNe)+s#yypy zbb2wN(V*mIOcMVo)Bi%U`-`V8iOXvU z>C@YBM{ioYxSJ>9zBttUDveamPUM24ok$!#44&I=6M;E$tcF)zA!*J*Xr)VZcBE! z3^G-f?FVOjKLDEp8^&iE-MRneT-o$B4bxCDK^(<&akxPf3FLS-FdwN8^oT81_9uSe zf(BG!5e%LNgG45;F4lGTIuaQBLfvDz{;1~VUS8a!b`ikJt$@^4<^ZsN@eH);RDRK& zT%=(XntFq%k9j5i$Tb+gXUY$Rcl?Nj05)s4$WekbH-jtWj0F(5nRoVyK3`PB#~*6S zxJKcO#>~3BOjw;7H7C~04S&j($Oh%BQxYP^pxES zQ{OnTQYg#extOIlrG(ja8{=%Zj~dat^7rx26-2wxX|Eh-kJ-PVr)nBGUrtmD#B!m` zY7f8<;`5LfJ&9AKh5%#%vRWku%LVV!QT~ABhXn(&XDRc2R6!|R1L{NCZ3%ogXVvKJ zW{_t)^EOK0Df|#Ke{1MAhJ9eoPy#cc{*s5O_P%pGz zJ7vV=!9Q)jBkWkTAMYXRNEQ!5qoS|OF2sYK5#=G8I~AzF+NF=&V8?UW1ld2Y%51$W z)R?7%g|#;MxYS_3uG0xczezNN&->7ThCr)+2k=kOGvX%ya922hT({h1tu}v_`umBV z&pP8x+sms?>{NG^G^bHsvFXjebVNO-33YCPI!E$nTG?Nw*3!9K1E)>hCOetx;lg1~?p0=UaA+ zP}ru0!02LK%n{Luq!6}k#^dI4M>?>s1Q0$SKW(|uyJya)pCt3JJIj>20nh>rQ6+$2 z_Q%J`_qs9I5DHWbg?Zx=?ijwf-G8X>h6YNY<&s`=EL#L92*v@^@dX(+=3Y+A;Lz$} zig13pj(nnxKapbk6bDN)SNFe)L&F5l%TzOmVj9R9D#S^r7mIN_;D*oIqWm`_-!>MR zB|k-@LD4@FR6@|2K#BvVMd64$0?I@ez((+KH%{J4)jr&)k)Tq)`o3tjp>pxfV4B@c z7J549Cuh}4*jYA-H~dtfsDT2dSIQNrMNb5JmPQ+w5;6eT#Q9ZD4adfDYh+Y`8LBkQ zn36FKO1BE})M;i1M~2*}XfccY?Hhlmx78S`7bGGRCHt7VRjO+jY4L3G>O{SO35!-o zV7wKKWT3Fx`r$s*dC5yBp6ES&nvKH=m^l1cP!Y@o%>GI;z9#C`f=iwu8FOfFE_=My zCMjOwHpLh|w>Q{^XQB9tytrV&OYv1<5x@EVxK7fv?p9>Ad?@FHGFtE@xhYoJ?@KrC z*5N=qH+t-Yg%va3g1>NF#WvrAG!G@?GL!Wu6GHz)P={-JWDzD(8a_-m=#x;|%h&?kR$&ObDp_6ajZU)Cp zbcPts=X-4oH~$oDS}8w#`x6f{Ihsn6l#BmNC;`+)-}AVv!MNdzz~6L{uZoyYQ0iGS z-oz@+S-_=aOcaN!NuOx12*G3mJFLru&_(vyjXAHjM6n=@Ih)Pkf4JG`zq!}2G%iNA z8@s90oA($F-HAZ11l*Gw%DZFTv80R%uuCFnJ=Gp-|Bm&`9zuwyZ(Oe^F9h4U&F%^E z&PlTJ0mZ%5k;zdp)^W7eFAHLTd2$Hd507Z4AFqdQP&pC3b>reRv*TM|Zl`Q;r(vX> zllbPdF!3Mv{q~yRk=(@%IWAqrD_R$4O&fQ2 zzx$~T-X=KC_)d!!n7ki+dPogCBvmg**{Twr${o-h$qh44Gw z$HRbOt4JgKaz|ySi^3{tyjS@r%kj~^87M-c~pNdp$>L-ZX zTa||Dnd;nAy=ccO0Zg9l;B%<$#i^d{hPxPyn{ZK-qMne!Ee7Z7j zyyC_*+0u@{AKSDjW&c5cX7W=JXJaQ6W$8{C`*hK1e2xFeic7gtWA`B(K}E)m9GEL zqt5^@C`P(DaH=Dd{ds@(T+=S`=AG@105j%{fmsuB-$IQw`}d+Vo5F0n-1!6)B3Xjx&jVNB8fTvOV|61){X%X-XShGlNL7F+dk- z=8m+kfqNyOc^B4dtHt<6D-He)+ZUgu-`lm5k$H^T)GquTVLr~aJyg6LDo&xL@;=UL zrUumng~qSW2LeTN@t|j}%-Sip*uwW^GRSq`Yz2tdYs|&hMltIOkl3fAuu_uQY`>?O z@wMeRyyakn{Igz6J8_h^l&_6h)L=BYxXB3IygH*ST-BYxokC6lCVFE;A+NAIPBJI9 z0^X2~M{_Vr#}O>IdbIrV7-23a;Z>~nDCS==)O>N)IVZaac}4AxD0F8ysY6~+BRjIJ z$t~>-qc$VM2gStlOuk0FRvM_EHShH(pTW1Z=IynBxcdhJ{wmtpVmxomF4#UM@BY)V zuKIaGjA+39iXhpW=HF38V-)9{n{B!BQ5j|UM3m=qsM-s_O!#EM4xqC zdVVfMm>&YD9R<2N16?=UDP)vO-Rw_*3UXq!0IE^^=n-=ZDUI#IrTv16FgEE+_lr-7 zBXqF)%Sy6XK&~_uh@Z3~`DQcd{+lUOWwjM$RPd!?6Z{C~Ge(p1w6lWKHR}|iQLA1D ze^tRfaEvXOL6FJZ95<4w1(a60qB$Ju=Pc~FTZ*EwUPX?JuWX&)?YA%6`Eye%@M9QQ z!TI}hN6G-D)!viR%3i!&5%tvD^pe?Dj=3n)EBBDcMm$eu_6=$|?6d_$hsNlGLLwcu z%hSlSM$GS{xq0~)>HQgeC5*M(^L!A;d4>ZqjH|H$89c8OR$e4Gp>^95j>#^bHJ#6M z?7qx`TjZSf&$jUFB;#k@MwlHILtXQ>11F6f{cLgAKP1cO7@H^QS-W`oIl~=&HG+`E zvx39vkmG4j@@YC;QLg(&u8Nu=nI*;i%%!v0MVvQ<&c%fU|Ay|<{ieSwGJMpK@Ouct zl_0Sz9`dYej z2E`m4^2hG-i-DD`-pZC##DbWoo*EHmX9lrm1vdOwtJ6RK>QSfZSfcy(t+0GfJ>ei<~p4VpzZ?b^7ObaMGx`08c;^Rd&C%lYc&=$G2~K>{#hcvN=AT`H~ZzB z$W3#Qg(x+~HH)_1FD3Z}8$RdH*un+#Rk241U=Z@hwEsFFzBZPAW~1!q%I6rcuiaeO(k8X(227bSy0BdS|x>!Hh|) zg5@5zkqV5y32v%}=l5YjZ+>fCBu$L5`Rx2-3^-h+`wbllAg*t=saht@0sVo&aQ^F& zJDl2@_)XVgbf35t$-wiSnk+9QY@t!5#6uzfVC7+RCHw|3kPVhPZG}>dXXjBK(Q8{} zGOyJ)<}XvMHlCf|Z3ajw)ad1&ICJJaJKyiTc;o}hgrCl!kk&%@c~dE4(x@1K6bUl~ zzB6qI#=QrWSZP&cEdXzx1$`RaRec{ERvt0z3V5#fD~%7qyzH6JOg>X8X?d`AEl=x) zeH~oO-Vz8cV>bs#>4=phJ7mGm+8H1Ao5wJhwe~IKi)cQY4zu(}4w@sZh+$}6-gb%y z0t4&Kerm9~`aN3>we)ke;P4_3SdwtP#~NMjG0>8$(-2nkQi_humA?nY+Bq>-wu-C= zj43n7d$a)h27s;<5Ug(q#}H5CLOR~A{Tk)T@;Snknj?QG?0i9E7G)7|(h(@i6&tLo z9Cw%OJqmqT{O72^TLT5W9as25tq<^I4C4qBML&Xrfj0>XpC&9ss z_{NgI+vmC$xckq~UL_K30H*EZ=&EpN>1%1|QK14+{DV#~ubt+*8*-?=U{v7C1bf)e zKkbC3{Lb*Ts?p3LZYwNeY3Z;8A5j-)UC<1T#hVl9A{XsUoPp#^kScW>Urw?dpDIg7 zpzBjgZUZ+5A4ve(60LW3A(w+5;#W{<-6ac+767iw>?h{X*)PJ^nWR~?+nL%sncolj z;d4NoHGCO{{p~SRSmd22v9W*9ma7noIkw+n+aGgG!q_<`{k5=+)_Z?vP7%oRjIAN)-Q`wKWEHnqTShJt2wZVa)Gc&A8Rj4^eVl^<&=p+O@YP;9Ks>``-Y%|rEbo9@{V z?JUbC!0vKD*yB7H4%PU`fjU9X_qsL}mzL4!W)AzK8>jGH*%VD8UOBzEbd7P>nPs?B z?+*iK)A?^zt*Y;V`$vDyHa?r&6neid>ii4opgqt|E-QUDL1B7)v(gJ5CW(^=s7*kf zHn`md39i{CsH8;?62ASeY4fRZDQKL9 zb8O=_*IPk#fc-W$c_o2jRDK4p$jER6Brh~pDkY*gt@z6s)G%mVfMw$=*G006YE|*G zjiBS34z#1v@sw<0d!~vghAaidWKk+M2F{AEVUkuLQbC&s{c5L`w{)!2?dmE3VzsU_ zN<0uob`p@|KzfG_go>BKB#Wi&4V**EVd97(@LX4~dsev3lH8UBc$SGb~pi^*+6b zEq-ul_m>>j^xjjM0Z|J37j^N>%<`a~?8M_HrJcJ*u6w^}^l$hNI7ts_&+U2o$W=QX zsNdL9d2Rne;x-bt>po$<$*x1)Y!`1HDb^9C6y{)bvWo#rth@aVd>9VciLMl*z~46u37pyxJpCxQQQa{0$9AA!r#C%tG9Zs#F2V$n3V_zfmcYCjwpkJa^aJ0UI{vmcoy6#Y#s2t?7U z2s$^LFbcS^BW$Tlr{vP-)%qykb?kI*T%L`dij39&OzYmAPkb?2O+rY^51$PQctY`u^#)-FG~4IqzN(Y zA_3?2fzw8+5SD|FM;2xuIi6U?in^I~=t<#-!c{3?Vj10_Le|;oa&g<1bjSF4`9Xh= zC^q>wpSdsu%2E_?N2{}jxMVyDMu7&@z`6Fy1Y?N|#|U1Fm+tc&`P`{$L4#AE zOaoKQH(p~8ltUD?M_{5;w@u;yU93bT-W2 z1ENAwue9dfqq*hyia#B@|kySs0(s7yrQjAl1QE z>#d2ctJwe#_{3}phnBNvLPP635W{(cdJ$@Ciag@Crxs2yc^;p>>})r1OWpHX?{N1J z7#eedse_bpyqSdZlqmm!SmyknjITJHA`k{>m~aGaToMgESo3}^h*iucU`^d2 z!NSo_L5{govX^^5(iFEjLjPP*8|@$hlO^{?U+&T4)M~~DMtAn{zS!d~#pCYs@y6q^ zf!O?h&}y9D!5HP#$gJseXJQozazS5z{qzmh6+af`kosd`RpR64SQMr}3_%4i zAc^k_qUVPxaNwUa4|V-lcynFSr?qk%d^~McrO!E%>^&}DLPT!CZaXQvJRbagyg9fd zCfwCe$*O)!XWfhzYE84Ep=*CGs#(KSz`&R1&9b;U6Hd4K_Yc496pYoyWzit0;)KI-)%9Q65 zJS3J?qWrw?r#$?Z?YzBksqGf)Qks5W+)3ztoLCicYq8;QwoB%G&U%vGgB-cZ-uJ(0 zr}ZZHJ-cMjLS4Ih6?ALQ*(xsEO=6(AN(HCln4#A#Y>X?vWVA z&IB);jrk-w+DKKvn(?{FiUh-q5He>A2bUURxB+cny-7vrbnvIjA>+X*J!dj2$0PT2 zwZ_zoz~U90@#>%t>>Ly5+o8Df`4F^i#T#z(hKbI2#e0r&6Y54mqUg6W0_KRkPWk=t!^I z)~j-STAG43ZdsW+i$5+X5apN5Qhm|#(4kPU-ov4-wY;#4mr3}p$&>TN*2Wf+Qnps{&7h=F&(CA(Plh6xI`RYRC4 zZjm7iPT3UCW2Nb+8`>fo5BODv}4e^ZJT8U0|~1X+=rfZP0GLnMB#}0~(Ig zgywe{qD)+DqdTCNV3qAibcL{0oN75g}rj0<;w**_0Bg zQyzFppUi%WrZUvv(DT;);98;6#nMb9KHh_H61d=R?hQ!5NdR_ z{K7@8H8Ugo4)JxKU+4|~)7^>^``3_%coh^o_3y-%JnHnfGxA*4{LG2{az}^o=9CA2jAs)< zs}_d8;PE!0r49LAIWBuTWP0RzG9Yn+@rMD4!ym-^@n_41+4bML60V-jPNGqF^~Tvi z49jY4aQ!WuaUce<3EZt~3o!udDSiVuEBuNJgWsD4xaKgQ5@-F>wjZ+F967yL-JaU^ zfLl=iwS8y|lv%fCyq7psZVL0sKj3B-?M@#BE%lj=D1*zh=G~r#&<%#r&OWS5XPge_ z4(bi6FzH@?0t3CD6lp*>^`RfXWyL&i58N^2f6|5Fnl0uUl6Rzit&W;Hx7ws@5j=G| zN|#!-U=Y>ndxc{xfy6fpT%X-5Yx*YA&4Jn>hJW7oeRr6SlPldhCJ!qj&YO-pe4><% z5FMK;g)_i=kw7J6=eN{sSn7s1iTJ`k|xc}y}(!hV6^ z4PravL3A5>olwjlbJs4Ff-+VeXM;dc^n%{?>KnaJKh+g45iYxTe$Hv_!h%iUXzE~1 z?&Mv=tQW;l;5P#&dN^P!Uc^<>faPTQ=yTi<0Tj~a+6U2upDubH*M4jpOCFyj;%4X! zd)6nkfd9r|jAn0I92Y%%x)BB|K|1U4)({85`~FLu+;}^0=z&u*CS$m{$7owTk@Mi0JC5iKUmHVp?7So1{GPY2~O0>r`Q9%k)v1@Sm7i^pL4KE5a+^?bXr2mF>@ zj>D%x)&r7o8`MszU*&OqhOV(aWn5hihyw66-~xDH47fgj$Z)9RiY352l1+2n$lY~M zA$VZ*;ZJoT-kZ%RbDJoaX!kH%kfkD+#Cq;qqbxe2{y0p}zjXRNsGaZfp!uHdATN3( z>0h@u*a2mbPX_?ZbTu67ii1yq@ZloO!}z)#tHz81^|Cts%SHr~H+>RHG?;5+iW$Ko z7I7gd8FK(N8Azng{Xr2BC#B_aZ>&kG`-P@mblWZ!qKcE$*DIX`MKLIL%)ZI{xNW;t zpzjLdKE8Crtd#trJaJk1ORcy_Ms!>0!e7Zo*l5&KML$GuWh~-cp?yDPuQ`~odbawX z*Mao{EMGHr+qhfg$c~!I%?0KCh28|e;%WR&m6_=0!hio5SJ!&`^${> z{|xelx<|^M!zt8yC*Kn{6eM4wb*%<)=aT~rd=v#%E2`1u2xjsH!C^WVCPR$zFekfi}op( z&fxEhDB{!K%B9|3_@f$VKbvnOrWGRX`Xpf2hH*0$8Gl$?*B%?hMuq~ujyj)g>Z`oOg!E@xNLQLwJpwB>*gqpo%!?rrIfth+rHa z*m5zU4BogU+JRmlGeF37u^mKa)jey$Cjmfseo)QI~jHFyubi0_ZJ8LDHvJB_& zLVk}WFCR|TE4k5=-|F_CcdXU!M%Pd_s2if#OLq~!Us2B>gsFA#=gv|LVamAL$+wma zMO74fGbU72?3M3j~W?+4bB5d9WacsNEv(YSEmg za|#P_Gg}tyNWqHVrLMin6V#kKNLx=;Y&5?NY$K}dQC}s5Fa7W15xrY{-n9<=V+x_C z%T;El%MCx1CgEu{VNcK9PGYYhb<&Qk`PeUKNyFLlYl;ZmFvep{_t`64HbpAeLB)DA z8^|E;I|rn9I!>2gE_xkOj#hBatE=IuW{>c{09XVJ7P-Uz5469pOT?{K91ng>DJuZK z?OOfarTazg&EyYMw(f$tJr6O|2+}tN2yq5%H842yBq&tqCPX9vfE z-}uw^lm24FRNs;p5@AIzHLSY{9%-f0x*}lFwA@Z&BEmXmx$v?U5y52qDL*A^i>jSY zraB~1n$706zvbQT6q2nsEUif&s#m5(d|_Y*`EvT;JG#zw^g>@K$WDNd?3O#tF6a{P zl0zh2v1%JfiV)%65pTvel=c##+r>BSh>FErX%0JWMaIXIlv;J{HYQF_>WH{{jomM% zjF!ozG+qcBXITPGAkDFur9wfxa3im^>`&x+GfTv#`$rEWx3V>?!(lnr43krfG)HC? zq{8~o6^iJkv4!D^BwGS=WbTnGd>?FB@(wk%%9ogsMAP3MlIG{T@^;5z9}eo7-xG@k zP6x{%hl7{M2&jE^Kof0yj;n?~BFG>oFl3CC5MG_OcFAkX@n9K9B_XZD6^X;p?FAnf zE*!3;MXILlbv69QUpJtlFXB@AGh|Q%S>bJ*t=FZDMLcf#{yUtTQ!`e7L`L!B zg&&S1J#8$)s??BMPHGNaSlET@6%EmXO-PK2?U^v;Koq2BW(kK2E0C;*+!bL(i~$vb zW*8@YXv3iXz73Zk%!Ve|+_HgYgSMB^*+XijGZ`-!6>90m!ETJevxd!+k z+OsrDuVXHg7|jWI+x$zs8+)z6t>$+X(%rqYcD_m|eyx8T7sMFgg_b2Uxgz)YYFk;F z%N1T(M|+d~++_$z+~e~*8l;VvotFE{RY>4wYjLD?*g|L!CB z=yvo#J(s57U*E=otALjlw+(xe3K~VCpOFeXGU)rhw=|bzNWN_<{$m$bAyeW0(FD^{ z%~4Du15Jw*aYS#k_9$4M^x@7_mJ_>4;?2;0@awkMUgU&Pl(PCx%K`i(&jrSc#XD)< zLjtRKr9G08>E(hKz9C(22g zs(K@}&!D^YMvt;l;x&oasw1bTqFcDNa?K@Ikb*x6?Tn8)?L_cuoDxb9w0i_i7n`zI z+Iq^F<6NpNA~5ocvPW$!-c75=#?pWF4iS)zNAxN!tS`~ORGY5r(tZwxE&^uYiUS>T zPUmGTtB<<_Z`((SZU^n;0Yxr;&t(O}Pxmd9vkK`BO6-Na$S(~`+ridD7K-w1{72Jh zRb04LwX%QpJTHH|M$T{*6K}w0-QiE?ZK}0vAHFzl4xb+hmMmkErt9Do$9Chp}(S{bpYiC^{M=#+gf=`7|G} zVc6OG(x1YJNGO{u17kCQ*O79>jEfm}VlPI#iD# zZX*~aa)nXxOKH@rN)4~|PY~pyu4C9M@GY%;bXu;?5N>H~LAeZ}dn#?jzD$y@dp)I? z437F~1~b&<(~A982*02pUd6ufocuUz$~80iTobae8vfbh$}`P^3f=i!Pd+o+8$?ff zcfoP3zsAcF!sLYy5k$z3PF`h8wzoU4KSjGpr*g^Tq*(~vI~przM*F9yGFGjKX-iQC zEoRB^%ExTh$>4CTrpCb*F5=F-S3KbJT#-~uzOV%Qr$7a{!5F*@ z#O~wf{NbRHNg9@mM!oDUODp?$rD2Q(`=Q8Ldi22S`L8jZC6y@Q@*T&^aOm?HhXS zMa)egY7kv#PtOI>s2LICB6E9LOG;_(A@24MY+vc0AncR%2}U~|njrj_DRi@BkjGWG zzGcpG#SE*$qx+}xH4+mjt%JBW@EwP&pB!sfk=Abq?DtggTWV>*+0WjE;n=hDB@eYb zTF6XkGH=me5sQW29Nk-AL*QrAhyLQyj`V!_=gTa6imCKM#L!TgqmUqokow*72Xt5h zFBH62rDM(y3f-5mLF^XjFF413J1B7!@?yQTD19ezEaIz|%NNqvvCruD@mb@WQPPss z$w@TD6X>F<`?QA26e^d3i1ZW^7<}~ts)lgo4<=t@hB?$Xcdk+PiD|wWQas}GT$WUq zDm-sM6*QS(Ko+V1dw2}|up9Wxw-L^d7DXqBd}gWX@O*ZY{wlfInwp6@!5XRm|E&@N z19ZKHl-tm(?yNn%yAMBGlhQe{Ius%{0*jjl%lt!+hg#eMy%F{tQwjQ48nY-lamA7d z8s#0!Kcdl}6W9fy>yI|sOoFSTU=UbOgIJSRRb_%Vq#D>{x0pU!uy5<%yT5O~a@F|! z*8xIgk;d;jXzSAvI`kJfr%CV;RaZ(G2CZKZ(;4Ixe==x@9r8TXUApg^S~4S5oV4&t zpPJA6V9*%PPbFBUs(MTs?|uO<a zOmTMlTmnIUg6>BILl?cPUnQjQN5|(4QTVdBh z-UZF?e3(~iWIW$h8r2+-E*}z@nx6TVxLcQKc`kPUBatDOZ}l8y?S;~1(2(Z<|WMvrd4(_qDa(+4rcQye+!PNgPpswb0fPPC#X z5|<(;{$@xQ!jd|(B#w=BF{wg)4E|ty{_Xf`x7DfVQ$?#CbN3y9p3jw_XIjAmMAV@J zlYpNQi)_7L!g1>Ll0LNmZT(t(Bo*_@RQAc9#D3@_BW}lLIFo4qJe$$n`e0zga8;#( zs?0Crs#HP~EJ12WW9s>n=|gF<*l z$R_R`>-0{)cXa-4Ixu@N+_$~2q}4jl-vx8y-WElq;kVpEXVQCf&Hh19K7suE$}N)* zwuvdq{MT@IA~e(Q4}a88CY!bZbXG;2cMHGcvoa}dBPxck62(H8NK^4$T!rWOz5*Zg zei=8{Mk%#_K!S;K>KgDGbJJ~Ht$vOMenvy9han=E* zHJb#9_bJX+6P~o;Mvcmp-ZV0Y#n9hBc^S#+jy>kixf!4Tli!eljRSS|So)QF_71ug zGn~VeVY(ls$#xaKYR7KM74)1_!9u$=s!qGALpY}s#U8^@%yYHtB)H19ss4e#w+^}9 zM~iAs@rDmQklnPDfHPqzOF8@r?WDRxllc>P5yV{$oL1(c?u)^#$EDpcL5f}xql_25 z`kFiWKbIeS$~`lJ+8(r)>fEZ3$Ll zL^-7Tq&+c+P(^=Fyoj5JwWxo5R#mG#hsJ3< zpq|a*U0k^JA&)(xOiAnSh(R{RzvL$fuQuUI=c&0BvK%bQOx8aEJcWPa#jsWDnFuqC z62GrBdWeCSyjaOQ1pkX4LlU-KM+A25Lw8;)$dJY(wy~-|w!Vj%2UVQNU}Zw&1hTA| z$z|`}dg_(yQ66x1eOS^QB1ArN$M|rYSE=x{)~4XoMtH%~5ipgpLYcn#XQNa;RKUnj zPwAGE`vtVK?VE{tL*?q}<|CT{l^M}OXw zs)lm4GgQvN{E}Z^dCB4Tf5>|Cc&OU{5BM&Lklj$o5*f_cDofd!C@siN)@)_TR`w#q zj7Vh($*5HJEnAktPEzl^JP_wHoZSG%(AW4-iB_OG$E*!U>d! zL})`U&VmQhWv>N#CFcz0Yl!is$LTusi$!gW_2OQI5Z*b1l z5cM`Ah2!UOvse$u3S#!Yvx4j%#$zY>8&fn2U+&uB2C_}FJ_VRw{I5IE(0d@tLTJOI zhXrChtXnZLQ=H8kkO`o(#|Ec6s%&83VtkJxiNw zp|1H6IIgMPz52tO2wU1-C}Fd!r{URsbjY^p-iD}B9=Ul(@5kw72NpXx;YjH}Y5zcN zGkQa<)FHsEfwhH-y!BuF64%RT%P%(!psEV?CTG9v6JvP>>xH!n$VM;!>hWYm`lgaV zy7&NbnW8r%!h#+d*w}$B!>cTMjv-U5`zJAat#2ds#B1E~e?P*3Cz=uF?X9*(7r$6& zpxHGV!fUvq+^asVt2=Wpoh&#lA<$Zgq9ZZUx{B<>rglFu8jeBJVJF`3wYM{e34DdU zM>uqUUbp?s4*Zl#LxIrGD3YRA!lMGe+Plo;6)^^lRA>-{Rm}XjZ6k|2Ii&-Kj|f+X zoD*4AxMa0k(e$Pn888}c-?~Xvs(8N+5glH0h6m*rH%Vn*E0Xp@a!jhp z7|pKowWzHm+_xHE&uj(V=Lg6D5dt;pR?ijx*(ZI7o9Uy6(z{_MFPCuj1uBg=$9t#h zr&>+QzzBV*Z5;9{IPA3U!htMfPlY_y9A z{U&|L2~ND=?sW1SU1P*3ioxvNrDzuO#I)@7fW4%s-0!8ItLQV5tGr~pn>F-GybfKU zQG8EGZJWLYut;TS(?jY;nKzto6YO7R;cG>tQa?7EE*3v)Q?Iqe*+%Bj#s{>lhs8a8 zcR%59z~BMeSF9~;9ozGP40xTj&kbd#jIvY%_tr&aSLLYEXEL6NaWM0)&+>A}sEeQ|JA1edcoA0_W{;LrO`<#{&$rE8*qL)0Z&2Ow!! zk7M(&sNioKCJWOs6wA^`Z)!#JM~LQ+bft9MQK9US+6-|V*9I@W+@rIOjI)B`%PaQ; zPlOKJmt&-=rW~_m7Lj+cSUF&jQn@_b#e4C7jDB~Lb8TbBt4x4tlN~JY|WuMs5 zXR01z1ywoC2ah;LTuflLffn%R^LUu|3p#rvNlm+cs(IxMU#!?B#RMVhrw`PBV|$mT z5fM7e&$Wu(hBy4dGCx3g6bT()Oeq|3eOZ1M?M#&nJ`$1@Sp1 z#xmwti(XW7EQ`Ko_S{5&8G3ZPLON_pXCCJh!Q3hn7pv1_c4zsnT;NX@-CT5x(%lx# z`+j7GYWBqIkU1S9zs19G{#5@usZ)(S;$Py{61*Y{Xst3n0s@3g%K3U`du&8wY4E5I z{#vc^Mran5iySG=^l$t)!Ubvs1EUVZrlQpd46kW?v$rn#R@}fEWAm*A<)W+_k*|xb z#=*x+|5zt7EpV9mUGLju&7ojLaCOqipw8P3-|vOT)28$2sSH=(&9oW?=l`USAd=}| zpXN$-eLsr?SE;(ytGyiSo?R6o=NCxa-C3K^JUECOksa9Tklx`NeqLCf(mS}o&i130 z(OH|@^MoJc^x7fctuAXCHr11SlPGR>M~y=lU-O<{XHEU@<=GiGv{C6d1jMx~SNCpI z9LUa^^I#BV9l4Y+l6vr$X=EEg`$c?l8q7sg&tc@-1!RyGz2$wjxT}`leISA8wpuibt1A2Px4HzBwYB|D$bC?=dQu-O+sxxE~tJlr3vWA3%R7^*o7$O zKGOLw!XA0okR)8WSYq6-USMg%C6)&*)8YZ~w+R~F;WgSDH4k`hSIrBtUhrwc@g(p7 z#B<&#r>O1NzF!;>0!4sBHjCu2!@nTRwG|i{hz&m)+~4S8V;@^ca=jeg3U&&2l?B-B zTx1s9cb+l zv{REm$Symp+}OHfE@dz$i5EAVgX!490D1X2I^Dv3Wj;OjV zZ|+i7%8X3jwixyP;fKmLSI23jro=xFoz0GOaYxT{*l6)|{TjZPMfZ%h;YooBjFvNTkA23Y$*w;GBt0=szxTPD&f*O@kO(siE4csQ z3C!O}gZP@#61;%KA(S96jvlbnp9%;PHSeKdOyghJI=807tgGv!zOWn%j54EK%z@e@ zZ$NsQvKn0Td`0_oXhm4F2LZfbvOPJVD!m*LF0(Xz4drg)lBsR=7sJvltg;bJ9%HKF zsGZ|{yOp&DDeM?l{o5IB1n>Hn002x1AS@51Ik{H*x*rGmBlkxVn&Ob9w{6*J`eb{S z0P6|Y{QKDK8UQZuwmS}rs-I<*?52Crf_atIMCl11TxH2Xunu>Y&#Sg-a>i_tMs;dC z7!So~&Mjkm1kQSoD&824_72!U=T29P@TZ<7=w4-!S)!{$(6oIk@VXHwrB}O*D1Xmu zJmz=hTlPsK$~*Bv5t{%9hwb?vO;C?uh91-lL8HIQ3BVG!qWgw_&?IZ8Ae$#<@8I` z%gVk5gg+Kv=d=5f5>OiGaysTg#=js9vhI{&^UTA$Pl|9pJpE%%^Cn%B@Z<=zr2gfM zb{Bw9@&BTQ6mXHHUVmvzHHOW2kjwPeTk(m*^=KxQo46D=dTDP2o+P0bh1h zd=S>kSt2B7K~M{}u1@ocM>W?ovAbT1VSY|oMAp)>jx9=agI(^?4XkyWN^iVV(7{VW z(aPs^Bs|O%FXvm_CxYwh3rQy2!aMi`UW+~~jU|%LDM;-l-PRZ#?hA`ditTZ}hhs@$yIWsM z%q`0Le6p{z$4i<6D%Fc_9$UB95Ni_y`x4>v@9$R{-iNG|bzMJKo@S7mr;7>GYtYJ)^MR7W7u zx&lLu_9RS>KWOPbf_C?+DHDQtzT0KoBolkN;g=PRpLc z@rbk&QJOh9MmRy^9{N5r5q_7I{@6}CI@v7}!FB{B0Nc3@*p41EX9#)YYzzIXmY8-^ z#ac<%;$LSk1oUKI%9*Nw-~Mt<*0^mI&CEI8^0^j)Q6dSdg{E&!w!iDq@1L;S&wuJ% zn(t1P7{|;SyNC0{ZO0OhvY#kR%$Z2SqP_V4{=k!Rm$iYL->>v=e=gYa$ieZfW;s8L zBt5A#@qiZ;OQQABt8p36H}^y1LyaeJe4g^ zf_d1(ovL$pRwPfM>8xt=Q+ACTG_F0C8m}3Va^-%TAfvr-|JYMQ0X{hn1$IZWXHp7j zqEky+-(9a}-_CQVYRDGlot%$wFBD#njTzFrHQr=jG0rXi*y4q5haWv1pdbTeb%o6( z8~5wCcL}3NxR23ePF?XvUZ}#BJ4cgY$SuRAahh+FI+&dH_0OKo^#brw0NLIJ9 z=;C7-Y7`fzof&QsdM6z5#-F}qCX=oI9K?~xXeOT|E4GLE2ZfsgJc}!ex%JEtyUrmc z;g3FF{Vw(dUedJH`30Dh*&c~rjW`N6&VQRSy`VM^bF!2xl6+X!`P5nd6h@Q2j3qo} zg)#CnW(tA=(6ri{tHK}uOJ9xDv%kW_xbIuzx<=pDuxBUFMVX{VAt-3LOEtE_TihLI z&MNeguXc)vrj56kfK;>3OHESoOaU*BW?`X1-rTb_+;izD3fInSE3rIA{1$knTct8t z^@Q0zE?+yj!+XdT6Noe)cS{EpTu?_nx5?2!#` z`aRvkc#ynKnoxpjuD3zy!I9KGId|}!;o7X;wSB&NH4Rc)X6SP&y?w0-^@1WGacUnNPyNai>~TTP1i zl^S$*JIg;%kX594ufFq9JzXmE>qMt&&Zy9$q|8f$EC?uB)XCYq#Mx`Z)xYeoTF$E< z%&BlOrbyQ%STo}#M!f_ug0cNa@ClBzWI%G`xX74g(j)3u$2ZUZ!*`AP^f4EpsIu0N z#Ny~0Zn5W{p|7M7gdq}kaY?oF)(l^&>WO5=MMu-mbxe%RR`O2@1;}EyDkmM=pB4x7 z#9UGZ(oekWGf(Boz309LFym3u<|~xfiqMLqdI9AgaZ6_kr(@jKV_P%c|6dhE!HU3st2U!T_Bma8!Os`retz zmGy6>mP7$=S;IweD$#gD7n{@BcF_6W%APB@5zq&`MJXDV6{4p?B1sP8Bmm1xr?rOs zDM&f66gb;`Ll@+e<*JI!Vo$zj%NVSY?@q8fwe@WAm?%4;4N5{MN@sT+h!cclI!4eg z*OLpf(b6P*GwgN*WW9m>MKI!dOJhbljUWmj@2urX@pEdsNH+k?J#+)ACai52a8wL_ zYo%2NIRMczIQH9)yWJ#S$-g)?(L~VHIk|6Tj8ln1j+Lj$X_DFGZ22W<4nL5*iEvW- z+aJALp4>sTQ_Df!1ey*>MwEW1(VlTNs8gBUZ1~%ouV|y6YrYtL(VFK@<8J(~4)<~V zMX446z`(7%!?+WBzGI(xvaFIg%p3U# z;J7r9&=gXkkuY828{>V_9qN= z!1P*7T~Ch`otUMbj!1|ka6HW})0GYuI<=Z${`woV^emjCIwykY5V{iz7#nYAt>^P4 z62p&4s#uZ4FE7u%7qRuolrY-dKb=n06w$l_yz!r~c)0JCDZ(<-7wVH!l9~OZPdfv! z2_dO6J{)Ro>qiJi+W?+RR{m@cZqLYb%54C0Q7b?qxn3i@CMkE;osSa@@$;h z1EDgPeo*KuC${2cN=yT;>{}LDhVA&waJ#m5wQoEDs6d2@QaC9ai${l^*E^Uh3G*OB60dh)K)1*-T2V@zOdttetj+}m5al+Zo!JQ)Q*L%u~y6c3?4A+ z^#y)6_>ha42&`#o^es&CN)4}Eu+u3q9xR#Hg$qV?mIFT{+7Fm~2^&6Cm;8>>NnfSX zU%5VE43AwX?}f?^DyGK<_?)^Y52ITWrR=4`Q!9@B{JB;^6d~Vzx=RTC;CZ>#12x^a zW>V5GXPrX4EBrCBdiD3cj_7Zgx;}VMi8!_EAF@3f(t>;Y^4tf=nhJ!Tnh!H`AN>)^a0L)_xYyV%*Z@;Vh+2+cBU4`FQu|l+z{%fLeOg2F_I`B zL2ub~lHW9L)Nf^30~GzF(4-6X+B1xwAs*Bd3xlynm-!FBWY~DOF9MKDVz3d!kZFi^{>k1R@KNLZ%XSYIUh|!JY7UuH@L) zaAh1>g6@Q}3+3TP)!LBXh-fE*J@8JHR7*Ndd;X-V_<>_`L$zGjYN~W|JRzPZ8KuPo zi|KR!)fQw-K0?R

w&=_1B;LoyHFJY#Vq|MRREErJqZjkfkv{iFWwUwZyOU?XRe zDf6D%^+iVNwnwP?xla$DW4^@B1!VNpge;rT!q^O>_5x+pE&|oU{M~!k9u(A_SDgw8 zDkVhiH)K{m67?k(L{;AAKRX8=g&Vb}0kjOlJR~#EY@E+l+txkVUz0y}7gM-Ako(n( zBMiYHAlUBSN{%x@L^r+zw#RXs0Ez?cNzPlx2-mv5IX(8c@iw+&QjF)nN1rsIg&Ic$ zzVK!6cE(>3T*u^W#8yKV&HmkJj_Rt^+=sBbh+4~ieplsR00(3-@B`2=Z5L(;%m#ly zv#{`9s*sS@A%-Al6WDjW1p>>{w{`CRx#t%vFsz<^Tp2+k9FJz}cB;smM!7DWH$QSZ z%~b8!!qRs7;x@EyTIR3$N0<_azRt(9?p%6kp-K_!JOGLZT?(Y1tat=@ZMD&9N5+Hp zZc3h;K$dr)AiHC~r@z$@z0o;>W%{>m-+_#mhi$*oL%aJe&<`z}?QyWsX6pM+$fnX6 z4UD6P!P#|6YG|9~;-d5C2Xz-}y62fY8rLKOkn;c()Rdk`&A|vI<#B;M{jt4gFMb+F z)JCLpBj4>4l}z?Q7*Gae0LNExO!W>V`E9b;xZP*N!?eRXs)ES88T9kx3J&$bxSzov?=XRu!OWU_8t`Y z6iE@*pM>l~W=CIkurW#hBXmwer&*tR6w>-*+|)4b1w=RklDrHgtKbG4S5H6Oqf!3z zdTN`~?c0hiitn}j?$@4??&UeNvtF1nOTu6?lb@M%Ay`k8<)y#tP`G~CsjsvA{!Hda zE-R(g?9h3bI7A7Z)hx-YVkiwndIpaUWxFpj$b*a-;f`USnmb};L3n!VL2&z|7|pZy z1@_BItl^df-4a0IxZtB@C6~Ak1ku>MOjy!1iCNq|L?wIUj62t{(wd>XjD-Wt=ISb{ zCAaYB!YwhA6HTYh;J0mD1}rZHFYrXtktSNq+uQIo%MlK&ievSG_XT@B)zY&0kz7|q zTLEgjd?>4?P+%!&3p9u)bq95=K71Kd-@YiuP}`!QW%`;=MX(DR{BFRg5iB#EqWsf? zVT@~7ZcwyFI)HocZO*kK*p!H{fF(oDAZaaTl7f~YTbQE8(HnHPiq;9vXx3-fl41$Q znP>3_^5CYMVogelPd;X>9)je9WN2W4J$2cQ$Sy>1C$)dMJ(eqEIVbTf)kpz}O=p&5 zIrT1()j=f%CFrZFPKw62v-dsoqZ!%d-jE$b*7E3vhb@$rwzG-_zY0Jfz;^$yVD7p4 zomJN_+p-L;6-k%$Apvr=8={E5N7-lpf3RHrVsL)W{oG8T& z2QRMIX8aWDa3@zlwAgNCs8u-9TBMF&r@2n6galzRZBGCoQ(be>)Fu9O;y&_!bQ5@6 z4hcJI9+eRITFj8eg9xAC=9ReyQ#d>u1z}$6x<-e~jupsZ$6t=8rh61Cvj5s3lj8Lc zC)fPdCbj9~rvft2FU-O+s{g(cN_s`srdQsq^QY-!vzy+GX<`kvS~__O(WjPy$8{kb z!?1wlxh`K}h4Mww^g%2d3`JYEDh+7b$KPXoyto<#9a%{Z%HFpYbVqOTSeKO_3;tJB zwSd4^V&d=8-Wh5vyd9cP8`nG57vv#%^Q(1|9gDFdKpLMt`$bFw+!$|9s3n=qAlsrO z>44-Az@T9z-d5+mhUL8dl$!?}ADbTeRZMN?+xggCII*v~|`_ z-o%BC@-{V3##B^yj$F}IS^hMso9fi!_lm0!om`-a0hQC1$C6cpJ_0g<@@RgCQH=)Ns5S_#LsgR~N+2zE} zM~r(Toh%w{vraP{@48+_+cDMglXn7OCEr&%Vdcx9M>XA~2(cOvp_lrFgG*O_KhUb9 zpCWuX%w_eEK#D52eg{8#Q9Q33&78P8Heph-jC_+)$+X-rhlX-?^U6IW2{hN_?WCj~ zpLp@DJn0_R*f_c zWb;H^eo4I|#T~)Y|5Bd8x0db*Gx!36Lwds!EoyOJ~EA5(Ss)6=i)R6-cR%HKQ-@(zG#RF ztQU$>Z?GhCy(uc;o0OsSY??-Jdu8{eXey4A%i%*^mDSM=Jf=j8?MTfOr|vqNHqaAw zeZ@7Hxt+OOo;+HF_W$~{zlB*OugE4k*zfcm?t;l3_0N0Miz-x8`RQPVqT@KPCa*@Ti zSH`>>Qx<5_40|RGvwxp*Nqn<%z5?{pKZQV}$8axLk%7A@NrtBd#3IOyjkDb%e|+0@ ze`!~>WzQ|<@+T(86b%0;6{dk;CEV=>6A!Tbb}Th6&t7e2ApacYY|2}&^SVO&VC;=( z=D{oJNn;rHXud+H{I|}dJe?(rwq1*SAK-`hV2W)Acx;zgf)@O=$8N*Y+l|kVN|}Xl z+_6v{QRZc6q=xTZ>Bh^(ZhxU%T_+jCaeslDt8?qrgBv5g}|7XlBQi?d`qX z%=})pR4B~&)p9I9zz$27Oppyb^$cOAE^Q5-AtaRR8VwI{rtnFks6{qD<(x}nvy}i4 zG&QMZ$td!8LxS#f{XhbtTe!&reRxE#gcVz^N=LGNp_H?A>Q4)HXN*Q$wI}fa5cdJtd!dpd)slcBO@U(pqBy+FRNrfvPvh@!(RsRK$Ldh0NRLVXc*9l{9nqzftathb54#eIk}Vi_Tm0{dYe3KwEx0UTG!GG^h|Zk zh_3xPbegoV3N>dU?p9oX5GcX=G|B&MrGC!hT$EI5N}c7Zr#bo?;kv1r-vH4XQ`%1% zTN@FSadIhQlInF5dH&ClmeSeh6!p&rzNZXg?0ib{(vgVS>ugrib6cATFccVGd%Zzm z{m3zA#o?r~EB)FDaVDRUm9i?NlcVYRQQgBs{{I*k(jzD0TyNW2ur&w$>PZzZ(iIJ! ziInX?fk5-*|3nrH8c^IfPaDivx=S7u%I`SEpyQTCTmK}h%7vTp<&SDa+XKuU)KcbB z$x{ZP3W4DTMp|O#(`d5>FS1wGm6I{PWMGTO~&iOZ~mZdBCanvY9!_q=tBwGecC=nQ-6?4 zkmZU^c5%gY4XCI-;W-@C=&c}CnnJ1^qReH$#GACcqsyL;^S+#1v)j>RZMrz^L_EK3 z(k|PJ%2+0QXU20zm0BJ}Q^w+Tx*#YCK6gWxCS&pnvR_ddt|`IJvnyPbeSl*m*}j$p z?MV)WLxAP9Q+ob;d&4?ml9#UV*NO!zf?3!VZpxyiBfyqsDMZ(3YYMQv4eKPsEhk%Db5f_Vty^wgC_8_klXyj@eJyWb6wDaa@Xfn9aQG(XV(~qb(cbBwFr6W zp-6KFUpgpS?XgiyD|Tso6AHSTH%ueF#qnhyPXD2eDDf236CD^(=gZ|`2#?n;bzb3& zx>aRc@qCDVdW>_XqPY(4E!yG`U=KE`ad=$a5)PJ@t*L937VCP6|Jbpp{Mzsm+jTct zpxYZ>bCB0(s<(0ely{A!SE|?s5SkO1l%T#>k0+=j1=?YDfxn~3j!gs~;#S@8)qD1A6MiPf}>8A0tK+nZr(2!qGjdET4- z_Lbf>!_TN2oQ`+DECAZhZXPpV_`?6eI>?`p70+0&sjlHd8SVc^{qkZ4=PP^-P`RyU zIh6+^^`~weD1@$>@}l~|3*!RBXYW=7|DfvbJZdaMUDt4gjjOb^Rx5*}Eab}-2uCDm zJVU0_**SeKS@AREn*k~(N_m-qIruvoQ3hFdVeUB5w2>mkh2cXxjlN3r?`@R@cDr%<93x{M z+96Q*s^bmSmt-%!h-bS;8K*p@3~DeDCAK{3k^z#q0L83LiwhYI0gDC^8gMyK6D2of zcp!NC#RzD`iU(MwuA;GNnr})(uL1LR=zhUt%X!xgK@35uWp$VbF_EblBre{v=xk3- zGi1c`>#FFuL@ndUG#_SWcN#PpUPQ{jkW_8eLp6Vl)h+o{%9_kcf9eYgK(&AEjes-C zJIOY-?D^4;i(m4-JpWE~f)@Co(-B7TdWtjdQP z*`leKq)G`_q=C?4QE5lT-MBXm7EU50;N<_odN`H&7<;XUk!i0#XO&lmg2o0wMpXbA zGY8H&4Nu6IcuKKTU@OD`CWAARW>rh;Y6)xNLeDq0Cm9pH4(JHC;-`|c>e;~bd0qu+?q9rPCQ`-XozV4I12!f>$ zRMzDlL6+pP^2QIZooHVvNgsEq7BQ-ZmxEoJD|8AsojaiwFGev{`Y>qmP#- z-g36h5L7WwC}XY36_lb99ra(hq#FQ0%B$35SqYJIn6x8=THX{j#l|M^I1@XOaxjz` z&iR+P`x{%pBWA2u34GXdE2656ezg2zS@z99iafo?^+w7OJY{63C~f{?2iw-e91+LB zmdlE(ax`64Ggv`NwKJ5CZ(8dASMvjml`n6EH+Pimi!?n2?B$aQOljG{@W>Skt<3|f zrl-PeUjsx1ONt0w5Mfi@!3!n^JDq!D10wCcOnL%i?X;uT3&p{#`SB}3ht^EtY|OlH zMCa77EwrD9>yce5SuO$RDwID=FXP}t?Z}!9jk>WrZEP|lL_Td zW3uy(fzgd43;;RrS&~_4zCl>%-F4Gz*k@aHx+h68jrAS8I)|>@Y(Yb8$n%LmY|_RU-YJKipRlZg_Z&9@iolNR1pMJpP~zrKHCsPrhi2@ zk)fWxem~cvLUiI9hW#F2i|LDgxS7fkd`6=!-(ROPjB#=#N%rm=es{N-y<+mG*9>Qs zju7B#bpOOH0TaoKj-Z|bIo$FA(*O@(Ji)ceVZ6K)A%ZbKWXh=2zUX<2gFulXYMPY3 zn{6WAsq!l6MAR&(fvK}r7cq!cGnKYwZI$|FlHuK=CC8!I%ZLUHo`251H4nY-m!VL^ z6o_x9DQFq^WWg6Ka4Lr7!qg>+h0g)*s43Cl1D;?nP841sg^ z8#8(LcuDI%r*O1%*a!SYg--N`0l#r1^6dDzGlPxw2%e};A9S#vM$KgEs{8uL*mC5I zh;fcTu%fSI=pCBP88qxs&8-0e;Ga6vw5VK}e5drXeSG|E!b^_C<3e~3%+Y$$r*k;| zn|=j9*-xX@S?|YzpQrAG3*fI&<-i4xTh~>bt{sxkm54dl03}_JCt&nRE|6n?o z>3fP%jV1c4Nv$mIz_H8q?h_g)6ca}8%A*4KQ@<^TTe61x>oaArWlKmW0g_uXmT%+A-$}ZyjUF*dU7DKvX5EDgG zEQX7{xtZp)&m)L|edZ*$9`%JKbpo#gZ{qN=^If+s(h;lSf?{{V$cP@wJFFBmHfr?! zw!<5Unh1a;u-_5&ST5^H3-4Q$MOR=Fz3?U9+dJGDv#!A3e#*4mL+&F1>G3tjjD{%NX%H;0r+%p*_ zhm)1Ldctm9V_vm584${!`Smz8%t;d@XRoQ>cqr4)R%EspZc{O8RWZ&dtG5L@xl4G|aKkn(4(a`BXhm^EK-(uNhK4qObiT^a; zNEA|kK?`o%KK}sWRp@Rq__5$L`uD&DJC7j#GPMuXP2NA>rj$Om9#)e;HC+P{)eDEI%=UF-UNzTwvwV-j7-iF>=AKp!(x+e`a1QjC&lo$xV}_WWL|S290@M#N8SGvk-M zPivXIw!FqC(7ym_g?P0|^?tXv9|$RX?EJ&8cQol3-f$HXyFt_( zLIpZ~Ce;5uU%_=F>`c(kEib?|8t|DMdJUP8|7b?6OLUjHi(%m-A!HNhKG(_%ync@6 z2pPza&78jHX8=+*XbK|9Y?o&`y)DiJ4P?#BU2LKO}xpJt2FE+QoJ)MWqz}#7l?1>u=6h=0rl!aV;y zzCaIAe`h$EvxWE`!?rM;CT|_cZCxpzkz;*ToAvLKzX8GN1rN0R94SNiFl9|LR$kf` zno6Ut-%mG4YxyyC8Sg*D>g5CezBTQV1B}>IWHWJ@B@0Btq``&nl@(;|IB46l7E<>U zKe;cmzT&3ru2qGdy%5HI;DDN_ADqDP5?U{$_(z0lA7Dt~xa3eZbY&rwXRQdQIeV-p zuX!%2N0AjJ?sd21#KUXd%>>O~z*PZ`#11+Y7o(jUjr5@z<7FyVuT9%hX0lcx_~Hso zh}?rV5xOwRfjS3KZRiCis(PSHw~r*{>{*Ond+2`tUQ$}SLZy1{y4|G`bvf%K0R|0k zc#@MaD|6w&LCY98zJxZGzZ88NE4#3TzQJS zPoNkS>{r8$f}~ZtPK0lzj5vytYMuE(bhU1qsUJwwb!?#I1|!N$fuGJQqwliQAU%7} ztv=NkhL=jaz8Ex$*4)VA7kqWHR*V@a^_=2VO_kUI8nJfJPa<5K$+OOKgmy#Fme~bT z{3<h=ol9go?@{lveM&zcHFM`PBGKL7n)v^^9_uZh?KjUxccsQW!6@jNvmDd zxc9^w3F8~`AwZ>r~S*M-D&48#Cg~?wk)gZ7$U}V8C15{-OHGVBA;Qn`(V-< zmIz`uXvf1xdI(fgLLm)DWm0biJ(g}MrI;q;3Qb;e4_1dGw?w*o1clH7L7`rein^W0?g%{(BR}@dIMfnqjUo8{%V!Al%T@q+9tghW`VJ|fOlB- z8ES+@V5eJ6*gA0Z^c^fPlrTh=RuB%kPuxT*5zua+lda_ByZ=<@JEHA78};AC-Tejy+w zpGAcCFf>C>7j$?*w&93z8XQYEIe9uQI~vH`xYEGcaao_APYpmy7m@MN_7k!J5l{Y# z0m~A*+<9%1L36Jx-&|DU9nnRto()hoAw+@B?i3up)QKAS1s*(#=9U)+D9L^6p%%p* zXQ7E^%$L+TuKqkbP$aHW4q?%(KsGC)lHskVW;|vLg>a*+xWOkIg?1C;I?>q zZ?(Z;90#>%HI9$s-Zuk0BQjaEvtIikEh>NDRBXe;Pa;x0qAeYzWZD&f4$5&eCX~bl`>h7^v1jhWRO`({aO4jFq$#aNt5y43i``c8uO=JKMdm7Qj z>;>*d-0`{iT*Nf0RcCAHCZ&C%BsnD>8VH-hD{0$^OcRWbY1dVxEpIkyvz|X#&;I)- z?jpzS*yq~0>%m1Pb%G}*Ng#o`oAEin43ZKfe)=88j>w^Bab(0t>nZhm_1a{~%QxU* zmEt)>=y6ob4iW;xuZ`2l8~UG}BmlsOFZc8=t`S6NbmujOHcnq|YwUZPA^s)O5x)*A z+eQos(NH-Okzh;POkqM`s1Cx;Syk<+d$zHS8Ru>E)@t<$*E?GC9L+k*d_G!_rR_6%c% zxoV-pem1(J&FKp)m8yp|X5LK?64G>AK;gH*rwfJj$$9NjZ<3QhVftsexKTp$`i7CU%27TZ) zJbI;&j479;vjH}ef%-~eyz&Bp4tQQ-eoa!Bn*KvgCzkJrTG|qlvI}&DcjZm;G*NbX z;hxp~tdv9TT76iMBXHm(tPEQFZqCg?Q1tw+{QU5^jwhI!1yNefwm7%5jJ9Ir}|5u$^S==KdsAL4lQYo6UMSk<@(YR6GW7FId+av-p?0FNaHQI7e^ z!Yr@cN~HR_c!nUt+E^j4V~9lh!-QY=59s3rEC6Hz#GnceHmuW*RCwLN+ENi@q>_G1 z@lX@}-sR&%2IsToK?50AOzVmw%H|`2(@`DvEvJ(fswSc|D_8c`uU=7N(DRhAS>Lfn ze_A#xg^@;zM+c++1c&YggBk)*ktoU)O1;AmOj5B%XmvR5`K~+C;mEPe6aD6vQV`paQ^4Ji^bCKu zeiL)A+{dD>FJWt(mxnEYMevaS%pu^a!Q%Jo2k7=S)Qnx;Q0y+=P=`k6LJ~Q{GrxKLYf$1X|X2G68zFkY~;eTvNm8L&mjxh z`|d$9=Oi8CJdo7;4d!dum;tig;6}6E*AVlmk^4x@7`BraFHC@L0UpKe=!*#;s4D?E zM79qJ*$8Q)+xTmT58F$hu}l7r99|W6)VR;3`<-sCF>^#Y09V6cS*1u8ez{)`Gg{?j zW8e>uoM;er&wB_+j#jDQW7n{2&=)rt0q8?Hh(N5cvv-8I9UoGa)c#K6@zTU9@f?^7 zcn-91@ChlXzlGos#HbGpdn(%qIS{itSOdsU@?DZzw)#}wFT5qYwmVwdy@5fKS*ua- zs%hBQs=Z(mhnsIYoeaw`tk*?)*O4C16ETka$iLprmA6)yb_&PCSM)>)v7qFxBQdjz zSV1tJkj#SL*?eN+k9`P{w9V~YFQm`9&p)$Iu5;QgOSS5yC%KJcC;zee2`4{Vm_9op zQ$)2{EU~rmr_zwhdXnnDaLP6;oxH=lI!DP~!$a zdP_vF<3Vn9l3gXgY4JnxmWU8fDKFW1sav7mc%Ecu{>}WVe2AFN+8EYJ&%UCl=64Ha zm_a@oWQgm_FxkXl)_+lZ-)#Hb_d)yeU<&B_(#%DWgP$hrCU;5vtR!jm> zC%Z?O-ik30(mWgk)ic}Gmufj&PqHdKCQu@( z9_W@lWq@!wc27xD{ombLXH?U9?cb4kwy^|hdzO`ESE_;XBn z!kW)lSt$m^t3QD!2UK8V0R#`ETQDXX*3Z9HGMZvae{w)c;Sy%XISqlV0Y?~L&i7*G z6fklBU3bSez4LQ!`{|tThzcBWV=fGGu38>RGKx8KNvhU9#IpaycwRGD(6q?5VXpIu z_afnBSMZ69_fPae(~8WD3l`3P8n5U$5;8{aOdprvoGeT_D#0E(=SUEiD7kSF{A%d_ zU2xlD1rUlQMiy}EX z_T_DFA6c3Dr4=}+Ly;{pC|q0mqJyGiNP^I{;3ev(Mho50O)S` z+Ro4kwa%NdCL7AppVG2GvyA9fG3M8+=%j!73=|UFzck<`G6Rt7@5>Z&7WBmhO;j4lCpt4nJw7|RradAS34>$AuKg}IVy z|4>eX1zQbToy7E$ymo&p(x>)=l(7d~<_Q{a4JL*Rnqyr9C8v0tFeAQ`qTG^v0 znws-@n?Bxt7kwAxo0NX>N34qc9bZF%&S&;8wZ*B;F_k>}p1{2>FnQ+B%#1@JSoc+3 zKHX-*rQ8sI;^$Ei<^-Ak+J8P$GViT7>A{@>5kfSQ$E2mN0@zcktBG{U6utp;4*$m$ zR*scmnkxJrH?kJBkhW!T@1HR%_{MY#yvV&@`f(GGakfTt_V4I319+nyOTAt%_0pAs zOlj_wV}?0da|;|7;?`rxJ_Y-CP1RR`{s5b{?e*LLIxS$@NU}T;NS58pl>W+Yrdj*9 z0dI2X67mP=NB!fVXM=Gw;d+P2pRn7177)TD*!Yt7%?5Jlyfbdj)!tZ&K$bHU_U{;4~^GesDa-U)gGVshu(!#_2txH{epEPLvj zIt=dS8_)9v*z+s7RN;6!rykvHRDbC7@gm&gMLpLi4*|+Mq7|q9cYQzjaZ})mn${%F z^JlN$cSgor7tCjT3Yo3#=2m`*h3fD9H!hZu0RL0ihk>0v9by{~xx#JFcmv?N(7ldM^rsGzk!z z1wkp&QIuYSG>IZzKu|h}QUikY-lc<7=|v!+cLAkKM-f5~5Xzm+dCz;l?|1Ki{>~Y8 zX3w5^W@bHWtw#awpHc;QNMF z(ki*oNHP4+SsFNpD!>fcRLx-T+?TuSdYgkR}C<~7S5n1ATsQ6G`2W8n#;Pn-nzkb zwZAAEj|{C==La~f3L$Xu0!PSyrxTs;qrw%Rnl$1oR5?~|vF>Mpb zvuZFN{`rCIpko)&L=9wJ6ytuULm=gYR9OrAj>(;l8VJ)i1foV*ZZP*d=-Qc%0>Xw` zk5VD8D)R5f>dwtNXb%(k6$tIGWg1~lgTda586|*iDorrv9Lhu+=AR3}t;-j_CkhAF z!5A|aZt}d-xd`$J>W&9j!5=W?*FZUR35U{yjD3(c4Xho|lXGZaNXoe2|D}J6=1pUY zBib)v(^GN+=`&Euhz4Y3{}`Lghdf>yrr^WxD9{{xN_2mV-O7HZbQjk<7Qrbv9D;>C zos%V9=6es;Tg*jbs4*}fuYn%@f$o5`W{=VOidjKoXdoW}n@K;^emu#*vVQ_-`~YKe zZCcRW^uj6fdd`7MP_dZMgxQKN6JkwP z|NO(|Nq(1FFNo;aDCT5PuW8YLX`Q)-OO0?xww~^mS)cA4`3g z5{-x4DUd1x3}O!HWcLmd&_b*PpbkCL^Bb-N^Iz4Yj?>u^xYmK_`yQURpXqpqk*^&H zCvm`#AQU=$K|0-kpS??(yd`_NljG27;L1^AcCZc^;6rmdh|JyHHz)podlJ$X^+4j6 z-;%P1)bUik10?5y4uT8+-=5%+dCNjnBF4IggI>aiW(mI+a!GgR@u73dp~ zk;R;G(*DnxFsX6 zS}}?K0Qj>sjGWV&#$5{Q=9Y^P#0)K7C@AghDj5nZ%-21a$=sw^<1Ww4A^|i_7Jfm+ zdP3NL=fHoHZh>>#n=+|W6d$Dx-(m0Kop=-wEL%({Ck2 zJEswp@;ok6U^q-Vl!5cz2W7Gm0|<|>2<8SI0f3MFY3{cbidYB9X=iK+ZR(jxEGOJRLKV$eorw+P_?r&B} ze7W+OPP?&;M)?Oiv|x6%a>f>Ge3OCuKgC!q7Fld|D7*-+Oz(kx_>H9*o2$|Ld+eHK z;dh>h|JCaOrTc0%1vh|F>`@1p!N=l{P}4MwXVjgHLnVG5h7o6dI?E9?#-3)`Vx~(? zX~)}3PI(DuBgvs{56r0$gAfgpOL+kJOQ@6Wq6Qh`ySF!wf88KhCKNcq(4-Ac!okU@23`do9d8y!3dVWsX zP&FZ~y_-|>k^tXx~t#Hg^NYP@Ak@(UVXQzR78P? z;%5ON_u^di3ZZ{*7^J?AIfCwJiNNOWr2?5W{(TN}eXnNM2A4`@Y{yD`>*BqK_j@-PZihtxay0UwM8NwP6 z^$BE87`?&A{$lo29DnAf54e@ux}tUII<59btOHJ`I_&Hr-&$aW`dqqCpv7!NO&}&a zYMcZ_G3vl|PxlDAU5?twa}67V=)tybkY2oP@(OG38&2g$ zrOQV}R7N1lU(ww-czusg&LrC>j+1{l9&3>Cq#~N8;U|sU3xGiuA6vY+^c3u1z^T*u zHa68HemoerP_3B+II{L098^lRTa|I6<+ToK>iD z1Dp3VScRQ$HUw`&>Do+VC@TU|3kqa<9UB}MRR~XUEwFW0l-5m9>?G88@p+^z5QhGf z3I24JFT4zG?8iO?Bb4*JA}Dgb^1F3NV4Nb;*q)(WrSMVVFCc)nm@to)1FI{j??AL2 zrT|~=C;FZk^Mkp?s+_(-LTfuL(N%Bl$m=m9+SH+mpH&xEm$79QG-LWEETgPUzt40; z3o}zL5sm)FGTG5r+@en0Q+VdNu1!QiN{Q)k_W-P@rdA-ll6^{HA_*!Vkd#ZMYo0S( z)Bh4!bSq$188~tmnK-({KQ|rSHU+7)wdJyr##R8+`c|+so&>VB$;k<9?!I^F{%cC5 zz51F7BBsm$W|Kd47cftsb4CEI8wot1m^A=a^h=NdCl)t>PfF?U;Bth8a>s;qRwY8p zLu+x6unh~VJ$-*uWFfy|lf@hY{xZw4jU%$h5k&5L?cl*M4Z+h4>-%?Ch`{j^3Dsf>|tj} zZod7dQq$bA9JrTIehxx(uS8=Ej$FMWk?^>yTJDXZ@J+5q=!k98+2kJTHt+_X0j(gv z*MguhN#13jl`;OnxW8K24={LhRO#*QrlYp5X%+^+!-CpQgSUk#GIh@uWCqeX@HP~z zy!&>vas_LU!rL~~qXt76-Uknoo_1p}GR^Ma-T=FtDHxi@6cIAkH+(ih+xq2HoDXe- zqQbDiYh`A^>n`iF>7RD4XFBJW>DW5D5EYT}o!i;A@rOY#1GiH2;TQ$~yd(FlYy zZhL!s>BoPYPF<)_Jo`(oGqc@6vR<){`2#2_WQr7DmWybp|GoO3$;FF&0&I<~GUweT z^d2+!na=(VgeY9bWk0#X0W%*c+8W)mKc+{;3C~2?Hk#KKD{Y|?-pPQTQO@`!HJ`E6L7x+xNb#h&?7}LN4kSX0f z12A+~RXC3P0V4yT;BGb&T%9N=G) zQW&rKhUISoR(LLmW|z-sq}qN!v1(xH+^#hxbn&-v(UNE6^g=Glo%t|{iAPguABCK5 zE1~#-@w&$uSgw;Uv>sG^)5R8_emZ+R7Q;LaG?nvp9CzmLqeMq^$G3sIYqTTCAy*C% z-PAap0=b9ipY)a^8b@iGsE2wF98#A8b;*%Q!X=rrNBgFg6_abUswpBdV!rDu&mI;A;l-y zD&#c6v*PBJe8hwqERff`6?e0_#EV_8FENoW&ebGsP7cX3xi-8XG}sl|A{ab#{p~xo z0O?5}PD+`%J}4L+ZFj{%!W~hoy^E6voEpL>QVgvxXMty7v3kO^sb|Y0Qy{AZ(YrZm zMwwFNsQ9*2&$K*8K##y;%r;Jp&h>3zk^1ZDU^xwhNKx{d8dD7@B!$-9HK8y&7rriQ7-JlZi;1%I<^w~ob`LvtHX;YI(7du)CZatR6L#jMB3`^ zZ4TXL4GaQG)rv^MT-#f%N}`itdq+Ran7H545R z%RdeXK)^&MTF+N#aGV#&+SywmQ1{95f6^#bBBO1;g>oc?=}tFu1a>T%teuwV_^pr>)z7=8@vmNJ zVDr?pi~_hNN>Nf9T24#Pl5C-Ou>t1v;`#S zjw@S!pB8N2`=`HQRIoC78`l;|vd-(7z@L;+Cf8h|U=wv+K*0dGHCO8$9oE=1fjBr;O?JlMLC{_Iu-)s1H`BIR`RRsbJX~9LWtCT7aw}7?%wA&YN2D)_M<#`u&WanyFwS7L0~M$OkPW=@7brOeC*hYlvKYCGF>$pHi({|CmbK!5^OAk6`QL&?r&rxNqTqj$6~-&SbXKx_-kQj{o@7tKEc zyTQ4KkZucj-_W|A{AF8_cy4jBbD%Kd1`P0SGVq(1UuK8sIsXxl03isFdwPmIe|ZUY z;r^vl8JKM!9#pOa7*i=ioY`qS;6F*=Gk+2=uN3G^!ekGwo1ye-w?!lPc2eGzWF5DU z*^Jl-yYXB!E#KOcTUGf4f;N%q_F{5x@1H0uIdck>fxD}a)D$HsO_mQb6%(jTZLES9 zZGH(%(jB6zcnd#S33$|9VS1m(ttlh}Hz0hMitmjpsj2eL$^O^v-p{$a z#3SxL>HC5{vuSQkrkDAajyws&;%aRiv9cKtp` z!YaFKoNkd9CR?hXtjv~=TcxDo*l*GzQ=wp+p~)i#E({HAkG74UG5sN?ZImxkvzX>5 z6WG)M3QFX@r-yy!N8`0+<^trnye2o4Q+$Cz{rsh1zknd-o#r>mXxD;)_G_RLT_U}R zhqQe(LXqLu@wul7*$#25M47A}VRDJTU=8ln0qm84a?rxFvrg~7#@Jt1&{CqG{H1yN z)86}Ev+u)$08*ag9QKpdS$@K{aA%zmb=JNQYLALh;H68{2}2(PlX8pOCyi3rrlvVf z$G*HNA?**?aEm&&sPQvfb%=f7JKoSTpveXVbYBW5p10nwxHjJ`8uAlDO8AH0gUjS9 zBI}Ki@PDkq)87;e(?utA|KL<&saR*+>Bvz0D%fG)b=p(>yr9~B&hW}_luqe+9m-^W zh%*s%Ne%pkI^D2RnoBnS>bB28sroa0`>+eqf0^kVoy>-713qQ+ew8Bx)|zGV%_26z zIRmwTd%EA1_NK|>`UTQ4s?XzWFeTL=j;gDy0xw%@UXNRc;TRujO@CSbGSNz`y!a+ z;GLiw0e4XV3y2~?D!Pta!=AfLtv7>$ekK&oe=jUrve(|mYWt+ z#rO;Pi9*^bn;molsO={rId7WONx^0ANmcY{!vhg*C_1lcJq0t0%Ir$TdiVK;^*HRw zlrbz;HP2D?H!FG7bFT3h2E0?GfO!b6loe|K{L@Umw!QiNJ+{8(A^RhR2dY6P^v+>C z;+dn8#-9`C^)07J0Zj_DN2~qwPqTqq6R!a$uXzWyop&iY-1eq$@pN~t$Iq{i+ZYf7 zne|zmC{+9BpJr6G0Vu)g^T!YL?N0+nN5+ojRF}Nc$vkGLY*#--)G$JZNf9+wt%gxO zU~Zrhm*Vleez>{nfwu6i7EcNLkSdFCpQ>K1V-i$=0^SUW>px3*V{&Tap1rH`y!%(2$0^MGXdYEb z?5w)+uKv%wWLH401{_N;2tzj5o(>_kc0N_uwrA;!?H~jTC^%P6Ogtq7VRPDY z!qV(c0BslqM?rqKR(>tI$3K=)uSa%lI_aM(rn^r*y}Abnq{>K0=MAibht)1m0diLb zxc9kpEHV5FU-u_sO8su)@Z|OD&wd$O?J2Md8(CSX-(&)VR6sl91ej*wK+ee0$U4J! zsbVhZ_uD@OjmS9uvJhFEbovWgn%eeHziP|fm{59f4y_8T7O1Pxb@TF6ruNsr_g{{B zim1%Lapq5qtTj>0!_QbUO^3e<_`FUj4VCh|?ge~X(sDh<-p|0^gjGEty-`b-)fger z*B_lUd@Rem-Mi^oraIkrIS_Kg4J1u26{=;r${uugNFA z3a9{`0cVE_!)h^B9B(N;n4TnOu8t11+$S7$G%)DPJE3lQ8a=ccb*STslZS=+F|fDM zMAWql(ux~KO*SO|&RBp#!_du+sd@W@xoN@?_l{O%TpK7{1UfE91=`4AsSFw zR1F{X6AHfiilun$g$Z{5@V{)63GYfD^Q!2efM8i!-}mzq<-cCUK~rqHax;62Te*eH zdrw+>3qisasI;BCGj6||5H^I(iL~{?Kjvdo9)RLeN^laL#7~pQXPHgQQY|CUGtmC& zr5e1g`@0gl!w@7XO-dr7XON+*8PGP4-zu-O;yBZdKt`8;Htp;U#&32_HL*LNwhfkf zn@lis=xlq6C>bi&;8O>_Qq%0de*nY~0Hw8=QtO0b5!&oioylO-htlFex#2NsX(#a~ zs*JZ?%)h)&aE<*Xe`VQ?+yODES4dId4lLOxu&#(HY;ad>uOmglbVI|ujlaZO*QCn+ z>5J_2{ib4#FAw5x$~dlk4Y&qLGpyoxyaE8Id)ZdpWAiG~`x(E-bKTn+r8}(ssuq(m zi|Q=F0=hr5#dkD9Dgt;-D*v1zGUKfj5cTANv}&@W-~DxFGL-~YsnpE}*45LM5%2)N zsszqn(}3cMS8BT--rg>*xLHLNpwl$v&ji0TZ_3h1RBu+xo}E(uY_#uk`ok(p;H+}l z3|M9Ls_FBRaYij}j0k&pvYO?R=3j7+9!cJF`y$?ef9b`-A9rm13Tlo~wu{>`v-au& z#W4Y}1uCANc+{Kl-%y_$b9`87r~Y_E|GJ|G3UiZcinT9>3PaV@8e5fe_P)=l9cQYr{y4ap@e(aSk7H06nJv~* zi|=cn7CiC#YQ@JTSHHX_17cAQ2Edc*x~V)pPcwIcY%__F0AnGMk}4oRf{)JM@Ki|T zYX4Vei(XgPNqg#LW@B@gO}!Z!1@=a zp?W=FRi^gz$~7WOv+XsRZPwy~2PxhzkKX0<4-_Q)`+P~7Cu#uRQUIEF*oeET8<>69uEOl8D2z;IB5a5KJ=s&DO;=j88$$5eOc&T~w zL+g1i+07#oHlvMS@2fG*kqi}}4EZ_PC-;evDkov;^%Q>&)WVm)@YaCZxllcWEv154 z3fZ)?%+dIJg4+v0&RmmXNjY+@X)7tM_FMwGaarO!cfrx}KH{RK$epG$XXVdXEt^>B z$Mk?!|AF-qlDfh4n?r)VUMX(>3B}i0cP=2HUl}2Ve7Y4hEK7$03-Zih`7oKIR#{ka zv;1_U&1zcx{Y_ax2Bq=Wwx$BBV?!6fZuCFP1{Tm-?C>J*PEK3k(e!JDLuI<9fRCAC z^7WSmHL-JGrTX(}u$B!Tj{v+embw86<)6Itc)Y7)VI$v{89!2OyV38P@-gI*&*)VE zi&NuXQ4R6h(98U-(RtR;@2QNTyE3Q`pR^7$>0}pCpAW( zoQ6HkH-)r&PP(_f2Rj`BU1ZXWs;8|)&*@vHioc~xA`I$grUq`y-A0VxkC{tN8Mt$ybJOA*<$ag|QH zLyzWYHiY*UquUVkj_2FMh+0 zlgm7I5;6>!>o6$1N1E`Cqux)~7%pmpTL}AdiK>^(b65Y5g{5?d`VrFNJ2TS0Izi#f z2S#fxDv))*xg~i<9J8vf7Y&$VrBEsq^Mcy!i1DKMpi2nfA4`I^-FHCF4Y8anh^weR zgPaTCVwuex7C|Ef?fMmjncEuTm}Dkny2+wt?GH?>rW7Mu{I4AmF00~OWeNv?9x}54 zZu0^l%X%^}6w)TV*SIgvqzvIUJi z95pg;KM=`QGga6023$i# z$`$xmZb_#WKozpWLB&#LvXn3`amY;UkmkSBS^0RQjbzAj!e77Kp!4lZtm+S_-qq$L zN96~G_!Ju8K7CD@*Erdw)FuZwa4Cx@sJPc0!WO7G)UT3a zzzTr5x~u%_c8)1GyKum7r@EUagm- zk~=ADa{P!4f5g~YJh0>&)9(uG)C1Mf$nFzLI<;i2_UJfJ@RR=@+Q8SXL12MD2(R+( zgz*_-{Pk4m$&aGL_llm+@UsWJRs&Vnm&Vb#JTo$ML7`bhBxowz*0IU}mb`bPzs zo=X66P)+)(o13TJ@*0{}w;c)QkTmhcY%3e6w%B-z2jmM3C*suUMj(C`R4~Pe0P}K$ zlX_y-npfyX^Z6W{sm{J+DEphaFlC^_V){xi4$u#%rre(#y7je`+{SLCa)IixF4-SU z&HL3g+~cbk5+d_3kHZ1%wtkAeZ}Bew5Oc*hX@^v?8LEQyGnk_OZ$sgy zFN{h+LZnQmwrSf0IaWM1;b=}55?op=%}~@)jHKKBd*C0pN|CATVbZ&&t9u!eCwPE0 z136y#zgG%UmZ%CaExSZOb=Ndd!BA#$(jQE2kc=~niEW6bZ-MF+L6sMqpm{_XMgyaG z6GOiJP`HFK;bLx5M)C{RVD2!0+H6I~+WMw_s;}br-v^n1M-W>W!dQX9vAYevZbRxb zOLS?^tLynt*Gv0e-zyKj9WJ{CkewEei|Ycbw2B4s;iw6wMAsm7mi?~1Z|+(WdmnD| zynLyqqN5`xb>6gl8J}&zzfYHi1D@0XXVcio91*7i-l?G69+^DcJW%`ro|J$#jzomK z)iB(TW0S~VH`6Q&FxGQBz;4qAc7#>{-U}hCk?+N)kx_eV%XlYhFJ7EZFV2NUlLJg0 zEn9X2V7&KKoEUhoba~278l%$>1?yrg708^U9#i|wWGPQQ&iPghVx!M20$F+w=%6OP z%cr|SQi~oBpoUBqTM)X3yIb*xDdTtHZJA$^1IcqkPw64iv@eF_pOUF9pfJiTLr=uN z`Av}MRx8*jA^AX7`9q4k@9g2qdaAxhVsW5K%hTR;0uUCJ&$MflwKg>ZI5Vltzly?8j5Gqyvtfjx z49s_RH9$Op#+y#yfVygYXC!F5F_V~RJ5;3e-1i^ZiAbGKpj3o1qOhlz&vzBbTL8rI zkCz!F3Jc_@jlaJ9?hlFaYGc!PVsRZAXdDqzd%B4cJAL$^w!tp^&p%QdNWS;uH;K8) zx95y>$!a+xizGRGkXVtpyFGgxXqqslxuBNgnRMnKH_-w#dCs z-kB+&2qY2yB&wjlcwsG!K481IpaN*+9yq*!>&{(jAp$Die^FLc`XQJt9|$=xFsC;6S}b1r6Fp>H^oG31*(_$8w0IuQ^_h*T*TC`CcI zKPf!NL8QG|mfn>|1)$IAu#%spf)HH^q_IZyXR65h+6#q277OHEpc0TT_)S0&%@Jq{ zLjdBxF{9W^EVG5k2oX|@KaEOODf?SnB)>d&UVB(+*SdP1u`%T|WIE%b6{+9nbN<&} ziKtJwDu+LOhdAmVC@TDhBPMS!5U)=2WF^ttC@(t+Ynjq?!zk9!={(*U%YED$VNi;w z^T#pH5)Px`aV2xfS0@4Q;14I6*2U%X`6Srmwv(@FiN4QCh zeVUbRW!}85E6G$u>Zq{*ZEf3JsK*yW797Ye0h(&VuP_9EE!eL$P3SBd*gN(@B6OGV zXu}=KCRLh2>57YN@sM$b64+J!iDXhC56K{V)_QtAQxs< zKJCyTQND@f=WtUzTsIrpS%QFTh11Xs^)f~8PVGN@`z)(bEq8@n1wAoISZ&;ki7=%u zl`rd8G=qAB8MmETuH_IfHe5<7p=uU>VkE=F)#%MgdA-mxKv16{II%YqN_Z%U@_D-2E}j*XB4c1_jR``LEQ%xSzICE6v;^v2Kr z=JJ=1v-7w9dAg;P$RkTI@(k9}(*qrT_ifi{nWN3BvwFtpb3s`bd3>#oq#CYVQk!{KQ58^d15Qd0lA?xg`9 z%}?Rs?XJEyO|vJ&_dcD_LJ?n`gxr%Pq#m1%ySdOVj-`Uf37a2WHU#c|@ z>8(6r9~PGW8Zpgt+ey){dn?zJw_KB8X2urOhfkuM7vdvVQN;LQ?YjMR4<^mLf9oi5 zCY4$m>tkv8UbTDTIi@cTxBc22V=Rp_pf2^D(#k~t&VHc{AKg=^mloV%O&8vjiwGkf z+dOz_*6>5;_U)rD$O=9J5uv_7Q9kh*n4zc5T!x|o{m2Y`|6S2P_PZNteBrDUJFVis z4@=6Xd92t8HJQu6Cl!hKvw3GNNgu*hDt z#Vy}8tzYgUv??)}4_F`C?`Xn9cJk}hU01BcaPKfB8+HDh)kZDTGb%svQhcSJB<@1D z!H)zqnBR%C$#%O~ITRcmnC`#EC0!cRV0O^(Gn2@Q*lkA1Y+HDoEPPpOz&+ixv&?11 z(|ssF$4701iN19@nSb6>R?ZG_u`k=Vewn)U@JYw*kQ;I@4(PUwExS1m_&C!z`CN`sWnCnQcQ)qE`n7b{haAGw3c}i*Q%H~^Bt;X-A zk&NMMnPTJdomqW(D{EcYyC>4x{=r;|xh=mlZ_+<6jGJFAy-CmC8O^$im5u2kd61>F zY9QP#H?h%bNXD}!#gQ0H-!$|*IYH_UtWrAKJ0=b_k@}Ub;^Ef$8%rlXlizX! zx}PijjFk?W*!bW(n>MBNf=0KiS-a5U;1G6~`a*@wiIJ>atdVQ!9mf_gZ@V?`HgUd` zls*CTSd$tBl5Pg!Rs4XGS-1Yr=w}CRg^w{Pqu-@_EcC^^V=={8$|o(#la!%i^{?W! z1}CCt)U@$15`*hfDEhh#Bxb8m3WQdqo?GE@=DtnMkpF&Kj3vzUK3WZKj@9MpSaP8M z@<6%zW7mxgd;av=@>ApuQ%ungqRzJndzqHnke@duGDvJu&n13QVkVvCn`WP<{GP_- z+zURGRmnz1@hLJJi7DPexc|OzQFYyZx)--O=M`1$-6Mjnn)NnN;Tqeajz=YWcWBdv za&zGigT+9 zLb*RD2LutG^NR}BxVwjfeIi&_%6T2K-#$_t8iR9OjvDrVGPhTo;&Nd~%#S(A(Nq*m zdCiA3)xABY8gq(yE%#uaX-BOC_u>Y^$CTpSa}6l&q*M&p#atPb>YHW+i&UoSvL5FZ zF;C6eo8He!?=O*$=w%T8QL=ytk@imd-hzX6$VdKJr@Q#R7K7WubN%yQC;L`* zZ-p*TIi`hemqeJPTU6Nxe^Rvj~FX^!9s?ioj4Q8&V6?W$Ddw?)slG320%p*E+cl?R+tcW$fknDQ2J(KVnvK9j za}!TKg9Dm+^!m8@e3!WiXq>T0Tkrlg1vA5~J*hb6W-K^-WqYR0`Bli~CcAlG855U#y13S+HI-X{`?1julaY`WX@6s7-t-F zs5tlZ@m~4qiQY{{daTs2>mDKq^#_aMyElEDT%&Aadt){lJD-j(5#lsGjsHH7qtDU) zk~GeXUZujF{%n7@1ivt_fncb-y@LseSW|n>`a1E}q`{I4PGZ4%>EOYYKjsR}xnX9e zlB%M#nS)eqY8&eqJ&j`{S%p~h8WSlA`OIhdCvDjKUh6%suU`osmX%yCA0t?DEPY`R zB9)4eZm!egBe~D5EpZs5D}di0>R%+`z~k>dJm}qA_t7}4VN#pqo`ymxX6;8`Z)(dF zk&WDGGx|t;Qj%z$@&#Y5VfxI6_Uz1r>~@xQFXj{Hg;16lDUPu;m1^(`T*q`+<6K0U z0E6{Jnn|h(zxQ?SCyW=2e^EYM8xHLiSXp~E6DA(sT1MEz=+s#aoA$izjdB(O=P<_h zI`2y3+#hPz-4m>Op6vk`@oC9Bzr1AHIh})`2}t^M;i+tvdD!ABv@ePqRNuvgF$-Z` z_1=77c3nO=-uH4>sQ>X1ET$Qb-X#r=X&l(WTB_0YsT|X!^eSx_YU~VYBI%P=i*rBS z9*XSGVPaX7vh;ddxkA#J*Y7?+i|J^00L*YI{Y3&Tu>z-dyhn7L61aZ|y*<9eEK z4!RzjSysce=Ks%M!1Wp1>%Ap}o>3lfh>|qdpd1yDZN4))A){I7`sGbiZo8~F+-$l@ zSag!p@C^6ToXFzPbVYrPx%(oKwFiQ*U+!bgZhi8pGx>AYCi(gqFVO}%S~xxquk{W8 zj%`Ic?x!?Sy|;(Z%wD}+`biw^ROj(umW>oe7cc7|rFEK(#`ew_$J@*;C&Y7}@}@f` zJ=zj44X2BE`Rr1##fu7iSl2R{eltUf&67P;)#Q6kp;s6i?d%L;!9nmB4VtAd3|Tf) z$~nR|#9uqsbfBa<<@!ai4?Ds144dR$0@E`O7NSz&{ob`0nmWN>uQ{G{MS6evx=e+u z4{^v)hf7`js@vCV87m0t96aY(mQQBIaa1qm67tbC zScjYRnN?J;)a=)zQQWD>QuyXk!v;y9;sG6xm*vS_DbxP&b>45gjv^@~DxbmGM@z(w zbawCJZ*ZyD;$^TFYEQ3?tR$@JPwZ>sG1y}HouVmOjFd2qRmMrGW+QhPTGoh%z)&n` zf&Z0p<0hQ#-UA6{w`Y@t1fp}pS0>8W?ZVB=2_H+AV>-1%77ao3QP5ob>dVr-n(7L# zvQ);xx0St-^peZkA^bDRkGRGq9u)I=@n(zZKDrovKrIhOmNXET*!rfmBK1YevO12E zvklA9cIz$i8?Z>VG{mA2t9c0A`I!N45v@%GbEToshER#SQ+rL-#gmf1E}?Hvg|B+( zi3}kQ0~ZEt!65o+1$n6ruKS_5bLK2?{dhluh`SiE3XcFb@wdJFzZz1IoAvI-7ux5M zyci^CuYlq{l*H6A^2HPe&rdl)L#)Uasu&d&%HT3#YH}H?p_<%rq=DHn? z^?SOxh?|0;F8$N8{-`bV-lMIDx4zBlIfS~XrCyU}oDVlp|AP^T`}kHzVZj_P9n})Sg4dm&i^Jge)rKt#(4?;j{+7_#i9%n8SNI-wdhiuR8} z@*E{Hlghk7bDfYeADUe+{Mc;_xDw%Wa{q~6!|qFS%<50)`86WHqtO124t>>3j9_EC z8cTyq$&|xXpFMjWJX@P(sT$`L*H9%DCDD>2VS~u3em%Xaq!8AyY{JU|b|&|Pcd=ZP zipjBVD7FwuH$3HyFBlE1gFr|7UXKAo!fq%t?@-z6S9k2pPo*e1Bg{zSjF*`5r6%-! zt%Iifk=J%L_+ku$F=g84*H+IxQ`f`Q55@Xgm(E(HxDU>i;;I5)%-(ZJelgak`KTK9 z*K@}U83ktkQ+SNg(uU&4e)h1Z9khkVFeSb|AFONGx~e!4{F%3sI>8KBwuaij@GHk? zSbD}`U54G}FZ|u(4C;;6OUFI}H>DPPe6RimaY94#8g*vp!Y$fABcSM4Kcjt<(KK4r_tMXApc??j zp*^X3iZ=}5dDUFVc$XU7>O|bgXHS5q5xq(y1{sfv`nLT<9>d&EQEU#ra1|39_`*vh zHu9AjDvjW#zOMF>EWV9{nn~P zP?Q6J97gi+*^^;dIx&X%Q)X8$^_=J8J|=O}>y5j6nuv%&!)eFg;n73;GZ%a@MEESOMC0rr z--ATt@SgHJW}I=1b@g*gqGYTwE!`8hr$F+<0{Pz(3mYaL-$_HUPrNxJ3+P2 zeXpUI*JcpVBSBh-$whXSGg(O`2Dg;UFPnpFno0@=Y0D?IJ{dSA^+{w@;^g6JN%VnK zdd4;!XDwk%YSWMqY>)(cub;794)#oN1qe3d^Ad`3tBjE=R+N`;l`FXk&4|98<2ojh zyAr&G`{tk*OwWn``}}xRShsBC*xe`RO_+i=dCmJunXY^N5FTCp8EE4zC=Nz{A6%y` z5wD?+FOMc)?03e?k|Jd(aLeM~kuLj2trkx&mKhf7? znqh(o9tiOCx9d&g(N=H646Z+HPtJ<67U}V*OsGOEkM_%oQ+uBjdLtsZzauk6sbiNH z9XWshxT@rc7SHoF3BtaiCZ}pr!?-jQJ=h+s_?as&PpyIgLkaD<7Ef)WuF+!|HD7zs zbRLT^7td1&@o$4en%&cjSqkTVW8POnhq?-M1-Kw-!{fo|C}HmVmG2Dg4{R|>T^dds zG87NQGD%01l%5Z4IyMIjC5*mFNyAxNMvlTFkDyh4)2G1K-Wb1QhoKGP8_c zz((aUULRb0BWz0%l^GM(cSzP&5L=#`x1W0;495WW<2J!zf0y8Kv942RB>*XV&ydG4 z=-Un@Y^k(`1>{Wj3%>i{KZALj0N|Ag2A?`=3%9H}1|BeXD&@D$o8t1$&eT-`w}pS@ zjS%q1Sg%J;nSppN*qG{caZ~ZzZxC+gLGa%;2e1G*?IOoqH@#+wg@(;v>lQO6hqZ&) zP|G9?Vna5`MKm~qlnDY{)B6n9G=j(ExzrUl7{R)I^gUi0=@J~FQG25bRCX6rKuZvj zhduvpAfX3O1EB2^mDdPr$Dpo0#5p$loQdAxyG>6)skxDibN#22mrmH!#H*~z7s@R% zL$}a8-pEfDk@<=W%gD^fa_^ib7t}FRhxXepuU(dmMK)=#x-D-70u#>r0EM|jNt^C_ zT&|&%6&*MsRI4G2wv7tvkh6Fi`QU~{7?0ZZg?f^t@icPA8J;UD?cjX_r%NwSeoXrj>q~RE!VAlDeppgdK2c z{&gJ^GWTmz3#?8)bf%u~EgW2+rB(r@^k_aBgBM16Tz90szr87NLN0o#fo=vj${u}k zm2i&%p5! zym@~6Kpq{TL!ItB3zmJL$n#9lyzA}KTR%C*2emQ-vY%pQGM|gum&DV~iF(+5R+}BO z`;s3fMqf&K@0257U1au$QgSd%$qzHBblfYiPEo|ny-~bcElEt&Ax_<(7xXDcAk902 zVQ)zaMk{g{{0rxno9T}mFwp&oqAd-(N!~=~|Mdj>uIbymkNtP)lzJ>z9O2Z1WrUCK z6e8&9`3@VT_-IR!I&OuY*4g5A1)|v+PT%ddzdn&$2anee@{JLA#)rfd=Bf{3cfkwL zmk&K(yH%>c__b`k4nD7YiEg?#2!Lp}Cb?dn*P`ZstKf$UEe`9FFZ>DMGduBZtE8yk zyBBUFd(t_-O7mn?@%P&gO%r5u=>A2K&`X-IMF$~D`6WlaJj?5IJx_;TC!XldNA8tn ztakGx+3wFQvGCsYlrIz4``^tQ#T{ckn;K!h)n3x`vG97w_{5;;FY+A##7PbUAegK| z{|_ZI@FR}^>;yM#d5wQi5tF z-s4sL`4k;;5t)UZ<)!p0erkJgCs;Pcd;z1d4MxG2z7C8+aEYt+u@0G~}V5SjBx4G8aQ!T&<} za#wL|>g|oJt;VU^abnh!D{Fej!!&fxyt3GC!A7di#`ZNekhA( zA+7eEDM!4bxkS;nce*`u+QlWYr1MJc>NjtnyB^Z&;&43g8pUhGq+>qmmc|F0w5f7-I>xh7%s-9H z1Y!$Om${EpK6brCcIG4Cv<@``)%6XPr^c=yL}E&={Q^s=506Ju;KqW}MuTl@UTDk1 zX7|AIGwsYp6dK8vb=!0O&OE*GYF+;qVEzLO8}zTA<>rs^3_A3K++8oRXV^W_o} z<2?>o#{PyTrRpba4qg?d!ku3VOK_V`7MV2E7+5%^{ZIF+JY~_}BtE+n3tduVYZd<^ zNuDxFUgb?TlfIH2rc4JHr^A;%hna3)+%ZnFM#HjxmF}^_RTE8Dc4(NQ!7)ix%JmRSHAzQBXq6tH0?n%b0$fZPjF|VIU{n4!I8C(bWEMovQ zW=M&Pzb5Lj$X@+Kq}(=m?R5jDgpPa-EAG1S;(kKa7q>dGxX8QO&lB$|)e#-8O9yPa z*mBZPt5ARsU@1vh$H)=Gg8DEuAI&x(S6wpOWh(yJN7DXPyc`5pZVR&(6Q&p!E=0&_ ze8vydXPrJAlK>;2RNM45!6qezoxK>|IKNe67V6YQE*;1etN_M^uTcC}Gw)^Am~I>x#4r5$O86506FFw(?GfC*-kk}(;aWw-?keI4ej*JaZY zB=x`}8NMhUh9ui>NpoGE5}R5qr%!TV{ctynovBk7*0 zdy6FI` zVnR^IakqMA$EGtjpg+)ZaKTS74xLLHxLDZ$h*4|-%69Cj11Af}8$7u5Z&~)ALRDb`!l&p-)y$DgVt{IhxOPQJ1 z3T0oRVJqWW4Vx${!WH5Y8QGcP+PP$B<=UHEd;QM!`TqX%@Hp>tU+0|Hd5!1m`F_0v zv{M!-;sSUx2-(|2f4}jOw4ltxirQwx0;~8Ag18Xdnm@*!M3E+G$hkh5aJtN(8M1bs=>Zni&!7P~hVy*`jn@D_@p~eLs>=+8{GmsfCk? zQ*>C7V-nKFvHlx&ja=nMBLd6lUj(h3*go=|@815EZAo;5X2|ATy0ivgIb*Sd=Ct>h zPag9izpeTnZ~}CTxZEh35{C6Hy|023mMTOH7IrTWRzvYWgDDGa=$_ZnPHe(sXK)Cc z&2u0QNUo0+*Hl99DO+?G=w(; z72?H@>Rg6r*Dw@4HsI*Pc!g+8)|KcSm<}hmc(*6(Zs;7eq(Pst6h-U9#aG%^f@H3I z7}nSHzDm-fEZ+mp21s3v+^8-FT7nl}@%ro{JMtG$_nDV>J-1ok?NpTm>yJFf{%aviS@nsDg|DEEGJY9W+91VGZ`F&%-13lS;c_PUN z!BYvI+no)!&%4|aEs5Hr_fuG_pCH+1FJSvHYi zqd#@wUR{}A~a}^M^Qew z$U?RmBd6^mVq~u#ikk2jc0n_DICPxaELQLB*g$q;G}CBKdgh^>r@vHi5jLgPq@yme z{p8GH?occv&ODIO3lRX(RdX|m#KlfiTYjh0L9{6V}3Pf}vkQ68^u#GRh5H6m} zpRlFg#HtYrikuT;t~QO>^?k=o3x2L9Y1<=0;?`&Q+mJW3Ty`j`hd|-l7O5WBFz6mW zR8R`e!|DYzG|qi{+2@dHF{N_et#u#we%1GRB*!N?g3ohN$kGFW5B*F{XkL`?RekQ` zSYNBoKaH{c>tQ>E*w(No1^y63(eaQt)*nJ%WZv2MpV?)qNAsq6zUL~3QRF$TeHq8| z$8o3bFy*c9hD%Z8SzMeSpR;P}CPLa|g>4L=1wbu+<{ELUoP3#`ELsi^%N($yGp&ES zEKqY0#1i6Gy^r%;tqR#-9#G-oj+YPwN$WC|X$_;UEDCaCV>cRuVU^9@{r755yZ?5< zxn}M%#lt4o?MrbX=?~B6n+wlHKLoB|f~f?}3zF^@f0J8G{t|-F%Vp`6YccUF{87JV zpKSS>*%Es4MxCnvrQOn5TBz4rUSG7$ECy^H zC&TzBhUYE8s`X|*qWW7pF_!&8R2ji~n;=c6Sa~jlnG!$xw&gJ!*lyh9E|<#CQx(U{ zSwe0F;xShtGDWV-si=O{vA}jr9ma)~`y1OXVH-ey+bv&j z5c-)f!?j7?;{gX2g9yx`bNX5|ne9Wjbx^nIoKNf!ST$3)mj-$41w7JMf&0FlEAack z@PgrAFkY>^@FO-kQ(UcD3?0)KxA*A9XJb1$mQT(bgB5=l4AbSxm4aoQ{vD}{mz_2v zZ>}*LE0l@4Ju9tpgu1xn%h$}}{YDprHxEI0^Wuqa2pV{%uTn|~8Sth{I)=UJdB;H1 zw|0(n3~1=aWL^5Uz8&m`p)zFA+8qf2Yv9cV-?CVOvpYW@Nl}Z%jLPQS-`|&+YnH#uYwB^IB%ty>+Q!?Vegx>oqC^;FP4^p{41Ae(UeVy2|0hg^|a5=zD&;FcT zP#f@H4cSej;-+^&-YMyvTwrM|U+Z|tQ7Yt!Yv8#}|3(I=z?)uG+gD^`sx5Pq#A^D_ zEN4~O-$#jM5!R{g>9mowVG~9KwJI{~niV6hvcQKarctjRX2T{%cUR z7Wu~HdX7(i;&eW#fp1dv?)C0mPkckT<$KeIOsfDY>;>qBA8pEohVWA(5<>Q(ijB8} zv#@p8q#%kRPh_4!HxZ%9m+pz=3YRATWF88XlX|Pa% zIA@8F^FAjZllYyBAO12q1FO*X4WyMy0{4lO=_cKCDuLUcd}aDRzIKw#hqcL5_p07t z`TL8HO&)Wr1YgzjAww1@k=Twl+_hrPJ{3WkJyI;8H_h0(nkt*Ziv9xFi1ftV`~ybt z3s!po&O4Rn`3axG3?XLs0tbbpo4}+gZNi7Y9#qy3UC6cZ8J*DqJr*J7Q(yH8QE8Zi z3PZp|$3gRE*g)Ue4haKNPHg!7wjf1eYc1A5y*`33_1iS7dIjQMmJu~b0&x3D0qX?M z`cCC#w?p%1k(x+v%DRjupSyvY6OE4k~rO6Y`)quNYUUFhAi@JJwE$KPu8o%3? z#f8@sgDyAWFi}K)PD!LP8_ylY&rBo6cs)Vp=c$3nPAc=LeKCv{EwjX8m+-X1uEj%l zo@^ondudi~-iLWxw?*RD-_5?yt!6s@4@v}l`||(3Rh&GR-bYa``ze5x= z`T!_4DgZBX+TI8FUYP1akv*p<^y@a7Z$`QT zuG5@^8Ept42Fw@lE-6L}POC2X$dBu1OJSH?WEOW2GnXR=P3jH6gt}?zJxdu95FyUY zmOlt;FDRP%_iUC+e=z5abjz2CAg*{fDa##5n`^!&6eh@a3WX7L(Rpx#s{A# zxjb6@Z=}*iWq?C*qi~r31DZcY!af2F$oob|Vl7=s#ifr>`MFm_!1a1F88W8ILv8pQ zE@(jX@-@h^#AOnVx9|CMTYRzkxNz&3aCWpU|997#O|GgJaMo+it*aojF26p!$A=Do zvb9fd%WdcbU@L!yp0%UO-&}8h*_hs?xaOm+WO2Afm-7nn7;+7g0eHsb?F%4FTn^Fr zK(o3Cdv{KuX?8D&e?^)HaYg?&)3r!Igi8(O<*A{DrLJ!2hK5i(aYgGP29EEY7gFBf zNV%-zoP0TO<<2dpcPHcc+aJCF0*P&Upjv$Kt6!>(c5hApY|d#M(JSAGi%{pXVUrVm z=c{2t#yN>HR*Cl_|1l^Zg1L`mTAdJ7QWqKs{RFZHBS;$R5%S1#nOaCS^eG4o{P~;2 z_-49b!kS-PCV#upsn^k-jC{JLk{NYh30(P~M-1&K?!>9w#a|u5&7pc zrUh=ibmaBM{tFFWrjGVI==VPnGHK&ZV*H0j7?#CqXydLen3xPTqo;LulQ7Zz-s&zMqQ&+3re8GMxa1I<&y*vYPM;?Cu6b#DXABfQI9&$)aM1O`^$GZ1W*BTR} zW6Kes>$c}#S@4q5kG_tkvD!XOgXM5n7rgGa_o|9Iox$Bo*lLqaR?l9K{+;STfLD=a zu33G@2{Q`HpFBIUX@T)6Eq4H$o<>Bziu%U99rk9WZCF4JpL9@}IR8yOpvV4fZ?P8e zM8duT(+2X{BvedF*w?$F1tH0+0e$u1G6Le#^_8`2Kj692lU&%_L zNfWpk`mB@p95Y+iOPiyZ{;1jzH)Hg^;157%zv{$J-Q_S4d3Dc}hPM%glDn8aI%X zFm)@xlJ|VjM>?VK8owVKmiyKR;B#7k$II8HOYRLTK?5zsYxom3DZ&o-?W_`jhyU#j zjF+0^n0WS+x$z5M@oig!97)GdD|UU6p;-H~u`Icju?D*K`#LSdCp;c**GPLcB(@sw zl|~(3cHYT-rZ_1{S`uLva+uGEABIW#(v#?^VoGRa(ZZ=?9U>+>8m=}3cd*MvOxcD^ z>&(f8j^`<{0bE#yn~hwKNjD*G{YAlvr6iMv?aW7Qd~q0_Lg{F$IcDDl?G&j^{}YdI z@=P*IBe|!LhP(i*H^Pt)xLBz>N0(cN5ANw}Te^&@(TYLl=85aIk&M5SMI3+iFolQ# zNd@ls61~H;#hBgNzJgv9tds1;Q3EUA_PQAnSHimltlM0;rb)Zs5*tp~psg3;eM{G6$lHI-FoiOb!pg>iJohJfdN)Okw z0JvID;h(I}IMrYVBJyOnL6ZDppO^b726o4nD>hW1y)1o7{{@IEPAhvx$fQl1}_m z6Xqr#VQvUwwDBvCpyGGPuA<<%>GMSKDo~h#06eK@Au-tFk(!&rmwN)v%`Oi&+O1QE z^pkZbA*#@#b^?FLYk8;cnL9t^;LfQe3vN;Nm@9lV88!lO71i&_*JWcUn}Lz25kT=4>!JF3^Y}S68m54U+vEFcjN;?e7oz`^ZC3ds zZAYr0fnICJhv`;rqW(x2f8Qn&%A3byD6IMYo4~dY(fbVoyap%a zS#Rs-uW`<<4CEkV*!0zjyCSR5PDl|On_qQIkwN24yP zt|MZiRMyT=HjVSsqvRl1G0N(DDTg9e3qR$SdIPJ58ptN+CNGy$AT}FwbF7^7F-4ue zbTtp84||Ct_OIhuger-o&^u#-2~2NQBB9`hB;ci`^yOpJFI~5_M{Auv%>pcdTax zwsrID`#yj1;W$JphA2=Td^3;insea3jO=hHI9ZHNosuvBq*u3uNkOk6*K=v|H1kfy zdsQr(-H^o`_rJH~oP=9%ymdvuLsG2!dVh`rtHt!3gwP`&fNILt`3&<=R*@*io(Hdf zUcmZJ#5_ElNx@y$m2b)Bx0>gT*eQgmkc_9Il}LIoQv}$Wh)hx~Rx&W@f@4vR^Bnm? z+sYtE#___XEsdt%r)lso6qg-kj=a7Y7t(`(LsF1?^-O*f1BN2IGaItW8ecsXju^cJ zH4htV?caX)stCN}fu`2s`t<2*ojTm}Cj2Iv?^e5A;p@8<2W%`dc>z278{O9yKZfXA z%^HzzI}n=)JV#r5L5#H1B%BUw>;9uky51y z^~ZgiuVSZYYor%GB7Cl1Z^@=O0_t_HqNb{bI-@qY&9JyNN2wgkx&`nc-HnKSYc zXY@_g9?3@Z)4|li{&gxU=K#8Am>w4|w5X|!Ah7jq3t=8F3!{@(_UQU+EP zoQlE&RpGVW2CKZ!Y7M6v(||E;Y)JYZrUl zD=XlXGcjfEj^gsxqIW;fOkK{SZD)2P%BE|Cr-PK#T6hcbLVM=i@FyS6h#Nl)KSZ;I zOMTSYh~1yL6ho8*HC0|L_wFX@mQ$K3TKX?bg%stn+SB=e@B)o%QLYZXkE`DujO){h z$LWfX9@Kuajx;|riTkfwP=WZF6E$X4ev*1j9?iJH58i&8*-A&KWHhHh?{vB>dTB2h zxm9=MJ1=@C*xx!e-GBRpNVeW((Wy_w*91o!toTPXrJCLtzP9B}fzO#!GZ{@5l8$4< z9zV?(o#hzA{Ov+4w` zyHI~M$rjH=?HR1~aaLe!^qe%+wGW^c$V0-DGUa|C?&W|yEH5mh?>6xoy}@6*hS4cz z_vDeJjK0mMTinfI?62daWQu9qo zVNx-Ab=3svDg_|B1O%rT=cxh(vYdJbUvWi7TTZ#VxT4!pj&+6&fDI;j#Yg^XErE1D zFjZe(-mzn(&SeC@pj*HL!p-o~wT^14&+oW+K9TDT0t=P_JdY1wLywB(3{Fq<5==Qe zDv@Ly8RL^d+}@$=iS9VN)Jzk%_7>HZ7UHIjd$oQb8v!EsjN_*yORyIN0@u5JG?Z_5 zbNkwwFTn~2N}8|qG+#>MdqxV~>yddjP6p;15oY)L+#E%|*;pA6st~}~uwK}(gG0(K z4m*1FIXe2!y`&Q}em4~=Wpp@mY3t^TI%5M+k>g*u3H3iL`s4IaE?KUKIF%uIVD@8Bl zWI!_4tzUt$5g-|x2jEKJXAWvVFn&HC@o#f-!4_OYV#NN;W&F_pCXY~m;T@(6YLk~> zz-~RGHLrmrwtdy{yK!fPWDi((G^GFE9_S3k&Tm1}K(zv(>N*pyg4cZoOn>x#OQ6`A z${eLzXVMs+GdLfcrI79a=d)(h6H%qhSYQoEmjkkx7(D9h;^LWEc}{JX2{M-l{5;ol z0rW4eK%JA4dx;@E<{`_N)6s@t^uqKGT6UElLuPR{H=-`sO$B#mbu631vEQA_`on`4 z-k>O^Q~f&15xKDt7?Al#43D)FcPpIyIrZ;H*~vv3mDyb)@>~UAFe|zTA^}9f?N~cG zx&jY5!q>6U^5O(u%F43A1|(8?@2VE3%gcUZTZIEqyK#Z#o<4-tQ=C>K(1HY@!dI(7 zXsj=ra@?ps8%KRW)xvI`b5Vl-=1b7L^4VFuBLbq@KRh{_e zbnrwPe|9mIQ-{6+(bW{DTSQ6%vX#p1iXzoGz;()&x{EM-#OV#FIvu>`^bzl%CNy$go}c%cQ1pNB-N!7lZ2$N(TfVU! z5~?lyemeg@_=|~GRp<o}SrC zm>qZ>sRwxmiQ+m*6>5rWlM|;o^4>>wQm{@R$?95g_qo|I#(#JAeikb=9-9|-txul$ zPRaW$i`)^5)RvdK(ZGw~y;Jb* zD?R}eZ^X9pt}a*~4%@Hsgfv2`0MwFKtp}j6{!=K7B=ek=Oux_g&8mlgucah2jI{Zn z^f%;Y7(7n6RTLQ*$!_L>kWQ|BFIJ~@m_(u?p;4i`e`!~znioEOM@i1wJlVdV))iWD zc~(SBc_hRl7tqP-LgMlOD8&VkrbO_kdX&uwwbmo45uGNr9$J>}3^4*n8m8vBA|dM6 z-ve2wz^)X%M!^2DHP#UT$Rb(?*>2^7sw=$1kI!2WkF%)8-fb5umfxnx#>!-Z46=xv z+WbjHpqivD1mL2ti;O_v2Ye6qDmm^7S5N%K%{B@bU|j!2?Arz85Fv9_pzNdsg~!cr z3&)#uoc;OvSlB92WI&WVGdNWUY4C{Y2YzVfm-wQW+JYU*^X?y2P_@YmL-hquo1gsD zB$@6Abl2^IQV+H{r*Ta{QgJiHU0FTAFd~OxDn}g`Y&Lax4AKGl02D2A39MUvqocdr zu4h#=$jahkNH6vM-+s&o7NvGq{bc#}A+xa(f=<~ik!qJ|u-w7CtSUbwaJU%)kI9Z{rA9XwP%@h0keKU>r|2fE2n zc`|kN)}%ehH1@ry5JIaSPP(;bjed6F zSw`{^OwDm_i?5lQLRro3jT)&Z?|-!HXu*2>_PP8Jgakf>a64wWL!P#)y}gK5*ZXn{ z6YOS-(_b9k%?JowEch7MHk)=`sAmljj}@#80|G__Vi+QX_ySuxS1yY>Pz6Ogfmq^S zEILmID*l5^oL<9L`OKPiQv5N$4L5Z$_i3WYkgJ+&pCGuxyzy8ZKvaztJBlc!-#VhD zem7#iRX#QXY0fU-8Kf5qk+`@HLN{BDN32Vl&&wKLx(u0P-4erR&qYdfHHJMbx;~6z zw&k4#8_Ng8QUIwX-(&+SidZrJt zRSv0m_ucEh`rrQd4b_dAA14x~_diYbIW_L2L(G#^AvC;}prQiQN5bcget{AcJm8sk z)>X|GUoMMjHdQxJqzrHl@bj^!5QQ@+V!hEnIgpsL zJiJxTdCfCxf&ZW0NvR*LxTPW8GybCV9fRZ6xF{`cv$`ZPddLUJp?GpSIW`!3?5{=; zONn4X>OVl-X6ndB`x zB6H=2`d=K1lMNIQ5y4t@u~U~uRCqKhV@agoLH!6)o|)~bf-s>rDoho-_^oP27k%Qo{>?AJe)dFWiLnhiCZ^(V*wT z<8b3=e8*V^R~W{wsM z0r@klPEM{{Pv|_34_+$=v2lJH%DF# zZ=HGDZnyURes+_d_E`Wp#ruT4D@7`n-MZfzP4*@&DTb}Fx7V&w4DJ$c_OrJf=$P^y zUgDFd)|C$)=E1=`lJeuf;a7w!=#C$mWcwPa_}FaSa z>sxJ~|7`xN5xlcmBwW<8ceU6#jr?-Tq4`d2h;;~Bj>eC#S<9VVcOpp5HeI1Fm{Y_d z?=fFkQB%V*!xsLwc3_ovr_voabF=TPihVKW>)g#+WtKlxtBe%-%7b~`kyjRRng${J z^9CI0$2Ik%2jGKC@)xvFV)+LVR~O|ub@BC94*6FWN3IWRMr5fDAzqW#8dhhYxAGDb zKSFPL_BVy_4X+%>XHFOO@GU=oHQmpA3K=xV1+^(fTu_X4rM?eXu#XMxg-w#`)(NXo zp_ubu9^4P#QRcWI#Hoi=O$PzBEot;<&<{zUH@aN#i!Xj;v&w~k$9bV#vT~;v_pVS| zm_FFMH@*~JjP2{aB=^jC1LW?A!8ihegEqCHFE!oy7fQG^-yV6WUvWO|Ir)K()XC!buX$Z zZIkyixZkGgBnkq6x89S0oIKC_N3Pds*3(TNXMNy`+xCzD*Q>V7`sH})>Q*6Qs>4(* zcj)@CEYzL%jvRRoK)05|g8S;uQ0j;^Jh090pC~YlE6NpAadE25Y0aEme6br5nrN(^ z3x0ve*~9Bd*NZJ$*0j4fisUNn~&A1AJZxYd4~#Om<}rsK=b$E5q>OgIMl zn#93q3VUbsd~;LfhSD7bvi?Hl*WdVos5EfM59qaSb~UM^NXGzWS(@5H1b#hTgknGP z#^v>=_UaiIcW>G||Dfq$&EpitCWU^`?B&j1QKvF22x7z_f0Eb{RRqf|gPsjZ{ARd$ zl_`DIyz=b}xhI~+h@l>E#O|5rbDijAj(sWdSW8lpZQ)Xg&c9hO-!Pbh?kb(SKe74l z=i*u0gGAQyKQLn;xpzu=u$i~qM)lA{31WSB=jt-JV1uu8%Zbm)ntAWH3YtwEzO5KU zq}uiO;(4^2)ZKyjH^C90is;cKD-pqM-_pG5@I6LxzK&(3gsF@#O zv>>JYF4VJ0^@XIN#`bv#W#tyTh#O-crQ&NU@+B_h@;{uKnDzW>_PQV3Ur3-;Xg*Sw zl)?#f8X}Kq%^()`P8DU}x6ds(GPPU`5AJV7hukPaffrTxqu%tZ4)UkRhVV9mZHZ-e zL)uKzap6L;hbJZ@el8{sfLmMRPHuR5YyTS?5JGLpMH~LMzxVshM&gj&hjdh`r%Sot zMI&V9cP4ShXOEh;B#Q!ajTA#QQ{T;O-@_JM3=ZX1&+DQyLQ-{I77(pTg)cXinszq5 zyk=pMf+)ni^YKg}w~@oh8|z3i{ZD+*JjTkU9=-i1*3Qiayq;ImG4BUyao*y1yz50$ z8dI?KF*>8hp%}1M5X=g38p)6OX{ijzPAvW7u)k<@vQdFHJOQu?X}^^s6v2xR);3~p zZTV8BM|bGcKK!0}Jxwh=cl<~8Lf)L=*VNFSe1lh84j@KIFo%#P{nexgPdGPFAgf1r zxv*2LHX``Oo8Gi-`X0G;1hFkn+Ob5a9=*8c@6VJNVA5AQrC+_@%E(Zs{tZ5IZs$B~ zrEXk|P17~$Qt@vsouj*}=h*g_)gR|N2bfFsif%0S>!e$SlMMj;#1_~52gzvF>_XPe z$^Lao*?oq>l?(81;-puLh3fZ_RK^qOrd#)4_)r^z8<+@Vf9~PKbVS7oI(x3-;~wr^ z{X?Nd#90;(ozkR?>4$F5CS^?N*OMH7o~PxJ>y*R|a{3oPTlQbBrM{b7gr9MKuFlg} z60&4gMjPY90Lurph``(^MnO~G5q&~klj#d5YHgEY_(uUzg2xv=Gk;ErcDfJ$Zq0DZRkG@ zZKf8V4oO)3TRmKAgMIZwKf7p*9HEq_pAxMi%X0!K4gd+qFsEArNJ&f59>(*ijp@TB z7fXedvS8~9G+Vvr*(tiJH@RRinacC4`Rj1g8U_-AlkfNIYo3UTl=3tk+g3=;wI;65 zHMN55_?muM%DYMzbkaPrEt52(8cf-H6)>Jf zDXT?ph`;qdIzSx4os4-yl3NW%k;`ATLtx~I(d7cJ^KXNmBaUdn5?-xeF6ez_@kowh~|Vi`xBiy?YEWmB=f znv8k(lGsPq{Mm8~S2C#Bmtwk{?^hGm zQB%zS^lxaH2n)?7Yd$6|vMOonb+;&rcuxG1C28lA7;7-^P6@S=1Aa(L9J8Mbqga^r zLB&3vHQ$eB;$ZKx#u;~r5U(Y#sjqi}Ohu?UzL`|ijJuJKnJpCVq&ZWmarYOP8)FPi zhOaxh(L_+2ZTfY=o|WNqH%K7ibl=O#uo{mw<0Jkjyt#X6yNqz z@BNh*s$Vv>{8OcZ^04wrzycH+yayJ*p0C*;Wc_q3_<~afw&wa^m$HC@520r@dx35$d zOa4std#8Gbf?-13@rb$B3EJkJf<=W(`Lz)YVs-B+!(de}%CeU#1yWfJ>~#$<9aJNp zs8PVCD@Wn=3th_jkO+=@_LCGa#VxIH&?<>pAzMIG!`W{oG%<`2s(=$|H~3<7wg&Yv zG`8L&EBAqz4gNPGi%_;xe(hnG&y*>2^-6Wmh_!FjN40E$-$Fn4OTpXFxHU#WM_vE? zkg1L#2l=acKF6}da~HA^oW0#%@iU|Sb6n4(LMI3@zmy=nH=$X)Qgz!+nu;4rX_L&Q zPj|Vi{-v0IR+GL}QE<0*)yQT3)6ot^m3%e+XGuv2c!m0EKv+ZR`K3Wj0@w}#gVFS5 z-`OYL4mbBQy~V#d?;ow#S?}(P9DYVQi->UR>n2d9r5B3)JSgS26%6C2hLsDLq)psD zZ~7#!3?t$v9Xa2HdnscvK35p(k%V*mxwfxDnazXQ_0T_Yq&xQQ;yZqW&QZ~_z%+!7>+40%%m@ikb2@{ zt{EgDWMTbljmC8yF2e2oT1Za$9iZBWw<|Frq3M_Pg?xL0B;YR^T-TQNq6)ii7M-x% z_(st9W2`Furc(bNzMI|fSU7EXyV(9LM!@MxltL79SIzw|M7rwio1ADf#1~A#vmt5_ zm=^L13fAUx3-WE1$u~^+U(RF}{K~Ha6M3e;v3Vh=Lw7Ky3hqpEFr#k3ADJ_+2S1aS z9)>#L!lqy84yP?8^abaQZ>zRG5FTc1$?@;WwE2~3x$Kt{JN{2bCFcaKj_~@FZPgY# zB(l;E*C(}~38yLhDSC$@x!vQN;iI;|r3EV^oVW==e5+%xW->!y=-;BR z=rAKLCu*Q8{rgVkV+{`M@zIxSgcL6+x3P=eeSu-tWlb@QA1PgVUG&901l2p0tu)dr z^|R69EE;7^X^V;BG?;*6)}6tolo}e?ps%N|;)?E;?6uk+)PY!S7_UjX@m9|$+>q2t zLxUq1Q@|bh_+bU+A3;I?gEm!j> zT0>4bg10k8*(UwLKmL_L{z%UVReVw84G)OqAMo9yy!CY#a^~#4C#%pXxfd<*@MG9N&Ab}di1*(4;vaDzFFKaXiq9rj0hVMr9O!GzlhyOla*}-xb)z-hC{fOL zZNy5ecR5~}`Z9L3U!XtSTsjjxc|g%S!{X-$Bf>b)hb$$Hb zSb5mGTpkdO9sTYI9Ey)EJrBi#7o2?yJ(xfEGHdKHiW`1eCfz>pG@{4AN9V0pjQj3S zA;X%g&`WQBpG(?Y+g!M8dJ8BXjCb7bgGlIXS2O66Hq6N0K}R>y$M3R`q$Dv&d?NP8 z!lcU2Yii#v*sYIi6#{}Vj{|J5tbq1?47qTWA zP9Cpv|0q>{yygZsL3QpmjZ#UFIuU7rG%n4eEPDTj{RjTMI?9Nr!Eh4Y_;?QtUQ3=E)JN??RBQQ_uRQ-$;k;r;w; zoNw_li1l>b<!zXA$eOk}d%h?;x zq(FG6myGgf&_`4=^jdd;5!(X)PS)Cz`i~WhiL=iT+mZ*SYE31M=d$K{FL--cfSx2b z89$_sW9B8_fc{)?7Y6hvTIaomQq+7U`R3w|^XyA5{hp7`$~sGumief*SOIF^68A6# zAbBDqGVNr1+y^Rshi8E}tM0s$FsW)@z}4?bhpDvzFr-_J|2>&bZZ zwkmDMB9!IjiD?dS7~;$b7V+vZameS-Yc7r7(g~(e-G)DKR^R0C#n7HiDb}N=8kNHS z7tXdV2p9i!6O+xxg29O=^clf_8OER7d(ZN`B#co|#qC-(CcOvrFk67CW~`xW#CK&W z-Q~9k7ff`=aeYEI$n!>WwRP0G{^-ATQng}cw~DgImV8&8?$GtJNB=txhk0F;$D8+p zuzQG_C<^%EQ1STe-%#D6!RH4vp@YRm%df)8PUhTdF5~wGKHz=-Se+CupK~hkfPWBC zFaFrR5g*3rRFJ(8_Ec6p{&(TDp1p8V=#?nuAMCE}?c?iyMU9+O_+n+(8{!ShCQ9Eq zN2J>5Oq7&;exTDy-K%6@9ao`p7i{ox(@)=M_0&9*X#PBR90vO9T__FF#OeK0is1e2 z$xUjo1E*7{<32l|*fGh1J+^leb-!gThZ9&?slOxfD2H39riV8bXrm)i2@cbc!|Bd% zbl&{)MflU~uftwd)ejz=uHG2MIwf^*{EOS?znL!=DjVLD06jD{A)KH)VX)Nh*gG)rw(9~|eh?BJ>AIX%P$1@r5vZ*uLt z;Z_gH!myf!ocpz-%wF|~Q29?|I-gANo6odCEQ3};WG@#nH3pr|`0eo>4&z#}KR+%N ztZm7vV)#OS=1@geJ3kXt4`EJu1lOwuPEF8WxY%~n z!9i5wj;MJ1OG}#ZvhOngm`aPeW$IGMa-T!}A-Sj-RM&Gi z!9(BEGY!>dASHX)Uh&qH#yehHpxWc8QYt8TZnV2Hf)eQfX&VkpFKwwwB{Zg{4_z0M z^D)pxh_719a`~&Ez+A9s!?a^A3;xlrZ<<>qx%iLn0vRNdbU7*WYgGZvL+x; zS>ss(*Ii9Msc-%lLz=~@SP(pesvMl?5ZlUl9-H6Be?AkEv=bd^xWo+QqFILI`arW?RJ%q`KQ z-lbGOjytHLof$Q~!db$Fc^tt3IhHO+R{xP9|GHl%uv;F>b*;f40l{c_LyI13o^LT* zrc;gu{t>hplHiM^I+^`SDI~5TF73A3y|p=Ot+`uZ*U)QpG3 zhe>UZtO;rBG`BeLx5NKcKPW%ib4&jBQT`(ODEp`4^d3!Jw(MOgLzz zBMnmxOtp<ARSE)`m+xdNzWI zrIMr5-XzLJ=}xNssxs4>&MFIe{MzvSc15G1@h&KV3ijwA#P^}Wf&RWu(Cp~4FFfIr zwpS{>Z@65f!8dISHn+3lETW!=1ItLA3Ud5+q$W)8NSYIKN2Wcu#XHV$4|V^lwausi zbN^HX3ug9-RUmVM`-WUDXb@JR!RQe#NovoHmAWpQg68W=m|$Rm7k&Af;Y}?W6l6Cy zJ&Ov!mW3?86WxO<`$py5A??l{6dyDh)Z}4IM`~Jwwmv+n=Z>~%4WSB?K3_fK5=qFoExj4-xsS*?8UJ zZ6;{wr>{tgRlJ}-<5|jc`HSD2^#>4}t_orVp zj|D;Zf-@a7SpFyL>dGMTA~4&~3RY~^^GuJ%h0`z7iy2N8*@sR3_JH9qwwq)?@1oM3 z*TswZo&SG)Y0hqn~zT!+RKH){ZBciHzHH5iA zLo$=|TZl*Ufk9CY4&+60Z+>}s`;SD8?IYwoehY$P)6O5tF9&y`V6Ua6WDDFzolTSg z;!#HVrRx4BT$W_RO*lA@H-FT?r;<~nVDKP^h3&ROt6mE8IR_KahO1&ylFl0}yNoic z>bidzJt@j#X789KBn^5U>reyIcc^lfcgXIwbo%THhoz*jIr^i0T`(rxRrMXJ@$$Ez z>;BXv^Ek{*5#m4Lz0Bs=|PhKl%c)*{Pq!z?I97cYO)FoI+p<&kF>rb75 zGQ(Wpk$wWnja5i#9v%z^CL+2tOC>4$kJE-&`8c^29z?W}k#6KemeyDU-nl;FF?@!?-KS^1Bbq&5!8ac1Sp;5t@^->p|Ymo2qXgsZC`q&YV)wekET z+p~^F@1&GFkPYKZHOwj^Q0HR3vt16mv0(Upk!*j`SEisf7VJG8GOXjQR zc4M^P3t=nlGFt9#_*jh5h}`98{s7vl1`NHRSz=t0B^NmD5=IB5FL^84*|%0)nCH2W zAGXOV;78*m@-s2qs&|Y#|AD>k-~}CJ<#R*nrsPVV88b|cruMq=M;S=k(W2YF$1Jwi zbzL#le_MwM%zirJ$!_IfFT-AflAxBy^;aUK-rr__jBU?w-DW$|ajtgXwfH!1ZZ;%) zA#9=nbduKc0DF{Y{84hC^sGJT{RUh5-Tb|Bx3b1oG^raKGjKWg|wCqIfb zCOOQi_;mVOFg+J;$ANp|d;OB9pMc-mzt2+M9Qt7RCM?#!yVQ{sO^_am6ja$6BuNAw zDa~n(KhScs+7929Uw_2n<{~hspwQ#Q{MP9yuq&{M-_5RqLDLq%{p3xX{xQzt4!DhO z{Z_Ry2%2&)Z_gEPMxB#ie?YwSu@wPCnwSf7cxrvQA5yR>85?t*c+EkvvL((8~M8LctB?qZult7RhGUWppez3rP=l(V0Rpe~5Fx|CEuq^{#{^3OyhwXZiy z6Bo=l6@~9!-SJmRu>G6&?%Knu|+lV1O zK+VmSlKQ5#MJtk3h|{^;!s$S%^M(=6p1fQ{LS254btH~q>08z`Z08U22U2TA>Z@G0 z!+#AmnEL64Ty62?BtzpZ4ee6x_9uZtob^ET(-6YR`$)4SngTNKvEkmF zhAP%mxA9vtLQM&`{~uXz9Ts)dJr0A22qIl7AxNhP5>nC#NDD~HBBgXUh~$ElbhAh) zwR9t}O9+UFN`oxj&C#|*b57-wnrLb%C(&1C6JdC{Z9SPEL&~7$^oSjl zO&h7lAP^t=yZ3`6_{e}e2iA@Uhud~{xEb9oaH^MwQ`6D&PA~Bjqda$tX#N8XrDh=i z$M?G*O|7nc(Q1D(Dqeo4$e}LCg?fHZgz^Te;61yEudDNwr-*?z$K!oI zq`r&)<^81ZWr=qbjg%Nf+=YKL8^m+r^U~M#OhiOU^L4lvp40OCBMDkX_5RuwXs+iv zOs@E@;VG4$igZUVb-A*`Q+Y6oR+@rO@FUrs z)Kzdp@BRlk%BmZ>^pRS47;|!5W`7AbX7hn8~Jbxrc^&^)3>|q@hrVt zf5-4>)-5{SCa=_rGl@F_g6yWifKvkSDpbS4bKUuamn~BWO?6?OM+%YIbUaxNG?=-I z!Q4;kr&z2@=E5-M`M5_dgSc_7>d7EzW&05=E|E%B^!)19aS52 zw0TCkr+^lX`ZYqBN;X#v%;4J~z{i=S+yIDeAcl;j#sqFMymlB2{Dvytm~09p+kvX# z6*H6+A!e?aM+DqghRs={p*m?e^aUA!gV`3>)?|S{{7<|;=UQ)s}cMj zsO@dPp)UIrk-OpoJe$QlOys5WD%Y1JEH@|!C9$Y1L$ayU?4z?+WnQW4^#Rj6aQ2#M)z}W`x-0K5YYY+FYCONqL}Nj^N@mE56Nl`)^-& zZEsg~b_!6=C5>Y?xJzz!*8DhMz9A+nB@+0I+uPM-g?tn@Qm-$6J6(1t1{U1Dx~q#* zv>|Yafv_ci=TGf*wZnJF@xqsMh4JGDKR)@x*F%v;e-PCR9-4IzJVHx2TV0vbDeO;u ztN-<0iAXLN?9b+&u{`E+g|qAjZG5iZ99iowB_fG|I;q%UENOEX)(*de6;Sd36RaJ+ zy9z;(-tSoFs}~)Wy}yZ^doF{l_qDR{%+%WzS>T|u5D+t*lvT4?-c170`OZDb0449% zcLjkFOZ{IA14Ki^vgzovQfY*LVUFInBzAFM!7F4cS%`=SrEwxx!oKX>+7cA3x&Y2x zDRDp2&Si$%8S%Ye!3ZdpF*50(7(YEsng(PSc2LKh{_HBSdz5tJ4D+a_#K-1X$79&@ zsgdx?E+#Tc6|CdK>QBtWUCKH->QPPdc(H}TcXjbizPScy-~H*A9e{OkqUtK{B^XtU zWO{TPc#0ZJ@=Sv;FkGSN{#clVP0u($zvno#xFG>x=6_7^L{M zogyoR@AA^pjc*x6UHbM4qluCJB!};i$UDi3uZMmt=U3M!L`B{oiJZ)O^^=RT4rB=I zH4hc|gd5F<*C9;LP6}qt0(herv}r$P{g$_N6?z{Fnrk^}s-9S1UjJxPOtbal%j5B> z=p~u2)n6Z-Itz6TKqEhVPjTY6Y`Q~jecj1zj8FzBmwI+qKhRL%+10WKkUVeFBPnQi zpS&ksqzO{MnR-iUF6Vt>c?4eQ&#DQR>m-WG4MjpV)q{>O?W8Vb0$p-XKH1M?OJHSJU^)DSp~wE$+hFe`VWj4L2FZ#D|V4 zadW_3CF~qMn4NmRoY>vOR%`QC#LnW11*Hc##1Xo?EJ4f+1DPEIVXdF4!pUxhbq=(( zQP9eW2XQHf9L&My9iybA#4f~{x?F&>4+^Zv+AFCS)ihR#JxT)Z`91Ohs?)QYW&TsW zL&xo+5%hsK-FGE*wo&sVue+Nkb;9`MA;>?yOXW|0hIpFx_t~!8*UO{g$#5LVPs&v_oMBJigugSwd|>7fb%)ZbHM2{Oj2cABGmK6w|GEyIcY= zJdjG({)t*mZDPRj>F060minFvmgcd3FDu}M#9riq zYD{Wp`|LJQqWmFYdDo7+&J{tG3Nk^i_iGOv%vr2%f~pJJ0Z~CJ96XsBL&W+{Y@0yi%8 zpFPBPZPwy%q#R<=###jVBL;~124~O@RHRbF!-#@NE>Ec^r-9H3`&J%a(z`gEb)O); zVeDUhU0tp{EXPDy@AXYDixtONoM_d47o~m{y4l+S(79dUllMTdK%RL9N%pGs!$NE1 z&;Y8+!+Bo++re+%c|q%CE=Sn-womaDsE@c3>1s;wf;J%_$}DAenp>u7R&Dx zcubw#{j{&G&6Hq{U2&C7tea%7O0i$NfhIna1>t)S^ywIxC1D(OR2-N4(lNC7@N$SHbdHRa1r`GZd;l zf*1lMXYM0EvF|}vM(M3px|h0eeaa-9d}SC}7fPl^+_)uB_`4+9s5OH}tOU#cch}nV z!e7T4Ft1X`NFFv&YIZP-1)o{L?t6O2FMo=hm?O~GH!mP=ZktX~g?Hrz$7bWwn%EXo z%1Dljg~Wm$8c`9lJ=%WTe#A)`g!_b@alWqm;#{vbN$!+sJyWI!@M0I@!Ms`8 zw(+Mp!t=wU!nS zfu5tR`^6HBAtok!y+t8|iQ2FdF%om8E*e)eETp{C#i8ytF~4~{_F%qP0ADPVOs4h2 z$6}g|$O?~t=Z=ZRV>&XSUZ&>xSWx36i~d00SL_)qj=Qai93W*rRQo!MR8y40_bDXV z1?fU<;OLdL+KI|50GwT!g{(TOM{ro1t;k1Cz=vx`VM%O`HX_%_hc0Ud<;$3&Eij%a}TCIZQ;Q>7(<6Emg(uXT3etL}~%+Zaf6qks{~A z)}@6ZG!7zQxf|=bckbNad9T^!>}r2;mK!0-;3s#}b`SAGp-@gxLFDTj zS*hYKkZ+_@^d8C>fmfJkrFz6?&Uq5`Y#wQgW4Uo?fa&wY+CE(jet^;(gmbF=KL>3w z&)YPh_1`u$j)`6Lr9dzId&O*v>#$uYx$=$dO7i{;O&7E0Ly#%ea8-F;hBoWP!2>({ zaVHB{GS0&8C}?PR*h)7l2kFOTep$e*&YU-V3Ss~J!Z_bp`*Ke9+FLJ4o9~Xfs91i+ z#09+GY@z>R2=lOQ{*p>hFv{T>pCUk+G7GH?W$60s{V#Hm(u~+C-Ym$x@yB#;0 zlJIvL&nkP3796EBu3yU(Fl#CCY+-2avO(k4Pt&+CE%wP{r>!l;b$DI4EvEu8tTm{k z_%!%bWK=-rE<(0_K-_02 zE2V^m=+fL^A1A`_;8{1_dwm2Sv2RtyjE@w`wM`!7B({g#A zh(Fkz*u^ISSSo*g*c^S6T&M%4hYN*r(H~p$e`gk;ULyX$-4n_*yEGxye&<2w@Z+8U z>i+^ONbQCD4+2vbOdL(9Ub7h1S~l77-l4&ycq6md`HfihFn_DJ3t;K0+bNr_6onEQ zy|r+%w!)t6!o-YoSU8kzf`@a={Bm^ZK%!Q$-6e%o>U}zP57ozbep2?Ml$JU6zvw5g zwaZ_tDzhJ15aPClTpq8=tEvJMjOz|Lc2vw*6DVA(ju+s_{ab9n;3LsO8X^7>(djy7&5^ zgyq&a`BXnS=ob@i)hV0$*ZaCXQRB52A;tu|kGeuZ*mTQbQaX8b+b-jE3#B+ZyzQ+y z3+K3guaT?m(cqlI1kYxh)b^9f_O#rhYsabj086wD!>W&be;xOI3?>aYq+HwC{&%J6>qWHttyq>osJFQ7No=>;`y2=zFkQFgXrDIJ5jv#ajHQS82>18Hx{P8gc7m5N2)^To5^3sG& z_V!89%)5YcUqtFD-@3QnPH7~-ex)~4y?h<^xGHzDdsUE5R74Yj9ya+qk8bbp4f!OKXO*G#F=43gA>A=6@LXuRUFY9@RcmN`+q|ole?lv4PHp>+Zimu7~ zvubA(S*ei_pxG~^;8FN!%5R^jxX>60Lazb1b;hdl@*Xh}LXpagFdCKO`g3oU6G{7D z1QzFKuo)=jg0oHps=LgdDzR6NO@oAq)YGxX_$a!z8*cLk(9^aS{$^93<=kgU>j z9|T8aikQ45Y|&8g=vqTy52gYE$e;FXP@=|wI4BuG;?bNZ_$E=uGF&z`76cf0;1Ol*ej2oPXdvWSEs>AfJdfn8?;_#384iWAV1c)+c0F8tpy7q zt0kTizC5!T1f^*4o|R%DLYl~L3#2DE&M$<>7z%qfoeJ4|HkFY_a`~v8U8RxusCSdp zIGI~2&&1hk!fjU*xg)-|L{g+D>3@mYR#7_8(LWe7_| zt(wym##w@xk10692;Ag1wd8u;-u`6BGz{yb8S%R{vzHWGFlPju6=_@QbQ+`^@W!q_ z(MHv8Ry=9@b?!IAm-5DRdiCO-L=xVPkq22H*vF6qd6J1lC=)SR=@nKg1qi39h9N8o z966@-QQ^8o-bEBBhBWeEr_%)Wb)*RQd{_mh^0TbzDq=Q5mvt-uBdv!Ma<2SZGE2^K8+K`SI_1O_(WS4xG>vP|Mk~ zIW+k)L2Msbu(lBongpn`1sgN|2;c3}sH69eT2qUwG0g?_{RlUQIp#lLkw|;~0y(Da zXA+tlF44ReFta;)3nIHHE;Zto$W)O4J8LMz2jvSkmA$ z@X^m-L~+Mc#ra;xg&A|RTpOaFF-1BN+rB@nBCA^x*(6C7nL+iOz~fi{w75HjNpv}x z61q-L&xA993{ZIf=~~bs?_?HMLwD_$4G#-WdTqAhezUcYJhP|ri3@8Gy}i7g5s%08 z^WRH1D}=e>IEi1p?l> z@Qg;_Byj4o!4d1lvAZ|pp=^|8RPHr2)!HORd$s28axoyLT6j~g$-pRBJ>q{Kq(b*T z#7nnEaf$jAWc#9d4{3V*_^nS`8xa!!U^BhE%p>PPf$a<`NI_Agg@fQP#EP`}0<|!l zpPSlYUq2!luR&_i$4|UU{g!~=FkRyEGRZf0NC5Qc}-g;EweF?p_LO_MpvRl zDT$qzpvt~^Dt9Q~S?f8tx>CcBS*(299BN!RA#K1M5c<>F_@mKwae4NlDSBAXiEu$p z?0H?I+@eaP(jo+ZVY%kOr|1`G&RXH0_|{ zgUe35m{q@V4cCpS7-JJ*)~xyOijRR;pED}bdyN|w+@j5n{^3s9`n*C$8EO7P@6wH~ zFWJ*s>_bYPb-r)JtIEhknYf9(bnSU}MuXn4QT@|_0rY8SjVn7cR;hu1MC>Jvc^zNY zl-gcL5ro5(sc2nUK8`QKUJX70gaiX>bmR5h7A=SBysGujo4Q3lr{*p_w!?@;VRj8u;Yzl6hAy>7b(6=d}i+=yLT_fCEzEHr=6;TB9?cv z@je!PXT$K{@lsyWou5o(=1fm9Vii- z_aaIyKiCCGlP{4S4T?U6j+-(_tx`SwrPi~fg#(tcxykG}`O4cVysX!Ll@*vYhtO|i zLRz~oayU6H#pwq87To-wO~IK1u#EjqebyNceV@u@^GC3g(VwQLkDVGe+C$9gmI^C8 zn3P@5UTDf7=iJB}n^}PZ>{y68 z+I;)?cGI12H1!EjZvxQC`DL{dYq1Ov^sGPr~5{ZB@sT~lR}`69xqY`Lz*R)S!>VazdP&LG zhIDM4%=o=QPKbI$VpHtULsQi19svp{z<@d!O!VzT zqV3hASaftWokHA!jWoP)?Yqu!k+UN?}c*3t47M(=C){L)l_y6S_{%t=LGtWR)wdE`JkM*0pj^0wPITW znRlnl+1P|dmLk1J_-n=Tw=cfXxZqI5wG8Gf^C8{AgL3lnGcXw9)JVA~3O+UV1vIa` z(S&dZ%%bh!GtOVkJ3T|g{ABl5O0&opCu@)CnVQk;FXSF+B4ryg-p2_HI4zq8;%n!} z(Wjuo!jA(!{Q_KmCpe^UI){QPRbjrWIx4J#i-K#xE_a`-9JEZw*h}WXJUF)MJB5tA z4=)r3@4G(+;4N7};EIC^FS32wsb4*|OT9I4525CB{4qoXgy%<3ArJR!J1-Cd>tyR- zTUveqVwM-a+5*g@?o@$ZdQW!<)p^ZYnPB92+uwEXUBmJo$2GK8Pzw%>xY+qIwY(&_ zI)&>4RI4xaeT$2FmIBv$Imi07v!8Z99)@~{k*_9EN&9f*PJ?iPDg`g*mAF-7Ykg-9 zqJsT>01QqGr}k z68L|eqBOVE|HTAcL|XKnAE>>0?sNz=zk@H#dns$<6i@5S`5=_eBcSfcJbgsZI>sB%mf{8Xf}U~_T0VCg zyt7vXyuc1dFFBy9J1-cA@GP^?9M`lVQ3v3*nL4SC2j~d=``lMCk}4lc!Z3J;n{(xI z??Dc6_+nr4^#VlfYN~bq+Z}=f`hm>AdM@B|A-cwSc9&$GIEA0nQnh;73WKv;L;2g6Vq^V=P0Hi%Tmv^(WLsSFeO~OSeX}i~1QDY(f4^6}-sVipd@R4j zB3~jWBiOMyfcQqn^E(xJL>tbmk#@>O;3UT#)k*F{MkKTrEoKAmV&R_k(L0G!ql@}# z^^VYFP_tUd{ekQoPoh1(QWaLW5U<5eDS+kU=lum+dK2E+dNfD*3Ri(bh; znrj;M4T=_=R8hT>NydI@_7vYFDth~Selw$&P(2|OS*X6|YZn$VAe3&DIMeOSnYA%R z4}t*?Sx@!dZ#z7f9r_vj7MqqHvh(ELy$h=~|CSTu>g`S=aff0ws3?-TX2AfWbG05m zpvhx<%DVa3-E-WCCqoIB2VZC6$d*%HUsyfR)8RhLL5bd}P<+ivMbLH#>TyCl1|&cY z(ks4l08-* zhBfDp&wsrkyy!NbJC6uX{U;XHT=~Un0u#HM(kIK0X(W zqZe&z@GDUjob?U1HcZ$sn64Ya;}i`kZJY3Hspgk5LxkxcWjjnV;?}!q&pfCs*4t)0 z)wf^G)OG`V0rP}uI7KQKOq0}Qn%f2RNo<=P0^fw7&X1ul!u@;yhGvBIYJ{d>x80Ue zy72^x$ZuY};JkjzUSI%|lj7~{a)xDUI~^(>otCQ#?#MBY!<`&~udF>=yxzS!U|emt zmj&_`l+bVf-(RA?NSRu52DX=u8#WP- z$#vY}HRR@*VjnANTyzdEZiz24iU((YD;GJaZSfX5P^Op384n@X_B6 zlebt~O9|U1IW%u8_A9VEy-=;@0d+O}$~1h~W_kk8#Wllg!7Zx~sW;yz7!LiW&}TmT zkfa*?^-Gv@u3yR9ISX03_cCqce(dNvW(IAg^_YtlV7w z{q?1;&!>EYa$?@$X!1v+=$(9~xq;MSq-ZZmaOOUEX0d)BUU<3sR;y$6~#T@}o45`2CB2t>8T z8xOVog!^o-YL(NYAvydPf+hrT`Wd|g#8I3}xLXdy6YlG7YJ3sxMLC+9bXRGYCOeaTh7fn z!KomTTQ;O9n${_u!8>^!@B>TVD|J=(HmP9G^P5xQb9TgJ+gxX)%1Q08%2ZbFP8)Z} z;Ii;lRJv{@AtmgII&BqUR*@v`IP`W(TULJad#2CXYt(2qOM%h^k0O%aUMT(PcJklv z(z`~g=I>nRGCtyNSB}AdgR3Kjj;N7ZDh7hKT~> z&3A^CfOXO_p|aOB*C`})_%_vD_tdQQ9}ylRTK)!Cver7@|5G?}sqKE_CDEz~#gwMn zwd?Ut=!=nGF)`nyWJY*wbjTT`O~(F)y`IXJ%rCp9Gtq3a^0E8Nmr736P6CrD-gMXc z4ox&Vp58N1kNCOyG1Q6vN%M&c?(1}@YOu1_ozfGd=ZzYio`V{m&5cZopymx0q*s^T zWEYZm3#a)B=(eh?UlJ#0G-%K4OxgEYTV{_La;u_jW8SW71o!^?qQ`dBUm4?c8PbQsrvXwiSi9?Jpb3hMbGaZt(!Y_m&d9V7g z;A$@Lg!~G9-Ris;O>g$k_KJ#ZgE5-J9xAus;76(`-qk!(VBVq%EbLrMC0XM)zY*>R z?!bw1>sI<);3-X~c#m|QJD_QQaDuxoqP9A9i`#WOWfdg<-s_jBAh*_XqkpKs zV7^teuQ}E`zg=%%ZKW)(Z9F=4M8mzn&#CIxBzmbF8g=L4?`NRc3a&qPZ4(sv=$LL9 zYgOl!J^bBWWQ*M?4PBlOELMXhLUX}0RO1a|4kdaHve$NHosAw;9vDDVxYi2^059$}d z1HACK8Q+gTgNoKnCN&*1Bk>383FQ0`&yu*Lf0=hh-Cxat^t$Vp*sdP@q{>QdEprMm zGAmH2*5aXvBETC|;9h$4v|4MtgMYY?;iF-=>r@V(A3;h2Ka*j3el;%RPcQ|^%6$WH z>qYLwsM=(_gM8@sT{hl9t$gi}aDO3*ViA5O*n4Xh;@R+JmmzPPTf|RIaj4{_s^=Lw zvGz&(TJFGhn6RFN>+X^e1xiwvTHT4lyRf=*4HtkU)D~@jSL9|h5HvxE;SB-S+xfl( zJD?5fPLK4tmu%I`h*eL^Ejc?d*Y#n05g@@R(tK=Goh7?1zW4Z{nxL4n&8_>1AFcnA zK!W-q77l~)vK9JcNqI{o;;lwX6C%UPM7cqHBU1NIZnDcU`-=8gYr#5Z`w9h0<+*0E zvUJO;K1f#+l=>o`3DtdNjQR5M;XIlu#kc*4f3*EuvT@+pU@bkHxA~>$y1U9|_2@z) z`4;gtPxB}xcRNb={&yoG;bNfRbCoIKUlUwEKa>&f+kZnjgv>pB} zQBUWdrjy)nCk>0mC*1sfj&=GACkK+ZNVi$c<$yZDtJrvBkF7X?o2P_~fA3i0>?TT9 zEtg2rH-@wKcp-UN5zQT9m014H6oNf`LXi{OnH8eHu+u{m#RSUNI;y_O*G-r@xV8Bf z|I)UBADdS!@r!>kI)6NOT0hUvS>k-Ic%O&#!yODXeJA`Ib#c*_jrT|2L3ghc0plTw z;n|l9%J=0dRSG&$uLqIV6gdf`kT@vI@T~sA+KxVYn9m>Ak9zCXXS2cgo>~poDVm+} zc-S{YORYrgu?=M%IV&E6Vn#(yJ%%B9+-qyO#NA;@$jg_5656O|3uy}=SjSNu`>iIc zJwc95G4Cm+jPcWMF^A2&-WQTvdsNtZKo;=L6QuYePK@XO&fgTJm(Sx;(9^>HxYp|Z zTa#Y3^+shVuIfx%AZ^$ILC;RTQG|-WgM+9&*Rxnr-s4?Ll>7&cK_UwoqUTj(JgNA-=cQNkIuM4JE;G_`)KS?Rgnjqb(_@KB0-0k@QQg&Y7DSmhE`%_sEo@AbETPfK?Lt`-lqiYkA{!z8NA_K-7bhm?!o5{T26kBZ18ofh&=@ zG87DOGB!*2s8FV5_rRI&LFTxfH8^r_SaqE22&mXZcio6A+oV>b&A zh_OdOyMUy)nGlyNavP-nDlNV~ruc(O&4|GG9vUtk!4S_S0v6%}UVRLg z94Jcxb$zt-Y^^IPF_M90OVSLPV(A_e_OpnCwU&jO6mk4a=@Be_*g1|5?XNnQi_)r{ z`5kiW4SF277`C2r%(?O&#^0r9FXIe7qUT_FjCDnk5F(? zdgrcNas5_d_Yv>bi|nT!11T%t%R+FHq0<(=bF8=(6Q7vYk7sl7cpp4@{Po?xUsw!W zFN^7wDWVMuEPA)^`tfJSJFa+-J2OY}#3a)|p<9p;iF^Ma@FmoXeV1BA?cQ#uKNtPO zvOKi)I{Oi~p9YPF&t%>J29NDi?(H{$IBb-&&>Ym`RjnO;hVg=8d_z+4i))?f5r`$(aB1vg6;rLlWePPX@7kCQkNq1l27gvJ)Ko zrzCLF-Up-#alt=&gzPEb;i{1Zv0>wvYaS|5+&as9gOj=OW5n7m;MC)-<0Kt3A?Bl*ZCeO6dRN3E#V0#OL z3MQVfN1+S8A4CFLgMIsdkr^82zY?{6%ok5uWx#Yg78-)JFG)*v(vT zFq`$lvHI`%GQ4Y-?u<8w>(9j%{bY1|ezLr_dYYFuyb`%E2@}?Fjg~Kkz@S5SVxcVE zGOxCmPj3+{>XvXyNxEpQS0SGJ9npq1N_c6(`OK}p<^X6SCL zSzRDC4bOLL$M@<-M(VSdTC;||Y5N{d9?oEWEuWcUr}txL-6KGDTa!b_I%J0o_$f+% zA00tt_zk_-aQ%ORm(&y|-Cm<}^uObfI%4x>*WNw4wAARbC;e3pu#HuR`3*mlUbs zcTnmo@RKN}bC@(-A0N8J^&H_vGiWX@%)QxVL>mStM!q)5JsR!v|0<SrxIE}vWaPVo8Z|tk2($~b3rhOvc zyCnxE8hGjk*@9_IiGzEagMJ`J=3;roZ zl;E_CVjpWxi(qbyc1%EZr$_t?&5nUi)SRrPek%7$Wb2jhn$tA@in$)`smM1AlU(b( zw{XWl0&pgmIQ^FeVX4j={6x_8NO#amHG^B=XxG1_FZ_pN$~es-e}bk!f>lZ60XG$W zC}XIMFKHH=E@s5k`dWgad}A`%%;?a3Jd}9oV6LH;pkkZtJaBDAyRp_*1Dr^$tdDMFyiNvorQN)Ow1^iq^;d2$rc7qQZBZ0hb#A2%-QoS**{va zLuQ}1kjOH&-Yg~OB1V$PlbUWVH1W*UvOvVkxq>KziQqqs77gcn7Pw^LlNOmo zILeGw1#ubrN!+(L|DFKRO@=-Apfj!4#>ge%PTXuYiZ(UWPyA15Dh^GzWVCat-mURT z6RWV9fSm%c%AS!orok%_iD6(I&KE)0JyNUQ_wL@h%ahP!+1I1pS_%VNtoKYY#Ap8a zi0nVG`QdzK-t{e6p?A2vQy1|?gdCc>B$L48E^&wv#v~?x+9a70TU`eV%qmn!O@mN7 z28&*X1$Hv-ThYr`|V zeesx6Pa-@mpVB;wRdbHy;}0lXhjL_xEI&dkoUdaa@eMOVjQB_lcDC}Z2@c?OWE^^D zo#~#a$8yb29p)k=mN~;3=khfHYKSp6V7cQiBzM+Nzg;j}Y>{y+(4# zo4oe!j*EU{makLOPazCWse(?ab|z0T(;@)u%d2Tgsdp9}7mH&GjTV3Zs7yy}6>no+@65Y&tw2G%l|hJ<{hE z{~^Qu{)oxwgFZhP$@+Tb9K=CiCR_0c(Sj|kCMmw(dh_uEAi+qANlfXj{b(V7e(Ohm z@JDt4+!z7b5LZ|-B9ie87%WmATj<;Fb5YS-Gk$yOONx>&p#h`&t>QmU*gG?1SL4Tv zRvVS^)Ac=bLfCP3Yzl<@|m?ww8xpG@}Z z@j9op9jPJ@o80<+O=#blNz0epp~w#DIDJxO>5`ko^Mh}9N6R!;uRxoF$2 zZN|1H3~;l|kU}?PHkSrQ0jO&US-^}?D&b>^Eb1XLJri`lqr85d^Y{_Y0j*WtUa9X% z6AGHJ$XloCdx25t(Kj0V&bJcO;>Ba?Bom~H&Mp#Y&1iY;=-SJToWsAp}irlm!J zhy3={>=~SHdD|ELNg3xK_}MMZyZ457+|dAYplO4Tm=Mn~ho4r?x7T~xWTLWv7)skE z$m1d=(6ApDkHJ>AB^-#2@B|Xm;NB|#E1-s6nYdEYG4QL-hSjF`I4wdt}pgAS-ykA@I4}Jz+WE zDLkf*JFU)^8AW6viBSLWII_N~eDIgL5eEswI~?mv^cY%Q5R-mx0I^W$TO5RPq@RB00^&lL&4-YVkCydRLucb?Rwy zs?>+kS~&I?i4;eLiqw$82$?vl9hSah&)e)c2(k3z>wYGZIire(x|uX%wqJR!cnd6Y z^=gshOs&%+Bx~Xi6c9e?a-R(G`oG!Y?-vk9U#ZI6o@{-fvw!lS)(U>CEDShD z4-3%`;1L@#I5hNcrFxEvD00&hMo`Z3Kn_w*Nb{_y88xIVEo>t7EpCp?h1{bRSIm$$ z{FCd|)^ZS+)Ca7l8={a|X&NRT1a+=Ivknw8j#DtiAnmmMSoNq8fV#r{Fv3@od zN-$zZ>0jh&0P-}JT=*1&Bh?QUgfsT4H$20Jqg_sv`U+fl3R^OOI7O1!0Q=?uW8a*~ z6C-U|*1Wv(&mCnA%djUc&I2I~He-S|qE{2M|94{c@#g1=+O4bF0n$b$Y=&u>bm%K{ z;1>EuJ1Kr?!pwxgpfc@d+7>QxY?ytA+|8%Py($*y#pjPELYY>pxcT^}A)R)js z$&3dIwJ)Ab@2CHgof)-ET0=HAIGrj1|Lh_Ym0*i;I)` z^LH_~YclS!qQY#yO)IMK?g>M!JM9hfF*en|vt!J!q;S5}$fdW^Tcwm9#&bbbMt0C< z^w5I+_WyGWOB?cg*_Xw%XcvRtuZF@|l!E^8>o;z}tQ zJ$At!ulM2KKqlviv>sBj2wN9F%qJ?{c;&eDfc1YL$f*L}-Mr#QRlw(h8P-tITq@ER z@Xm+6o|{!@(vKL49(sic6D!~Y4-6tqtSA-UmcWc`(>BkuZMd-|? zOB9NLA|@dV+X;CXONr9@aM2WFO<>Hkd7vSF$N(z{SQ9((vcElZ8wy^Q{!899frbYP za%+M5@CaQi6%G6FHE!lj4a4_gl%DxN8>z)l0BQXP#bGaitO#c`0MZrw+MC6fH8>EY zMoATm=gh~?p>=j@r8a7lRYIEUc`RX^y9U{FPj8X@zqG}}^L~Yzl+bVz$$eVvRp))B z_R=DRGJqM-Nx}h6c?~$_U<9mqk{g3`PMr_$Zw2>ji|IN4t06;&F@$9Tgw1uo+D{gg zO2-NhB^w?2T{w*F9*lGVHjo#SaAkI4rbYgDS}-Ctt<0*UZsd~Pz?w?ctR0hldBWbC z>&&v|V)991DE{_OSeick>5mhjRDt^y${ukhse}(OY^QXv8Vq?4*No8RB3d98u zGb#gmmh>1g>jsZT zB)(^YGigTa_mS&IjROakj~aMlXyue22#rm9Hea+7-KEY?<-@ z*fLS5P=m`)e}X;4tWNY!<8#%I=$1Trif=xnXyPzvl)m<6XD&Dh0oAw?261FN2dpeV zPM+2TtbPgp66$WYmY+sHI3Ce+a%v67c2hc{BfU7No6S|S?wyMyDR|wn_1Y98LJqPb zJqHrq+4sQ)7Y^sCFvKWEIdZ->{=Fk9JU*+1Hp2SbR;rG%d z>|zHfB@Xc0toGiR3Szs#kD{2WDyxHEzI4_rku0~iw3y3)u*?$?@5XP68R+ExWq6TV zs)YBpg^f6SfQ#aJurW{s52}NYE9c9AR#LEjw!?CY+8uYa2!C^960&7I`=fC;jdyB) z8~S}w7^-`i544zqC&8XYuQ3pQQQHY9l|?39*EgbppD!&sF5e&~HN02TBhqTc(*rIl zUWaAn&{RRz#^Uew9YFbrJC#1@?Hr`8Jx^Bh?nVqZd>jiN9m^#)2ri;9WoV#3&Dn_S zGlEN0_*+-<{PordRT}L6G4e%pqfzPGC;1t@FREJH&H48jIr@|Hz5AikG>ja`WZOWE z-mxFxj&wnJjt_&-SKjzX5nbN+gjVDx)PR+4J~M0O{%x*H2u2*O!4KyWdsEWv3OYn+ z`v$h})>Zf!=CE>9;g?s*20T`L#P-k3xtN)O7uofDpaS(+IALZG{)V|+C8SWLpv7GA zS1LIdI-c%?Mu8n{L)7;P`@6ok#)Y2aEb`_IgjCGFyh3s|t zmDV|9bcFcwiwmi|$uBP%g8If=8;Urf;Ke4Q^9GvUl82_ji==*@U~U-u0@_BgF7kFG zK|RjBozI1I3;WO7&Aq?)4Jr3f+;T7_bfRLSQ4goDdbCPlVf6UY>=g3cz#wMlM}4lt z>cXqG>!Z52S;Y1WO~MOEo_p-kl>2JLndcFHEupdW(%*WUhGRUx{hH|+w{PR4=IcXq z@LVeYZbWTU&H{riH6fe=_D8LrUfZ9IJyaihL`G3aJ5jr+<;M4KV52W)AKNT= zRDY>*;n$jdapuYP=vlrvSp>d(-&5f6X2B9>sI74V3rm(8uP^i*cOJN9?~1Lu>>E{& zjL4N{mM)X1n|<4gNskVOS==W?lBjEb+gg?iN$TW4baiy86dgi0hy4|CLu7_lzXJlx zhgHv{@^u^{l3f}}xGG_Zho_-WL?;ajr?q=xAvVI=9Dd_l;S-OcLR~~QZPZfVp_b!Y zF_||f^_W79@XA*(8H(>7Q{udr3j`2V%{6@FE9O}I!%*P%f=1!)B7?oMf>8>A17NTa}^ zI}Rlx-AG8cbeD8VcX~Ja#rwVY{SWTl{DkA#vu3S5Yi8E0XJ+rc^ZkQmjl6&Z*-iX1 zeaT&m){VwkL(XG-2S!o%s~e>P1PgqQQ1e?JNx-yuo*X_9;O{uyEQXZFv_92;SUj20 zhp`xwX4M#%hG-%=^LrwsBEI!+Z_qjUtq{TRwPb#3DbhU^pf@469j!dvvl)k_ob_+@pySZke4^5~QxK!Ip_+|OzajNPJ9Vcu9 zEhp9Dp$l&FI750WZ$Aq8lq4GT-krbuB-?M(u@|%J{qQ9umDy0<9BGnz?g*_*g;BDv zDN}OLo*f?_bu>Pf@UClp-^$cFqCVhB~K8;awsm zr}F)0YBYOpWcZZV6{iZv?f zAH-FQwFFRbd2^Zj0-sSJI#UyqI~k|VSYkhSJus^RRwP*=|0#6a9Xk7wcqE4b@j>1v`NHCI1^S-YO zTJmn9<<^;q?{UDFKB|p?jv|h)YA>jvtC6BsOu&>Q& zoK~6<=BYrL_A8GP${Np2AmE2`EC0?|FxGf#Ol21_!E5JO2klN5W$^E)r-g>)G@jRS za6E?QXK2AiLwa7)6PrX4dH6poLhoWFh|BzibVkAJuSmCZ<4@@LQQ~zm9Y(1?MJ`!4 zdvdw4chcnXWC`y{g(^5MjqyWyhE=L*LWjBDiDujTXNq!oIJhEG#jVOidtU$uNX z-<3NFZb zUELgjzB5E1td!kxwrNimy*qw5H^kA4B%mGH|=6I!iX|t)MZGj2!aHu5w z74Oc794lXl0-0*MB?HLFOfocCF6)~>V|_7#k8S}7l+5bagrKk;obMs=CKEs~8Rg!z z$R4shIgd&yz@gJxgFoaF<7AF!qm(-~YlwxGN|(cV_2ljJNfd|oWkzARLZqoYlhvEO zyc;eykEEO}$RwyRg`Sz+4eCA-H~?7^^2wtz(n+ zbrw`$$z=Q7whTo);-o9O$}a`ll<%RSx&)Q+f#ixm?_Y~Wjb0tyF$vgJ-^E;dA-$^} zpH&UL(+!q3X9>aGh$P39l&1qai1^)sYQzaRHuO>%JonMPfy9Ak9w8(t$je7>BI4xb zif1lie==4H*cpaX(o_C087d$O?2+Q30;MIdC-BG+$!4LNa0O<2bB)UCWNz+WlJhZq zdC!gRz%CFog%{1HseigiJ+Wn!5t9V1SB+FYiL2CU@eBXtF(|31z}R}&V1*3BYY?U@ zkQO;b(6_{{=ox4dmdv1d?FL4tC!Tq;*3QYmpi+0)C1e79xiJdAj1?DCAOMQQjveIM zVE#)N6n!%|J#e5T8R`599~a3&nI=;fMw7r*a3sJFs|L(XP{MlhBHFqiSW-GAtosWY zX(jlQgE?>F{ZulOS+LJBaPidJjvZCeHZ*|Yh(pRz)4udPK~fv0&qh&RQV|iVwnL_K02SFH0$%(=fFL&gwd}x-3x-%=p;}x-&quw0fGoT zE?EB6*|~*m;nod5fR{9(0@I}@@a)^ZmJb&6VD7BiOn{h!2unQn_`2$V`Q@?XHmU0S zx;dPmkwM7}Ltu?U@d$j@ED^cy@^8xBT+fQSbN9tb(M;DtA0V{eWiTtvfM+nX$#yFAS`Cv*o@;4~C*ScN8E{{oLp^`j+2NLct) zKZ6pYS#K=-ESj&G23{>Eo$u%4IIWL8<{^@!AQ+GujvzRMijppI348J}hDrud(nyu2 zq6`n$64|14OiH&u{2|2E06XN-_kh)ynCy`Lo4wc88s-^ZGKBjbv7CwLkZ(nFvehJ5 zfDafw2b~0FCSYvXUeX$CL=luG08%QM&I0a7Viq{_Qq;NTj_~c2DDHL~#=#l768n5S z7?C3r77ke6%go=q`@Q=M&q?UFdrhQ2RPL znP7<_?sWqRggP7W7ouVc@p1x{9!2W+uN){ssVfCMkj}xUp2+zhi1H>zj zKe>KbKm##}9VR<8?Kljz^s0Q32ObP13@t%I(SJfyjWEB#>Dzd`eViSQL82Dwx)w>k zRhqTat>@L_x`qtM4(ht*G~EiNa*%qTfROXj;fo$H^A2};ZfmMP%(fn-AqlR{MrN-# zi`wb_Z3a8txq`hY2QLf!U}Z@g4AN%{NumpVY*=(fQ%qEcIm8abr~pgA!U1gG;g9V* z2pH0sQoo9wxPHnUeGd0HJA2)ZTIv=kJKdasuphOAeG&ZIh!?^7WMVnkv&tI5*=Efp zU<}#O7HVBf1>oIBW7wpCXZcNE`zr3gqtQes>!#ihn!{=`y&EZOQReX*Ydp3hv%Lg- zA!Wo7nxxpHUuisk^o^>3EJ!fACJ{E!tW+Yb3NxDb#y@6NI|vBq0VW>A(9Vfl!W#P* z>Q6R=Ief_2n0bOn)wkkJGfC>BjmPtG*m-!Q{Y62pYbmU1>3zRNSpbNVox$Dr&z27qP#1^WO` z9tp4>ZsvqhHJN)W=r7e$nrp#xp0C3XO^wn$vyrW*4tgui-`;XIxn`%l*SnUGg}@#! zLTPhqh6%tH``Aolpd=k^+&y~wTX?4tYC&%?|4n(A{4VvIpmkB-i5uEEdP>S`0K94q zDg<(EfU-O$1UK}Z48(Rw)s@F%XA^yaO9CEzT0oXVo49~&1SGMTA4QhUOC=s;nT_LM zJ0e=8gE3wB3^PYh@cbv!NQ^TO11buN$-iizHZlVc-*B}kX>OsS+e~m=6LtARkO?{S zX#a6Na^?IR02fpQd@3??#d_40$sgL~r@dDBYJ0w+F|fNOX&Y~S`vET!6{_-45&$N{ z=m1R`0Re#vE8s#B88)uWy=a5wUH9br%+4CDHV;?;we(eO7M%@?P>|s#+rWpQYMV$B z7Xw=INWXtYYc6tphOx-$g#49E0K}&0%>Gb$-owM`;VX9hYU!m_ zVq#FxDzHX!seLeScor8&pB;?yRZQ%7``dKaMzxASW^md$F(_g<%dO!VB%n?|YNvBv zC|AKHQ=vLF_{WuCc1;{eiQ(qkM}JP-CiuxU#8CVw=_+qj69zr8^^Lw``7cOGru~ol zi;8WeRtnQb)f<(A{;E=`zQ|$u^wPSAx?W>1sW&aC8Zd9idyuY4!?W=yV7q7)J{Dz^ zG|jte!KaDj$gdQAq5pFr(P|&Oj>+Z39;@0th-vmx@QwWCeq1Frte<+fL5m;X8Ny z32~fZF3=5n`F#^gi)_aNjHZwDKcnf{iFNf=ooYlMh{$e~J-De}bFwZ_t)_I7;cvM* z-lXbF)OPNve)LSb(CW2u?{$s&2BC>N=ex3(mAHh* z^wZnl@7G2&t5NIiUU9-l;}*M*Y8A03TAh(rS-gdx72-{EMgjPN*ncb(YQd<*O~-%M z`2NvZ850z+Rb2sC5L?3P-f788?tY;WSL11$WI6Zk_ZH)fXpdU3Z2+Xbb@R%LuTDhd zpLi#y3?{|jH5-;>%T^ps6>lK_E^jdgsIuD!t8*+nOa@;1A1#uy&?dj_YPH7+G8TF< zja%js@F_fek~4+%*wcSgC-s{MP3O8OgTEM_C{RfB0V{4*etYpjOyDE;L(#nHuI&^R z9fR360Fz%rzrPd67UZMqr178nYW+q_Yp3ZE?Fs|XE(YlKngn7YTrlT$+d_fI$k?v9 zDo&gnorod1W$1%`HP^c#z+a*@GZ~%z>B5i|YD>4oa=`mUxX$s0a(X5dN=K*l9FT|V z2P~Y;0hT%>jfm(xASfBvD3B_ZEuFzp@v4Ng@m-zjPsTN|(ON1xEye! z=!17ORKVHAH%Kec&7HAAB4OL1t0B2ATvWua7z=?@a~F@Ueo=@IM1oI_Q)ovgaE|o( zDf|lbLh@pvkSJ^D@a<>>K#bYdi-WrvIl@2cdHERnt2G<1PN? z42I988>2)y2%^COFpCsoAYy)mQ%xAsx{lTtKmHUe&!aA!}O-yJ4 zdzuoLv9{Adw3ONpkNw(7tGP*@Q!e(|p`kbK+rt9#I5UG)THM~RTeQW8fbE#ii+iA? z3KWaxc-~7)W!M96a%`trRlL9m&tavz0q1)< zT$-^%v)5?Iom*`?I`(Owaeo&OL_$GfS(Z%+0qcwIB6K9I1oL|h$`oWDlg{+HUZQuk zr$kEJWTcNMP%Nz<^b-i;B4blAOtC7_oeeHf=mo7W@}(m#QP-8)-oUR~GGn#sJPs>| zP%fM89hig}^bTCf1Tv!0)qp)-i$K-Z=g8*ce6BM^Ds~=O(V$Evl(vjJ`>1LFTp<-K z(>Xsjcp~Y@63ZIfpk1i(=$x?0=E%IB=w*s&EX;Bfu0B`5UW(<1A`KUiT~j?Np<}ck zxdDey1xQN<3G&2gP;dWiMHZgs7_xySysp7&uujh_x#*#anGV356X^2^{SQ+NE~ z0j`{h-4vNlXc1A_l2h`gmyuYQHNj8?(IGj8li~yc5E(-_D99+rN$egQ-fGm8^&Hb` zQ|<O4tpdAlAHi}s0Br>cy; z39pt|;L&xJ16d|=yPio8{0ehOC&4UE0)@cbQwq_*E#3eub~imZ+(-b*Qj}HF8xoON z=<~urZr^Ha0P;S(g`B``a#rya0yoR6#MmVq1bH!`%5AP$6ao$-z#F|0+Pm6yhABnC zn~`Sw$rJ)9kwNf4T(m_`;ST1i!2HxFp-368k*4}RKY*a;Gg43Bv)?Z0+U5Lh??u?ZyQDLV5dY*ch+c79`4)jxRg=JtEp|R- zuLpQ$X0c6ZC8`!d{B*w@I4{|` zD@qK2aUU(0vTT{4>U&Q=aKWp~A|Uf}`Eu6*!8jrq=Q?=e5!YQ0yWtmw_>MXR%H5MT z;~xc>NsePyTL;HT?1{B1-}+9?9jcIVrBRI9(hnqQL#mMgoK_*qs+AK>+!{OoJ}T5W zES={pRhpk%WOfN2Ky{B#RLd;Ac}U0nfUxU)uzK14o_5p^*meSS&Cc9yLPZ`}gy^a0 ziQ;O9UD1fC`N%=_*hL;sb;F#)$~-o(OW}{TNCpXN5*Y~XSbnsG((3N|)}64{3vq5e z)BBcFjTGlq2cdv?0Sp;JuVQoIHxhRY=+9I=^e^go+HjM5ObMCKTJ(+ys;MxHWv&A9 z?SWLDy~K8gF`)^?>VRI0FPa3`1S(vGE~|^obcRkvT*nEn*nC(iKitQ+WX?BJxOuYk zH?E7Ld4~I*TJZoVPbnYfg4N|-sa?e+$*WMT!u`V8@w3@t5B9F@C+wiWMP&E$m3y_U z@BHyNTwxa%ym_aFcVNWO-xu~0%!DV}Dq~@2kcv=*Bw(U!-)|A*&pDBriYUFDhRu8y z_1%a21l1`fW`v(iBNaWqDc@A=(2NBT7RbwJKg5^gP(@~nb5v@W zGUal$Cna2AbMfRA-KhgRO-rwr7RTYYw_RO?y5A|hM7`=s+t2!%av`x&|0Wk7J$c%` zmORKW%HAm&KgchULrLBMO-DjOl&RveG+7>SXFiFNYOr{TW9$xk8X~bonQ#5NY6HMJ zGu{A_lUpa%QLO&>?Q1^GNz~jIcgUNp-rFydOF0>DRI77@#o6N#CzfJii6 zBD{&mG*w8^qsdEcV0yO+7nQyM*}9R_D5!^B(F0aD>>Y!m#{~Q;Qn&fp6ZC$2TXGs; zUYHu!cA7H4lutowuF@J7NoDA0DNz14woZ|1=6ep6*rHn2Yk_aFUa@y9}6$@>h}U8!cuxaLU_xo}#&70wL|d{ar#PV9mA=Jqk&`{Lli%NhjEA|~F?JeqU)Q#!W{gZlOlX0A zHyQOh_>|b48D<$uubW6pzI$tK8BvubtS1EnBfHR6!BpxcNj9|&+I4O%k8q;7A0+j3Z@ zD+6g}bG*w$G?r&!zZ(_fI4AD2^LmV5B%w!2&RoA&@~T(}pSNAuV*&c`b% zS-yfT2T|%4BrCaL-{8zPNa|5ER`N5cpWY>J%9V$Gr@RiVI_=t%V=}#u*K%aks1q01 z1wDYZ1}Lm?+&$kAtM<8l_^|tJL+oUKFyh5)x@CtdXN8uG0Q^--xfW99kND}8R`f{5 zGfr13<5Tp@7;?Gz<(h8K@i8!w`|1F}4zXXm1MET>t}pf^dxI>BQ5!KRvnE>BcHroPhz-s zU#K&YmY_alHW>SHu6UYQtlZ)|b+SL>&K!|rNlkfuaZ=;e4dU+k&tNDgf?8gY#4_{* zCh?FTSh6&qTnz;#IxZ9F@#4ff_PAvC4#8GXHSp6dD*?sUD)BzvoGY`1(@`M4X7ssa zzPxE*{^luEP3r8T`fBt9;W%DA0`dDUS^D1Gt*03h_B1pA2{;*>%tYNyV#1dyI~W{x z)h_@~n>2HlZPQrzH0d^!CCArmziP_1C!=omER2N@`Kh=UAZ13vYC9#Lq7x%vUXp>5 zc;t;N$BP=$?la{2+;r_> zNdPJQLoB9XH=F-EkEReGaNF2C;tk8JGSap7?Ta14M3SMWK-lQX{vJ^ z8D;xV8qMFgBIrUl=11Roi9~9;GK%w+?1tTJGtAI0McGIQV%on_%$oq53>C|2B+YJT ze+Qd{Ts4?KQ$pktca9jb(RQV_Mn@Q+#>gQ22t_XxU)}GDu1oFf0D*{sJ5GQqA=TCW z{kM&*ywJ%x59klE>hakLR41&0j)&g-9A_X$HzM{-qNrvqgvr336zge+|4Lx74*z)J zR8$g-*RGBlNX-@qEsA<@r|b=tzC6PXXKj}gv)Olv*5xlfHhvhIJ}gZ{_J7?v0$YB= zsSvhHKuoRGIx;Wx%j;*WAZXi~CXTr8p6TQILxPm^FFEQ4|9pskYbLd%r_n;|m`thBFeL9sK3tsUKcB+NHl<1X$L@CO zyA}NTzy9Mn*-%wPy3l_-wezYUL|6C^r_N*&$v!TSQ}l3YP#5{%e|q>yf4*mVzL#{e z#~b^V&HV=7WpjD$WY2N=VNF>0;i1Q|^%u@b-3V*|Oeq&sheF=_u;TpjqKC)+BY?eW z7a9*;9Ko(VUZg6(O%7lWp#q@=ts`|Ja{*H2TN9 zB~t<IMr(Q1hLI=D^4;-r>#v8c=+=kp zz3E?w|G6PfB5EF`oeeF-D1Y5>nqBglJ4N7X@DC=9B>S$}4=+vwt`6MCjaeSn2j1N= z%_G|zF!x4Ou$4!RFTT*RQNdK=Pa*;D7%7NFd5`)NSvrPQYN|t|owlfp1Dk$!IHAd9 zq7xrgQ~HZLkX{>4CxXlM`IbKuqq1igqN@DtxA_vBn!}IiMupuNrI9lP@s^|TfxdY_ zZ4-B;Z>h$Envws$XAcfWMIH_w7v?Dp67YXY7)9ZL7eP;9U;>C?U~qtMz;>*zR_6BR ze}1#GfbH#m>I^&1bK|yP&ita-KG3bDoyFM2t+k`KjjOdHqo)I-8*$Kx2aAb{DQ^lE|>V$51yRK#;xEM#o_m--&?y?@dbODYoAW{ZJ~FAEG_y6IptT zt%j1`!Pt2Gb}zr?+%Ww_dB0xN*#tP<(rHP zOXG0%^uhUi*y495@{xyrRm-1vVu*4>wZQyGp+jCPO%kFn?LN!9vzr#pXxbus zbyiZs(8N%Qs!4yr2Zzw9bYVly(y_O~22B~75aftLEy@nio1fk#2^$#{g=GtV;Ii#Q zO+G52>+4yG`|v`dry1GM0R80c&NDXbIN8x+XA|vbP}fqurkm?|f4{r_>)X04L%~~? z!!xSPy{n{Fzx$JDVWTawZr!IVU6;c~4-a=;n9FzP1DuhU4k6uqJ`$k8jQOjDh8_@M zYOpYwrK2T#1UCB%6|@%SQbep*q+svhPw3Mf7D(h~*{s|@+D;&C=QQ6F1GT+c@aUf& zC9}_&ro4oyGc(c!HrK8WhXQ)L_1;Jsvo}#HHZ*|32b6=-1{k`Z4t>SVhaGZ?2&T@e zQ2)T3>n#(8J#EU-M`+2pB0!^riES$*WoCz~X6-5i>K*8p!cw#E+`M`>ZaTreD;Qg1~_*MRe!SM{oR623cwa`!rg^$AG zpTDd*Pd{xxPv=MNHZCbkounf)Ix2BT8(2?PJQcJlp}Wj&0KcW_cW0deo^>X+zFn?! z@Rj>uQpvsaqQ@3??808gDkFu47#}p`Ms2p@Ml;b{8YDU24y9czQ?L2HN{1_AABU&S z!COAETmLTW@ZOjtm?5)YV50|S?`QwIzIA~ILfkGM%wDJ)LVSvYd?|k@y>wF29EsT+ zMNUGRtQb=Vp`w#&#*J~G=4PWo!9+H_c!eK{9P!XpXm1TJ1a(eC-Qc;^Zo0=pIq)*a z8F*&C1GI~qw<#~O-Rt#tsrW+Ss=|9u2p$aL%^^;s(gNS%N+sxDzYju6=Z&W!eU0I- zj?Se?46UEsr9dKQZs6|i?~%fYb)X{}wg6L|RogOij9oJ`_>|&J753jPwvELyQ!dgU zpDhrLo4aF*h9eB6qWy?}`|6}BJ~s$&63N5`Llu0YKFTwFp^+>;*~Nr=Z8!LrG# zAHE$D)i~cN@|CF@?a3m>3Ydj{<}%ZbHS4qaI@jonEd_DS^)J>0XC*CB; z*UfPssqM|#mb@3M`3_1NY;4h1(1X-}idfiG}1SLIA&u7`rsx zT#D}ihjsf#RWOXvkcm*^F62;CSF}!b=c#BeHuY}$oXD^^-57r>qIPUt1H3Ff zjoT^uBs1TBcl-8OM)fR>&k^-E>uSY^59_N~SQ+%jzB3Nh-p2Wh+ZarKt<#^YbIV`v zgsN5v|CAOj7WCQWvqbsnK>{Ji&`l`Em&6l3qq*1mI(D{qUb_6V_Nm{cb;XBwkgXHd zPT2Nh@cu8<$Hn#Ui`MgLN5&X<7#LP;7#Q?Fi;;)jnzW1&5zW` z1TvScK%b9%RU3d%EBd5RvaDKSwo@%nbTquFYsqUT?rh=~nk1xH;rjvOd0#Bg_c_@{ zgKQy*RjaGxH1h8;M9qvY7K_+ADs4U7jYoEoM-qAQW!)E;brs`vC1S>zCUmce=K0lT z%IKE;I2k9RAYG|g=1<^RyDxN$hDO{X%@m#W^~2Ctj%aZ>dQd`~H}FYh-|BA5^1ApA zM=m^5cX*A|A;61B({r-#zD7f}{6;b=kc>^Z2ovCkea0MVT{leg`)%f31-BEnz)s_h>c_dwJ8T z6om;}uaqL)8=~p!nGtxzQCY&sI8B8h?1|;-!Qs5u7pF=Gj%0^%si6*aY^dL=c+hxa zThjLlY`JLs6qLsfY745yrGfcDjNHRp#3l6bu-7B}FriEI{rSA6nh~l(|AY;1FbuKK zBy8va+F6o%vY@#6NAo9rIu&~e%6f=IN!6$ZvJtbY2M$S_Y9!;}&KK9nL2mE!aW4xY z4+gkvwT!QvYo5(JEpWDqJ3h#mMezE;@*Yyzbr@26Z+e9J6x?1&k^dl5E+okbs`=Jb zjz68n!JXoPl`L4FQWXHpt{2`m;Ih7|GT(Mv4qWv4? zw+v8CXOd=|Ii_%rULF>ACe70Twnts0&c_Y37+t?p-a z`P?Y_@mSsXbIDniaKZ^Xg@j^X$8TRxOWvc178x;RX}}##n#SD?D>{G0WfUW2kwM(w zG5f;7{~e{K;Ec=|fpO9f+?Iwuzi4+#3#qUk<@=VPL2&KP^|ELxs*KR!IGyVrfk)hy znZ_4U-Hl)72miZFRC-j&xYBjQ)gk8RL$G|v13x7VKM@M|@3h@WCwR(y+P!bq+$FLf?i3?F^2q1Hah=K>tD!#&rH_nA$O@rgpL(fcHe zkmL=$H&ZR93C{KTBcJ>;rKTBDhfRC-x>|6uJ(p81a=#r@<03o^W(}*JEv3t+QH!A8 zPKRRp|GJQ;=N@1-k2vuax98_pQ5?AV@un0`ZkcHaVxmmulNyWWn9{iD|7?E$)`KR0 z{@IsLH7$i6VnksTi!H6AmM?Lw`$veDao>ImEUZbu{gR=~ide-3Kcts5 z!vyFLJbRI;7CdfLSzTALuB#Z5zjwU)`4%(F6Y)Yn=<|Fv9XSTx5OQu)xF$|cDQ+Nq z$1vdvH{|X7FIR7vrhr7JO7foCvQ}-*zm=cYeo>(JJGbAh#YF9# zxWQz%UHVRzhnXh;P=rtvk2&S@hYf66qYWXfZywSn_rmZYKErfv)Of2b=uiBtqEJ zvyP{b&%`!0kspFtYjRDii%;-RojwH1j`o|LvF`0EG5UT8w^nDZLlyj`*>A&})joO0 zlxXF5?;<>|BKm?N1Zf$o5NqXi(w!miQ8>)9zMOjN*A7v$0M<5ozi!TfQp_ejb@E2^ zqn@FM4S`$kWCTy$^g(_1rTjz<3gY*K_qlM#KD+zO%$-r?Lk8qEA9t6_-yJ<4%*a~A zm#A}Er5cZ(8l4@lI9(Lx3q)W$mVhY?_*@Sks*Zup5}$2E-AUa?`r!TMw#8e7<3~M4 z$Ow}@#(N@2hs3FN(>Ix0(St&^UK~#lfTKWU5nrj}IFGY$?3~t^32)ud&R9y#fEVp} zgHn-14BkcRQoMf!MbD6QP4+1hcoO-m^Ns>ec=#zNnch%5{RG*_+&;u8dkL2WD@nX; zuNXCS=@KLeO8F$;S5kajZd8thzc7f~nF+EgDm%^^CmE4ScN$>Qwoc*t#r73+#Ohlu z?f50*mzu$<}i?Yo&a+h!wksGBPc%AVQ<~f^L%=^lx;{mJWWe6E*!mM0!!zqRC zdW}k<^}e;ZB{K~?bo(DPH4!Gu!&gq!@2cyY9Y5a`efFS*>B_-Xu=-Iua=hhmdEWAC z&OfN0r-Ucj$&|ohb^v#MlH0)`^EC9TZZ}5peIn|ZON#dVk52*^Wsb4D5r{!iJ{%W) z3^_Pbg?8X$ePKW6*6%YIC*nc=;qOCaq>)oN#E2SxVn+wm_xawit~&msKJpoyVW$Gh zh6M-w`>Q@0zkU1v)$zY3m>%2aFb{xPhX^+qAqOAR9T-AHv{dt$yk^Q_WW8RS^Wihf z2X+$oWtu0EKK+UAJ~!|BjV!nNV}a7Vx+9JM>a?`ls?Yq)m*C_N9tdel&WZ2FTeuw5 z%g=Xr*Kr0wjISB@NOA2D!!rifs4K#RL?qyzQ==xKXi*Jw_+xc8IaDoA1W4g7(clgOs=hpV3w9ZtR#sqbKFYkMW2mDxCii7O zU@7V7H)0Gbv?s&ws?@@kZgx@&7j7F}nG)#9Ln*N>d zGxEb-!Upw!^!W)aGY~!e8^bd>$gna5H2QH%^oZg8@1ZB|=-_JZ;A*7d`PSUU;8Dd@ zn}$VU?fJ)=5{WuERVnc-ta4p4~z&vVgq2HkKX^kM-zZ4 z0$8GlqqD7vqoeIV_4-aIKIs%_(FpVkppSn!o}>J3{HNo;viQeA_(HR0GkM@(w({^` z2>-$b21Xz4KZ9_wGB&yfOZhtx`d>zD5^( Date: Sun, 13 Jul 2025 21:13:31 +0800 Subject: [PATCH 26/32] update smc_reac.md --- docs/zh/examples/smc_reac.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/zh/examples/smc_reac.md b/docs/zh/examples/smc_reac.md index 7a9ead9599..b4c3b34b8d 100644 --- a/docs/zh/examples/smc_reac.md +++ b/docs/zh/examples/smc_reac.md @@ -3,7 +3,7 @@ !!! note 1. 开始训练、评估前,请先下载数据文件[data_set.xlsx](https://paddle-org.bj.bcebos.com/paddlescience/datasets/SMCReac/data_set.xlsx),并对应修改 yaml 配置文件中的 `data_dir` 为数据文件路径。 - 2. 如果需要使用预训练模型进行评估,请先下载预训练模型[smc_reac_model.pdparams](https://paddle-org.bj.bcebos.com/paddlescience/models/TADF/Est/Est_pretrained.pdparams), 并对应修改 yaml 配置文件中的 `load_model_path` 为模型参数路径。 + 2. 如果需要使用预训练模型进行评估,请先下载预训练模型[smc_reac_model.pdparams](https://paddle-org.bj.bcebos.com/paddlescience/models/smc_reac/smc_reac_model.pdparams), 并对应修改 yaml 配置文件中的 `load_model_path` 为模型参数路径。 3. 开始训练、评估前,请安装 `rdkit` 等,相关依赖请执行`pip install -r requirements.txt`安装。 === "模型训练命令" @@ -15,7 +15,7 @@ === "模型评估命令" ``` sh - python smc_reac.py mode=eval + python smc_reac.py mode=eval Eval.load_model_path=https://paddle-org.bj.bcebos.com/paddlescience/models/smc_reac/smc_reac_model.pdparams ``` ## 1. 背景简介 From a0a399517da4734d0aeebac203908bd26a19539d Mon Sep 17 00:00:00 2001 From: Dubhe Chang Date: Sun, 13 Jul 2025 21:44:03 +0800 Subject: [PATCH 27/32] update smc_reac.md --- docs/zh/examples/smc_reac.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/zh/examples/smc_reac.md b/docs/zh/examples/smc_reac.md index 7a9ead9599..b4c3b34b8d 100644 --- a/docs/zh/examples/smc_reac.md +++ b/docs/zh/examples/smc_reac.md @@ -3,7 +3,7 @@ !!! note 1. 开始训练、评估前,请先下载数据文件[data_set.xlsx](https://paddle-org.bj.bcebos.com/paddlescience/datasets/SMCReac/data_set.xlsx),并对应修改 yaml 配置文件中的 `data_dir` 为数据文件路径。 - 2. 如果需要使用预训练模型进行评估,请先下载预训练模型[smc_reac_model.pdparams](https://paddle-org.bj.bcebos.com/paddlescience/models/TADF/Est/Est_pretrained.pdparams), 并对应修改 yaml 配置文件中的 `load_model_path` 为模型参数路径。 + 2. 如果需要使用预训练模型进行评估,请先下载预训练模型[smc_reac_model.pdparams](https://paddle-org.bj.bcebos.com/paddlescience/models/smc_reac/smc_reac_model.pdparams), 并对应修改 yaml 配置文件中的 `load_model_path` 为模型参数路径。 3. 开始训练、评估前,请安装 `rdkit` 等,相关依赖请执行`pip install -r requirements.txt`安装。 === "模型训练命令" @@ -15,7 +15,7 @@ === "模型评估命令" ``` sh - python smc_reac.py mode=eval + python smc_reac.py mode=eval Eval.load_model_path=https://paddle-org.bj.bcebos.com/paddlescience/models/smc_reac/smc_reac_model.pdparams ``` ## 1. 背景简介 From 3b16af50abef46d532ac53857ff638101147e2a6 Mon Sep 17 00:00:00 2001 From: Dubhe-Chang Date: Sun, 13 Jul 2025 23:05:13 +0800 Subject: [PATCH 28/32] update smc_reac.yaml --- examples/smc_reac/config/smc_reac.yaml | 14 ++++---------- 1 file changed, 4 insertions(+), 10 deletions(-) diff --git a/examples/smc_reac/config/smc_reac.yaml b/examples/smc_reac/config/smc_reac.yaml index 68cd4300d9..5232ee9cb6 100644 --- a/examples/smc_reac/config/smc_reac.yaml +++ b/examples/smc_reac/config/smc_reac.yaml @@ -44,19 +44,13 @@ MODEL: # TRAIN: # epochs: 1500 # iters_per_epoch: 20 # - # save_freq: 100 # - # eval_during_train: False # batch_size: 8 # learning_rate: 0.0001 - save_model_path: './smc_reac_model.pdparams' - # weight_decay: 1e-5 - # pretrained_model_path: null # - # checkpoint_path: null # - # k: 9 - # i: 2 + pretrained_model_path: null + checkpoint_path: null # evaluation settings EVAL: - test_size: 0.1 - load_model_path: './smc_reac_model.pdparams' + pretrained_model_path: null + batch_size: 8 # seed: 20 From 8126b2ae3fcc3fab228f508673023536f3a5e6c1 Mon Sep 17 00:00:00 2001 From: Dubhe-Chang Date: Sun, 13 Jul 2025 23:13:51 +0800 Subject: [PATCH 29/32] fix(smc_reac): correct evaluate unpacking logic --- examples/smc_reac/smc_reac.py | 73 +++++++++++++++++++++++------------ 1 file changed, 49 insertions(+), 24 deletions(-) diff --git a/examples/smc_reac/smc_reac.py b/examples/smc_reac/smc_reac.py index cf81115181..30356e617b 100644 --- a/examples/smc_reac/smc_reac.py +++ b/examples/smc_reac/smc_reac.py @@ -8,11 +8,11 @@ import rdkit.Chem as Chem from omegaconf import DictConfig from rdkit.Chem import rdFingerprintGenerator -from sklearn.metrics import r2_score from sklearn.model_selection import train_test_split import ppsci +paddle.set_device("gpu:5") os.environ["HYDRA_FULL_ERROR"] = "1" os.environ["KMP_DUPLICATE_LIB_OK"] = "True" plt.rcParams["axes.unicode_minus"] = False @@ -107,37 +107,62 @@ def train(cfg: DictConfig): def evaluate(cfg: DictConfig): global x_test, y_test + x_test, y_test = data_processed(x_test, y_test) - # Reformat data for evaluation - x_test = {"v": x_test} - y_test = {"u": y_test} + + test_validator = ppsci.validate.SupervisedValidator( + dataloader_cfg={ + "dataset": { + "input": {"v": x_test}, + "label": {"u": y_test}, + "name": "IterableNamedArrayDataset", + }, + "batch_size": cfg.EVAL.batch_size, + "shuffle": False, + }, + loss=ppsci.loss.MSELoss("mean"), + metric={ + "MAE": ppsci.metric.MAE(), + "RMSE": ppsci.metric.RMSE(), + "R2": ppsci.metric.R2Score(), + }, + name="test_eval", + ) + validators = {"test_eval": test_validator} + model = ppsci.arch.SuzukiMiyauraModel(**cfg.MODEL) - model.set_state_dict(paddle.load(cfg.EVAL.load_model_path)) - ypred = model(x_test) - - # Calculate evaluation metrics - loss = ppsci.metric.MAE() - MAE = loss(ypred, y_test).get("u").numpy() - loss = ppsci.metric.RMSE() - RMSE = loss(ypred, y_test).get("u").numpy() - ypred = ypred.get("u").numpy() - ytest = y_test.get("u").numpy() - R2 = r2_score(ytest, ypred) - print("MAE", MAE) - print("RMSE", RMSE) - print("R2", R2) - - # Visualization - plt.scatter(ytest, ypred, s=15, color="royalblue", marker="s", linewidth=1) - plt.plot([ytest.min(), ytest.max()], [ytest.min(), ytest.max()], "r-", lw=1) - plt.legend(title="R²={:.3f}\n\nMAE={:.3f}".format(R2, MAE)) + solver = ppsci.solver.Solver( + model, + validator=validators, + cfg=cfg, + ) + + loss_val, metric_dict = solver.eval() + + ypred = model({"v": x_test})["u"].numpy() + ytrue = y_test.numpy() + + mae = metric_dict["MAE"]["u"] + rmse = metric_dict["RMSE"]["u"] + r2 = metric_dict["R2"]["u"] + + plt.figure() + plt.scatter(ytrue, ypred, s=15, color="royalblue", marker="s", linewidth=1) + plt.plot([ytrue.min(), ytrue.max()], [ytrue.min(), ytrue.max()], "r-", lw=1) + plt.legend(title="R²={:.3f}\n\nMAE={:.3f}".format(r2, mae)) plt.xlabel("Test Yield(%)") plt.ylabel("Predicted Yield(%)") save_path = "smc_reac.png" plt.savefig(save_path) - print(f"Iamge saved to: {save_path}") + print(f"Image saved to: {save_path}") plt.show() + print("Evaluation metrics:") + print(f"Loss: {loss_val:.4f}") + print(f"MAE : {mae:.4f}") + print(f"RMSE: {rmse:.4f}") + print(f"R2 : {r2:.4f}") + @hydra.main(version_base=None, config_path="./config", config_name="smc_reac.yaml") def main(cfg: DictConfig): From 150d12c39ca63bf17d240159ac1c49f560bbee31 Mon Sep 17 00:00:00 2001 From: Dubhe Chang Date: Mon, 28 Jul 2025 21:21:50 +0800 Subject: [PATCH 30/32] update smc_reac.md --- docs/zh/examples/smc_reac.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/zh/examples/smc_reac.md b/docs/zh/examples/smc_reac.md index b4c3b34b8d..d2590c2c6a 100644 --- a/docs/zh/examples/smc_reac.md +++ b/docs/zh/examples/smc_reac.md @@ -15,7 +15,7 @@ === "模型评估命令" ``` sh - python smc_reac.py mode=eval Eval.load_model_path=https://paddle-org.bj.bcebos.com/paddlescience/models/smc_reac/smc_reac_model.pdparams + python smc_reac.py mode=eval EVAL.load_model_path=https://paddle-org.bj.bcebos.com/paddlescience/models/smc_reac/smc_reac_model.pdparams ``` ## 1. 背景简介 From 16aae512126244383cf811d8769d9f4dd014d754 Mon Sep 17 00:00:00 2001 From: Dubhe Chang Date: Mon, 28 Jul 2025 21:26:17 +0800 Subject: [PATCH 31/32] merge code of upstream --- examples/smc_reac/ppsci/__init__.py | 78 + examples/smc_reac/ppsci/arch/__init__.py | 136 ++ examples/smc_reac/ppsci/arch/activation.py | 160 ++ examples/smc_reac/ppsci/arch/afno.py | 687 ++++++ examples/smc_reac/ppsci/arch/amgnet.py | 649 ++++++ examples/smc_reac/ppsci/arch/base.py | 279 +++ examples/smc_reac/ppsci/arch/cfdgcn.py | 350 +++ .../smc_reac/ppsci/arch/chip_deeponets.py | 214 ++ .../ppsci/arch/crystalgraphconvnet.py | 167 ++ .../smc_reac/ppsci/arch/cuboid_transformer.py | 958 ++++++++ .../ppsci/arch/cuboid_transformer_decoder.py | 1245 +++++++++++ .../ppsci/arch/cuboid_transformer_encoder.py | 1515 +++++++++++++ .../ppsci/arch/cuboid_transformer_utils.py | 347 +++ examples/smc_reac/ppsci/arch/cvit.py | 1095 +++++++++ examples/smc_reac/ppsci/arch/deeponet.py | 154 ++ examples/smc_reac/ppsci/arch/dgmr.py | 1151 ++++++++++ .../smc_reac/ppsci/arch/embedding_koopman.py | 544 +++++ examples/smc_reac/ppsci/arch/epnn.py | 126 ++ .../ppsci/arch/extformer_moe_cuboid.py | 996 +++++++++ .../arch/extformer_moe_cuboid_decoder.py | 1475 ++++++++++++ .../arch/extformer_moe_cuboid_encoder.py | 1992 +++++++++++++++++ .../ppsci/arch/extformer_moe_cuboid_utils.py | 350 +++ .../ppsci/arch/extformer_moe_utils.py | 563 +++++ examples/smc_reac/ppsci/arch/fno_block.py | 1269 +++++++++++ examples/smc_reac/ppsci/arch/gan.py | 400 ++++ examples/smc_reac/ppsci/arch/geofno.py | 205 ++ examples/smc_reac/ppsci/arch/graphcast.py | 492 ++++ examples/smc_reac/ppsci/arch/he_deeponets.py | 197 ++ examples/smc_reac/ppsci/arch/ifm_mlp.py | 540 +++++ examples/smc_reac/ppsci/arch/kan.py | 385 ++++ examples/smc_reac/ppsci/arch/lno.py | 312 +++ examples/smc_reac/ppsci/arch/mlp.py | 828 +++++++ examples/smc_reac/ppsci/arch/model_list.py | 72 + examples/smc_reac/ppsci/arch/moflow_basic.py | 297 +++ examples/smc_reac/ppsci/arch/moflow_glow.py | 477 ++++ examples/smc_reac/ppsci/arch/moflow_net.py | 335 +++ examples/smc_reac/ppsci/arch/nowcastnet.py | 639 ++++++ .../ppsci/arch/paddle_harmonics/legendre.py | 176 ++ .../ppsci/arch/paddle_harmonics/quadrature.py | 156 ++ .../arch/paddle_harmonics/random_fields.py | 148 ++ .../ppsci/arch/paddle_harmonics/sht.py | 461 ++++ examples/smc_reac/ppsci/arch/phycrnet.py | 540 +++++ examples/smc_reac/ppsci/arch/phylstm.py | 239 ++ .../smc_reac/ppsci/arch/physx_transformer.py | 407 ++++ examples/smc_reac/ppsci/arch/regdgcnn.py | 250 +++ examples/smc_reac/ppsci/arch/regpointnet.py | 146 ++ examples/smc_reac/ppsci/arch/sfnonet.py | 568 +++++ examples/smc_reac/ppsci/arch/smc_reac.py | 107 + examples/smc_reac/ppsci/arch/spinn.py | 180 ++ examples/smc_reac/ppsci/arch/tfnonet.py | 514 +++++ examples/smc_reac/ppsci/arch/tgcn.py | 200 ++ examples/smc_reac/ppsci/arch/transformer.py | 417 ++++ examples/smc_reac/ppsci/arch/unetex.py | 290 +++ examples/smc_reac/ppsci/arch/unonet.py | 289 +++ examples/smc_reac/ppsci/arch/uscnn.py | 124 + examples/smc_reac/ppsci/arch/vae.py | 103 + examples/smc_reac/ppsci/arch/velocitygan.py | 354 +++ examples/smc_reac/ppsci/autodiff/__init__.py | 17 + examples/smc_reac/ppsci/autodiff/ad.py | 341 +++ .../smc_reac/ppsci/constraint/__init__.py | 86 + examples/smc_reac/ppsci/constraint/base.py | 62 + .../ppsci/constraint/boundary_constraint.py | 163 ++ .../ppsci/constraint/initial_constraint.py | 172 ++ .../ppsci/constraint/integral_constraint.py | 178 ++ .../ppsci/constraint/interior_constraint.py | 174 ++ .../ppsci/constraint/periodic_constraint.py | 169 ++ .../ppsci/constraint/supervised_constraint.py | 92 + examples/smc_reac/ppsci/data/__init__.py | 205 ++ examples/smc_reac/ppsci/data/dataloader.py | 47 + .../smc_reac/ppsci/data/dataset/__init__.py | 118 + .../ppsci/data/dataset/airfoil_dataset.py | 241 ++ .../ppsci/data/dataset/array_dataset.py | 390 ++++ .../ppsci/data/dataset/atmospheric_dataset.py | 1781 +++++++++++++++ .../ppsci/data/dataset/cgcnn_dataset.py | 312 +++ .../ppsci/data/dataset/csv_dataset.py | 287 +++ .../ppsci/data/dataset/cylinder_dataset.py | 215 ++ .../ppsci/data/dataset/darcyflow_dataset.py | 296 +++ .../ppsci/data/dataset/dgmr_dataset.py | 95 + .../ppsci/data/dataset/drivaernet_dataset.py | 316 +++ .../dataset/drivaernetplusplus_dataset.py | 321 +++ .../ppsci/data/dataset/enso_dataset.py | 405 ++++ .../ppsci/data/dataset/era5_dataset.py | 249 +++ .../data/dataset/ext_moe_enso_dataset.py | 406 ++++ .../ppsci/data/dataset/fwi_dataset.py | 103 + .../ppsci/data/dataset/ifm_moe_dataset.py | 462 ++++ .../ppsci/data/dataset/mat_dataset.py | 287 +++ .../ppsci/data/dataset/moflow_dataset.py | 437 ++++ .../ppsci/data/dataset/mrms_dataset.py | 251 +++ .../ppsci/data/dataset/npz_dataset.py | 279 +++ .../ppsci/data/dataset/pems_dataset.py | 151 ++ .../ppsci/data/dataset/radar_dataset.py | 146 ++ .../ppsci/data/dataset/sevir_dataset.py | 814 +++++++ .../data/dataset/spherical_swe_dataset.py | 104 + .../ppsci/data/dataset/trphysx_dataset.py | 326 +++ .../ppsci/data/dataset/vtu_dataset.py | 106 + .../smc_reac/ppsci/data/process/__init__.py | 21 + .../data/process/batch_transform/__init__.py | 135 ++ .../process/batch_transform/preprocess.py | 74 + .../ppsci/data/process/transform/__init__.py | 72 + .../data/process/transform/preprocess.py | 331 +++ examples/smc_reac/ppsci/equation/__init__.py | 76 + .../smc_reac/ppsci/equation/fpde/__init__.py | 19 + .../ppsci/equation/fpde/fractional_poisson.py | 196 ++ .../smc_reac/ppsci/equation/ide/__init__.py | 19 + .../smc_reac/ppsci/equation/ide/volterra.py | 127 ++ .../smc_reac/ppsci/equation/pde/__init__.py | 43 + .../smc_reac/ppsci/equation/pde/allen_cahn.py | 64 + examples/smc_reac/ppsci/equation/pde/base.py | 243 ++ .../smc_reac/ppsci/equation/pde/biharmonic.py | 74 + .../ppsci/equation/pde/heat_exchanger.py | 94 + .../smc_reac/ppsci/equation/pde/helmholtz.py | 119 + .../smc_reac/ppsci/equation/pde/laplace.py | 55 + .../ppsci/equation/pde/linear_elasticity.py | 184 ++ .../ppsci/equation/pde/navier_stokes.py | 151 ++ .../smc_reac/ppsci/equation/pde/nls_m_b.py | 101 + .../ppsci/equation/pde/normal_dot_vec.py | 59 + .../smc_reac/ppsci/equation/pde/poisson.py | 53 + examples/smc_reac/ppsci/equation/pde/viv.py | 64 + .../smc_reac/ppsci/experimental/__init__.py | 37 + .../ppsci/experimental/math_module.py | 646 ++++++ examples/smc_reac/ppsci/externals/__init__.py | 20 + examples/smc_reac/ppsci/geometry/__init__.py | 83 + examples/smc_reac/ppsci/geometry/csg.py | 337 +++ examples/smc_reac/ppsci/geometry/geometry.py | 696 ++++++ .../smc_reac/ppsci/geometry/geometry_1d.py | 119 + .../smc_reac/ppsci/geometry/geometry_2d.py | 706 ++++++ .../smc_reac/ppsci/geometry/geometry_3d.py | 203 ++ .../smc_reac/ppsci/geometry/geometry_nd.py | 196 ++ examples/smc_reac/ppsci/geometry/inflation.py | 192 ++ examples/smc_reac/ppsci/geometry/mesh.py | 1392 ++++++++++++ .../smc_reac/ppsci/geometry/pointcloud.py | 312 +++ examples/smc_reac/ppsci/geometry/sampler.py | 92 + examples/smc_reac/ppsci/geometry/sdf.py | 198 ++ .../smc_reac/ppsci/geometry/timedomain.py | 793 +++++++ examples/smc_reac/ppsci/loss/__init__.py | 67 + examples/smc_reac/ppsci/loss/base.py | 38 + examples/smc_reac/ppsci/loss/chamfer.py | 92 + examples/smc_reac/ppsci/loss/func.py | 94 + examples/smc_reac/ppsci/loss/integral.py | 112 + examples/smc_reac/ppsci/loss/kl.py | 51 + examples/smc_reac/ppsci/loss/l1.py | 219 ++ examples/smc_reac/ppsci/loss/l2.py | 310 +++ examples/smc_reac/ppsci/loss/mae.py | 109 + examples/smc_reac/ppsci/loss/mse.py | 355 +++ examples/smc_reac/ppsci/loss/mtl/__init__.py | 49 + examples/smc_reac/ppsci/loss/mtl/agda.py | 161 ++ examples/smc_reac/ppsci/loss/mtl/base.py | 68 + examples/smc_reac/ppsci/loss/mtl/grad_norm.py | 145 ++ examples/smc_reac/ppsci/loss/mtl/ntk.py | 118 + examples/smc_reac/ppsci/loss/mtl/pcgrad.py | 124 + examples/smc_reac/ppsci/loss/mtl/relobralo.py | 127 ++ examples/smc_reac/ppsci/loss/mtl/sum.py | 60 + examples/smc_reac/ppsci/metric/__init__.py | 63 + .../smc_reac/ppsci/metric/anomaly_coef.py | 122 + examples/smc_reac/ppsci/metric/base.py | 25 + examples/smc_reac/ppsci/metric/func.py | 66 + examples/smc_reac/ppsci/metric/l2_rel.py | 139 ++ examples/smc_reac/ppsci/metric/mae.py | 73 + examples/smc_reac/ppsci/metric/max_ae.py | 77 + examples/smc_reac/ppsci/metric/mse.py | 73 + examples/smc_reac/ppsci/metric/r2_score.py | 97 + examples/smc_reac/ppsci/metric/rmse.py | 155 ++ examples/smc_reac/ppsci/optimizer/__init__.py | 84 + .../smc_reac/ppsci/optimizer/lr_scheduler.py | 911 ++++++++ .../smc_reac/ppsci/optimizer/optimizer.py | 649 ++++++ examples/smc_reac/ppsci/optimizer/soap.py | 558 +++++ .../smc_reac/ppsci/probability/__init__.py | 19 + examples/smc_reac/ppsci/probability/hmc.py | 175 ++ examples/smc_reac/ppsci/solver/__init__.py | 25 + examples/smc_reac/ppsci/solver/eval.py | 316 +++ examples/smc_reac/ppsci/solver/printer.py | 161 ++ examples/smc_reac/ppsci/solver/solver.py | 1219 ++++++++++ examples/smc_reac/ppsci/solver/train.py | 324 +++ examples/smc_reac/ppsci/solver/visu.py | 98 + examples/smc_reac/ppsci/utils/__init__.py | 66 + examples/smc_reac/ppsci/utils/callbacks.py | 136 ++ examples/smc_reac/ppsci/utils/checker.py | 287 +++ examples/smc_reac/ppsci/utils/config.py | 457 ++++ examples/smc_reac/ppsci/utils/download.py | 285 +++ examples/smc_reac/ppsci/utils/ema.py | 172 ++ examples/smc_reac/ppsci/utils/expression.py | 212 ++ examples/smc_reac/ppsci/utils/initializer.py | 498 +++++ examples/smc_reac/ppsci/utils/logger.py | 264 +++ examples/smc_reac/ppsci/utils/misc.py | 684 ++++++ examples/smc_reac/ppsci/utils/reader.py | 266 +++ examples/smc_reac/ppsci/utils/save_load.py | 300 +++ examples/smc_reac/ppsci/utils/symbolic.py | 981 ++++++++ examples/smc_reac/ppsci/utils/writer.py | 225 ++ examples/smc_reac/ppsci/validate/__init__.py | 81 + examples/smc_reac/ppsci/validate/base.py | 69 + .../smc_reac/ppsci/validate/geo_validator.py | 161 ++ .../smc_reac/ppsci/validate/sup_validator.py | 103 + examples/smc_reac/ppsci/visualize/__init__.py | 82 + examples/smc_reac/ppsci/visualize/base.py | 65 + examples/smc_reac/ppsci/visualize/plot.py | 580 +++++ examples/smc_reac/ppsci/visualize/radar.py | 124 + .../smc_reac/ppsci/visualize/visualizer.py | 409 ++++ examples/smc_reac/ppsci/visualize/vtu.py | 278 +++ 198 files changed, 60861 insertions(+) create mode 100644 examples/smc_reac/ppsci/__init__.py create mode 100644 examples/smc_reac/ppsci/arch/__init__.py create mode 100644 examples/smc_reac/ppsci/arch/activation.py create mode 100644 examples/smc_reac/ppsci/arch/afno.py create mode 100644 examples/smc_reac/ppsci/arch/amgnet.py create mode 100644 examples/smc_reac/ppsci/arch/base.py create mode 100644 examples/smc_reac/ppsci/arch/cfdgcn.py create mode 100644 examples/smc_reac/ppsci/arch/chip_deeponets.py create mode 100644 examples/smc_reac/ppsci/arch/crystalgraphconvnet.py create mode 100644 examples/smc_reac/ppsci/arch/cuboid_transformer.py create mode 100644 examples/smc_reac/ppsci/arch/cuboid_transformer_decoder.py create mode 100644 examples/smc_reac/ppsci/arch/cuboid_transformer_encoder.py create mode 100644 examples/smc_reac/ppsci/arch/cuboid_transformer_utils.py create mode 100644 examples/smc_reac/ppsci/arch/cvit.py create mode 100644 examples/smc_reac/ppsci/arch/deeponet.py create mode 100644 examples/smc_reac/ppsci/arch/dgmr.py create mode 100644 examples/smc_reac/ppsci/arch/embedding_koopman.py create mode 100644 examples/smc_reac/ppsci/arch/epnn.py create mode 100644 examples/smc_reac/ppsci/arch/extformer_moe_cuboid.py create mode 100644 examples/smc_reac/ppsci/arch/extformer_moe_cuboid_decoder.py create mode 100644 examples/smc_reac/ppsci/arch/extformer_moe_cuboid_encoder.py create mode 100644 examples/smc_reac/ppsci/arch/extformer_moe_cuboid_utils.py create mode 100644 examples/smc_reac/ppsci/arch/extformer_moe_utils.py create mode 100644 examples/smc_reac/ppsci/arch/fno_block.py create mode 100644 examples/smc_reac/ppsci/arch/gan.py create mode 100644 examples/smc_reac/ppsci/arch/geofno.py create mode 100644 examples/smc_reac/ppsci/arch/graphcast.py create mode 100644 examples/smc_reac/ppsci/arch/he_deeponets.py create mode 100644 examples/smc_reac/ppsci/arch/ifm_mlp.py create mode 100644 examples/smc_reac/ppsci/arch/kan.py create mode 100644 examples/smc_reac/ppsci/arch/lno.py create mode 100644 examples/smc_reac/ppsci/arch/mlp.py create mode 100644 examples/smc_reac/ppsci/arch/model_list.py create mode 100644 examples/smc_reac/ppsci/arch/moflow_basic.py create mode 100644 examples/smc_reac/ppsci/arch/moflow_glow.py create mode 100644 examples/smc_reac/ppsci/arch/moflow_net.py create mode 100644 examples/smc_reac/ppsci/arch/nowcastnet.py create mode 100644 examples/smc_reac/ppsci/arch/paddle_harmonics/legendre.py create mode 100644 examples/smc_reac/ppsci/arch/paddle_harmonics/quadrature.py create mode 100644 examples/smc_reac/ppsci/arch/paddle_harmonics/random_fields.py create mode 100644 examples/smc_reac/ppsci/arch/paddle_harmonics/sht.py create mode 100644 examples/smc_reac/ppsci/arch/phycrnet.py create mode 100644 examples/smc_reac/ppsci/arch/phylstm.py create mode 100644 examples/smc_reac/ppsci/arch/physx_transformer.py create mode 100644 examples/smc_reac/ppsci/arch/regdgcnn.py create mode 100644 examples/smc_reac/ppsci/arch/regpointnet.py create mode 100644 examples/smc_reac/ppsci/arch/sfnonet.py create mode 100644 examples/smc_reac/ppsci/arch/smc_reac.py create mode 100644 examples/smc_reac/ppsci/arch/spinn.py create mode 100644 examples/smc_reac/ppsci/arch/tfnonet.py create mode 100644 examples/smc_reac/ppsci/arch/tgcn.py create mode 100644 examples/smc_reac/ppsci/arch/transformer.py create mode 100644 examples/smc_reac/ppsci/arch/unetex.py create mode 100644 examples/smc_reac/ppsci/arch/unonet.py create mode 100644 examples/smc_reac/ppsci/arch/uscnn.py create mode 100644 examples/smc_reac/ppsci/arch/vae.py create mode 100644 examples/smc_reac/ppsci/arch/velocitygan.py create mode 100644 examples/smc_reac/ppsci/autodiff/__init__.py create mode 100644 examples/smc_reac/ppsci/autodiff/ad.py create mode 100644 examples/smc_reac/ppsci/constraint/__init__.py create mode 100644 examples/smc_reac/ppsci/constraint/base.py create mode 100644 examples/smc_reac/ppsci/constraint/boundary_constraint.py create mode 100644 examples/smc_reac/ppsci/constraint/initial_constraint.py create mode 100644 examples/smc_reac/ppsci/constraint/integral_constraint.py create mode 100644 examples/smc_reac/ppsci/constraint/interior_constraint.py create mode 100644 examples/smc_reac/ppsci/constraint/periodic_constraint.py create mode 100644 examples/smc_reac/ppsci/constraint/supervised_constraint.py create mode 100644 examples/smc_reac/ppsci/data/__init__.py create mode 100644 examples/smc_reac/ppsci/data/dataloader.py create mode 100644 examples/smc_reac/ppsci/data/dataset/__init__.py create mode 100644 examples/smc_reac/ppsci/data/dataset/airfoil_dataset.py create mode 100644 examples/smc_reac/ppsci/data/dataset/array_dataset.py create mode 100644 examples/smc_reac/ppsci/data/dataset/atmospheric_dataset.py create mode 100644 examples/smc_reac/ppsci/data/dataset/cgcnn_dataset.py create mode 100644 examples/smc_reac/ppsci/data/dataset/csv_dataset.py create mode 100644 examples/smc_reac/ppsci/data/dataset/cylinder_dataset.py create mode 100644 examples/smc_reac/ppsci/data/dataset/darcyflow_dataset.py create mode 100644 examples/smc_reac/ppsci/data/dataset/dgmr_dataset.py create mode 100644 examples/smc_reac/ppsci/data/dataset/drivaernet_dataset.py create mode 100644 examples/smc_reac/ppsci/data/dataset/drivaernetplusplus_dataset.py create mode 100644 examples/smc_reac/ppsci/data/dataset/enso_dataset.py create mode 100644 examples/smc_reac/ppsci/data/dataset/era5_dataset.py create mode 100644 examples/smc_reac/ppsci/data/dataset/ext_moe_enso_dataset.py create mode 100644 examples/smc_reac/ppsci/data/dataset/fwi_dataset.py create mode 100644 examples/smc_reac/ppsci/data/dataset/ifm_moe_dataset.py create mode 100644 examples/smc_reac/ppsci/data/dataset/mat_dataset.py create mode 100644 examples/smc_reac/ppsci/data/dataset/moflow_dataset.py create mode 100644 examples/smc_reac/ppsci/data/dataset/mrms_dataset.py create mode 100644 examples/smc_reac/ppsci/data/dataset/npz_dataset.py create mode 100644 examples/smc_reac/ppsci/data/dataset/pems_dataset.py create mode 100644 examples/smc_reac/ppsci/data/dataset/radar_dataset.py create mode 100644 examples/smc_reac/ppsci/data/dataset/sevir_dataset.py create mode 100644 examples/smc_reac/ppsci/data/dataset/spherical_swe_dataset.py create mode 100644 examples/smc_reac/ppsci/data/dataset/trphysx_dataset.py create mode 100644 examples/smc_reac/ppsci/data/dataset/vtu_dataset.py create mode 100644 examples/smc_reac/ppsci/data/process/__init__.py create mode 100644 examples/smc_reac/ppsci/data/process/batch_transform/__init__.py create mode 100644 examples/smc_reac/ppsci/data/process/batch_transform/preprocess.py create mode 100644 examples/smc_reac/ppsci/data/process/transform/__init__.py create mode 100644 examples/smc_reac/ppsci/data/process/transform/preprocess.py create mode 100644 examples/smc_reac/ppsci/equation/__init__.py create mode 100644 examples/smc_reac/ppsci/equation/fpde/__init__.py create mode 100644 examples/smc_reac/ppsci/equation/fpde/fractional_poisson.py create mode 100644 examples/smc_reac/ppsci/equation/ide/__init__.py create mode 100644 examples/smc_reac/ppsci/equation/ide/volterra.py create mode 100644 examples/smc_reac/ppsci/equation/pde/__init__.py create mode 100644 examples/smc_reac/ppsci/equation/pde/allen_cahn.py create mode 100644 examples/smc_reac/ppsci/equation/pde/base.py create mode 100644 examples/smc_reac/ppsci/equation/pde/biharmonic.py create mode 100644 examples/smc_reac/ppsci/equation/pde/heat_exchanger.py create mode 100644 examples/smc_reac/ppsci/equation/pde/helmholtz.py create mode 100644 examples/smc_reac/ppsci/equation/pde/laplace.py create mode 100644 examples/smc_reac/ppsci/equation/pde/linear_elasticity.py create mode 100644 examples/smc_reac/ppsci/equation/pde/navier_stokes.py create mode 100644 examples/smc_reac/ppsci/equation/pde/nls_m_b.py create mode 100644 examples/smc_reac/ppsci/equation/pde/normal_dot_vec.py create mode 100644 examples/smc_reac/ppsci/equation/pde/poisson.py create mode 100644 examples/smc_reac/ppsci/equation/pde/viv.py create mode 100644 examples/smc_reac/ppsci/experimental/__init__.py create mode 100644 examples/smc_reac/ppsci/experimental/math_module.py create mode 100644 examples/smc_reac/ppsci/externals/__init__.py create mode 100644 examples/smc_reac/ppsci/geometry/__init__.py create mode 100644 examples/smc_reac/ppsci/geometry/csg.py create mode 100644 examples/smc_reac/ppsci/geometry/geometry.py create mode 100644 examples/smc_reac/ppsci/geometry/geometry_1d.py create mode 100644 examples/smc_reac/ppsci/geometry/geometry_2d.py create mode 100644 examples/smc_reac/ppsci/geometry/geometry_3d.py create mode 100644 examples/smc_reac/ppsci/geometry/geometry_nd.py create mode 100644 examples/smc_reac/ppsci/geometry/inflation.py create mode 100644 examples/smc_reac/ppsci/geometry/mesh.py create mode 100644 examples/smc_reac/ppsci/geometry/pointcloud.py create mode 100644 examples/smc_reac/ppsci/geometry/sampler.py create mode 100644 examples/smc_reac/ppsci/geometry/sdf.py create mode 100644 examples/smc_reac/ppsci/geometry/timedomain.py create mode 100644 examples/smc_reac/ppsci/loss/__init__.py create mode 100644 examples/smc_reac/ppsci/loss/base.py create mode 100644 examples/smc_reac/ppsci/loss/chamfer.py create mode 100644 examples/smc_reac/ppsci/loss/func.py create mode 100644 examples/smc_reac/ppsci/loss/integral.py create mode 100644 examples/smc_reac/ppsci/loss/kl.py create mode 100644 examples/smc_reac/ppsci/loss/l1.py create mode 100644 examples/smc_reac/ppsci/loss/l2.py create mode 100644 examples/smc_reac/ppsci/loss/mae.py create mode 100644 examples/smc_reac/ppsci/loss/mse.py create mode 100644 examples/smc_reac/ppsci/loss/mtl/__init__.py create mode 100644 examples/smc_reac/ppsci/loss/mtl/agda.py create mode 100644 examples/smc_reac/ppsci/loss/mtl/base.py create mode 100644 examples/smc_reac/ppsci/loss/mtl/grad_norm.py create mode 100644 examples/smc_reac/ppsci/loss/mtl/ntk.py create mode 100644 examples/smc_reac/ppsci/loss/mtl/pcgrad.py create mode 100644 examples/smc_reac/ppsci/loss/mtl/relobralo.py create mode 100644 examples/smc_reac/ppsci/loss/mtl/sum.py create mode 100644 examples/smc_reac/ppsci/metric/__init__.py create mode 100644 examples/smc_reac/ppsci/metric/anomaly_coef.py create mode 100644 examples/smc_reac/ppsci/metric/base.py create mode 100644 examples/smc_reac/ppsci/metric/func.py create mode 100644 examples/smc_reac/ppsci/metric/l2_rel.py create mode 100644 examples/smc_reac/ppsci/metric/mae.py create mode 100644 examples/smc_reac/ppsci/metric/max_ae.py create mode 100644 examples/smc_reac/ppsci/metric/mse.py create mode 100644 examples/smc_reac/ppsci/metric/r2_score.py create mode 100644 examples/smc_reac/ppsci/metric/rmse.py create mode 100644 examples/smc_reac/ppsci/optimizer/__init__.py create mode 100644 examples/smc_reac/ppsci/optimizer/lr_scheduler.py create mode 100644 examples/smc_reac/ppsci/optimizer/optimizer.py create mode 100644 examples/smc_reac/ppsci/optimizer/soap.py create mode 100644 examples/smc_reac/ppsci/probability/__init__.py create mode 100644 examples/smc_reac/ppsci/probability/hmc.py create mode 100644 examples/smc_reac/ppsci/solver/__init__.py create mode 100644 examples/smc_reac/ppsci/solver/eval.py create mode 100644 examples/smc_reac/ppsci/solver/printer.py create mode 100644 examples/smc_reac/ppsci/solver/solver.py create mode 100644 examples/smc_reac/ppsci/solver/train.py create mode 100644 examples/smc_reac/ppsci/solver/visu.py create mode 100644 examples/smc_reac/ppsci/utils/__init__.py create mode 100644 examples/smc_reac/ppsci/utils/callbacks.py create mode 100644 examples/smc_reac/ppsci/utils/checker.py create mode 100644 examples/smc_reac/ppsci/utils/config.py create mode 100644 examples/smc_reac/ppsci/utils/download.py create mode 100644 examples/smc_reac/ppsci/utils/ema.py create mode 100644 examples/smc_reac/ppsci/utils/expression.py create mode 100644 examples/smc_reac/ppsci/utils/initializer.py create mode 100644 examples/smc_reac/ppsci/utils/logger.py create mode 100644 examples/smc_reac/ppsci/utils/misc.py create mode 100644 examples/smc_reac/ppsci/utils/reader.py create mode 100644 examples/smc_reac/ppsci/utils/save_load.py create mode 100644 examples/smc_reac/ppsci/utils/symbolic.py create mode 100644 examples/smc_reac/ppsci/utils/writer.py create mode 100644 examples/smc_reac/ppsci/validate/__init__.py create mode 100644 examples/smc_reac/ppsci/validate/base.py create mode 100644 examples/smc_reac/ppsci/validate/geo_validator.py create mode 100644 examples/smc_reac/ppsci/validate/sup_validator.py create mode 100644 examples/smc_reac/ppsci/visualize/__init__.py create mode 100644 examples/smc_reac/ppsci/visualize/base.py create mode 100644 examples/smc_reac/ppsci/visualize/plot.py create mode 100644 examples/smc_reac/ppsci/visualize/radar.py create mode 100644 examples/smc_reac/ppsci/visualize/visualizer.py create mode 100644 examples/smc_reac/ppsci/visualize/vtu.py diff --git a/examples/smc_reac/ppsci/__init__.py b/examples/smc_reac/ppsci/__init__.py new file mode 100644 index 0000000000..bde7aa49a5 --- /dev/null +++ b/examples/smc_reac/ppsci/__init__.py @@ -0,0 +1,78 @@ +# 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. + +from ppsci import arch # isort:skip +from ppsci import autodiff # isort:skip +from ppsci import constraint # isort:skip +from ppsci import data # isort:skip +from ppsci import equation # isort:skip +from ppsci import geometry # isort:skip +from ppsci import loss # isort:skip +from ppsci import metric # isort:skip +from ppsci import optimizer # isort:skip +from ppsci import utils # isort:skip +from ppsci import visualize # isort:skip +from ppsci import validate # isort:skip +from ppsci import solver # isort:skip +from ppsci import experimental # isort:skip + +from ppsci.utils.checker import run_check # isort:skip +from ppsci.utils.checker import run_check_mesh # isort:skip +from ppsci.utils import lambdify # isort:skip + + +try: + # import auto-generated version information from '._version' file, using + # setuptools_scm via 'pip install'. Details of versioning rule can be referd to: + # https://peps.python.org/pep-0440/#public-version-identifiers + from ._version import version as __version__ +except ImportError: + __version__ = "unknown version" + +__all__ = [ + "arch", + "autodiff", + "constraint", + "data", + "equation", + "geometry", + "loss", + "metric", + "optimizer", + "utils", + "visualize", + "validate", + "solver", + "experimental", + "run_check", + "run_check_mesh", + "lambdify", +] + + +# NOTE: Register custom solvers for parsing values from omegaconf more flexible +def _register_config_solvers(): + import numpy as np + from omegaconf import OmegaConf + + # register solver for "${numpy:xxx}" item, e.g. pi: "${numpy:pi}" + if not OmegaConf.has_resolver("numpy"): + OmegaConf.register_new_resolver("numpy", lambda x: getattr(np, x)) + + # register solver for "${sum:xxx}" item, e.g. pi: "${sum:[10, 20, 30]}" + if not OmegaConf.has_resolver("sum"): + OmegaConf.register_new_resolver("sum", lambda x: sum(x)) + + +_register_config_solvers() diff --git a/examples/smc_reac/ppsci/arch/__init__.py b/examples/smc_reac/ppsci/arch/__init__.py new file mode 100644 index 0000000000..16387994c7 --- /dev/null +++ b/examples/smc_reac/ppsci/arch/__init__.py @@ -0,0 +1,136 @@ +# 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. + +from __future__ import annotations + +import copy + +from ppsci.arch.afno import AFNONet # isort:skip +from ppsci.arch.afno import PrecipNet # isort:skip +from ppsci.arch.amgnet import AMGNet # isort:skip +from ppsci.arch.base import Arch # isort:skip +from ppsci.arch.cfdgcn import CFDGCN # isort:skip +from ppsci.arch.smc_reac import SuzukiMiyauraModel # isort:skip +from ppsci.arch.chip_deeponets import ChipDeepONets # isort:skip +from ppsci.arch.crystalgraphconvnet import CrystalGraphConvNet # isort:skip +from ppsci.arch.cuboid_transformer import CuboidTransformer # isort:skip +from ppsci.arch.cvit import CVit # isort:skip +from ppsci.arch.cvit import CVit1D # isort:skip +from ppsci.arch.deeponet import DeepONet # isort:skip +from ppsci.arch.dgmr import DGMR # isort:skip +from ppsci.arch.embedding_koopman import CylinderEmbedding # isort:skip +from ppsci.arch.embedding_koopman import LorenzEmbedding # isort:skip +from ppsci.arch.embedding_koopman import RosslerEmbedding # isort:skip +from ppsci.arch.epnn import Epnn # isort:skip +from ppsci.arch.extformer_moe_cuboid import ExtFormerMoECuboid # isort:skip +from ppsci.arch.gan import Discriminator # isort:skip +from ppsci.arch.gan import Generator # isort:skip +from ppsci.arch.geofno import FNO1d # isort:skip +from ppsci.arch.graphcast import GraphCastNet # isort:skip +from ppsci.arch.he_deeponets import HEDeepONets # isort:skip +from ppsci.arch.lno import LNO # isort:skip +from ppsci.arch.mlp import MLP # isort:skip +from ppsci.arch.mlp import ModifiedMLP # isort:skip +from ppsci.arch.mlp import PirateNet # isort:skip +from ppsci.arch.model_list import ModelList # isort:skip +from ppsci.arch.nowcastnet import NowcastNet # isort:skip +from ppsci.arch.phycrnet import PhyCRNet # isort:skip +from ppsci.arch.phylstm import DeepPhyLSTM # isort:skip +from ppsci.arch.physx_transformer import PhysformerGPT2 # isort:skip +from ppsci.arch.sfnonet import SFNONet # isort:skip +from ppsci.arch.spinn import SPINN # isort:skip +from ppsci.arch.tfnonet import TFNO1dNet, TFNO2dNet, TFNO3dNet # isort:skip +from ppsci.arch.transformer import Transformer # isort:skip +from ppsci.arch.unetex import UNetEx # isort:skip +from ppsci.arch.unonet import UNONet # isort:skip +from ppsci.arch.uscnn import USCNN # isort:skip +from ppsci.arch.vae import AutoEncoder # isort:skip +from ppsci.arch.velocitygan import VelocityDiscriminator # isort:skip +from ppsci.arch.velocitygan import VelocityGenerator # isort:skip +from ppsci.arch.moflow_net import MoFlowNet, MoFlowProp # isort:skip +from ppsci.utils import logger # isort:skip +from ppsci.arch.regdgcnn import RegDGCNN # isort:skip +from ppsci.arch.regpointnet import RegPointNet # isort:skip +from ppsci.arch.ifm_mlp import IFMMLP # isort:skip + +__all__ = [ + "MoFlowNet", + "MoFlowProp", + "AFNONet", + "AMGNet", + "Arch", + "AutoEncoder", + "build_model", + "CFDGCN", + "SuzukiMiyauraModel", + "ChipDeepONets", + "CrystalGraphConvNet", + "CuboidTransformer", + "CVit", + "CVit1D", + "CylinderEmbedding", + "DeepONet", + "DeepPhyLSTM", + "DGMR", + "Discriminator", + "Epnn", + "ExtFormerMoECuboid", + "FNO1d", + "Generator", + "GraphCastNet", + "HEDeepONets", + "LorenzEmbedding", + "LNO", + "MLP", + "ModelList", + "ModifiedMLP", + "NowcastNet", + "PhyCRNet", + "PhysformerGPT2", + "PirateNet", + "PrecipNet", + "RosslerEmbedding", + "SFNONet", + "SPINN", + "TFNO1dNet", + "TFNO2dNet", + "TFNO3dNet", + "Transformer", + "UNetEx", + "UNONet", + "USCNN", + "VelocityDiscriminator", + "VelocityGenerator", + "RegDGCNN", + "RegPointNet", + "IFMMLP", +] + + +def build_model(cfg): + """Build model + + Args: + cfg (DictConfig): Arch config. + + Returns: + nn.Layer: Model. + """ + cfg = copy.deepcopy(cfg) + arch_cls = cfg.pop("name") + arch = eval(arch_cls)(**cfg) + + logger.debug(str(arch)) + + return arch diff --git a/examples/smc_reac/ppsci/arch/activation.py b/examples/smc_reac/ppsci/arch/activation.py new file mode 100644 index 0000000000..3f78eb1a61 --- /dev/null +++ b/examples/smc_reac/ppsci/arch/activation.py @@ -0,0 +1,160 @@ +# 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. + +from __future__ import annotations + +from typing import Callable + +import numpy as np +import paddle +import paddle.nn.functional as F +from paddle import nn + +from ppsci.utils import initializer +from ppsci.utils import misc + + +class Stan(nn.Layer): + """Self-scalable Tanh. + paper: https://arxiv.org/abs/2204.12589v1 + + Args: + out_features (int, optional): Output features. Defaults to 1. + """ + + def __init__(self, out_features: int = 1): + super().__init__() + self.beta = self.create_parameter( + shape=(out_features,), + default_initializer=nn.initializer.Constant(1), + ) + + def forward(self, x): + # TODO: manually broadcast beta to x.shape for preventing backward error yet. + return F.tanh(x) * (1 + paddle.broadcast_to(self.beta, x.shape) * x) + # return F.tanh(x) * (1 + self.beta * x) + + +class Swish(nn.Layer): + def __init__(self, beta: float = 1.0): + super().__init__() + self.beta = self.create_parameter( + shape=[], + default_initializer=nn.initializer.Constant(beta), + ) + + def forward(self, x): + return x * F.sigmoid(self.beta * x) + + +class Cos(nn.Layer): + def __init__(self): + super().__init__() + + def forward(self, x): + return paddle.cos(x) + + +class Sin(nn.Layer): + def __init__(self): + super().__init__() + + def forward(self, x): + return paddle.sin(x) + + +class Siren(nn.Layer): + """Implicit Neural Representations with Periodic Activation Functions. + paper link: https://arxiv.org/abs/2006.09661 + code ref: https://github.com/vsitzmann/siren/tree/master + """ + + def __init__(self, w0: float = 30): + super().__init__() + self.w0 = w0 + + def forward(self, x): + return paddle.sin(self.w0 * x) + + @staticmethod + def init_for_first_layer(layer: nn.Linear): + """Initialization only for first hidden layer. + ref: https://github.com/vsitzmann/siren/blob/master/modules.py#L630 + """ + if not isinstance(layer, nn.Linear): + raise TypeError( + "Siren initialization only support Linear layer now, " + f"but got {misc.typename(layer)}" + ) + in_features = layer.weight.shape[0] + with paddle.no_grad(): + initializer.uniform_(layer.weight, -1 / in_features, 1 / in_features) + initializer.zeros_(layer.bias) + + @staticmethod + def init_for_hidden_layer(layer: nn.Linear, w0: float = 30): + """Initialization for hidden layer except first layer. + ref: https://github.com/vsitzmann/siren/blob/master/modules.py#L622 + """ + if not isinstance(layer, nn.Linear): + raise TypeError( + "Siren initialization only support Linear layer now, " + f"but got {misc.typename(layer)}" + ) + in_features = layer.weight.shape[0] + with paddle.no_grad(): + initializer.uniform_( + layer.weight, + -np.sqrt(6 / in_features) / w0, + np.sqrt(6 / in_features) / w0, + ) + initializer.zeros_(layer.bias) + + +act_func_dict = { + "elu": nn.ELU(), + "relu": nn.ReLU(), + "selu": nn.SELU(), + "gelu": nn.GELU(), + "leaky_relu": nn.LeakyReLU(), + "sigmoid": nn.Sigmoid(), + "silu": nn.Silu(), + "sin": Sin(), + "cos": Cos(), + "swish": Swish, + "tanh": nn.Tanh(), + "identity": nn.Identity(), + "siren": Siren(), + "stan": Stan, +} + + +def get_activation(act_name: str) -> Callable: + """Get activation function according to act_name. + + Args: + act_name (str): Name of activation, such as "tanh". + + Returns: + Callable: Paddle activation function. + """ + if act_name.lower() not in act_func_dict: + raise ValueError(f"act_name({act_name}) not found in act_func_dict") + + act_layer = act_func_dict[act_name.lower()] + if isinstance(act_layer, type) and act_name != "stan": + # Is a activation class but not a instance of it, instantiate manually(except for 'Stan') + return act_layer() + + return act_layer diff --git a/examples/smc_reac/ppsci/arch/afno.py b/examples/smc_reac/ppsci/arch/afno.py new file mode 100644 index 0000000000..62ec6fdd7c --- /dev/null +++ b/examples/smc_reac/ppsci/arch/afno.py @@ -0,0 +1,687 @@ +# 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. + +""" +Code below is heavily based on [FourCastNet](https://github.com/NVlabs/FourCastNet) +""" +from __future__ import annotations + +from functools import partial +from typing import Optional +from typing import Tuple + +import paddle +import paddle.fft +import paddle.nn.functional as F +from paddle import nn + +from ppsci.arch import activation as act_mod +from ppsci.arch import base +from ppsci.utils import initializer + + +def drop_path( + x: paddle.Tensor, + drop_prob: float = 0.0, + training: bool = False, + scale_by_keep: bool = True, +) -> paddle.Tensor: + """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks). + the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper... + See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... + + Args: + x (paddle.Tensor): The tensor to apply. + drop_prob (float, optional): Drop paths probability. Defaults to 0.0. + training (bool, optional): Whether at training mode. Defaults to False. + scale_by_keep (bool, optional): Whether upscale the output. Defaults to True. + + Returns: + paddle.Tensor: Output tensor after apply dropout. + """ + if drop_prob == 0.0 or not training: + return x + keep_prob = 1 - drop_prob + shape = (x.shape[0],) + (1,) * (x.ndim - 1) + random_tensor = paddle.full(shape, keep_prob, x.dtype) + random_tensor = paddle.bernoulli(random_tensor) + if keep_prob > 0.0 and scale_by_keep: + random_tensor = random_tensor / keep_prob + return x * random_tensor + + +class DropPath(nn.Layer): + """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks). + + Args: + drop_prob (float, optional): Drop paths probability. Defaults to 0.0. + scale_by_keep (bool, optional): Whether upscale the output. Defaults to True. + """ + + def __init__(self, drop_prob: float = 0.0, scale_by_keep: bool = True): + super().__init__() + self.drop_prob = drop_prob + self.scale_by_keep = scale_by_keep + + def forward(self, x): + return drop_path(x, self.drop_prob, self.training, self.scale_by_keep) + + def extra_repr(self): + return f"drop_prob={round(self.drop_prob,3):0.3f}" + + +class PeriodicPad2d(nn.Layer): + """Pad longitudinal (left-right) circular and pad latitude (top-bottom) with zeros. + + Args: + pad (int): Number of pad. + """ + + def __init__(self, pad: int): + super(PeriodicPad2d, self).__init__() + self.pad = pad + + def forward(self, x): + # pad left and right circular + out = F.pad(x, (self.pad, self.pad, 0, 0), mode="circular") + # pad top and bottom zeros + out = F.pad( + out, + (0, 0, 0, 0, self.pad, self.pad, 0, 0), + mode="constant", + value=0, + ) + return out + + +class MLP(nn.Layer): + """Multi layer perceptron module used in Transformer. + + Args: + in_features (int): Number of the input features. + hidden_features (Optional[int]): Number of the hidden size. Defaults to None. + out_features (Optional[int]): Number of the output features. Defaults to None. + activation (str, optional): Name of activation function. Defaults to "gelu". + drop (float, optional): Probability of dropout the units. Defaults to 0.0. + """ + + def __init__( + self, + in_features: int, + hidden_features: Optional[int] = None, + out_features: Optional[int] = None, + activation: str = "gelu", + drop: float = 0.0, + ): + super().__init__() + out_features = out_features or in_features + hidden_features = hidden_features or in_features + self.fc1 = nn.Linear(in_features, hidden_features) + self.act = act_mod.get_activation(activation) + self.fc2 = nn.Linear(hidden_features, out_features) + self.drop = nn.Dropout(drop) + + def forward(self, x): + x = self.fc1(x) + x = self.act(x) + x = self.drop(x) + x = self.fc2(x) + x = self.drop(x) + return x + + +class AFNO2D(nn.Layer): + """2D Adaptive Fourier Neural Operators. + + Args: + hidden_size (int): Number of hidden size. + num_blocks (int, optional): Number of blocks. Defaults to 8. + sparsity_threshold (float, optional): The value of threshold for softshrink. Defaults to 0.01. + hard_thresholding_fraction (float, optional): The value of threshold for keep mode. Defaults to 1.0. + hidden_size_factor (int, optional): The factor of hidden size. Defaults to 1. + scale (float, optional): The scale factor of the parameter when initialization. Defaults to 0.02. + """ + + def __init__( + self, + hidden_size: int, + num_blocks: int = 8, + sparsity_threshold: float = 0.01, + hard_thresholding_fraction: float = 1.0, + hidden_size_factor: int = 1, + scale: float = 0.02, + ): + super().__init__() + if hidden_size % num_blocks != 0: + raise ValueError( + f"hidden_size({hidden_size}) should be divisible by num_blocks({num_blocks})." + ) + + self.hidden_size = hidden_size + self.sparsity_threshold = sparsity_threshold + self.num_blocks = num_blocks + self.block_size = self.hidden_size // self.num_blocks + self.hard_thresholding_fraction = hard_thresholding_fraction + self.hidden_size_factor = hidden_size_factor + self.scale = scale + + self.w1 = self.create_parameter( + shape=( + 2, + self.num_blocks, + self.block_size, + self.block_size * self.hidden_size_factor, + ), + default_initializer=nn.initializer.Normal(std=self.scale), + ) + self.b1 = self.create_parameter( + shape=(2, self.num_blocks, self.block_size * self.hidden_size_factor), + default_initializer=nn.initializer.Normal(std=self.scale), + ) + self.w2 = self.create_parameter( + shape=( + 2, + self.num_blocks, + self.block_size * self.hidden_size_factor, + self.block_size, + ), + default_initializer=nn.initializer.Normal(std=self.scale), + ) + self.b2 = self.create_parameter( + shape=(2, self.num_blocks, self.block_size), + default_initializer=nn.initializer.Normal(std=self.scale), + ) + + def forward(self, x): + bias = x + + B, H, W, C = x.shape + + x = paddle.fft.rfft2(x, axes=(1, 2), norm="ortho") + x = x.reshape((B, H, W // 2 + 1, self.num_blocks, self.block_size)) + + o1_shape = ( + B, + H, + W // 2 + 1, + self.num_blocks, + self.block_size * self.hidden_size_factor, + ) + o1_real = paddle.zeros(o1_shape) + o1_imag = paddle.zeros(o1_shape) + o2_real = paddle.zeros(x.shape) + o2_imag = paddle.zeros(x.shape) + + total_modes = H // 2 + 1 + kept_modes = int(total_modes * self.hard_thresholding_fraction) + + st, end = total_modes - kept_modes, total_modes + kept_modes + + o1_real[:, st:end, :kept_modes] = F.relu( + paddle.einsum( + "xyzbi,bio->xyzbo", + x[:, st:end, :kept_modes].real(), + self.w1[0], + ) + - paddle.einsum( + "xyzbi,bio->xyzbo", + x[:, st:end, :kept_modes].imag(), + self.w1[1], + ) + + self.b1[0] + ) + + o1_imag[:, st:end, :kept_modes] = F.relu( + paddle.einsum( + "xyzbi,bio->xyzbo", + x[:, st:end, :kept_modes].imag(), + self.w1[0], + ) + + paddle.einsum( + "xyzbi,bio->xyzbo", + x[:, st:end, :kept_modes].real(), + self.w1[1], + ) + + self.b1[1] + ) + + o2_real[:, st:end, :kept_modes] = ( + paddle.einsum( + "xyzbi,bio->xyzbo", + o1_real[:, st:end, :kept_modes], + self.w2[0], + ) + - paddle.einsum( + "xyzbi,bio->xyzbo", + o1_imag[:, st:end, :kept_modes], + self.w2[1], + ) + + self.b2[0] + ) + + o2_imag[:, st:end, :kept_modes] = ( + paddle.einsum( + "xyzbi,bio->xyzbo", + o1_imag[:, st:end, :kept_modes], + self.w2[0], + ) + + paddle.einsum( + "xyzbi,bio->xyzbo", + o1_real[:, st:end, :kept_modes], + self.w2[1], + ) + + self.b2[1] + ) + + x = paddle.stack([o2_real, o2_imag], axis=-1) + x = F.softshrink(x, threshold=self.sparsity_threshold) + x = paddle.as_complex(x) + x = x.reshape((B, H, W // 2 + 1, C)) + x = paddle.fft.irfft2(x, s=(H, W), axes=(1, 2), norm="ortho") + + return x + bias + + +class Block(nn.Layer): + """AFNO network block. + + Args: + dim (int): The input tensor dimension. + mlp_ratio (float, optional): The ratio used in MLP. Defaults to 4.0. + drop (float, optional): The drop ratio used in MLP. Defaults to 0.0. + drop_path (float, optional): The drop ratio used in DropPath. Defaults to 0.0. + activation (str, optional): Name of activation function. Defaults to "gelu". + norm_layer (nn.Layer, optional): Class of norm layer. Defaults to nn.LayerNorm. + double_skip (bool, optional): Whether use double skip. Defaults to True. + num_blocks (int, optional): The number of blocks. Defaults to 8. + sparsity_threshold (float, optional): The value of threshold for softshrink. Defaults to 0.01. + hard_thresholding_fraction (float, optional): The value of threshold for keep mode. Defaults to 1.0. + """ + + def __init__( + self, + dim: int, + mlp_ratio: float = 4.0, + drop: float = 0.0, + drop_path: float = 0.0, + activation: str = "gelu", + norm_layer: nn.Layer = nn.LayerNorm, + double_skip: bool = True, + num_blocks: int = 8, + sparsity_threshold: float = 0.01, + hard_thresholding_fraction: float = 1.0, + ): + super().__init__() + self.norm1 = norm_layer(dim) + self.filter = AFNO2D( + dim, num_blocks, sparsity_threshold, hard_thresholding_fraction + ) + self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() + + self.norm2 = norm_layer(dim) + mlp_hidden_dim = int(dim * mlp_ratio) + self.mlp = MLP( + in_features=dim, + hidden_features=mlp_hidden_dim, + activation=activation, + drop=drop, + ) + self.double_skip = double_skip + + def forward(self, x): + residual = x + x = self.norm1(x) + x = self.filter(x) + + if self.double_skip: + x = x + residual + residual = x + + x = self.norm2(x) + x = self.mlp(x) + x = self.drop_path(x) + x = x + residual + return x + + +class PatchEmbed(nn.Layer): + """Patch embedding module. + + Args: + img_size (Tuple[int, ...], optional): Image size. Defaults to (224, 224). + patch_size (Tuple[int, ...], optional): Patch size. Defaults to (16, 16). + in_channels (int, optional): The input tensor channels. Defaults to 3. + embed_dim (int, optional): The output tensor channels. Defaults to 768. + """ + + def __init__( + self, + img_size: Tuple[int, ...] = (224, 224), + patch_size: Tuple[int, ...] = (16, 16), + in_channels: int = 3, + embed_dim: int = 768, + ): + super().__init__() + num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0]) + self.img_size = img_size + self.patch_size = patch_size + self.num_patches = num_patches + self.proj = nn.Conv2D( + in_channels, embed_dim, kernel_size=patch_size, stride=patch_size + ) + + def forward(self, x): + _, _, H, W = x.shape + if not (H == self.img_size[0] and W == self.img_size[1]): + raise ValueError( + f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})." + ) + x = self.proj(x).flatten(2).transpose((0, 2, 1)) + return x + + +class AFNONet(base.Arch): + """Adaptive Fourier Neural Network. + + Args: + input_keys (Tuple[str, ...]): Name of input keys, such as ("input",). + output_keys (Tuple[str, ...]): Name of output keys, such as ("output",). + img_size (Tuple[int, ...], optional): Image size. Defaults to (720, 1440). + patch_size (Tuple[int, ...], optional): Path. Defaults to (8, 8). + in_channels (int, optional): The input tensor channels. Defaults to 20. + out_channels (int, optional): The output tensor channels. Defaults to 20. + embed_dim (int, optional): The embedding dimension for PatchEmbed. Defaults to 768. + depth (int, optional): Number of transformer depth. Defaults to 12. + mlp_ratio (float, optional): Number of ratio used in MLP. Defaults to 4.0. + drop_rate (float, optional): The drop ratio used in MLP. Defaults to 0.0. + drop_path_rate (float, optional): The drop ratio used in DropPath. Defaults to 0.0. + num_blocks (int, optional): Number of blocks. Defaults to 8. + sparsity_threshold (float, optional): The value of threshold for softshrink. Defaults to 0.01. + hard_thresholding_fraction (float, optional): The value of threshold for keep mode. Defaults to 1.0. + num_timestamps (int, optional): Number of timestamp. Defaults to 1. + + Examples: + >>> import ppsci + >>> model = ppsci.arch.AFNONet(("input", ), ("output", )) + >>> input_data = {"input": paddle.randn([1, 20, 720, 1440])} + >>> output_data = model(input_data) + >>> for k, v in output_data.items(): + ... print(k, v.shape) + output [1, 20, 720, 1440] + """ + + def __init__( + self, + input_keys: Tuple[str, ...], + output_keys: Tuple[str, ...], + img_size: Tuple[int, ...] = (720, 1440), + patch_size: Tuple[int, ...] = (8, 8), + in_channels: int = 20, + out_channels: int = 20, + embed_dim: int = 768, + depth: int = 12, + mlp_ratio: float = 4.0, + drop_rate: float = 0.0, + drop_path_rate: float = 0.0, + num_blocks: int = 8, + sparsity_threshold: float = 0.01, + hard_thresholding_fraction: float = 1.0, + num_timestamps: int = 1, + ): + super().__init__() + self.input_keys = input_keys + self.output_keys = output_keys + + self.img_size = img_size + self.patch_size = patch_size + self.in_channels = in_channels + self.out_channels = out_channels + self.embed_dim = embed_dim + self.num_blocks = num_blocks + self.num_timestamps = num_timestamps + norm_layer = partial(nn.LayerNorm, epsilon=1e-6) + + self.patch_embed = PatchEmbed( + img_size=img_size, + patch_size=self.patch_size, + in_channels=self.in_channels, + embed_dim=embed_dim, + ) + num_patches = self.patch_embed.num_patches + + data = paddle.zeros((1, num_patches, embed_dim)) + data = initializer.trunc_normal_(data, std=0.02) + self.pos_embed = paddle.create_parameter( + shape=data.shape, + dtype=data.dtype, + default_initializer=nn.initializer.Assign(data), + ) + self.pos_drop = nn.Dropout(p=drop_rate) + + dpr = [x.item() for x in paddle.linspace(0, drop_path_rate, depth)] + + self.h = img_size[0] // self.patch_size[0] + self.w = img_size[1] // self.patch_size[1] + + self.blocks = nn.LayerList( + [ + Block( + dim=embed_dim, + mlp_ratio=mlp_ratio, + drop=drop_rate, + drop_path=dpr[i], + norm_layer=norm_layer, + num_blocks=self.num_blocks, + sparsity_threshold=sparsity_threshold, + hard_thresholding_fraction=hard_thresholding_fraction, + ) + for i in range(depth) + ] + ) + + self.norm = norm_layer(embed_dim) + self.head = nn.Linear( + embed_dim, + self.out_channels * self.patch_size[0] * self.patch_size[1], + bias_attr=False, + ) + + self.apply(self._init_weights) + + def _init_weights(self, m): + if isinstance(m, nn.Linear): + initializer.trunc_normal_(m.weight, std=0.02) + if m.bias is not None: + initializer.zeros_(m.bias) + elif isinstance(m, nn.LayerNorm): + initializer.ones_(m.weight) + initializer.zeros_(m.bias) + elif isinstance(m, nn.Conv2D): + initializer.conv_init_(m) + + def forward_tensor(self, x): + B = x.shape[0] + x = self.patch_embed(x) + x = x + self.pos_embed + x = self.pos_drop(x) + + x = x.reshape((B, self.h, self.w, self.embed_dim)) + for block in self.blocks: + x = block(x) + + x = self.head(x) + + b = x.shape[0] + p1 = self.patch_size[0] + p2 = self.patch_size[1] + h = self.img_size[0] // self.patch_size[0] + w = self.img_size[1] // self.patch_size[1] + c_out = x.shape[3] // (p1 * p2) + x = x.reshape((b, h, w, p1, p2, c_out)) + x = x.transpose((0, 5, 1, 3, 2, 4)) + x = x.reshape((b, c_out, h * p1, w * p2)) + + return x + + @staticmethod + def split_to_dict(data_tensors: Tuple[paddle.Tensor, ...], keys: Tuple[str, ...]): + return {key: data_tensors[i] for i, key in enumerate(keys)} + + def forward(self, x): + if self._input_transform is not None: + x = self._input_transform(x) + + x_tensor = self.concat_to_tensor(x, self.input_keys) + + y = [] + input = x_tensor + for _ in range(self.num_timestamps): + out = self.forward_tensor(input) + y.append(out) + input = out + y = self.split_to_dict(y, self.output_keys) + + if self._output_transform is not None: + y = self._output_transform(x, y) + return y + + +class PrecipNet(base.Arch): + """Precipitation Network. + + Args: + input_keys (Tuple[str, ...]): Name of input keys, such as ("input",). + output_keys (Tuple[str, ...]): Name of output keys, such as ("output",). + wind_model (base.Arch): Wind model. + img_size (Tuple[int, ...], optional): Image size. Defaults to (720, 1440). + patch_size (Tuple[int, ...], optional): Path. Defaults to (8, 8). + in_channels (int, optional): The input tensor channels. Defaults to 20. + out_channels (int, optional): The output tensor channels. Defaults to 1. + embed_dim (int, optional): The embedding dimension for PatchEmbed. Defaults to 768. + depth (int, optional): Number of transformer depth. Defaults to 12. + mlp_ratio (float, optional): Number of ratio used in MLP. Defaults to 4.0. + drop_rate (float, optional): The drop ratio used in MLP. Defaults to 0.0. + drop_path_rate (float, optional): The drop ratio used in DropPath. Defaults to 0.0. + num_blocks (int, optional): Number of blocks. Defaults to 8. + sparsity_threshold (float, optional): The value of threshold for softshrink. Defaults to 0.01. + hard_thresholding_fraction (float, optional): The value of threshold for keep mode. Defaults to 1.0. + num_timestamps (int, optional): Number of timestamp. Defaults to 1. + + Examples: + >>> import ppsci + >>> wind_model = ppsci.arch.AFNONet(("input", ), ("output", )) + >>> model = ppsci.arch.PrecipNet(("input", ), ("output", ), wind_model) + >>> data = paddle.randn([1, 20, 720, 1440]) + >>> data_dict = {"input": data} + >>> output = model.forward(data_dict) + >>> print(output['output'].shape) + [1, 1, 720, 1440] + """ + + def __init__( + self, + input_keys: Tuple[str, ...], + output_keys: Tuple[str, ...], + wind_model: base.Arch, + img_size: Tuple[int, ...] = (720, 1440), + patch_size: Tuple[int, ...] = (8, 8), + in_channels: int = 20, + out_channels: int = 1, + embed_dim: int = 768, + depth: int = 12, + mlp_ratio: float = 4.0, + drop_rate: float = 0.0, + drop_path_rate: float = 0.0, + num_blocks: int = 8, + sparsity_threshold: float = 0.01, + hard_thresholding_fraction: float = 1.0, + num_timestamps=1, + ): + super().__init__() + self.input_keys = input_keys + self.output_keys = output_keys + + self.img_size = img_size + self.patch_size = patch_size + self.in_channels = in_channels + self.out_channels = out_channels + self.embed_dim = embed_dim + self.num_blocks = num_blocks + self.num_timestamps = num_timestamps + self.backbone = AFNONet( + ("input",), + ("output",), + img_size=img_size, + patch_size=patch_size, + in_channels=in_channels, + out_channels=out_channels, + embed_dim=embed_dim, + depth=depth, + mlp_ratio=mlp_ratio, + drop_rate=drop_rate, + drop_path_rate=drop_path_rate, + num_blocks=num_blocks, + sparsity_threshold=sparsity_threshold, + hard_thresholding_fraction=hard_thresholding_fraction, + ) + self.ppad = PeriodicPad2d(1) + self.conv = nn.Conv2D( + self.out_channels, self.out_channels, kernel_size=3, stride=1, padding=0 + ) + self.act = nn.ReLU() + self.apply(self._init_weights) + self.wind_model = wind_model + self.wind_model.eval() + + def _init_weights(self, m): + if isinstance(m, nn.Linear): + initializer.trunc_normal_(m.weight, std=0.02) + if m.bias is not None: + initializer.zeros_(m.bias) + elif isinstance(m, nn.LayerNorm): + initializer.ones_(m.weight) + initializer.zeros_(m.bias) + elif isinstance(m, nn.Conv2D): + initializer.conv_init_(m) + + def forward_tensor(self, x): + x = self.backbone.forward_tensor(x) + x = self.ppad(x) + x = self.conv(x) + x = self.act(x) + return x + + @staticmethod + def split_to_dict(data_tensors: Tuple[paddle.Tensor, ...], keys: Tuple[str, ...]): + return {key: data_tensors[i] for i, key in enumerate(keys)} + + def forward(self, x): + if self._input_transform is not None: + x = self._input_transform(x) + + x_tensor = self.concat_to_tensor(x, self.input_keys) + + input_wind = x_tensor + y = [] + for _ in range(self.num_timestamps): + with paddle.no_grad(): + out_wind = self.wind_model.forward_tensor(input_wind) + out = self.forward_tensor(out_wind) + y.append(out) + input_wind = out_wind + y = self.split_to_dict(y, self.output_keys) + + if self._output_transform is not None: + y = self._output_transform(x, y) + return y diff --git a/examples/smc_reac/ppsci/arch/amgnet.py b/examples/smc_reac/ppsci/arch/amgnet.py new file mode 100644 index 0000000000..ce728317d6 --- /dev/null +++ b/examples/smc_reac/ppsci/arch/amgnet.py @@ -0,0 +1,649 @@ +# 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. + +from __future__ import annotations + +import functools +from typing import Callable +from typing import Dict +from typing import List +from typing import Optional +from typing import Tuple + +import numpy as np +import paddle +import paddle.nn as nn +from typing_extensions import Literal + +try: + import pgl +except ModuleNotFoundError: + pass + +try: + import pyamg +except ModuleNotFoundError: + pass + +from paddle import sparse as pd_sparse +from scipy import sparse as sci_sparse + + +def _knn_interpolate( + features: paddle.Tensor, coarse_nodes: paddle.Tensor, fine_nodes: paddle.Tensor +) -> paddle.Tensor: + coarse_nodes_input = paddle.repeat_interleave( + coarse_nodes.unsqueeze(0), fine_nodes.shape[0], axis=0 + ) # [6684,352,2] + fine_nodes_input = paddle.repeat_interleave( + fine_nodes.unsqueeze(1), coarse_nodes.shape[0], axis=1 + ) # [6684,352,2] + dist_w = 1.0 / ( + paddle.norm(x=coarse_nodes_input - fine_nodes_input, p=2, axis=-1) + 1e-9 + ) # [6684,352] + knn_value, knn_index = paddle.topk(dist_w, k=3, largest=True) # [6684,3],[6684,3] + weight = knn_value.unsqueeze(-2) + features_input = features[knn_index] + output = paddle.bmm(weight, features_input).squeeze(-2) / paddle.sum( + knn_value, axis=-1, keepdim=True + ) + return output + + +def _get_corse_node(latent_graph: "pgl.Graph") -> paddle.Tensor: + row = latent_graph.edge_index[0].numpy() + col = latent_graph.edge_index[1].numpy() + data = paddle.ones(shape=[row.size]).numpy() + A = sci_sparse.coo_matrix((data, (row, col))).tocsr() + splitting = pyamg.classical.split.RS(A) + index = np.array(np.nonzero(splitting)) + b = paddle.to_tensor(index) + b = paddle.squeeze(b) + return b + + +def StAS( + index_A: paddle.Tensor, + value_A: paddle.Tensor, + index_S: paddle.Tensor, + value_S: paddle.Tensor, + N: int, + kN: int, + norm_layer: nn.Layer, +) -> Tuple[paddle.Tensor, paddle.Tensor]: + """ASAP: Adaptive Structure Aware Pooling for Learning Hierarchical Graph Representations. + Ranjan, E., Sanyal, S., Talukdar, P. (2020, April). AAAI(2020) + + Args: + index_A (paddle.Tensor): Indices of sparse matrix A. + value_A (paddle.Tensor): Values of sparse matrix A. + index_S (paddle.Tensor): Indices of sparse matrix S. + value_S (paddle.Tensor): Values of sparse matrix S. + N (int): Dimension N. + kN (int): Dimension kN. + norm_layer (nn.Layer): Normalization layer. + + Returns: + Tuple[paddle.Tensor, paddle.Tensor]: Indices and values of result matrix E. + """ + sp_x = pd_sparse.sparse_coo_tensor(index_A, value_A) + sp_x = pd_sparse.coalesce(sp_x) + index_A = sp_x.indices() + value_A = sp_x.values() + + sp_s = pd_sparse.sparse_coo_tensor(index_S, value_S) + sp_s = pd_sparse.coalesce(sp_s) + index_S = sp_s.indices() + value_S = sp_s.values() + + indices_A = index_A.numpy() + values_A = value_A.numpy() + coo_A = sci_sparse.coo_matrix( + (values_A, (indices_A[0], indices_A[1])), shape=(N, N) + ) + + indices_S = index_S.numpy() + values_S = value_S.numpy() + coo_S = sci_sparse.coo_matrix( + (values_S, (indices_S[0], indices_S[1])), shape=(N, kN) + ) + + ans = coo_A.dot(coo_S).tocoo() + row = paddle.to_tensor(ans.row) + col = paddle.to_tensor(ans.col) + index_B = paddle.stack([row, col], axis=0) + value_B = paddle.to_tensor(ans.data) + + indices_A = index_S + values_A = value_S + coo_A = pd_sparse.sparse_coo_tensor(indices_A, values_A) + out = pd_sparse.transpose(coo_A, [1, 0]) + index_St = out.indices() + value_St = out.values() + + sp_x = pd_sparse.sparse_coo_tensor(index_B, value_B) + sp_x = pd_sparse.coalesce(sp_x) + index_B = sp_x.indices() + value_B = sp_x.values() + + indices_A = index_St.numpy() + values_A = value_St.numpy() + coo_A = sci_sparse.coo_matrix( + (values_A, (indices_A[0], indices_A[1])), shape=(kN, N) + ) + + indices_S = index_B.numpy() + values_S = value_B.numpy() + coo_S = sci_sparse.coo_matrix( + (values_S, (indices_S[0], indices_S[1])), shape=(N, kN) + ) + + ans = coo_A.dot(coo_S).tocoo() + row = paddle.to_tensor(ans.row) + col = paddle.to_tensor(ans.col) + index_E = paddle.stack([row, col], axis=0) + value_E = paddle.to_tensor(ans.data) + + # index_E排序 + sp_x = pd_sparse.sparse_coo_tensor(index_E, value_E) + sp_x = pd_sparse.coalesce(sp_x) + index_E = sp_x.indices() + value_E = sp_x.values() + + return index_E.astype("int64"), value_E + + +def FillZeros( + index_E: paddle.Tensor, value_E: paddle.Tensor, standard_index, kN: int +) -> Tuple[paddle.Tensor, paddle.Tensor]: + shape = [kN, kN] + row_E = index_E[0] + col_E = index_E[1] + DenseMatrix_E = sci_sparse.coo_matrix( + (paddle.ones_like(value_E), (row_E, col_E)), shape + ).toarray() + + row_S = standard_index[0] + col_S = standard_index[1] + DenseMatrix_S = sci_sparse.coo_matrix( + (paddle.ones([row_S.shape[0]]), (row_S, col_S)), shape + ).toarray() + + diff = DenseMatrix_S - DenseMatrix_E + rows, cols = np.nonzero(diff) + rows = paddle.to_tensor(rows, dtype="int32") + cols = paddle.to_tensor(cols, dtype="int32") + index = paddle.stack([rows, cols], axis=0) + value = paddle.zeros([index.shape[1]]) + index_E = paddle.concat([index_E, index], axis=1) + value_E = paddle.concat([value_E, value], axis=-1) + + sp_x = pd_sparse.sparse_coo_tensor(index_E, value_E) + sp_x = pd_sparse.coalesce(sp_x) + index_E = sp_x.indices() + value_E = sp_x.values() + + return index_E.astype("int64"), value_E + + +def remove_self_loops( + edge_index: paddle.Tensor, edge_attr: Optional[paddle.Tensor] = None +) -> Tuple[paddle.Tensor, Optional[paddle.Tensor]]: + # remove self-loop + mask = edge_index[0] != edge_index[1] + mask = mask.tolist() + edge_index = edge_index.t() + edge_index = edge_index[mask] + edge_index = edge_index.t() + if edge_attr is None: + return edge_index, None + else: + return edge_index, edge_attr[mask] + + +def faster_graph_connectivity(perm, edge_index, edge_weight, score, pos, N, norm_layer): + """ + Adapted from Ranjan, E., Sanyal, S., Talukdar, P. (2020, April). Asap: Adaptive structure aware pooling + for learning hierarchical graph representations. AAAI(2020) + """ + + kN = perm.shape[0] + perm2 = perm.reshape((-1, 1)) + mask = (edge_index[0] == perm2).sum(axis=0).astype("bool") + + S0 = edge_index[1][mask].reshape((1, -1)) + S1 = edge_index[0][mask].reshape((1, -1)) + index_S = paddle.concat([S0, S1], axis=0) + value_S = score[mask].detach().squeeze() + n_idx = paddle.zeros([N], dtype=paddle.int64) + n_idx[perm] = paddle.arange(perm.shape[0]) + index_S = index_S.astype("int64") + index_S[1] = n_idx[index_S[1]] + subgraphnode_pos = pos[perm] + index_A = edge_index.clone() + if edge_weight is None: + value_A = value_S.new_ones(edge_index[0].shape[0]) + else: + value_A = edge_weight.clone() + + value_A = paddle.squeeze(value_A) + model_1 = nn.Sequential( + ("l1", nn.Linear(128, 256)), + ("act1", nn.ReLU()), + ("l2", nn.Linear(256, 256)), + ("act2", nn.ReLU()), + ("l4", nn.Linear(256, 128)), + ("act4", nn.ReLU()), + ("l5", nn.Linear(128, 1)), + ) + model_2 = nn.Sequential( + ("l1", nn.Linear(1, 64)), + ("act1", nn.ReLU()), + ("l2", nn.Linear(64, 128)), + ("act2", nn.ReLU()), + ("l4", nn.Linear(128, 128)), + ) + + val_A = model_1(value_A) + val_A = paddle.squeeze(val_A) + index_E, value_E = StAS(index_A, val_A, index_S, value_S, N, kN, norm_layer) + value_E = paddle.reshape(value_E, shape=[-1, 1]) + edge_weight = model_2(value_E) + + return index_E, edge_weight, subgraphnode_pos + + +def norm_graph_connectivity(perm, edge_index, edge_weight, score, pos, N, norm_layer): + """ + Come from Ranjan, E., Sanyal, S., Talukdar, P. (2020, April). Asap: Adaptive + structure aware pooling for learning hierarchical graph representations. AAAI(2020) + """ + + kN = perm.shape[0] + perm2 = perm.reshape((-1, 1)) + mask = (edge_index[0] == perm2).sum(axis=0).astype("bool") + S0 = edge_index[1][mask].reshape((1, -1)) + S1 = edge_index[0][mask].reshape((1, -1)) + + index_S = paddle.concat([S0, S1], axis=0) + value_S = score[mask].detach().squeeze() + n_idx = paddle.zeros([N], dtype=paddle.int64) + n_idx[perm] = paddle.arange(perm.shape[0]) + + index_S = index_S.astype("int64") + index_S[1] = n_idx[index_S[1]] + subgraphnode_pos = pos[perm] + index_A = edge_index.clone() + + if edge_weight is None: + value_A = value_S.new_ones(edge_index[0].shape[0]) + else: + value_A = edge_weight.clone() + + value_A = paddle.squeeze(value_A) + eps_mask = (value_S == 0).astype(paddle.get_default_dtype()) + value_S = paddle.full_like(value_S, 1e-4) * eps_mask + (1 - eps_mask) * value_S + attrlist = [] + standard_index, _ = StAS( + index_A, + paddle.ones_like(value_A[:, 0]), + index_S, + paddle.ones_like(value_S), + N, + kN, + norm_layer, + ) + for i in range(128): + mask = (value_A[:, i] == 0).astype(paddle.get_default_dtype()) + val_A = paddle.full_like(mask, 1e-4) * mask + (1 - mask) * value_A[:, i] + index_E, value_E = StAS(index_A, val_A, index_S, value_S, N, kN, norm_layer) + + if index_E.shape[1] != standard_index.shape[1]: + index_E, value_E = FillZeros(index_E, value_E, standard_index, kN) + + index_E, value_E = remove_self_loops(edge_index=index_E, edge_attr=value_E) + attrlist.append(value_E) + edge_weight = paddle.stack(attrlist, axis=1) + + return index_E, edge_weight, subgraphnode_pos + + +class GraphNetBlock(nn.Layer): + """Multi-Edge Interaction Network with residual connections.""" + + def __init__( + self, model_fn, output_dim, message_passing_aggregator, attention=False + ): + super().__init__() + self.edge_model = model_fn(output_dim, 384) + self.node_model = model_fn(output_dim, 256) + self.message_passing_aggregator = message_passing_aggregator + + def _update_edge_features(self, graph): + """Aggregates node features, and applies edge function.""" + senders = graph.edge_index[0] + receivers = graph.edge_index[1] + sender_features = paddle.index_select(x=graph.x, index=senders, axis=0) + receiver_features = paddle.index_select(x=graph.x, index=receivers, axis=0) + features = [sender_features, receiver_features, graph.edge_attr] + features = paddle.concat(features, axis=-1) + return self.edge_model(features) + + def unsorted_segment_operation(self, data, segment_ids, num_segments, operation): + """Computes the sum along segments of a tensor. Analogous to tf.unsorted_segment_sum. + + Args: + data (paddle.Tensor): A tensor whose segments are to be summed. + segment_ids (paddle.Tensor): The segment indices tensor. + num_segments (int): The number of segments. + operation (str): _description_ + + Returns: + paddle.Tensor: A tensor of same data type as the data argument. + """ + if not all([i in data.shape for i in segment_ids.shape]): + raise ValueError("segment_ids.shape should be a prefix of data.shape") + + if not (data.shape[0] == segment_ids.shape[0]): + raise ValueError("data.shape and segment_ids.shape should be equal") + + shape = [num_segments] + list(data.shape[1:]) + result_shape = paddle.zeros(shape) + if operation == "sum": + result = paddle.scatter(result_shape, segment_ids, data, overwrite=False) + return result + + def _update_node_features(self, node_features, edge_attr, edge_index): + """Aggregates edge features, and applies node function.""" + num_nodes = node_features.shape[0] + features = [node_features] + features.append( + self.unsorted_segment_operation( + edge_attr, + edge_index[1], + num_nodes, + operation=self.message_passing_aggregator, + ) + ) + features = paddle.concat(features, axis=-1) + return self.node_model(features) + + def forward(self, graph): + """Applies GraphNetBlock and returns updated MultiGraph.""" + new_edge_features = self._update_edge_features(graph) + new_node_features = self._update_node_features( + graph.x, graph.edge_attr, graph.edge_index + ) + + new_node_features += graph.x + new_edge_features += graph.edge_attr + latent_graph = pgl.Graph( + num_nodes=new_node_features.shape[0], edges=graph.edge_index + ) + latent_graph.x = new_node_features + latent_graph.edge_attr = new_edge_features + latent_graph.pos = graph.pos + latent_graph.edge_index = graph.edge_index + return latent_graph + + +class Processor(nn.Layer): + """This class takes the nodes with the most influential feature (sum of square) + The the chosen numbers of nodes in each ripple will establish connection(features and distances) with the most influential nodes and this connection will be learned + Then the result is add to output latent graph of encoder and the modified latent graph will be feed into original processor + + Args: + make_mlp (Callable): Function to make MLP. + output_dim (int): Number of dimension of output. + message_passing_steps (int): Message passing steps. + message_passing_aggregator (str): Message passing aggregator. + attention (bool, optional): Whether use attention. Defaults to False. + use_stochastic_message_passing (bool, optional): Whether use stochastic message passing. Defaults to False. + """ + + # Each mesh can be coarsened to have no fewer points than this value + min_nodes = 2000 + + def __init__( + self, + make_mlp: Callable, + output_dim: int, + message_passing_steps: int, + message_passing_aggregator: str, + attention: bool = False, + use_stochastic_message_passing: bool = False, + ): + super().__init__() + self.use_stochastic_message_passing = use_stochastic_message_passing + self.graphnet_blocks = nn.LayerList() + self.cofe_edge_blocks = nn.LayerList() + self.pool_blocks = nn.LayerList() + self.latent_dim = output_dim + self.normalization = nn.LayerNorm(128) + for index in range(message_passing_steps): + self.graphnet_blocks.append( + GraphNetBlock( + model_fn=make_mlp, + output_dim=output_dim, + message_passing_aggregator=message_passing_aggregator, + attention=attention, + ) + ) + + self.pool_blocks.append( + GraphNetBlock( + model_fn=make_mlp, + output_dim=output_dim, + message_passing_aggregator=message_passing_aggregator, + attention=attention, + ) + ) + + def forward(self, latent_graph, speed, normalized_adj_mat=None): + x = [] + pos = [] + new = [] + for graphnet_block, pool in zip(self.graphnet_blocks, self.pool_blocks): + if latent_graph.x.shape[0] > self.min_nodes: + pre_matrix = graphnet_block(latent_graph) + x.append(pre_matrix) + cofe_graph = pool(pre_matrix) + coarsenodes = _get_corse_node(pre_matrix) + nodesfeatures = cofe_graph.x[coarsenodes] + if speed == "fast": + subedge_index, edge_weight, subpos = faster_graph_connectivity( + perm=coarsenodes, + edge_index=cofe_graph.edge_index, + edge_weight=cofe_graph.edge_attr, + score=cofe_graph.edge_attr[:, 0], + pos=cofe_graph.pos, + N=cofe_graph.x.shape[0], + norm_layer=self.normalization, + ) + elif speed == "norm": + subedge_index, edge_weight, subpos = norm_graph_connectivity( + perm=coarsenodes, + edge_index=cofe_graph.edge_index, + edge_weight=cofe_graph.edge_attr, + score=cofe_graph.edge_attr[:, 0], + pos=cofe_graph.pos, + N=cofe_graph.x.shape[0], + norm_layer=self.normalization, + ) + else: + raise ValueError( + f"Argument 'speed' should be 'sum' or 'fast', bot got {speed}." + ) + edge_weight = self.normalization(edge_weight) + pos.append(subpos) + latent_graph = pgl.Graph( + num_nodes=nodesfeatures.shape[0], edges=subedge_index + ) + latent_graph.x = nodesfeatures + latent_graph.edge_attr = edge_weight + latent_graph.pos = subpos + latent_graph.edge_index = subedge_index + else: + latent_graph = graphnet_block(latent_graph) + new.append(latent_graph) + if len(new): + x.append(new[-1]) + return x, pos + + +class FullyConnectedLayer(nn.Layer): + def __init__(self, input_dim: int, hidden_size: Tuple[int, ...]): + super(FullyConnectedLayer, self).__init__() + num_layers = len(hidden_size) + self._layers_ordered_dict = {} + self.in_dim = input_dim + for index, output_dim in enumerate(hidden_size): + self._layers_ordered_dict["linear_" + str(index)] = nn.Linear( + self.in_dim, output_dim + ) + if index < (num_layers - 1): + self._layers_ordered_dict["relu_" + str(index)] = nn.ReLU() + self.in_dim = output_dim + + self.layers = nn.LayerDict(self._layers_ordered_dict) + + def forward(self, input): + for key in self.layers: + layer = self.layers[key] + output = layer(input) + input = output + return input + + +class Encoder(nn.Layer): + """Encodes node and edge features into latent features.""" + + def __init__(self, input_dim, make_mlp, latent_dim): + super(Encoder, self).__init__() + self._make_mlp = make_mlp + self._latent_dim = latent_dim + self.node_model = self._make_mlp(latent_dim, input_dim=input_dim) + self.mesh_edge_model = self._make_mlp(latent_dim, input_dim=1) + + def forward(self, graph): + node_latents = self.node_model(graph.x) + edge_latent = self.mesh_edge_model(graph.edge_attr) + + graph.x = node_latents + graph.edge_attr = edge_latent + return graph + + +class Decoder(nn.Layer): + """Decodes node features from graph. + Encodes node and edge features into latent features. + """ + + def __init__(self, make_mlp, output_dim): + super(Decoder, self).__init__() + self.model = make_mlp(output_dim, 128) + + def forward(self, node_features): + return self.model(node_features) + + +class AMGNet(nn.Layer): + """A Multi-scale Graph neural Network model + based on Encoder-Process-Decoder structure for flow field prediction. + + https://doi.org/10.1080/09540091.2022.2131737 + + Code reference: https://github.com/baoshiaijhin/amgnet + + Args: + input_keys (Tuple[str, ...]): Name of input keys, such as ("input", ). + output_keys (Tuple[str, ...]): Name of output keys, such as ("pred", ). + input_dim (int): Number of input dimension. + output_dim (int): Number of output dimension. + latent_dim (int): Number of hidden(feature) dimension. + num_layers (int): Number of layer(s). + message_passing_aggregator (Literal["sum"]): Message aggregator method in graph. + Only "sum" available now. + message_passing_steps (int): Message passing steps in graph. + speed (str): Whether use vanilla method or fast method for graph_connectivity + computation. + + Examples: + >>> import ppsci + >>> model = ppsci.arch.AMGNet( + ... ("input", ), ("pred", ), 5, 3, 64, 2, "sum", 6, "norm", + ... ) + """ + + def __init__( + self, + input_keys: Tuple[str, ...], + output_keys: Tuple[str, ...], + input_dim: int, + output_dim: int, + latent_dim: int, + num_layers: int, + message_passing_aggregator: Literal["sum"], + message_passing_steps: int, + speed: Literal["norm", "fast"], + ): + super().__init__() + self.input_keys = input_keys + self.output_keys = output_keys + self._latent_dim = latent_dim + self.speed = speed + self._output_dim = output_dim + self._num_layers = num_layers + + self.encoder = Encoder(input_dim, self._make_mlp, latent_dim=self._latent_dim) + self.processor = Processor( + make_mlp=self._make_mlp, + output_dim=self._latent_dim, + message_passing_steps=message_passing_steps, + message_passing_aggregator=message_passing_aggregator, + use_stochastic_message_passing=False, + ) + self.post_processor = self._make_mlp(self._latent_dim, 128) + self.decoder = Decoder( + make_mlp=functools.partial(self._make_mlp, layer_norm=False), + output_dim=self._output_dim, + ) + + def forward(self, x: Dict[str, "pgl.Graph"]) -> Dict[str, paddle.Tensor]: + graphs = x[self.input_keys[0]] + latent_graph = self.encoder(graphs) + x, p = self.processor(latent_graph, speed=self.speed) + node_features = self._spa_compute(x, p) + pred_field = self.decoder(node_features) + return {self.output_keys[0]: pred_field} + + def _make_mlp(self, output_dim: int, input_dim: int = 5, layer_norm: bool = True): + widths = (self._latent_dim,) * self._num_layers + (output_dim,) + network = FullyConnectedLayer(input_dim, widths) + if layer_norm: + network = nn.Sequential(network, nn.LayerNorm(normalized_shape=widths[-1])) + return network + + def _spa_compute(self, x: List["pgl.Graph"], p): + j = len(x) - 1 + node_features = x[j].x + + for k in range(1, j + 1): + pos = p[-k] + fine_nodes = x[-(k + 1)].pos + feature = _knn_interpolate(node_features, pos, fine_nodes) + node_features = x[-(k + 1)].x + feature + node_features = self.post_processor(node_features) + + return node_features diff --git a/examples/smc_reac/ppsci/arch/base.py b/examples/smc_reac/ppsci/arch/base.py new file mode 100644 index 0000000000..5b51efec22 --- /dev/null +++ b/examples/smc_reac/ppsci/arch/base.py @@ -0,0 +1,279 @@ +# 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. + +from __future__ import annotations + +from typing import Callable +from typing import Dict +from typing import Tuple + +import numpy as np +import paddle +from paddle import nn + +from ppsci.utils import logger + + +class Arch(nn.Layer): + """Base class for Network.""" + + input_keys: Tuple[str, ...] + output_keys: Tuple[str, ...] + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._input_transform: Callable[ + [Dict[str, paddle.Tensor]], Dict[str, paddle.Tensor] + ] = None + + self._output_transform: Callable[ + [Dict[str, paddle.Tensor], Dict[str, paddle.Tensor]], + Dict[str, paddle.Tensor], + ] = None + + def forward(self, *args, **kwargs): + raise NotImplementedError("Arch.forward is not implemented") + + @property + def num_params(self) -> int: + """Return number of parameters within network. + + Returns: + int: Number of parameters. + """ + num = 0 + for name, param in self.named_parameters(): + if hasattr(param, "shape"): + num += np.prod(list(param.shape), dtype="int") + else: + logger.warning(f"{name} has no attribute 'shape'") + return num + + @property + def num_buffers(self) -> int: + """Return number of buffers within network. + + Returns: + int: Number of buffers. + """ + num = 0 + for name, buffer in self.named_buffers(): + if hasattr(buffer, "shape"): + num += np.prod(list(buffer.shape), dtype="int") + else: + logger.warning(f"{name} has no attribute 'shape'") + return num + + @staticmethod + def concat_to_tensor( + data_dict: Dict[str, paddle.Tensor], keys: Tuple[str, ...], axis=-1 + ) -> Tuple[paddle.Tensor, ...]: + """Concatenate tensors from dict in the order of given keys. + + Args: + data_dict (Dict[str, paddle.Tensor]): Dict contains tensor. + keys (Tuple[str, ...]): Keys tensor fetched from. + axis (int, optional): Axis concatenate at. Defaults to -1. + + Returns: + Tuple[paddle.Tensor, ...]: Concatenated tensor. + + Examples: + >>> import paddle + >>> import ppsci + >>> model = ppsci.arch.Arch() + >>> # fetch one tensor + >>> out = model.concat_to_tensor({'x':paddle.rand([64, 64, 1])}, ('x',)) + >>> print(out.dtype, out.shape) + paddle.float32 [64, 64, 1] + >>> # fetch more tensors + >>> out = model.concat_to_tensor( + ... {'x1':paddle.rand([64, 64, 1]), 'x2':paddle.rand([64, 64, 1])}, + ... ('x1', 'x2'), + ... axis=2) + >>> print(out.dtype, out.shape) + paddle.float32 [64, 64, 2] + + """ + if len(keys) == 1: + return data_dict[keys[0]] + data = [data_dict[key] for key in keys] + return paddle.concat(data, axis) + + @staticmethod + def split_to_dict( + data_tensor: paddle.Tensor, keys: Tuple[str, ...], axis=-1 + ) -> Dict[str, paddle.Tensor]: + """Split tensor and wrap into a dict by given keys. + + Args: + data_tensor (paddle.Tensor): Tensor to be split. + keys (Tuple[str, ...]): Keys tensor mapping to. + axis (int, optional): Axis split at. Defaults to -1. + + Returns: + Dict[str, paddle.Tensor]: Dict contains tensor. + + Examples: + >>> import paddle + >>> import ppsci + >>> model = ppsci.arch.Arch() + >>> # split one tensor + >>> out = model.split_to_dict(paddle.rand([64, 64, 1]), ('x',)) + >>> for k, v in out.items(): + ... print(f"{k} {v.dtype} {v.shape}") + x paddle.float32 [64, 64, 1] + >>> # split more tensors + >>> out = model.split_to_dict(paddle.rand([64, 64, 2]), ('x1', 'x2'), axis=2) + >>> for k, v in out.items(): + ... print(f"{k} {v.dtype} {v.shape}") + x1 paddle.float32 [64, 64, 1] + x2 paddle.float32 [64, 64, 1] + + """ + if len(keys) == 1: + return {keys[0]: data_tensor} + data = paddle.split(data_tensor, len(keys), axis=axis) + return {key: data[i] for i, key in enumerate(keys)} + + def register_input_transform( + self, + transform: Callable[[Dict[str, paddle.Tensor]], Dict[str, paddle.Tensor]], + ): + """Register input transform. + + Args: + transform (Callable[[Dict[str, paddle.Tensor]], Dict[str, paddle.Tensor]]): + Input transform of network, receive a single tensor dict and return a single tensor dict. + + Examples: + >>> import ppsci + >>> def transform_in(in_): + ... x = in_["x"] + ... # transform input + ... x_ = 2.0 * x + ... input_trans = {"2x": x_} + ... return input_trans + >>> # `MLP` inherits from `Arch` + >>> model = ppsci.arch.MLP( + ... input_keys=("2x",), + ... output_keys=("y",), + ... num_layers=5, + ... hidden_size=32) + >>> model.register_input_transform(transform_in) + >>> out = model({"x":paddle.rand([64, 64, 1])}) + >>> for k, v in out.items(): + ... print(f"{k} {v.dtype} {v.shape}") + y paddle.float32 [64, 64, 1] + + """ + self._input_transform = transform + + def register_output_transform( + self, + transform: Callable[ + [Dict[str, paddle.Tensor], Dict[str, paddle.Tensor]], + Dict[str, paddle.Tensor], + ], + ): + """Register output transform. + + Args: + transform (Callable[[Dict[str, paddle.Tensor], Dict[str, paddle.Tensor]], Dict[str, paddle.Tensor]]): + Output transform of network, receive two single tensor dict(raw input + and raw output) and return a single tensor dict(transformed output). + + Examples: + >>> import ppsci + >>> def transform_out(in_, out): + ... x = in_["x"] + ... y = out["y"] + ... u = 2.0 * x * y + ... output_trans = {"u": u} + ... return output_trans + >>> # `MLP` inherits from `Arch` + >>> model = ppsci.arch.MLP( + ... input_keys=("x",), + ... output_keys=("y",), + ... num_layers=5, + ... hidden_size=32) + >>> model.register_output_transform(transform_out) + >>> out = model({"x":paddle.rand([64, 64, 1])}) + >>> for k, v in out.items(): + ... print(f"{k} {v.dtype} {v.shape}") + u paddle.float32 [64, 64, 1] + + """ + self._output_transform = transform + + def freeze(self): + """Freeze all parameters. + + Examples: + >>> import ppsci + >>> model = ppsci.arch.Arch() + >>> # freeze all parameters and make model `eval` + >>> model.freeze() + >>> assert not model.training + >>> for p in model.parameters(): + ... assert p.stop_gradient + + """ + for param in self.parameters(): + param.stop_gradient = True + + self.eval() + + def unfreeze(self): + """Unfreeze all parameters. + + Examples: + >>> import ppsci + >>> model = ppsci.arch.Arch() + >>> # unfreeze all parameters and make model `train` + >>> model.unfreeze() + >>> assert model.training + >>> for p in model.parameters(): + ... assert not p.stop_gradient + + """ + for param in self.parameters(): + param.stop_gradient = False + + self.train() + + def __str__(self): + num_fc = 0 + num_conv = 0 + num_bn = 0 + for layer in self.sublayers(include_self=True): + if isinstance(layer, nn.Linear): + num_fc += 1 + elif isinstance(layer, (nn.Conv2D, nn.Conv3D, nn.Conv1D)): + num_conv += 1 + elif isinstance(layer, (nn.BatchNorm, nn.BatchNorm2D, nn.BatchNorm3D)): + num_bn += 1 + + return ", ".join( + [ + self.__class__.__name__, + f"input_keys = {self.input_keys}", + f"output_keys = {self.output_keys}", + f"num_fc = {num_fc}", + f"num_conv = {num_conv}", + f"num_bn = {num_bn}", + f"num_params = {self.num_params}", + f"num_buffers = {self.num_buffers}", + ] + ) diff --git a/examples/smc_reac/ppsci/arch/cfdgcn.py b/examples/smc_reac/ppsci/arch/cfdgcn.py new file mode 100644 index 0000000000..da43d3dc8f --- /dev/null +++ b/examples/smc_reac/ppsci/arch/cfdgcn.py @@ -0,0 +1,350 @@ +# 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. + +import os +from typing import Callable +from typing import Dict +from typing import List +from typing import Optional +from typing import Sequence +from typing import Tuple +from typing import TypeVar +from typing import Union + +import numpy as np +import paddle +from paddle import nn +from paddle.nn import functional as F + +from ppsci.data.dataset import airfoil_dataset + +try: + import pgl +except ModuleNotFoundError: + pass + +GenTensor = TypeVar("GenTensor", paddle.Tensor, np.ndarray) + +SU2_SHAPE_IDS = { + "line": 3, + "triangle": 5, + "quad": 9, +} + + +def _knn_interpolate( + features: paddle.Tensor, coarse_nodes: paddle.Tensor, fine_nodes: paddle.Tensor +) -> paddle.Tensor: + coarse_nodes_input = paddle.repeat_interleave( + coarse_nodes.unsqueeze(0), fine_nodes.shape[0], axis=0 + ) # [6684,352,2] + fine_nodes_input = paddle.repeat_interleave( + fine_nodes.unsqueeze(1), coarse_nodes.shape[0], axis=1 + ) # [6684,352,2] + dist_w = 1.0 / ( + paddle.norm(x=coarse_nodes_input - fine_nodes_input, p=2, axis=-1) + 1e-9 + ) # [6684,352] + knn_value, knn_index = paddle.topk(dist_w, k=3, largest=True) # [6684,3],[6684,3] + weight = knn_value.unsqueeze(-2) + features_input = features[knn_index] + output = paddle.bmm(weight, features_input).squeeze(-2) / paddle.sum( + knn_value, axis=-1, keepdim=True + ) + return output + + +def is_cw( + points: paddle.Tensor, triangles: paddle.Tensor, ret_val=False +) -> Union[bool, paddle.Tensor]: + tri_pts = points[triangles] + a = tri_pts[:, 0] - tri_pts[:, 1] + b = tri_pts[:, 1] - tri_pts[:, 2] + cross = b[:, 0] * a[:, 1] - b[:, 1] * a[:, 0] + + if not ret_val: + return cross > 0 + else: + return cross + + +def left_orthogonal(v: paddle.Tensor) -> paddle.Tensor: + return paddle.stack([-v[..., 1], v[..., 0]], axis=-1) + + +def signed_dist_graph( + nodes: paddle.Tensor, marker_inds, with_sign=False +) -> paddle.Tensor: + # assumes shape is convex + # approximate signed distance by distance to closest point on surface + signed_dists = paddle.zeros([nodes.shape[0]], dtype=paddle.float32) + marker_nodes = nodes[marker_inds] + if type(marker_inds) is paddle.Tensor: + marker_inds = marker_inds.tolist() + marker_inds = set(marker_inds) + + if with_sign: + marker_surfaces = marker_nodes[:-1] - marker_nodes[1:] + last_surface = marker_nodes[-1] - marker_nodes[0] + marker_surfaces = paddle.concat([marker_surfaces, last_surface.unsqueeze(0)]) + normals = left_orthogonal(marker_surfaces) / marker_surfaces.norm( + axis=1 + ).unsqueeze(1) + for i, x in enumerate(nodes): + if i not in marker_inds: + vecs = marker_nodes - x + dists = paddle.linalg.norm(vecs, axis=1) + min_dist = dists.min() + + if with_sign: + # if sign is requested, check if inside marker shape + # dot product with normals to find if inside shape + surface_dists = (vecs * normals).sum(axis=1) + if (surface_dists < 0).unique().shape[0] == 1: + # if all point in same direction it is inside + min_dist *= -1 + + signed_dists[i] = min_dist + return signed_dists + + +def quad2tri(elems: np.array) -> Tuple[List[int], Union[List[int], paddle.Tensor]]: + new_elems = [] + new_edges = [] + for e in elems: + if len(e) <= 3: + new_elems.append(e) + else: + new_elems.append([e[0], e[1], e[2]]) + new_elems.append([e[0], e[2], e[3]]) + new_edges.append(paddle.to_tensor([[e[0]], [e[2]]], dtype=paddle.int64)) + new_edges = ( + paddle.concat(new_edges, axis=1) + if new_edges + else paddle.to_tensor([], dtype=paddle.int64) + ) + return new_elems, new_edges + + +def write_graph_mesh( + output_filename: str, + points: GenTensor, + elems_list: Sequence[Sequence[Sequence[int]]], + marker_dict: Dict[str, Sequence[Sequence[int]]], + dims: int = 2, +) -> None: + def seq2str(s: Sequence[int]) -> str: + return " ".join(str(x) for x in s) + + with open(output_filename, "w") as f: + f.write(f"NDIME={dims}\n") + + num_points = points.shape[0] + f.write(f"NPOIN={num_points}\n") + for i, p in enumerate(points): + f.write(f"{seq2str(p.tolist())} {i}\n") + f.write("\n") + + num_elems = sum([len(elems) for elems in elems_list]) + f.write(f"NELEM={num_elems}\n") + for elems in elems_list: + for e in elems: + if len(e) != 3 and len(e) != 4: + raise ValueError( + f"Meshes only support triangles and quadrilaterals, " + f"passed element had {len(e)} vertices." + ) + elem_id = ( + SU2_SHAPE_IDS["triangle"] if len(e) == 3 else SU2_SHAPE_IDS["quad"] + ) + f.write(f"{elem_id} {seq2str(e)}\n") + f.write("\n") + + num_markers = len(marker_dict) + f.write(f"NMARK={num_markers}\n") + for marker_tag in marker_dict: + f.write(f"MARKER_TAG={marker_tag}\n") + marker_elems = marker_dict[marker_tag] + f.write(f"MARKER_ELEMS={len(marker_elems)}\n") + for m in marker_elems: + f.write(f'{SU2_SHAPE_IDS["line"]} {seq2str(m)}\n') + f.write("\n") + + +class CFDGCN(nn.Layer): + """Graph Neural Networks for Fluid Flow Prediction. + + [Filipe De Avila Belbute-Peres, Thomas Economon, Zico Kolter Proceedings of the 37th International Conference on Machine Learning, PMLR 119:2402-2411, 2020.](https://proceedings.mlr.press/v119/de-avila-belbute-peres20a.html) + + Code reference: https://github.com/locuslab/cfd-gcn + + Args: + input_keys (Tuple[str, ...]): Name of input keys, such as ("input", ). + output_keys (Tuple[str, ...]): Name of output keys, such as ("pred", ). + config_file (str): Name of configuration file for su2 module. + coarse_mesh (str): Path of coarse mesh file. + fine_marker_dict (Dict[str, List[List[int]]]): Dict of fine marker. + process_sim (Callable, optional): Preprocess function. Defaults to `lambda x, y: x`. + freeze_mesh (bool, optional): Whether set `stop_gradient=True` for nodes. Defaults to False. + num_convs (int, optional): Number of conv layers. Defaults to 6. + num_end_convs (int, optional): Number of end conv layers. Defaults to 3. + hidden_channel (int, optional): Number of channels of hidden layer. Defaults to 512. + out_channel (int, optional): Number of channels of output. Defaults to 3. + su2_module (Optional[Callable]): SU2Module Object. Defaults to None. + + Examples: + >>> import ppsci + >>> import su2paddle # doctest: +SKIP + >>> model = ppsci.arch.CFDGCN( + ... input_keys=("input"), + ... output_keys=("pred"), + ... config_file="/path/to/file.cfg", + ... coarse_mesh="/path/to/file.su2", + ... process_sim=None, + ... freeze_mesh=False, + ... num_convs=6, + ... num_end_convs=3, + ... hidden_channel=512, + ... out_channel=3, + ... su2_module=su2paddle.SU2Module, + ... ) # doctest: +SKIP + """ + + def __init__( + self, + input_keys: Tuple[str, ...], + output_keys: Tuple[str, ...], + config_file: str, + coarse_mesh: str, + fine_marker_dict: Dict[str, List[List[int]]], + process_sim: Callable = lambda x, y: x, + freeze_mesh: bool = False, + num_convs: int = 6, + num_end_convs: int = 3, + hidden_channel: int = 512, + out_channel: int = 3, + su2_module: Optional[Callable] = None, + ): + + super().__init__() + self.input_keys = input_keys + self.output_keys = output_keys + meshes_temp_dir = "temp_meshes" + os.makedirs(meshes_temp_dir, exist_ok=True) + self.mesh_file = os.path.join(meshes_temp_dir, f"{str(os.getpid())}_mesh.su2") + + if not coarse_mesh: + raise ValueError("Need to provide a coarse mesh for CFD-GCN.") + nodes, edges, self.elems, self.marker_dict = airfoil_dataset._get_mesh_graph( + coarse_mesh + ) + if not freeze_mesh: + self.nodes = paddle.to_tensor(nodes, stop_gradient=False) + else: + self.nodes = paddle.to_tensor(nodes, stop_gradient=True) + + self.elems, new_edges = quad2tri(sum(self.elems, [])) + self.elems = [self.elems] + + if is_cw(self.nodes, paddle.to_tensor(self.elems[0])).nonzero().shape[0] != 0: + raise ("Mesh has flipped elems.") + + self.edges = paddle.to_tensor(edges) + self.edges = paddle.concat([self.edges, new_edges], axis=1) + self.marker_inds = paddle.to_tensor(sum(self.marker_dict.values(), [])).unique() + self.fine_marker_dict = paddle.to_tensor(fine_marker_dict["airfoil"]).unique() + self.process_sim = process_sim + + self.write_mesh_file( + self.nodes, self.elems, self.marker_dict, filename=self.mesh_file + ) + self.su2 = su2_module(config_file, mesh_file=self.mesh_file) + self.sdf = None + + self.num_convs = num_end_convs + self.convs = [] + if self.num_convs > 0: + self.convs = nn.LayerList() + in_channels = out_channel + hidden_channel + for i in range(self.num_convs - 1): + self.convs.append(pgl.nn.GCNConv(in_channels, hidden_channel)) + in_channels = hidden_channel + self.convs.append(pgl.nn.GCNConv(in_channels, out_channel)) + + self.num_pre_convs = num_convs - num_end_convs + self.pre_convs = [] + if self.num_pre_convs > 0: + in_channels = 5 + 1 # one extra channel for sdf + self.pre_convs = nn.LayerList() + for i in range(self.num_pre_convs - 1): + self.pre_convs.append(pgl.nn.GCNConv(in_channels, hidden_channel)) + in_channels = hidden_channel + self.pre_convs.append(pgl.nn.GCNConv(in_channels, hidden_channel)) + + def forward(self, x: Dict[str, "pgl.Graph"]) -> Dict[str, paddle.Tensor]: + graph = x[self.input_keys[0]] + batch_size = graph.shape[0] + x_list = paddle.split(graph.x, batch_size) + fine_x_list = [] + + for idx in range(batch_size): + x = x_list[idx] + if self.sdf is None: + with paddle.no_grad(): + self.sdf = signed_dist_graph( + x[:, :2], self.fine_marker_dict + ).unsqueeze(1) + fine_x = paddle.concat([x, self.sdf], axis=1) + for conv in self.pre_convs: + fine_x = F.relu(conv(graph, fine_x)) + fine_x_list.append(fine_x) + nodes_input = self.get_nodes().tile([batch_size, 1, 1]) + + batch_y = self.su2( + nodes_input[..., 0], + nodes_input[..., 1], + graph.aoa[..., None], + graph.mach_or_reynolds[..., None], + ) + batch_y = self.process_sim(batch_y, False) + + pred_fields = [] + for idx in range(batch_size): + coarse_y = paddle.stack([y[idx].flatten() for y in batch_y], axis=1).astype( + "float32" + ) + nodes = self.get_nodes() + x = x_list[idx] + fine_y = _knn_interpolate( + features=coarse_y, coarse_nodes=nodes[:, :2], fine_nodes=x[:, :2] + ) + fine_y = paddle.concat([fine_y, fine_x_list[idx]], axis=1) + + for conv in self.convs[:-1]: + fine_y = F.relu(conv(graph, fine_y)) + fine_y = self.convs[-1](graph, fine_y) + pred_fields.append(fine_y) + pred_fields = paddle.concat(pred_fields, axis=0) + return {self.output_keys[0]: pred_fields} + + def get_nodes(self) -> paddle.Tensor: + return self.nodes + + @staticmethod + def write_mesh_file( + x: paddle.Tensor, + elems: paddle.Tensor, + marker_dict: Dict[str, Sequence[Sequence[int]]], + filename: str = "mesh.su2", + ) -> None: + write_graph_mesh(filename, x[:, :2], elems, marker_dict) diff --git a/examples/smc_reac/ppsci/arch/chip_deeponets.py b/examples/smc_reac/ppsci/arch/chip_deeponets.py new file mode 100644 index 0000000000..30c87c5656 --- /dev/null +++ b/examples/smc_reac/ppsci/arch/chip_deeponets.py @@ -0,0 +1,214 @@ +# 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. + +from __future__ import annotations + +from typing import Tuple +from typing import Union + +import paddle +import paddle.nn as nn + +from ppsci.arch import activation as act_mod +from ppsci.arch import base +from ppsci.arch import mlp + + +class ChipDeepONets(base.Arch): + """Multi-branch physics-informed deep operator neural network. The network consists of three branch networks: random heat source, boundary function, and boundary type, as well as a trunk network. + + Args: + branch_input_keys (Tuple[str, ...]): Name of input data for internal heat source on branch nets. + BCtype_input_keys (Tuple[str, ...]): Name of input data for boundary types on branch nets. + BC_input_keys (Tuple[str, ...]): Name of input data for boundary on branch nets. + trunk_input_keys (Tuple[str, ...]): Name of input data for trunk net. + output_keys (Tuple[str, ...]): Output name of predicted temperature. + num_loc (int): Number of sampled input data for internal heat source. + bctype_loc (int): Number of sampled input data for boundary types. + BC_num_loc (int): Number of sampled input data for boundary. + num_features (int): Number of features extracted from trunk net, same for all branch nets. + branch_num_layers (int): Number of hidden layers of internal heat source on branch nets. + BC_num_layers (int): Number of hidden layers of boundary on branch nets. + trunk_num_layers (int): Number of hidden layers of trunk net. + branch_hidden_size (Union[int, Tuple[int, ...]]): Number of hidden size of internal heat source on branch nets. + An integer for all layers, or list of integer specify each layer's size. + BC_hidden_size (Union[int, Tuple[int, ...]]): Number of hidden size of boundary on branch nets. + An integer for all layers, or list of integer specify each layer's size. + trunk_hidden_size (Union[int, Tuple[int, ...]]): Number of hidden size of trunk net. + An integer for all layers, or list of integer specify each layer's size. + branch_skip_connection (bool, optional): Whether to use skip connection for internal heat source on branch net. Defaults to False. + BC_skip_connection (bool, optional): Whether to use skip connection for boundary on branch net. Defaults to False. + trunk_skip_connection (bool, optional): Whether to use skip connection for trunk net. Defaults to False. + branch_activation (str, optional): Name of activation function for internal heat source on branch net. Defaults to "tanh". + BC_activation (str, optional): Name of activation function for boundary on branch net. Defaults to "tanh". + trunk_activation (str, optional): Name of activation function for trunk net. Defaults to "tanh". + branch_weight_norm (bool, optional): Whether to apply weight norm on parameter(s) for internal heat source on branch net. Defaults to False. + BC_weight_norm (bool, optional): Whether to apply weight norm on parameter(s) for boundary on branch net. Defaults to False. + trunk_weight_norm (bool, optional): Whether to apply weight norm on parameter(s) for trunk net. Defaults to False. + use_bias (bool, optional): Whether to add bias on predicted G(u)(y). Defaults to True. + + Examples: + >>> import ppsci + >>> model = ppsci.arch.ChipDeepONets( + ... ('u',), + ... ('bc',), + ... ('bc_data',), + ... ("x",'y'), + ... ("T",), + ... 324, + ... 1, + ... 76, + ... 400, + ... 9, + ... 9, + ... 6, + ... 256, + ... 256, + ... 128, + ... branch_activation="swish", + ... BC_activation="swish", + ... trunk_activation="swish", + ... use_bias=True, + ... ) + """ + + def __init__( + self, + branch_input_keys: Tuple[str, ...], + BCtype_input_keys: Tuple[str, ...], + BC_input_keys: Tuple[str, ...], + trunk_input_keys: Tuple[str, ...], + output_keys: Tuple[str, ...], + num_loc: int, + bctype_loc: int, + BC_num_loc: int, + num_features: int, + branch_num_layers: int, + BC_num_layers: int, + trunk_num_layers: int, + branch_hidden_size: Union[int, Tuple[int, ...]], + BC_hidden_size: Union[int, Tuple[int, ...]], + trunk_hidden_size: Union[int, Tuple[int, ...]], + branch_skip_connection: bool = False, + BC_skip_connection: bool = False, + trunk_skip_connection: bool = False, + branch_activation: str = "tanh", + BC_activation: str = "tanh", + trunk_activation: str = "tanh", + branch_weight_norm: bool = False, + BC_weight_norm: bool = False, + trunk_weight_norm: bool = False, + use_bias: bool = True, + ): + super().__init__() + self.trunk_input_keys = trunk_input_keys + self.branch_input_keys = branch_input_keys + self.BCtype_input_keys = BCtype_input_keys + self.BC_input_keys = BC_input_keys + self.input_keys = ( + self.trunk_input_keys + + self.branch_input_keys + + self.BC_input_keys + + self.BCtype_input_keys + ) + self.output_keys = output_keys + + self.branch_net = mlp.MLP( + self.branch_input_keys, + ("b",), + branch_num_layers, + branch_hidden_size, + branch_activation, + branch_skip_connection, + branch_weight_norm, + input_dim=num_loc, + output_dim=num_features, + ) + + self.BCtype_net = mlp.MLP( + self.BCtype_input_keys, + ("bctype",), + BC_num_layers, + BC_hidden_size, + BC_activation, + BC_skip_connection, + BC_weight_norm, + input_dim=bctype_loc, + output_dim=num_features, + ) + + self.BC_net = mlp.MLP( + self.BC_input_keys, + ("bc",), + BC_num_layers, + BC_hidden_size, + BC_activation, + BC_skip_connection, + BC_weight_norm, + input_dim=BC_num_loc, + output_dim=num_features, + ) + + self.trunk_net = mlp.MLP( + self.trunk_input_keys, + ("t",), + trunk_num_layers, + trunk_hidden_size, + trunk_activation, + trunk_skip_connection, + trunk_weight_norm, + input_dim=len(self.trunk_input_keys), + output_dim=num_features, + ) + self.trunk_act = act_mod.get_activation(trunk_activation) + self.bc_act = act_mod.get_activation(BC_activation) + self.branch_act = act_mod.get_activation(branch_activation) + + self.use_bias = use_bias + if use_bias: + # register bias to parameter for updating in optimizer and storage + self.b = self.create_parameter( + shape=(1,), + attr=nn.initializer.Constant(0.0), + ) + + def forward(self, x): + + if self._input_transform is not None: + x = self._input_transform(x) + + # Branch net to encode the input function + u_features = self.branch_net(x)[self.branch_net.output_keys[0]] + bc_features = self.BC_net(x)[self.BC_net.output_keys[0]] + bctype_features = self.BCtype_net(x)[self.BCtype_net.output_keys[0]] + # Trunk net to encode the domain of the output function + y_features = self.trunk_net(x)[self.trunk_net.output_keys[0]] + y_features = self.trunk_act(y_features) + # Dot product + G_u = paddle.sum( + u_features * y_features * bc_features * bctype_features, + axis=1, + keepdim=True, + ) + # Add bias + if self.use_bias: + G_u += self.b + + result_dict = { + self.output_keys[0]: G_u, + } + if self._output_transform is not None: + result_dict = self._output_transform(x, result_dict) + + return result_dict diff --git a/examples/smc_reac/ppsci/arch/crystalgraphconvnet.py b/examples/smc_reac/ppsci/arch/crystalgraphconvnet.py new file mode 100644 index 0000000000..bb82aa0b81 --- /dev/null +++ b/examples/smc_reac/ppsci/arch/crystalgraphconvnet.py @@ -0,0 +1,167 @@ +import paddle +import paddle.nn as nn + +from ppsci.arch import base + + +class ConvLayer(nn.Layer): + def __init__(self, atom_fea_len, nbr_fea_len): + super(ConvLayer, self).__init__() + self.atom_fea_len = atom_fea_len + self.nbr_fea_len = nbr_fea_len + self.fc_full = nn.Linear( + 2 * self.atom_fea_len + self.nbr_fea_len, 2 * self.atom_fea_len + ) + self.sigmoid = nn.Sigmoid() + self.softplus1 = nn.Softplus() + self.bn1 = nn.BatchNorm1D(2 * self.atom_fea_len) + self.bn2 = nn.BatchNorm1D(self.atom_fea_len) + self.softplus2 = nn.Softplus() + + def forward(self, atom_in_fea, nbr_fea, nbr_fea_idx): + # TODO will there be problems with the index zero padding? + N, M = nbr_fea_idx.shape + atom_nbr_fea = atom_in_fea[nbr_fea_idx, :] + total_nbr_fea = paddle.concat( + [ + paddle.expand( + atom_in_fea.unsqueeze(1), shape=[N, M, self.atom_fea_len] + ), + atom_nbr_fea, + nbr_fea, + ], + axis=2, + ) + total_gated_fea = self.fc_full(total_nbr_fea) + total_gated_fea = paddle.reshape( + self.bn1(paddle.reshape(total_gated_fea, [-1, self.atom_fea_len * 2])), + [N, M, self.atom_fea_len * 2], + ) + nbr_filter, nbr_core = paddle.chunk(total_gated_fea, chunks=2, axis=2) + nbr_filter = self.sigmoid(nbr_filter) + nbr_core = self.softplus1(nbr_core) + nbr_sumed = paddle.sum(nbr_filter * nbr_core, axis=1) + nbr_sumed = self.bn2(nbr_sumed) + out = self.softplus2(atom_in_fea + nbr_sumed) + return out + + +class CrystalGraphConvNet(base.Arch): + """ + Create a crystal graph convolutional neural network for predicting total + material properties. + + Args: + orig_atom_fea_len (int): Number of atom features in the input. + nbr_fea_len (int): Number of bond features. + atom_fea_len (int): Number of hidden atom features in the convolutional layers. + n_conv (int): Number of convolutional layers. + h_fea_len (int): Number of hidden features after pooling. + n_h (int): Number of hidden layers after pooling. + + Examples: + >>> import paddle + >>> import ppsci + >>> model = ppsci.arch.CrystalGraphConvNet( + ... orig_atom_fea_len=92, + ... nbr_fea_len=41, + ... atom_fea_len=64, + ... n_conv=3, + ... h_fea_len=128, + ... n_h=1, + ... ) + >>> input_dict = { + ... "i": [ + ... paddle.rand(shape=[45, 92]), paddle.rand(shape=[45, 12, 41]), + ... paddle.randint(high=45, shape=[45, 12]), + ... [ + ... paddle.randint(high=32, shape=[32]), paddle.randint(high=8, shape=[8]), + ... paddle.randint(high=2, shape=[2]), paddle.randint(high=3, shape=[3]) + ... ] + ... ] + ... } + >>> output_dict = model(input_dict) + >>> print(output_dict["out"].shape) + [4, 1] + """ + + def __init__( + self, + orig_atom_fea_len: int, + nbr_fea_len: int, + atom_fea_len: int, + n_conv: int, + h_fea_len: int, + n_h: int, + ): + + super().__init__() + self.embedding = nn.Linear(orig_atom_fea_len, atom_fea_len) + self.convs = nn.LayerList( + [ + ConvLayer(atom_fea_len=atom_fea_len, nbr_fea_len=nbr_fea_len) + for _ in range(n_conv) + ] + ) + self.conv_to_fc = nn.Linear(atom_fea_len, h_fea_len) + self.conv_to_fc_softplus = nn.Softplus() + if n_h > 1: + self.fcs = nn.LayerList( + [nn.Linear(h_fea_len, h_fea_len) for _ in range(n_h - 1)] + ) + self.softpluses = nn.LayerList([nn.Softplus() for _ in range(n_h - 1)]) + + self.fc_out = nn.Linear(h_fea_len, 1) + + def forward(self, input) -> paddle.Tensor: + """ + Forward pass. + + N: Total number of atoms in the batch. + M: Max number of neighbors. + N0: Total number of crystals in the batch. + + Args: + input (list): List of input, which includes the following elements: + atom_fea (paddle.Tensor): Shape (N, orig_atom_fea_len). Atom features from atom type. + nbr_fea (paddle.Tensor): Shape (N, M, nbr_fea_len). Bond features of each atom's M neighbors. + nbr_fea_idx (paddle.Tensor): Shape (N, M). Indices of M neighbors of each atom. + crystal_atom_idx (list): List of paddle.Tensor of length N0. Mapping from the crystal idx to atom idx. + + Returns: + paddle.Tensor: Shape (N,). Atom hidden features after convolution. + """ + atom_fea, nbr_fea, nbr_fea_idx, crystal_atom_idx = input["i"] + atom_fea = self.embedding(atom_fea) + for conv_func in self.convs: + atom_fea = conv_func(atom_fea, nbr_fea, nbr_fea_idx) + crys_fea = self.pooling(atom_fea, crystal_atom_idx) + crys_fea = self.conv_to_fc(self.conv_to_fc_softplus(crys_fea)) + crys_fea = self.conv_to_fc_softplus(crys_fea) + if hasattr(self, "fcs") and hasattr(self, "softpluses"): + for fc, softplus in zip(self.fcs, self.softpluses): + crys_fea = softplus(fc(crys_fea)) + out = self.fc_out(crys_fea) + out_dict = {"out": out} + return out_dict + + def pooling(self, atom_fea, crystal_atom_idx): + """ + Pooling the atom features to crystal features + + N: Total number of atoms in the batch + N0: Total number of crystals in the batch + + Args: + atom_fea (paddle.Tensor): Shape (N, atom_fea_len). Atom feature vectors of the batch. + crystal_atom_idx (List[paddle.Tensor]): Length N0. Mapping from the crystal idx to atom idx + """ + assert ( + sum([len(idx_map) for idx_map in crystal_atom_idx]) + == atom_fea.data.shape[0] + ) + summed_fea = [ + paddle.mean(atom_fea[idx_map], axis=0, keepdim=True) + for idx_map in crystal_atom_idx + ] + return paddle.concat(summed_fea, axis=0) diff --git a/examples/smc_reac/ppsci/arch/cuboid_transformer.py b/examples/smc_reac/ppsci/arch/cuboid_transformer.py new file mode 100644 index 0000000000..08a80104bc --- /dev/null +++ b/examples/smc_reac/ppsci/arch/cuboid_transformer.py @@ -0,0 +1,958 @@ +from typing import Sequence +from typing import Tuple +from typing import Union + +import paddle +from paddle import nn + +import ppsci.arch.cuboid_transformer_decoder as cuboid_decoder +import ppsci.arch.cuboid_transformer_encoder as cuboid_encoder +import ppsci.arch.cuboid_transformer_utils as cuboid_utils +from ppsci.arch import activation as act_mod +from ppsci.arch import base +from ppsci.arch.cuboid_transformer_encoder import NEGATIVE_SLOPE +from ppsci.utils import initializer + +"""A space-time Transformer with Cuboid Attention""" + + +class InitialEncoder(nn.Layer): + def __init__( + self, + dim, + out_dim, + downsample_scale: Union[int, Sequence[int]], + num_conv_layers: int = 2, + activation: str = "leaky", + padding_type: str = "nearest", + conv_init_mode: str = "0", + linear_init_mode: str = "0", + norm_init_mode: str = "0", + ): + super(InitialEncoder, self).__init__() + self.num_conv_layers = num_conv_layers + self.conv_init_mode = conv_init_mode + self.linear_init_mode = linear_init_mode + self.norm_init_mode = norm_init_mode + conv_block = [] + for i in range(num_conv_layers): + if i == 0: + conv_block.append( + nn.Conv2D( + kernel_size=(3, 3), + padding=(1, 1), + in_channels=dim, + out_channels=out_dim, + ) + ) + conv_block.append(nn.GroupNorm(num_groups=16, num_channels=out_dim)) + conv_block.append( + act_mod.get_activation(activation) + if activation != "leaky_relu" + else nn.LeakyReLU(NEGATIVE_SLOPE) + ) + else: + conv_block.append( + nn.Conv2D( + kernel_size=(3, 3), + padding=(1, 1), + in_channels=out_dim, + out_channels=out_dim, + ) + ) + conv_block.append(nn.GroupNorm(num_groups=16, num_channels=out_dim)) + conv_block.append( + act_mod.get_activation(activation) + if activation != "leaky_relu" + else nn.LeakyReLU(NEGATIVE_SLOPE) + ) + self.conv_block = nn.Sequential(*conv_block) + if isinstance(downsample_scale, int): + patch_merge_downsample = (1, downsample_scale, downsample_scale) + elif len(downsample_scale) == 2: + patch_merge_downsample = (1, *downsample_scale) + elif len(downsample_scale) == 3: + patch_merge_downsample = tuple(downsample_scale) + else: + raise NotImplementedError( + f"downsample_scale {downsample_scale} format not supported!" + ) + self.patch_merge = cuboid_encoder.PatchMerging3D( + dim=out_dim, + out_dim=out_dim, + padding_type=padding_type, + downsample=patch_merge_downsample, + linear_init_mode=linear_init_mode, + norm_init_mode=norm_init_mode, + ) + self.reset_parameters() + + def reset_parameters(self): + for m in self.children(): + cuboid_utils.apply_initialization( + m, + conv_mode=self.conv_init_mode, + linear_mode=self.linear_init_mode, + norm_mode=self.norm_init_mode, + ) + + def forward(self, x): + """x --> [K x Conv2D] --> PatchMerge + + Args: + x: (B, T, H, W, C) + + Returns: + out: (B, T, H_new, W_new, C_out) + """ + + B, T, H, W, C = x.shape + + if self.num_conv_layers > 0: + x = x.reshape([B * T, H, W, C]).transpose(perm=[0, 3, 1, 2]) + x = self.conv_block(x).transpose(perm=[0, 2, 3, 1]) + x = self.patch_merge(x.reshape([B, T, H, W, -1])) + else: + x = self.patch_merge(x) + return x + + +class FinalDecoder(nn.Layer): + def __init__( + self, + target_thw: Tuple[int, ...], + dim: int, + num_conv_layers: int = 2, + activation: str = "leaky", + conv_init_mode: str = "0", + linear_init_mode: str = "0", + norm_init_mode: str = "0", + ): + super(FinalDecoder, self).__init__() + self.target_thw = target_thw + self.dim = dim + self.num_conv_layers = num_conv_layers + self.conv_init_mode = conv_init_mode + self.linear_init_mode = linear_init_mode + self.norm_init_mode = norm_init_mode + conv_block = [] + for i in range(num_conv_layers): + conv_block.append( + nn.Conv2D( + kernel_size=(3, 3), + padding=(1, 1), + in_channels=dim, + out_channels=dim, + ) + ) + conv_block.append(nn.GroupNorm(num_groups=16, num_channels=dim)) + conv_block.append( + act_mod.get_activation(activation) + if activation != "leaky_relu" + else nn.LeakyReLU(NEGATIVE_SLOPE) + ) + self.conv_block = nn.Sequential(*conv_block) + self.upsample = cuboid_decoder.Upsample3DLayer( + dim=dim, + out_dim=dim, + target_size=target_thw, + kernel_size=3, + conv_init_mode=conv_init_mode, + ) + self.reset_parameters() + + def reset_parameters(self): + for m in self.children(): + cuboid_utils.apply_initialization( + m, + conv_mode=self.conv_init_mode, + linear_mode=self.linear_init_mode, + norm_mode=self.norm_init_mode, + ) + + def forward(self, x): + """x --> Upsample --> [K x Conv2D] + + Args: + x: (B, T, H, W, C) + + Returns: + out: (B, T, H_new, W_new, C) + """ + + x = self.upsample(x) + if self.num_conv_layers > 0: + B, T, H, W, C = x.shape + x = x.reshape([B * T, H, W, C]).transpose(perm=[0, 3, 1, 2]) + x = ( + self.conv_block(x) + .transpose(perm=[0, 2, 3, 1]) + .reshape([B, T, H, W, -1]) + ) + return x + + +class InitialStackPatchMergingEncoder(nn.Layer): + def __init__( + self, + num_merge: int, + in_dim: int, + out_dim_list: Tuple[int, ...], + downsample_scale_list: Tuple[float, ...], + num_conv_per_merge_list: Tuple[int, ...] = None, + activation: str = "leaky", + padding_type: str = "nearest", + conv_init_mode: str = "0", + linear_init_mode: str = "0", + norm_init_mode: str = "0", + ): + super(InitialStackPatchMergingEncoder, self).__init__() + self.conv_init_mode = conv_init_mode + self.linear_init_mode = linear_init_mode + self.norm_init_mode = norm_init_mode + self.num_merge = num_merge + self.in_dim = in_dim + self.out_dim_list = out_dim_list[:num_merge] + self.downsample_scale_list = downsample_scale_list[:num_merge] + self.num_conv_per_merge_list = num_conv_per_merge_list + self.num_group_list = [max(1, out_dim // 4) for out_dim in self.out_dim_list] + self.conv_block_list = nn.LayerList() + self.patch_merge_list = nn.LayerList() + for i in range(num_merge): + if i == 0: + in_dim = in_dim + else: + in_dim = self.out_dim_list[i - 1] + out_dim = self.out_dim_list[i] + downsample_scale = self.downsample_scale_list[i] + conv_block = [] + for j in range(self.num_conv_per_merge_list[i]): + if j == 0: + conv_in_dim = in_dim + else: + conv_in_dim = out_dim + conv_block.append( + nn.Conv2D( + kernel_size=(3, 3), + padding=(1, 1), + in_channels=conv_in_dim, + out_channels=out_dim, + ) + ) + conv_block.append( + nn.GroupNorm( + num_groups=self.num_group_list[i], num_channels=out_dim + ) + ) + conv_block.append( + act_mod.get_activation(activation) + if activation != "leaky_relu" + else nn.LeakyReLU(NEGATIVE_SLOPE) + ) + conv_block = nn.Sequential(*conv_block) + self.conv_block_list.append(conv_block) + patch_merge = cuboid_encoder.PatchMerging3D( + dim=out_dim, + out_dim=out_dim, + padding_type=padding_type, + downsample=(1, downsample_scale, downsample_scale), + linear_init_mode=linear_init_mode, + norm_init_mode=norm_init_mode, + ) + self.patch_merge_list.append(patch_merge) + self.reset_parameters() + + def reset_parameters(self): + for m in self.children(): + cuboid_utils.apply_initialization( + m, + conv_mode=self.conv_init_mode, + linear_mode=self.linear_init_mode, + norm_mode=self.norm_init_mode, + ) + + def get_out_shape_list(self, input_shape): + out_shape_list = [] + for patch_merge in self.patch_merge_list: + input_shape = patch_merge.get_out_shape(input_shape) + out_shape_list.append(input_shape) + return out_shape_list + + def forward(self, x): + """x --> [K x Conv2D] --> PatchMerge --> ... --> [K x Conv2D] --> PatchMerge + + Args: + x: (B, T, H, W, C) + + Returns: + out: (B, T, H_new, W_new, C_out) + """ + + for i, (conv_block, patch_merge) in enumerate( + zip(self.conv_block_list, self.patch_merge_list) + ): + B, T, H, W, C = x.shape + if self.num_conv_per_merge_list[i] > 0: + x = x.reshape([B * T, H, W, C]).transpose(perm=[0, 3, 1, 2]) + x = conv_block(x).transpose(perm=[0, 2, 3, 1]).reshape([B, T, H, W, -1]) + x = patch_merge(x) + return x + + +class FinalStackUpsamplingDecoder(nn.Layer): + def __init__( + self, + target_shape_list: Tuple[Tuple[int, ...]], + in_dim: int, + num_conv_per_up_list: Tuple[int, ...] = None, + activation: str = "leaky", + conv_init_mode: str = "0", + linear_init_mode: str = "0", + norm_init_mode: str = "0", + ): + super(FinalStackUpsamplingDecoder, self).__init__() + self.conv_init_mode = conv_init_mode + self.linear_init_mode = linear_init_mode + self.norm_init_mode = norm_init_mode + self.target_shape_list = target_shape_list + self.out_dim_list = [ + target_shape[-1] for target_shape in self.target_shape_list + ] + self.num_upsample = len(target_shape_list) + self.in_dim = in_dim + self.num_conv_per_up_list = num_conv_per_up_list + self.num_group_list = [max(1, out_dim // 4) for out_dim in self.out_dim_list] + self.conv_block_list = nn.LayerList() + self.upsample_list = nn.LayerList() + for i in range(self.num_upsample): + if i == 0: + in_dim = in_dim + else: + in_dim = self.out_dim_list[i - 1] + out_dim = self.out_dim_list[i] + upsample = cuboid_decoder.Upsample3DLayer( + dim=in_dim, + out_dim=in_dim, + target_size=target_shape_list[i][:-1], + kernel_size=3, + conv_init_mode=conv_init_mode, + ) + self.upsample_list.append(upsample) + conv_block = [] + for j in range(num_conv_per_up_list[i]): + if j == 0: + conv_in_dim = in_dim + else: + conv_in_dim = out_dim + conv_block.append( + nn.Conv2D( + kernel_size=(3, 3), + padding=(1, 1), + in_channels=conv_in_dim, + out_channels=out_dim, + ) + ) + conv_block.append( + nn.GroupNorm( + num_groups=self.num_group_list[i], num_channels=out_dim + ) + ) + conv_block.append( + act_mod.get_activation(activation) + if activation != "leaky_relu" + else nn.LeakyReLU(NEGATIVE_SLOPE) + ) + conv_block = nn.Sequential(*conv_block) + self.conv_block_list.append(conv_block) + self.reset_parameters() + + def reset_parameters(self): + for m in self.children(): + cuboid_utils.apply_initialization( + m, + conv_mode=self.conv_init_mode, + linear_mode=self.linear_init_mode, + norm_mode=self.norm_init_mode, + ) + + @staticmethod + def get_init_params(enc_input_shape, enc_out_shape_list, large_channel=False): + dec_target_shape_list = list(enc_out_shape_list[:-1])[::-1] + [ + tuple(enc_input_shape) + ] + if large_channel: + dec_target_shape_list_large_channel = [] + for i, enc_out_shape in enumerate(enc_out_shape_list[::-1]): + dec_target_shape_large_channel = list(dec_target_shape_list[i]) + dec_target_shape_large_channel[-1] = enc_out_shape[-1] + dec_target_shape_list_large_channel.append( + tuple(dec_target_shape_large_channel) + ) + dec_target_shape_list = dec_target_shape_list_large_channel + dec_in_dim = enc_out_shape_list[-1][-1] + return dec_target_shape_list, dec_in_dim + + def forward(self, x): + """x --> Upsample --> [K x Conv2D] --> ... --> Upsample --> [K x Conv2D] + + Args: + x: Shape (B, T, H, W, C) + + Returns: + out: Shape (B, T, H_new, W_new, C) + """ + for i, (conv_block, upsample) in enumerate( + zip(self.conv_block_list, self.upsample_list) + ): + x = upsample(x) + if self.num_conv_per_up_list[i] > 0: + B, T, H, W, C = x.shape + x = x.reshape([B * T, H, W, C]).transpose(perm=[0, 3, 1, 2]) + x = conv_block(x).transpose(perm=[0, 2, 3, 1]).reshape([B, T, H, W, -1]) + return x + + +class CuboidTransformer(base.Arch): + """Cuboid Transformer for spatiotemporal forecasting + + We adopt the Non-autoregressive encoder-decoder architecture. + The decoder takes the multi-scale memory output from the encoder. + + The initial downsampling / upsampling layers will be + Downsampling: [K x Conv2D --> PatchMerge] + Upsampling: [Nearest Interpolation-based Upsample --> K x Conv2D] + + x --> downsample (optional) ---> (+pos_embed) ---> enc --> mem_l initial_z (+pos_embed) ---> FC + | | + |------------| + | + | + y <--- upsample (optional) <--- dec <---------- + + Args: + input_keys (Tuple[str, ...]): Name of input keys, such as ("input",). + output_keys (Tuple[str, ...]): Name of output keys, such as ("output",). + input_shape (Tuple[int, ...]): The shape of the input data. + target_shape (Tuple[int, ...]): The shape of the target data. + base_units (int, optional): The base units. Defaults to 128. + block_units (int, optional): The block units. Defaults to None. + scale_alpha (float, optional): We scale up the channels based on the formula: + - round_to(base_units * max(downsample_scale) ** units_alpha, 4). Defaults to 1.0. + num_heads (int, optional): The number of heads. Defaults to 4. + attn_drop (float, optional): The attention dropout. Defaults to 0.0. + proj_drop (float, optional): The projection dropout. Defaults to 0.0. + ffn_drop (float, optional): The ffn dropout. Defaults to 0.0. + downsample (int, optional): The rate of downsample. Defaults to 2. + downsample_type (str, optional): The type of downsample. Defaults to "patch_merge". + upsample_type (str, optional): The rate of upsample. Defaults to "upsample". + upsample_kernel_size (int, optional): The kernel size of upsample. Defaults to 3. + enc_depth (list, optional): The depth of encoder. Defaults to [4, 4, 4]. + enc_attn_patterns (str, optional): The pattern of encoder attention. Defaults to None. + enc_cuboid_size (list, optional): The cuboid size of encoder. Defaults to [(4, 4, 4), (4, 4, 4)]. + enc_cuboid_strategy (list, optional): The cuboid strategy of encoder. Defaults to [("l", "l", "l"), ("d", "d", "d")]. + enc_shift_size (list, optional): The shift size of encoder. Defaults to [(0, 0, 0), (0, 0, 0)]. + enc_use_inter_ffn (bool, optional): Whether to use intermediate FFN for encoder. Defaults to True. + dec_depth (list, optional): The depth of decoder. Defaults to [2, 2]. + dec_cross_start (int, optional): The cross start of decoder. Defaults to 0. + dec_self_attn_patterns (str, optional): The partterns of decoder. Defaults to None. + dec_self_cuboid_size (list, optional): The cuboid size of decoder. Defaults to [(4, 4, 4), (4, 4, 4)]. + dec_self_cuboid_strategy (list, optional): The strategy of decoder. Defaults to [("l", "l", "l"), ("d", "d", "d")]. + dec_self_shift_size (list, optional): The shift size of decoder. Defaults to [(1, 1, 1), (0, 0, 0)]. + dec_cross_attn_patterns (_type_, optional): The cross attention patterns of decoder. Defaults to None. + dec_cross_cuboid_hw (list, optional): The cuboid_hw of decoder. Defaults to [(4, 4), (4, 4)]. + dec_cross_cuboid_strategy (list, optional): The cuboid strategy of decoder. Defaults to [("l", "l", "l"), ("d", "l", "l")]. + dec_cross_shift_hw (list, optional): The shift_hw of decoder. Defaults to [(0, 0), (0, 0)]. + dec_cross_n_temporal (list, optional): The cross_n_temporal of decoder. Defaults to [1, 2]. + dec_cross_last_n_frames (int, optional): The cross_last_n_frames of decoder. Defaults to None. + dec_use_inter_ffn (bool, optional): Whether to use intermediate FFN for decoder. Defaults to True. + dec_hierarchical_pos_embed (bool, optional): Whether to use hierarchical pos_embed for decoder. Defaults to False. + num_global_vectors (int, optional): The num of global vectors. Defaults to 4. + use_dec_self_global (bool, optional): Whether to use global vector for decoder. Defaults to True. + dec_self_update_global (bool, optional): Whether to update global vector for decoder. Defaults to True. + use_dec_cross_global (bool, optional): Whether to use cross global vector for decoder. Defaults to True. + use_global_vector_ffn (bool, optional): Whether to use global vector FFN. Defaults to True. + use_global_self_attn (bool, optional): Whether to use global attentions. Defaults to False. + separate_global_qkv (bool, optional): Whether to separate global qkv. Defaults to False. + global_dim_ratio (int, optional): The ratio of global dim. Defaults to 1. + self_pattern (str, optional): The pattern. Defaults to "axial". + cross_self_pattern (str, optional): The self cross pattern. Defaults to "axial". + cross_pattern (str, optional): The cross pattern. Defaults to "cross_1x1". + z_init_method (str, optional): How the initial input to the decoder is initialized. Defaults to "nearest_interp". + initial_downsample_type (str, optional): The downsample type of initial. Defaults to "conv". + initial_downsample_activation (str, optional): The downsample activation of initial. Defaults to "leaky". + initial_downsample_scale (int, optional): The downsample scale of initial. Defaults to 1. + initial_downsample_conv_layers (int, optional): The conv layer of downsample of initial. Defaults to 2. + final_upsample_conv_layers (int, optional): The conv layer of final upsample. Defaults to 2. + initial_downsample_stack_conv_num_layers (int, optional): The num of stack conv layer of initial downsample. Defaults to 1. + initial_downsample_stack_conv_dim_list (list, optional): The dim list of stack conv of initial downsample. Defaults to None. + initial_downsample_stack_conv_downscale_list (list, optional): The downscale list of stack conv of initial downsample. Defaults to [1]. + initial_downsample_stack_conv_num_conv_list (list, optional): The num of stack conv list of initial downsample. Defaults to [2]. + ffn_activation (str, optional): The activation of FFN. Defaults to "leaky". + gated_ffn (bool, optional): Whether to use gate FFN. Defaults to False. + norm_layer (str, optional): The type of normilize. Defaults to "layer_norm". + padding_type (str, optional): The type of padding. Defaults to "ignore". + pos_embed_type (str, optional): The type of pos embedding. Defaults to "t+hw". + checkpoint_level (bool, optional): Whether to use checkpoint. Defaults to True. + use_relative_pos (bool, optional): Whether to use relative pose. Defaults to True. + self_attn_use_final_proj (bool, optional): Whether to use final projection. Defaults to True. + dec_use_first_self_attn (bool, optional): Whether to use first self attention for decoder. Defaults to False. + attn_linear_init_mode (str, optional): The mode of attention linear init. Defaults to "0". + ffn_linear_init_mode (str, optional): The mode of FFN linear init. Defaults to "0". + conv_init_mode (str, optional): The mode of conv init. Defaults to "0". + down_up_linear_init_mode (str, optional): The mode of downsample and upsample linear init. Defaults to "0". + norm_init_mode (str, optional): The mode of normalization init. Defaults to "0". + """ + + def __init__( + self, + input_keys: Tuple[str, ...], + output_keys: Tuple[str, ...], + input_shape: Tuple[int, ...], + target_shape: Tuple[int, ...], + base_units: int = 128, + block_units: int = None, + scale_alpha: float = 1.0, + num_heads: int = 4, + attn_drop: float = 0.0, + proj_drop: float = 0.0, + ffn_drop: float = 0.0, + downsample: int = 2, + downsample_type: str = "patch_merge", + upsample_type: str = "upsample", + upsample_kernel_size: int = 3, + enc_depth: Tuple[int, ...] = [4, 4, 4], + enc_attn_patterns: str = None, + enc_cuboid_size: Tuple[Tuple[int, ...], ...] = [(4, 4, 4), (4, 4, 4)], + enc_cuboid_strategy: Tuple[Tuple[str, ...], ...] = [ + ("l", "l", "l"), + ("d", "d", "d"), + ], + enc_shift_size: Tuple[Tuple[int, ...], ...] = [(0, 0, 0), (0, 0, 0)], + enc_use_inter_ffn: bool = True, + dec_depth: Tuple[int, ...] = [2, 2], + dec_cross_start: int = 0, + dec_self_attn_patterns: str = None, + dec_self_cuboid_size: Tuple[Tuple[int, ...], ...] = [(4, 4, 4), (4, 4, 4)], + dec_self_cuboid_strategy: Tuple[Tuple[str, ...], ...] = [ + ("l", "l", "l"), + ("d", "d", "d"), + ], + dec_self_shift_size: Tuple[Tuple[int, ...], ...] = [(1, 1, 1), (0, 0, 0)], + dec_cross_attn_patterns: str = None, + dec_cross_cuboid_hw: Tuple[Tuple[int, ...], ...] = [(4, 4), (4, 4)], + dec_cross_cuboid_strategy: Tuple[Tuple[str, ...], ...] = [ + ("l", "l", "l"), + ("d", "l", "l"), + ], + dec_cross_shift_hw: Tuple[Tuple[int, ...], ...] = [(0, 0), (0, 0)], + dec_cross_n_temporal: Tuple[int, ...] = [1, 2], + dec_cross_last_n_frames: int = None, + dec_use_inter_ffn: bool = True, + dec_hierarchical_pos_embed: bool = False, + num_global_vectors: int = 4, + use_dec_self_global: bool = True, + dec_self_update_global: bool = True, + use_dec_cross_global: bool = True, + use_global_vector_ffn: bool = True, + use_global_self_attn: bool = False, + separate_global_qkv: bool = False, + global_dim_ratio: int = 1, + self_pattern: str = "axial", + cross_self_pattern: str = "axial", + cross_pattern: str = "cross_1x1", + z_init_method: str = "nearest_interp", + initial_downsample_type: str = "conv", + initial_downsample_activation: str = "leaky", + initial_downsample_scale: int = 1, + initial_downsample_conv_layers: int = 2, + final_upsample_conv_layers: int = 2, + initial_downsample_stack_conv_num_layers: int = 1, + initial_downsample_stack_conv_dim_list: Tuple[int, ...] = None, + initial_downsample_stack_conv_downscale_list: Tuple[int, ...] = [1], + initial_downsample_stack_conv_num_conv_list: Tuple[int, ...] = [2], + ffn_activation: str = "leaky", + gated_ffn: bool = False, + norm_layer: str = "layer_norm", + padding_type: str = "ignore", + pos_embed_type: str = "t+hw", + checkpoint_level: bool = True, + use_relative_pos: bool = True, + self_attn_use_final_proj: bool = True, + dec_use_first_self_attn: bool = False, + attn_linear_init_mode: str = "0", + ffn_linear_init_mode: str = "0", + conv_init_mode: str = "0", + down_up_linear_init_mode: str = "0", + norm_init_mode: str = "0", + ): + super().__init__() + self.input_keys = input_keys + self.output_keys = output_keys + self.attn_linear_init_mode = attn_linear_init_mode + self.ffn_linear_init_mode = ffn_linear_init_mode + self.conv_init_mode = conv_init_mode + self.down_up_linear_init_mode = down_up_linear_init_mode + self.norm_init_mode = norm_init_mode + assert len(enc_depth) == len(dec_depth) + self.base_units = base_units + self.num_global_vectors = num_global_vectors + + num_blocks = len(enc_depth) + if isinstance(self_pattern, str): + enc_attn_patterns = [self_pattern] * num_blocks + + if isinstance(cross_self_pattern, str): + dec_self_attn_patterns = [cross_self_pattern] * num_blocks + + if isinstance(cross_pattern, str): + dec_cross_attn_patterns = [cross_pattern] * num_blocks + + if global_dim_ratio != 1: + assert ( + separate_global_qkv is True + ), "Setting global_dim_ratio != 1 requires separate_global_qkv == True." + self.global_dim_ratio = global_dim_ratio + self.z_init_method = z_init_method + assert self.z_init_method in ["zeros", "nearest_interp", "last", "mean"] + self.input_shape = input_shape + self.target_shape = target_shape + T_in, H_in, W_in, C_in = input_shape + T_out, H_out, W_out, C_out = target_shape + assert H_in == H_out and W_in == W_out + if self.num_global_vectors > 0: + init_data = paddle.zeros( + (self.num_global_vectors, global_dim_ratio * base_units) + ) + self.init_global_vectors = paddle.create_parameter( + shape=init_data.shape, + dtype=init_data.dtype, + default_initializer=nn.initializer.Constant(0.0), + ) + + self.init_global_vectors.stop_gradient = not True + new_input_shape = self.get_initial_encoder_final_decoder( + initial_downsample_scale=initial_downsample_scale, + initial_downsample_type=initial_downsample_type, + activation=initial_downsample_activation, + initial_downsample_conv_layers=initial_downsample_conv_layers, + final_upsample_conv_layers=final_upsample_conv_layers, + padding_type=padding_type, + initial_downsample_stack_conv_num_layers=initial_downsample_stack_conv_num_layers, + initial_downsample_stack_conv_dim_list=initial_downsample_stack_conv_dim_list, + initial_downsample_stack_conv_downscale_list=initial_downsample_stack_conv_downscale_list, + initial_downsample_stack_conv_num_conv_list=initial_downsample_stack_conv_num_conv_list, + ) + T_in, H_in, W_in, _ = new_input_shape + self.encoder = cuboid_encoder.CuboidTransformerEncoder( + input_shape=(T_in, H_in, W_in, base_units), + base_units=base_units, + block_units=block_units, + scale_alpha=scale_alpha, + depth=enc_depth, + downsample=downsample, + downsample_type=downsample_type, + block_attn_patterns=enc_attn_patterns, + block_cuboid_size=enc_cuboid_size, + block_strategy=enc_cuboid_strategy, + block_shift_size=enc_shift_size, + num_heads=num_heads, + attn_drop=attn_drop, + proj_drop=proj_drop, + ffn_drop=ffn_drop, + gated_ffn=gated_ffn, + ffn_activation=ffn_activation, + norm_layer=norm_layer, + use_inter_ffn=enc_use_inter_ffn, + padding_type=padding_type, + use_global_vector=num_global_vectors > 0, + use_global_vector_ffn=use_global_vector_ffn, + use_global_self_attn=use_global_self_attn, + separate_global_qkv=separate_global_qkv, + global_dim_ratio=global_dim_ratio, + checkpoint_level=checkpoint_level, + use_relative_pos=use_relative_pos, + self_attn_use_final_proj=self_attn_use_final_proj, + attn_linear_init_mode=attn_linear_init_mode, + ffn_linear_init_mode=ffn_linear_init_mode, + conv_init_mode=conv_init_mode, + down_linear_init_mode=down_up_linear_init_mode, + norm_init_mode=norm_init_mode, + ) + self.enc_pos_embed = cuboid_decoder.PosEmbed( + embed_dim=base_units, typ=pos_embed_type, maxH=H_in, maxW=W_in, maxT=T_in + ) + mem_shapes = self.encoder.get_mem_shapes() + self.z_proj = nn.Linear( + in_features=mem_shapes[-1][-1], out_features=mem_shapes[-1][-1] + ) + self.dec_pos_embed = cuboid_decoder.PosEmbed( + embed_dim=mem_shapes[-1][-1], + typ=pos_embed_type, + maxT=T_out, + maxH=mem_shapes[-1][1], + maxW=mem_shapes[-1][2], + ) + self.decoder = cuboid_decoder.CuboidTransformerDecoder( + target_temporal_length=T_out, + mem_shapes=mem_shapes, + cross_start=dec_cross_start, + depth=dec_depth, + upsample_type=upsample_type, + block_self_attn_patterns=dec_self_attn_patterns, + block_self_cuboid_size=dec_self_cuboid_size, + block_self_shift_size=dec_self_shift_size, + block_self_cuboid_strategy=dec_self_cuboid_strategy, + block_cross_attn_patterns=dec_cross_attn_patterns, + block_cross_cuboid_hw=dec_cross_cuboid_hw, + block_cross_shift_hw=dec_cross_shift_hw, + block_cross_cuboid_strategy=dec_cross_cuboid_strategy, + block_cross_n_temporal=dec_cross_n_temporal, + cross_last_n_frames=dec_cross_last_n_frames, + num_heads=num_heads, + attn_drop=attn_drop, + proj_drop=proj_drop, + ffn_drop=ffn_drop, + upsample_kernel_size=upsample_kernel_size, + ffn_activation=ffn_activation, + gated_ffn=gated_ffn, + norm_layer=norm_layer, + use_inter_ffn=dec_use_inter_ffn, + max_temporal_relative=T_in + T_out, + padding_type=padding_type, + hierarchical_pos_embed=dec_hierarchical_pos_embed, + pos_embed_type=pos_embed_type, + use_self_global=num_global_vectors > 0 and use_dec_self_global, + self_update_global=dec_self_update_global, + use_cross_global=num_global_vectors > 0 and use_dec_cross_global, + use_global_vector_ffn=use_global_vector_ffn, + use_global_self_attn=use_global_self_attn, + separate_global_qkv=separate_global_qkv, + global_dim_ratio=global_dim_ratio, + checkpoint_level=checkpoint_level, + use_relative_pos=use_relative_pos, + self_attn_use_final_proj=self_attn_use_final_proj, + use_first_self_attn=dec_use_first_self_attn, + attn_linear_init_mode=attn_linear_init_mode, + ffn_linear_init_mode=ffn_linear_init_mode, + conv_init_mode=conv_init_mode, + up_linear_init_mode=down_up_linear_init_mode, + norm_init_mode=norm_init_mode, + ) + self.reset_parameters() + + def get_initial_encoder_final_decoder( + self, + initial_downsample_type, + activation, + initial_downsample_scale, + initial_downsample_conv_layers, + final_upsample_conv_layers, + padding_type, + initial_downsample_stack_conv_num_layers, + initial_downsample_stack_conv_dim_list, + initial_downsample_stack_conv_downscale_list, + initial_downsample_stack_conv_num_conv_list, + ): + T_in, H_in, W_in, C_in = self.input_shape + T_out, H_out, W_out, C_out = self.target_shape + self.initial_downsample_type = initial_downsample_type + if self.initial_downsample_type == "conv": + if isinstance(initial_downsample_scale, int): + initial_downsample_scale = ( + 1, + initial_downsample_scale, + initial_downsample_scale, + ) + elif len(initial_downsample_scale) == 2: + initial_downsample_scale = 1, *initial_downsample_scale + elif len(initial_downsample_scale) == 3: + initial_downsample_scale = tuple(initial_downsample_scale) + else: + raise NotImplementedError( + f"initial_downsample_scale {initial_downsample_scale} format not supported!" + ) + self.initial_encoder = InitialEncoder( + dim=C_in, + out_dim=self.base_units, + downsample_scale=initial_downsample_scale, + num_conv_layers=initial_downsample_conv_layers, + padding_type=padding_type, + activation=activation, + conv_init_mode=self.conv_init_mode, + linear_init_mode=self.down_up_linear_init_mode, + norm_init_mode=self.norm_init_mode, + ) + + self.final_decoder = FinalDecoder( + dim=self.base_units, + target_thw=(T_out, H_out, W_out), + num_conv_layers=final_upsample_conv_layers, + activation=activation, + conv_init_mode=self.conv_init_mode, + linear_init_mode=self.down_up_linear_init_mode, + norm_init_mode=self.norm_init_mode, + ) + new_input_shape = self.initial_encoder.patch_merge.get_out_shape( + self.input_shape + ) + self.dec_final_proj = nn.Linear( + in_features=self.base_units, out_features=C_out + ) + elif self.initial_downsample_type == "stack_conv": + if initial_downsample_stack_conv_dim_list is None: + initial_downsample_stack_conv_dim_list = [ + self.base_units + ] * initial_downsample_stack_conv_num_layers + self.initial_encoder = InitialStackPatchMergingEncoder( + num_merge=initial_downsample_stack_conv_num_layers, + in_dim=C_in, + out_dim_list=initial_downsample_stack_conv_dim_list, + downsample_scale_list=initial_downsample_stack_conv_downscale_list, + num_conv_per_merge_list=initial_downsample_stack_conv_num_conv_list, + padding_type=padding_type, + activation=activation, + conv_init_mode=self.conv_init_mode, + linear_init_mode=self.down_up_linear_init_mode, + norm_init_mode=self.norm_init_mode, + ) + initial_encoder_out_shape_list = self.initial_encoder.get_out_shape_list( + self.target_shape + ) + ( + dec_target_shape_list, + dec_in_dim, + ) = FinalStackUpsamplingDecoder.get_init_params( + enc_input_shape=self.target_shape, + enc_out_shape_list=initial_encoder_out_shape_list, + large_channel=True, + ) + self.final_decoder = FinalStackUpsamplingDecoder( + target_shape_list=dec_target_shape_list, + in_dim=dec_in_dim, + num_conv_per_up_list=initial_downsample_stack_conv_num_conv_list[::-1], + activation=activation, + conv_init_mode=self.conv_init_mode, + linear_init_mode=self.down_up_linear_init_mode, + norm_init_mode=self.norm_init_mode, + ) + self.dec_final_proj = nn.Linear( + in_features=dec_target_shape_list[-1][-1], out_features=C_out + ) + new_input_shape = self.initial_encoder.get_out_shape_list(self.input_shape)[ + -1 + ] + else: + raise NotImplementedError(f"{self.initial_downsample_type} is invalid.") + self.input_shape_after_initial_downsample = new_input_shape + T_in, H_in, W_in, _ = new_input_shape + return new_input_shape + + def reset_parameters(self): + if self.num_global_vectors > 0: + self.init_global_vectors = initializer.trunc_normal_( + self.init_global_vectors, std=0.02 + ) + if hasattr(self.initial_encoder, "reset_parameters"): + self.initial_encoder.reset_parameters() + else: + cuboid_utils.apply_initialization( + self.initial_encoder, + conv_mode=self.conv_init_mode, + linear_mode=self.down_up_linear_init_mode, + norm_mode=self.norm_init_mode, + ) + if hasattr(self.final_decoder, "reset_parameters"): + self.final_decoder.reset_parameters() + else: + cuboid_utils.apply_initialization( + self.final_decoder, + conv_mode=self.conv_init_mode, + linear_mode=self.down_up_linear_init_mode, + norm_mode=self.norm_init_mode, + ) + cuboid_utils.apply_initialization( + self.dec_final_proj, linear_mode=self.down_up_linear_init_mode + ) + self.encoder.reset_parameters() + self.enc_pos_embed.reset_parameters() + self.decoder.reset_parameters() + self.dec_pos_embed.reset_parameters() + cuboid_utils.apply_initialization(self.z_proj, linear_mode="0") + + def get_initial_z(self, final_mem, T_out): + B = final_mem.shape[0] + if self.z_init_method == "zeros": + z_shape = list((1, T_out)) + final_mem.shape[2:] + initial_z = paddle.zeros(shape=z_shape, dtype=final_mem.dtype) + initial_z = self.z_proj(self.dec_pos_embed(initial_z)).expand( + shape=[B, -1, -1, -1, -1] + ) + elif self.z_init_method == "nearest_interp": + initial_z = nn.functional.interpolate( + x=final_mem.transpose(perm=[0, 4, 1, 2, 3]), + size=(T_out, final_mem.shape[2], final_mem.shape[3]), + ).transpose(perm=[0, 2, 3, 4, 1]) + initial_z = self.z_proj(initial_z) + elif self.z_init_method == "last": + initial_z = paddle.broadcast_to( + x=final_mem[:, -1:, :, :, :], shape=(B, T_out) + final_mem.shape[2:] + ) + initial_z = self.z_proj(initial_z) + elif self.z_init_method == "mean": + initial_z = paddle.broadcast_to( + x=final_mem.mean(axis=1, keepdims=True), + shape=(B, T_out) + final_mem.shape[2:], + ) + initial_z = self.z_proj(initial_z) + else: + raise NotImplementedError + return initial_z + + def forward(self, x: "paddle.Tensor", verbose: bool = False) -> "paddle.Tensor": + """ + Args: + x (paddle.Tensor): Tensor with shape (B, T, H, W, C). + verbose (bool): If True, print intermediate shapes. + + Returns: + out (paddle.Tensor): The output Shape (B, T_out, H, W, C_out) + """ + + x = self.concat_to_tensor(x, self.input_keys) + flag_ndim = x.ndim + if flag_ndim == 6: + x = x.reshape([-1, *x.shape[2:]]) + B, _, _, _, _ = x.shape + + T_out = self.target_shape[0] + x = self.initial_encoder(x) + x = self.enc_pos_embed(x) + + if self.num_global_vectors > 0: + init_global_vectors = self.init_global_vectors.expand( + shape=[ + B, + self.num_global_vectors, + self.global_dim_ratio * self.base_units, + ] + ) + mem_l, mem_global_vector_l = self.encoder(x, init_global_vectors) + else: + mem_l = self.encoder(x) + + if verbose: + for i, mem in enumerate(mem_l): + print(f"mem[{i}].shape = {mem.shape}") + initial_z = self.get_initial_z(final_mem=mem_l[-1], T_out=T_out) + + if self.num_global_vectors > 0: + dec_out = self.decoder(initial_z, mem_l, mem_global_vector_l) + else: + dec_out = self.decoder(initial_z, mem_l) + + dec_out = self.final_decoder(dec_out) + + out = self.dec_final_proj(dec_out) + if flag_ndim == 6: + out = out.reshape([-1, *out.shape]) + return {key: out for key in self.output_keys} diff --git a/examples/smc_reac/ppsci/arch/cuboid_transformer_decoder.py b/examples/smc_reac/ppsci/arch/cuboid_transformer_decoder.py new file mode 100644 index 0000000000..8cbaf5fee9 --- /dev/null +++ b/examples/smc_reac/ppsci/arch/cuboid_transformer_decoder.py @@ -0,0 +1,1245 @@ +from functools import lru_cache +from typing import Tuple + +import numpy as np +import paddle +import paddle.nn.functional as F +from paddle import nn +from paddle.distributed import fleet + +import ppsci.arch.cuboid_transformer_encoder as cuboid_encoder +import ppsci.arch.cuboid_transformer_utils as cuboid_utils +from ppsci.utils import initializer + + +class PosEmbed(nn.Layer): + """Pose embedding + + Args: + embed_dim (int): The dimension of embedding. + maxT (int): The embedding max time. + maxH (int): The embedding max height. + maxW (int): The embedding max width. + typ (str): + The type of the positional embedding. + - t+h+w: + Embed the spatial position to embeddings + - t+hw: + Embed the spatial position to embeddings + """ + + def __init__(self, embed_dim, maxT, maxH, maxW, typ: str = "t+h+w"): + super(PosEmbed, self).__init__() + self.typ = typ + assert self.typ in ["t+h+w", "t+hw"] + self.maxT = maxT + self.maxH = maxH + self.maxW = maxW + self.embed_dim = embed_dim + if self.typ == "t+h+w": + self.T_embed = nn.Embedding(num_embeddings=maxT, embedding_dim=embed_dim) + self.H_embed = nn.Embedding(num_embeddings=maxH, embedding_dim=embed_dim) + self.W_embed = nn.Embedding(num_embeddings=maxW, embedding_dim=embed_dim) + elif self.typ == "t+hw": + self.T_embed = nn.Embedding(num_embeddings=maxT, embedding_dim=embed_dim) + self.HW_embed = nn.Embedding( + num_embeddings=maxH * maxW, embedding_dim=embed_dim + ) + else: + raise NotImplementedError(f"{self.typ} is invalid.") + self.reset_parameters() + + def reset_parameters(self): + for m in self.children(): + cuboid_utils.apply_initialization(m, embed_mode="0") + + def forward(self, x): + """ + Args: + x : Shape (B, T, H, W, C) + + Returns: + out : the x + positional embeddings + """ + + _, T, H, W, _ = x.shape + t_idx = paddle.arange(end=T) + h_idx = paddle.arange(end=H) + w_idx = paddle.arange(end=W) + if self.typ == "t+h+w": + return ( + x + + self.T_embed(t_idx).reshape([T, 1, 1, self.embed_dim]) + + self.H_embed(h_idx).reshape([1, H, 1, self.embed_dim]) + + self.W_embed(w_idx).reshape([1, 1, W, self.embed_dim]) + ) + elif self.typ == "t+hw": + spatial_idx = h_idx.unsqueeze(axis=-1) * self.maxW + w_idx + return ( + x + + self.T_embed(t_idx).reshape([T, 1, 1, self.embed_dim]) + + self.HW_embed(spatial_idx) + ) + else: + raise NotImplementedError(f"{self.typ} is invalid.") + + +@lru_cache() +def compute_cuboid_cross_attention_mask( + T_x, T_mem, H, W, n_temporal, cuboid_hw, shift_hw, strategy, padding_type, device +): + pad_t_mem = (n_temporal - T_mem % n_temporal) % n_temporal + pad_t_x = (n_temporal - T_x % n_temporal) % n_temporal + pad_h = (cuboid_hw[0] - H % cuboid_hw[0]) % cuboid_hw[0] + pad_w = (cuboid_hw[1] - W % cuboid_hw[1]) % cuboid_hw[1] + mem_cuboid_size = ((T_mem + pad_t_mem) // n_temporal,) + cuboid_hw + x_cuboid_size = ((T_x + pad_t_x) // n_temporal,) + cuboid_hw + if pad_t_mem > 0 or pad_h > 0 or pad_w > 0: + if padding_type == "ignore": + mem_mask = paddle.ones(shape=(1, T_mem, H, W, 1), dtype="bool") + mem_mask = F.pad( + mem_mask, [0, 0, 0, pad_w, 0, pad_h, pad_t_mem, 0], data_format="NDHWC" + ) + else: + mem_mask = paddle.ones( + shape=(1, T_mem + pad_t_mem, H + pad_h, W + pad_w, 1), dtype="bool" + ) + if pad_t_x > 0 or pad_h > 0 or pad_w > 0: + if padding_type == "ignore": + x_mask = paddle.ones(shape=(1, T_x, H, W, 1), dtype="bool") + x_mask = F.pad( + x_mask, [0, 0, 0, pad_w, 0, pad_h, 0, pad_t_x], data_format="NDHWC" + ) + else: + x_mask = paddle.ones( + shape=(1, T_x + pad_t_x, H + pad_h, W + pad_w, 1), dtype="bool" + ) + if any(i > 0 for i in shift_hw): + if padding_type == "ignore": + x_mask = paddle.roll( + x=x_mask, shifts=(-shift_hw[0], -shift_hw[1]), axis=(2, 3) + ) + mem_mask = paddle.roll( + x=mem_mask, shifts=(-shift_hw[0], -shift_hw[1]), axis=(2, 3) + ) + x_mask = cuboid_encoder.cuboid_reorder(x_mask, x_cuboid_size, strategy=strategy) + x_mask = x_mask.squeeze(axis=-1).squeeze(axis=0) + num_cuboids, x_cuboid_volume = x_mask.shape + mem_mask = cuboid_encoder.cuboid_reorder( + mem_mask, mem_cuboid_size, strategy=strategy + ) + mem_mask = mem_mask.squeeze(axis=-1).squeeze(axis=0) + _, mem_cuboid_volume = mem_mask.shape + shift_mask = np.zeros(shape=(1, n_temporal, H + pad_h, W + pad_w, 1)) + cnt = 0 + for h in ( + slice(-cuboid_hw[0]), + slice(-cuboid_hw[0], -shift_hw[0]), + slice(-shift_hw[0], None), + ): + for w in ( + slice(-cuboid_hw[1]), + slice(-cuboid_hw[1], -shift_hw[1]), + slice(-shift_hw[1], None), + ): + shift_mask[:, :, h, w, :] = cnt + cnt += 1 + shift_mask = paddle.to_tensor(shift_mask) + shift_mask = cuboid_encoder.cuboid_reorder( + shift_mask, (1,) + cuboid_hw, strategy=strategy + ) + shift_mask = shift_mask.squeeze(axis=-1).squeeze(axis=0) + shift_mask = shift_mask.unsqueeze(axis=1) - shift_mask.unsqueeze(axis=2) == 0 + bh_bw = cuboid_hw[0] * cuboid_hw[1] + attn_mask = ( + shift_mask.reshape((num_cuboids, 1, bh_bw, 1, bh_bw)) + * x_mask.reshape((num_cuboids, -1, bh_bw, 1, 1)) + * mem_mask.reshape([num_cuboids, 1, 1, -1, bh_bw]) + ) + attn_mask = attn_mask.reshape([num_cuboids, x_cuboid_volume, mem_cuboid_volume]) + return attn_mask + + +class CuboidCrossAttentionLayer(nn.Layer): + """Implements the cuboid cross attention. + + The idea of Cuboid Cross Attention is to extend the idea of cuboid self attention to work for the + encoder-decoder-type cross attention. + + Assume that there is a memory tensor with shape (T1, H, W, C) and another query tensor with shape (T2, H, W, C), + + Here, we decompose the query tensor and the memory tensor into the same number of cuboids and attend the cuboid in + the query tensor with the corresponding cuboid in the memory tensor. + + For the height and width axes, we reuse the grid decomposition techniques described in the cuboid self-attention. + For the temporal axis, the layer supports the "n_temporal" parameter, that controls the number of cuboids we can + get after cutting the tensors. For example, if the temporal dilation is 2, both the query and + memory will be decomposed into 2 cuboids along the temporal axis. Like in the Cuboid Self-attention, + we support "local" and "dilated" decomposition strategy. + + The complexity of the layer is O((T2 / n_t * Bh * Bw) * (T1 / n_t * Bh * Bw) * n_t (H / Bh) (W / Bw)) = O(T2 * T1 / n_t H W Bh Bw) + + Args: + dim (int): The dimension of input tensor. + num_heads (int): The number of head. + n_temporal (int, optional): The num of temporal. Defaults to 1. + cuboid_hw (tuple, optional): The height and width of cuboid. Defaults to (7, 7). + shift_hw (tuple, optional): The height and width of shift. Defaults to (0, 0). + strategy (tuple, optional): The strategy. Defaults to ("d", "l", "l"). + padding_type (str, optional): The type of padding. Defaults to "ignore". + cross_last_n_frames (int, optional): The cross_last_n_frames of decoder. Defaults to None. + qkv_bias (bool, optional): Whether to enable bias in calculating qkv attention. Defaults to False. + qk_scale (float, optional): Whether to enable scale factor when calculating the attention. Defaults to None. + attn_drop (float, optional): The attention dropout. Defaults to 0.0. + proj_drop (float, optional): The projrction dropout. Defaults to 0.0. + max_temporal_relative (int, optional): The max temporal. Defaults to 50. + norm_layer (str, optional): The normalization layer. Defaults to "layer_norm". + use_global_vector (bool, optional): Whether to use the global vector or not. Defaults to True. + separate_global_qkv (bool, optional): Whether to use different network to calc q_global, k_global, v_global. Defaults to False. + global_dim_ratio (int, optional): The dim (channels) of global vectors is `global_dim_ratio*dim`. Defaults to 1. + checkpoint_level (int, optional): Whether to enable gradient checkpointing. Defaults to 1. + use_relative_pos (bool, optional): Whether to use relative pos. Defaults to True. + attn_linear_init_mode (str, optional): The mode of attention linear initialization. Defaults to "0". + ffn_linear_init_mode (str, optional): The mode of FFN linear initialization. Defaults to "0". + norm_init_mode (str, optional): The mode of normalization initialization. Defaults to "0". + """ + + def __init__( + self, + dim: int, + num_heads: int, + n_temporal: int = 1, + cuboid_hw: Tuple[int, ...] = (7, 7), + shift_hw: Tuple[int, ...] = (0, 0), + strategy: Tuple[str, ...] = ("d", "l", "l"), + padding_type: str = "ignore", + cross_last_n_frames: int = None, + qkv_bias: bool = False, + qk_scale: float = None, + attn_drop: float = 0.0, + proj_drop: float = 0.0, + max_temporal_relative: int = 50, + norm_layer: str = "layer_norm", + use_global_vector: bool = True, + separate_global_qkv: bool = False, + global_dim_ratio: int = 1, + checkpoint_level: int = 1, + use_relative_pos: bool = True, + attn_linear_init_mode: str = "0", + ffn_linear_init_mode: str = "0", + norm_init_mode: str = "0", + ): + super(CuboidCrossAttentionLayer, self).__init__() + self.attn_linear_init_mode = attn_linear_init_mode + self.ffn_linear_init_mode = ffn_linear_init_mode + self.norm_init_mode = norm_init_mode + self.dim = dim + self.num_heads = num_heads + self.n_temporal = n_temporal + assert n_temporal > 0 + head_dim = dim // num_heads + self.scale = qk_scale or head_dim**-0.5 + shift_hw = list(shift_hw) + if strategy[1] == "d": + shift_hw[0] = 0 + if strategy[2] == "d": + shift_hw[1] = 0 + self.cuboid_hw = cuboid_hw + self.shift_hw = tuple(shift_hw) + self.strategy = strategy + self.padding_type = padding_type + self.max_temporal_relative = max_temporal_relative + self.cross_last_n_frames = cross_last_n_frames + self.use_relative_pos = use_relative_pos + self.use_global_vector = use_global_vector + self.separate_global_qkv = separate_global_qkv + if global_dim_ratio != 1 and separate_global_qkv is False: + raise ValueError( + "Setting global_dim_ratio != 1 requires separate_global_qkv == True." + ) + self.global_dim_ratio = global_dim_ratio + if self.padding_type not in ["ignore", "zeros", "nearest"]: + raise ValueError('padding_type should be ["ignore", "zeros", "nearest"]') + if use_relative_pos: + init_data = paddle.zeros( + ( + (2 * max_temporal_relative - 1) + * (2 * cuboid_hw[0] - 1) + * (2 * cuboid_hw[1] - 1), + num_heads, + ) + ) + self.relative_position_bias_table = paddle.create_parameter( + shape=init_data.shape, + dtype=init_data.dtype, + default_initializer=nn.initializer.Constant(0.0), + ) + self.relative_position_bias_table.stop_gradient = not True + self.relative_position_bias_table = initializer.trunc_normal_( + self.relative_position_bias_table, std=0.02 + ) + + coords_t = paddle.arange(end=max_temporal_relative) + coords_h = paddle.arange(end=self.cuboid_hw[0]) + coords_w = paddle.arange(end=self.cuboid_hw[1]) + coords = paddle.stack(x=paddle.meshgrid(coords_t, coords_h, coords_w)) + coords_flatten = paddle.flatten(x=coords, start_axis=1) + relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] + relative_coords = relative_coords.transpose(perm=[1, 2, 0]) + relative_coords[:, :, 0] += max_temporal_relative - 1 + relative_coords[:, :, 1] += self.cuboid_hw[0] - 1 + relative_coords[:, :, 2] += self.cuboid_hw[1] - 1 + relative_position_index = ( + relative_coords[:, :, 0] + * (2 * self.cuboid_hw[0] - 1) + * (2 * self.cuboid_hw[1] - 1) + + relative_coords[:, :, 1] * (2 * self.cuboid_hw[1] - 1) + + relative_coords[:, :, 2] + ) + self.register_buffer( + name="relative_position_index", tensor=relative_position_index + ) + self.q_proj = nn.Linear(in_features=dim, out_features=dim, bias_attr=qkv_bias) + self.kv_proj = nn.Linear( + in_features=dim, out_features=dim * 2, bias_attr=qkv_bias + ) + self.attn_drop = nn.Dropout(p=attn_drop) + self.proj = nn.Linear(in_features=dim, out_features=dim) + self.proj_drop = nn.Dropout(p=proj_drop) + if self.use_global_vector: + if self.separate_global_qkv: + self.l2g_q_net = nn.Linear( + in_features=dim, out_features=dim, bias_attr=qkv_bias + ) + self.l2g_global_kv_net = nn.Linear( + in_features=global_dim_ratio * dim, + out_features=dim * 2, + bias_attr=qkv_bias, + ) + self.norm = cuboid_utils.get_norm_layer(norm_layer, in_channels=dim) + self._checkpoint_level = checkpoint_level + self.reset_parameters() + + def reset_parameters(self): + cuboid_utils.apply_initialization( + self.q_proj, linear_mode=self.attn_linear_init_mode + ) + cuboid_utils.apply_initialization( + self.kv_proj, linear_mode=self.attn_linear_init_mode + ) + cuboid_utils.apply_initialization( + self.proj, linear_mode=self.ffn_linear_init_mode + ) + cuboid_utils.apply_initialization(self.norm, norm_mode=self.norm_init_mode) + if self.use_global_vector: + if self.separate_global_qkv: + cuboid_utils.apply_initialization( + self.l2g_q_net, linear_mode=self.attn_linear_init_mode + ) + cuboid_utils.apply_initialization( + self.l2g_global_kv_net, linear_mode=self.attn_linear_init_mode + ) + + def forward(self, x, mem, mem_global_vectors=None): + """Calculate the forward + + Along the temporal axis, we pad the mem tensor from the left and the x tensor from the right so that the + relative position encoding can be calculated correctly. For example: + + mem: 0, 1, 2, 3, 4 + x: 0, 1, 2, 3, 4, 5 + + n_temporal = 1 + mem: 0, 1, 2, 3, 4 x: 0, 1, 2, 3, 4, 5 + + n_temporal = 2 + mem: pad, 1, 3 x: 0, 2, 4 + mem: 0, 2, 4 x: 1, 3, 5 + + n_temporal = 3 + mem: pad, 2 dec: 0, 3 + mem: 0, 3 dec: 1, 4 + mem: 1, 4 dec: 2, 5 + + Args: + x (paddle.Tensor): The input of the layer. It will have shape (B, T, H, W, C) + mem (paddle.Tensor): The memory. It will have shape (B, T_mem, H, W, C) + mem_global_vectors (paddle.Tensor): The global vectors from the memory. It will have shape (B, N, C) + + Returns: + out (paddle.Tensor): Output tensor should have shape (B, T, H, W, C_out) + """ + + if self.cross_last_n_frames is not None: + cross_last_n_frames = int(min(self.cross_last_n_frames, mem.shape[1])) + mem = mem[:, -cross_last_n_frames:, ...] + if self.use_global_vector: + _, num_global, _ = mem_global_vectors.shape + x = self.norm(x) + B, T_x, H, W, C_in = x.shape + B_mem, T_mem, H_mem, W_mem, C_mem = mem.shape + assert T_x < self.max_temporal_relative and T_mem < self.max_temporal_relative + cuboid_hw = self.cuboid_hw + n_temporal = self.n_temporal + shift_hw = self.shift_hw + assert ( + B_mem == B and H == H_mem and W == W_mem and C_in == C_mem + ), f"Shape of memory and the input tensor does not match. x.shape={x.shape}, mem.shape={mem.shape}" + pad_t_mem = (n_temporal - T_mem % n_temporal) % n_temporal + pad_t_x = (n_temporal - T_x % n_temporal) % n_temporal + pad_h = (cuboid_hw[0] - H % cuboid_hw[0]) % cuboid_hw[0] + pad_w = (cuboid_hw[1] - W % cuboid_hw[1]) % cuboid_hw[1] + mem = cuboid_utils.generalize_padding( + mem, pad_t_mem, pad_h, pad_w, self.padding_type, t_pad_left=True + ) + + x = cuboid_utils.generalize_padding( + x, pad_t_x, pad_h, pad_w, self.padding_type, t_pad_left=False + ) + + if any(i > 0 for i in shift_hw): + shifted_x = paddle.roll( + x=x, shifts=(-shift_hw[0], -shift_hw[1]), axis=(2, 3) + ) + shifted_mem = paddle.roll( + x=mem, shifts=(-shift_hw[0], -shift_hw[1]), axis=(2, 3) + ) + else: + shifted_x = x + shifted_mem = mem + mem_cuboid_size = (mem.shape[1] // n_temporal,) + cuboid_hw + x_cuboid_size = (x.shape[1] // n_temporal,) + cuboid_hw + reordered_mem = cuboid_encoder.cuboid_reorder( + shifted_mem, cuboid_size=mem_cuboid_size, strategy=self.strategy + ) + reordered_x = cuboid_encoder.cuboid_reorder( + shifted_x, cuboid_size=x_cuboid_size, strategy=self.strategy + ) + _, num_cuboids_mem, mem_cuboid_volume, _ = reordered_mem.shape + _, num_cuboids, x_cuboid_volume, _ = reordered_x.shape + assert ( + num_cuboids_mem == num_cuboids + ), f"Number of cuboids do not match. num_cuboids={num_cuboids}, num_cuboids_mem={num_cuboids_mem}" + attn_mask = compute_cuboid_cross_attention_mask( + T_x, + T_mem, + H, + W, + n_temporal, + cuboid_hw, + shift_hw, + strategy=self.strategy, + padding_type=self.padding_type, + device=x.place, + ) + head_C = C_in // self.num_heads + kv = ( + self.kv_proj(reordered_mem) + .reshape([B, num_cuboids, mem_cuboid_volume, 2, self.num_heads, head_C]) + .transpose(perm=[3, 0, 4, 1, 2, 5]) + ) + k, v = kv[0], kv[1] + q = ( + self.q_proj(reordered_x) + .reshape([B, num_cuboids, x_cuboid_volume, self.num_heads, head_C]) + .transpose(perm=[0, 3, 1, 2, 4]) + ) + q = q * self.scale + perm_4 = list(range(k.ndim)) + perm_4[-2] = -1 + perm_4[-1] = -2 + attn_score = q @ k.transpose(perm=perm_4) + if self.use_relative_pos: + relative_position_bias = self.relative_position_bias_table[ + self.relative_position_index[ + :x_cuboid_volume, :mem_cuboid_volume + ].reshape([-1]) + ].reshape([x_cuboid_volume, mem_cuboid_volume, -1]) + relative_position_bias = relative_position_bias.transpose( + perm=[2, 0, 1] + ).unsqueeze(axis=1) + attn_score = attn_score + relative_position_bias + if self.use_global_vector: + if self.separate_global_qkv: + l2g_q = ( + self.l2g_q_net(reordered_x) + .reshape([B, num_cuboids, x_cuboid_volume, self.num_heads, head_C]) + .transpose(perm=[0, 3, 1, 2, 4]) + ) + l2g_q = l2g_q * self.scale + l2g_global_kv = ( + self.l2g_global_kv_net(mem_global_vectors) + .reshape([B, 1, num_global, 2, self.num_heads, head_C]) + .transpose(perm=[3, 0, 4, 1, 2, 5]) + ) + l2g_global_k, l2g_global_v = l2g_global_kv[0], l2g_global_kv[1] + else: + kv_global = ( + self.kv_proj(mem_global_vectors) + .reshape([B, 1, num_global, 2, self.num_heads, head_C]) + .transpose(perm=[3, 0, 4, 1, 2, 5]) + ) + l2g_global_k, l2g_global_v = kv_global[0], kv_global[1] + l2g_q = q + perm_5 = list(range(l2g_global_k.ndim)) + perm_5[-2] = -1 + perm_5[-1] = -2 + l2g_attn_score = l2g_q @ l2g_global_k.transpose(perm=perm_5) + attn_score_l2l_l2g = paddle.concat(x=(attn_score, l2g_attn_score), axis=-1) + if attn_mask.ndim == 5: + attn_mask_l2l_l2g = F.pad( + attn_mask, [0, num_global], "constant", 1, data_format="NDHWC" + ) + else: + attn_mask_l2l_l2g = F.pad(attn_mask, [0, num_global], "constant", 1) + v_l_g = paddle.concat( + x=( + v, + l2g_global_v.expand( + shape=[B, self.num_heads, num_cuboids, num_global, head_C] + ), + ), + axis=3, + ) + attn_score_l2l_l2g = cuboid_encoder.masked_softmax( + attn_score_l2l_l2g, mask=attn_mask_l2l_l2g + ) + attn_score_l2l_l2g = self.attn_drop(attn_score_l2l_l2g) + reordered_x = ( + (attn_score_l2l_l2g @ v_l_g) + .transpose(perm=[0, 2, 3, 1, 4]) + .reshape(B, num_cuboids, x_cuboid_volume, self.dim) + ) + else: + attn_score = cuboid_encoder.masked_softmax(attn_score, mask=attn_mask) + attn_score = self.attn_drop(attn_score) + reordered_x = ( + (attn_score @ v) + .transpose(perm=[0, 2, 3, 1, 4]) + .reshape([B, num_cuboids, x_cuboid_volume, self.dim]) + ) + reordered_x = paddle.cast(reordered_x, dtype="float32") + reordered_x = self.proj_drop(self.proj(reordered_x)) + shifted_x = cuboid_encoder.cuboid_reorder_reverse( + reordered_x, + cuboid_size=x_cuboid_size, + strategy=self.strategy, + orig_data_shape=(x.shape[1], x.shape[2], x.shape[3]), + ) + if any(i > 0 for i in shift_hw): + x = paddle.roll(x=shifted_x, shifts=(shift_hw[0], shift_hw[1]), axis=(2, 3)) + else: + x = shifted_x + x = cuboid_utils.generalize_unpadding( + x, pad_t=pad_t_x, pad_h=pad_h, pad_w=pad_w, padding_type=self.padding_type + ) + return x + + +class StackCuboidCrossAttentionBlock(nn.Layer): + """A stack of cuboid cross attention layers. + + The advantage of cuboid attention is that we can combine cuboid attention building blocks with different + hyper-parameters to mimic a broad range of space-time correlation patterns. + + - "use_inter_ffn" is True + x, mem --> attn1 -----+-------> ffn1 ---+---> attn2 --> ... --> ffn_k --> out + | ^ | ^ + | | | | + |-------------|----|-------------| + - "use_inter_ffn" is False + x, mem --> attn1 -----+------> attn2 --> ... attnk --+----> ffnk ---+---> out, mem + | ^ | ^ ^ | ^ + | | | | | | | + |-------------|----|------------|-- ----------|--|-----------| + + Args: + dim (int): The dimension of the input. + num_heads (int): The number of head. + block_cuboid_hw (list, optional): The height and width of block cuboid.Defaults to [(4, 4), (4, 4)]. + block_shift_hw (list, optional): The height and width of shift cuboid . Defaults to [(0, 0), (2, 2)]. + block_n_temporal (list, optional): The length of block temporal. Defaults to [1, 2]. + block_strategy (list, optional): The strategy of block. Defaults to [("d", "d", "d"), ("l", "l", "l")]. + padding_type (str, optional): The type of paddling. Defaults to "ignore". + cross_last_n_frames (int, optional): The num of cross_last_n_frames. Defaults to None. + qkv_bias (bool, optional): Whether to enable bias in calculating qkv attention. Defaults to False. + qk_scale (float, optional): Whether to enable scale factor when calculating the attention. Defaults to None. + attn_drop (float, optional): The attention dropout. Defaults to 0.0. + proj_drop (float, optional): The projection dropout. Defaults to 0.0. + ffn_drop (float, optional): The ratio of FFN dropout. Defaults to 0.0. + activation (str, optional): The activation. Defaults to "leaky". + gated_ffn (bool, optional): Whether to use gate FFN. Defaults to False. + norm_layer (str, optional): The normalization layer. Defaults to "layer_norm". + use_inter_ffn (bool, optional): Whether to use inter FFN. Defaults to True. + max_temporal_relative (int, optional): The max temporal. Defaults to 50. + checkpoint_level (int, optional): Whether to enable gradient checkpointing. Defaults to 1. + use_relative_pos (bool, optional): Whether to use relative pos. Defaults to True. + use_global_vector (bool, optional): Whether to use the global vector or not. Defaults to False. + separate_global_qkv (bool, optional): Whether to use different network to calc q_global, k_global, v_global. Defaults to False. + global_dim_ratio (int, optional): The dim (channels) of global vectors is `global_dim_ratio*dim`. Defaults to 1. + attn_linear_init_mode (str, optional): The mode of attention linear initialization. Defaults to "0". + ffn_linear_init_mode (str, optional): The mode of FFN linear initialization. Defaults to "0". + norm_init_mode (str, optional): The mode of normalization. Defaults to "0". + """ + + def __init__( + self, + dim: int, + num_heads: int, + block_cuboid_hw: Tuple[Tuple[int, ...], ...] = [(4, 4), (4, 4)], + block_shift_hw: Tuple[Tuple[int, ...], ...] = [(0, 0), (2, 2)], + block_n_temporal: Tuple[int, ...] = [1, 2], + block_strategy: Tuple[Tuple[str, ...], ...] = [ + ("d", "d", "d"), + ("l", "l", "l"), + ], + padding_type: str = "ignore", + cross_last_n_frames: int = None, + qkv_bias: bool = False, + qk_scale: float = None, + attn_drop: float = 0.0, + proj_drop: float = 0.0, + ffn_drop: float = 0.0, + activation: str = "leaky", + gated_ffn: bool = False, + norm_layer: str = "layer_norm", + use_inter_ffn: bool = True, + max_temporal_relative: int = 50, + checkpoint_level: int = 1, + use_relative_pos: bool = True, + use_global_vector: bool = False, + separate_global_qkv: bool = False, + global_dim_ratio: int = 1, + attn_linear_init_mode: str = "0", + ffn_linear_init_mode: str = "0", + norm_init_mode: str = "0", + ): + super(StackCuboidCrossAttentionBlock, self).__init__() + self.attn_linear_init_mode = attn_linear_init_mode + self.ffn_linear_init_mode = ffn_linear_init_mode + self.norm_init_mode = norm_init_mode + if ( + len(block_cuboid_hw[0]) <= 0 + or len(block_shift_hw) <= 0 + or len(block_strategy) <= 0 + ): + raise ValueError( + "Incorrect format.The lengths of block_cuboid_hw[0], block_shift_hw, and block_strategy must be greater than zero." + ) + if len(block_cuboid_hw) != len(block_shift_hw) and len(block_shift_hw) == len( + block_strategy + ): + raise ValueError( + "The lengths of block_cuboid_size, block_shift_size, and block_strategy must be equal." + ) + + self.num_attn = len(block_cuboid_hw) + self.checkpoint_level = checkpoint_level + self.use_inter_ffn = use_inter_ffn + self.use_global_vector = use_global_vector + if self.use_inter_ffn: + self.ffn_l = nn.LayerList( + sublayers=[ + cuboid_encoder.PositionwiseFFN( + units=dim, + hidden_size=4 * dim, + activation_dropout=ffn_drop, + dropout=ffn_drop, + gated_proj=gated_ffn, + activation=activation, + normalization=norm_layer, + pre_norm=True, + linear_init_mode=ffn_linear_init_mode, + norm_init_mode=norm_init_mode, + ) + for _ in range(self.num_attn) + ] + ) + else: + self.ffn_l = nn.LayerList( + sublayers=[ + cuboid_encoder.PositionwiseFFN( + units=dim, + hidden_size=4 * dim, + activation_dropout=ffn_drop, + dropout=ffn_drop, + gated_proj=gated_ffn, + activation=activation, + normalization=norm_layer, + pre_norm=True, + linear_init_mode=ffn_linear_init_mode, + norm_init_mode=norm_init_mode, + ) + ] + ) + self.attn_l = nn.LayerList( + sublayers=[ + CuboidCrossAttentionLayer( + dim=dim, + num_heads=num_heads, + cuboid_hw=ele_cuboid_hw, + shift_hw=ele_shift_hw, + strategy=ele_strategy, + n_temporal=ele_n_temporal, + cross_last_n_frames=cross_last_n_frames, + padding_type=padding_type, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + attn_drop=attn_drop, + proj_drop=proj_drop, + norm_layer=norm_layer, + max_temporal_relative=max_temporal_relative, + use_global_vector=use_global_vector, + separate_global_qkv=separate_global_qkv, + global_dim_ratio=global_dim_ratio, + checkpoint_level=checkpoint_level, + use_relative_pos=use_relative_pos, + attn_linear_init_mode=attn_linear_init_mode, + ffn_linear_init_mode=ffn_linear_init_mode, + norm_init_mode=norm_init_mode, + ) + for ele_cuboid_hw, ele_shift_hw, ele_strategy, ele_n_temporal in zip( + block_cuboid_hw, block_shift_hw, block_strategy, block_n_temporal + ) + ] + ) + + def reset_parameters(self): + for m in self.ffn_l: + m.reset_parameters() + for m in self.attn_l: + m.reset_parameters() + + def forward(self, x, mem, mem_global_vector=None): + """ + Args: + x (paddle.Tensor): Shape (B, T_x, H, W, C) + mem (paddle.Tensor): Shape (B, T_mem, H, W, C) + mem_global_vector (paddle.Tensor): Shape (B, N_global, C) + + Returns: + out (paddle.Tensor): (B, T_x, H, W, C_out) + """ + + if self.use_inter_ffn: + for attn, ffn in zip(self.attn_l, self.ffn_l): + if self.checkpoint_level >= 2 and self.training: + x = x + fleet.utils.recompute(attn, x, mem, mem_global_vector) + else: + x = x + attn(x, mem, mem_global_vector) + if self.checkpoint_level >= 1 and self.training: + x = fleet.utils.recompute(ffn, x) + else: + x = ffn(x) + return x + else: + for attn in self.attn_l: + if self.checkpoint_level >= 2 and self.training: + x = x + fleet.utils.recompute(attn, x, mem, mem_global_vector) + else: + x = x + attn(x, mem, mem_global_vector) + if self.checkpoint_level >= 1 and self.training: + x = fleet.utils.recompute(self.ffn_l[0], x) + else: + x = self.ffn_l[0](x) + return x + + +class Upsample3DLayer(nn.Layer): + """Upsampling based on nn.UpSampling and Conv3x3. + + If the temporal dimension remains the same: + x --> interpolation-2d (nearest) --> conv3x3(dim, out_dim) + Else: + x --> interpolation-3d (nearest) --> conv3x3x3(dim, out_dim) + + Args: + dim (int): The dimension of the input tensor. + out_dim (int): The dimension of the output tensor. + target_size (Tuple[int,...]): The size of output tensor. + temporal_upsample (bool, optional): Whether the temporal axis will go through upsampling. Defaults to False. + kernel_size (int, optional): The kernel size of the Conv2D layer. Defaults to 3. + layout (str, optional): The layout of the inputs. Defaults to "THWC". + conv_init_mode (str, optional): The mode of conv initialization. Defaults to "0". + """ + + def __init__( + self, + dim: int, + out_dim: int, + target_size: Tuple[int, ...], + temporal_upsample: bool = False, + kernel_size: int = 3, + layout: str = "THWC", + conv_init_mode: str = "0", + ): + super(Upsample3DLayer, self).__init__() + self.conv_init_mode = conv_init_mode + self.target_size = target_size + self.out_dim = out_dim + self.temporal_upsample = temporal_upsample + if temporal_upsample: + self.up = nn.Upsample(size=target_size, mode="nearest") + else: + self.up = nn.Upsample(size=(target_size[1], target_size[2]), mode="nearest") + self.conv = nn.Conv2D( + in_channels=dim, + out_channels=out_dim, + kernel_size=(kernel_size, kernel_size), + padding=(kernel_size // 2, kernel_size // 2), + ) + assert layout in ["THWC", "CTHW"] + self.layout = layout + self.reset_parameters() + + def reset_parameters(self): + for m in self.children(): + cuboid_utils.apply_initialization(m, conv_mode=self.conv_init_mode) + + def forward(self, x): + """ + + Args: + x : (B, T, H, W, C) or (B, C, T, H, W) + + Returns: + out : (B, T, H_new, W_out, C_out) or (B, C, T, H_out, W_out) + """ + + if self.layout == "THWC": + B, T, H, W, C = x.shape + if self.temporal_upsample: + x = x.transpose(perm=[0, 4, 1, 2, 3]) + return self.conv(self.up(x)).transpose(perm=[0, 2, 3, 4, 1]) + else: + assert self.target_size[0] == T + x = x.reshape([B * T, H, W, C]).transpose(perm=[0, 3, 1, 2]) + x = self.up(x) + return ( + self.conv(x) + .transpose(perm=[0, 2, 3, 1]) + .reshape(list((B,) + self.target_size + (self.out_dim,))) + ) + elif self.layout == "CTHW": + B, C, T, H, W = x.shape + if self.temporal_upsample: + return self.conv(self.up(x)) + else: + assert self.output_size[0] == T + x = x.transpose(perm=[0, 2, 1, 3, 4]) + x = x.reshape([B * T, C, H, W]) + return ( + self.conv(self.up(x)) + .reshape( + [ + B, + self.target_size[0], + self.out_dim, + self.target_size[1], + self.target_size[2], + ] + ) + .transpose(perm=[0, 2, 1, 3, 4]) + ) + + +class CuboidTransformerDecoder(nn.Layer): + """Decoder of the CuboidTransformer. + + For each block, we first apply the StackCuboidSelfAttention and then apply the StackCuboidCrossAttention + + Repeat the following structure K times + + x --> StackCuboidSelfAttention --> | + |----> StackCuboidCrossAttention (If used) --> out + mem --> | + + Args: + target_temporal_length (int): The temporal length of the target. + mem_shapes (Tuple[int,...]): The mem shapes of the decoder. + cross_start (int, optional): The block to start cross attention. Defaults to 0. + depth (list, optional): The number of layers for each block. Defaults to [2, 2]. + upsample_type (str, optional): The type of upsample. Defaults to "upsample". + upsample_kernel_size (int, optional): The kernel size of upsample. Defaults to 3. + block_self_attn_patterns (str, optional): The patterns of block attention. Defaults to None. + block_self_cuboid_size (list, optional): The size of cuboid block. Defaults to [(4, 4, 4), (4, 4, 4)]. + block_self_cuboid_strategy (list, optional): The strategy of cuboid. Defaults to [("l", "l", "l"), ("d", "d", "d")]. + block_self_shift_size (list, optional): The size of shift. Defaults to [(1, 1, 1), (0, 0, 0)]. + block_cross_attn_patterns (str, optional): The patterns of cross attentions. Defaults to None. + block_cross_cuboid_hw (list, optional): The height and width of cross cuboid. Defaults to [(4, 4), (4, 4)]. + block_cross_cuboid_strategy (list, optional): The strategy of cross cuboid. Defaults to [("l", "l", "l"), ("d", "l", "l")]. + block_cross_shift_hw (list, optional): The height and width of cross shift. Defaults to [(0, 0), (0, 0)]. + block_cross_n_temporal (list, optional): The cross temporal of block. Defaults to [1, 2]. + cross_last_n_frames (int, optional): The num of cross last frames. Defaults to None. + num_heads (int, optional): The num of head. Defaults to 4. + attn_drop (float, optional): The ratio of attention dropout. Defaults to 0.0. + proj_drop (float, optional): The ratio of projection dropout. Defaults to 0.0. + ffn_drop (float, optional): The ratio of FFN dropout. Defaults to 0.0. + ffn_activation (str, optional): The activation layer of FFN. Defaults to "leaky". + gated_ffn (bool, optional): Whether to use gate FFN. Defaults to False. + norm_layer (str, optional): The normalization layer. Defaults to "layer_norm". + use_inter_ffn (bool, optional): Whether to use inter FFN. Defaults to False. + hierarchical_pos_embed (bool, optional): Whether to use hierarchical pos_embed. Defaults to False. + pos_embed_type (str, optional): The type of pos embedding. Defaults to "t+hw". + max_temporal_relative (int, optional): The max number of teemporal relative. Defaults to 50. + padding_type (str, optional): The type of padding. Defaults to "ignore". + checkpoint_level (bool, optional): Whether to enable gradient checkpointing. Defaults to True. + use_relative_pos (bool, optional): Whether to use relative pos. Defaults to True. + self_attn_use_final_proj (bool, optional): Whether to use self attention for final projection. Defaults to True. + use_first_self_attn (bool, optional): Whether to use first self attention. Defaults to False. + use_self_global (bool, optional): Whether to use self global vector. Defaults to False. + self_update_global (bool, optional): Whether to update global vector. Defaults to True. + use_cross_global (bool, optional): Whether to use cross global vector. Defaults to False. + use_global_vector_ffn (bool, optional): Whether to use FFN global vectors. Defaults to True. + use_global_self_attn (bool, optional): Whether to use global self attention. Defaults to False. + separate_global_qkv (bool, optional): Whether to use different network to calc q_global, k_global, v_global. Defaults to False. + global_dim_ratio (int, optional): The dim (channels) of global vectors is `global_dim_ratio*dim`. Defaults to 1. + attn_linear_init_mode (str, optional): The mode of attention linear initialization. Defaults to "0". + ffn_linear_init_mode (str, optional): The mode of FFN linear initialization. Defaults to "0". + conv_init_mode (str, optional): The mode of conv initialization. Defaults to "0". + up_linear_init_mode (str, optional): The mode of up linear initialization. Defaults to "0". + norm_init_mode (str, optional): The mode of normalization initialization. Defaults to "0". + """ + + def __init__( + self, + target_temporal_length: int, + mem_shapes: Tuple[int, ...], + cross_start: int = 0, + depth: Tuple[int, ...] = [2, 2], + upsample_type: str = "upsample", + upsample_kernel_size: int = 3, + block_self_attn_patterns: str = None, + block_self_cuboid_size: Tuple[Tuple[int, ...], ...] = [(4, 4, 4), (4, 4, 4)], + block_self_cuboid_strategy: Tuple[Tuple[str, ...], ...] = [ + ("l", "l", "l"), + ("d", "d", "d"), + ], + block_self_shift_size: Tuple[Tuple[int, ...], ...] = [(1, 1, 1), (0, 0, 0)], + block_cross_attn_patterns: str = None, + block_cross_cuboid_hw: Tuple[Tuple[int, ...], ...] = [(4, 4), (4, 4)], + block_cross_cuboid_strategy: Tuple[Tuple[str, ...], ...] = [ + ("l", "l", "l"), + ("d", "l", "l"), + ], + block_cross_shift_hw: Tuple[Tuple[int, ...], ...] = [(0, 0), (0, 0)], + block_cross_n_temporal: Tuple[int, ...] = [1, 2], + cross_last_n_frames: int = None, + num_heads: int = 4, + attn_drop: float = 0.0, + proj_drop: float = 0.0, + ffn_drop: float = 0.0, + ffn_activation: str = "leaky", + gated_ffn: bool = False, + norm_layer: str = "layer_norm", + use_inter_ffn: bool = False, + hierarchical_pos_embed: bool = False, + pos_embed_type: str = "t+hw", + max_temporal_relative: int = 50, + padding_type: str = "ignore", + checkpoint_level: bool = True, + use_relative_pos: bool = True, + self_attn_use_final_proj: bool = True, + use_first_self_attn: bool = False, + use_self_global: bool = False, + self_update_global: bool = True, + use_cross_global: bool = False, + use_global_vector_ffn: bool = True, + use_global_self_attn: bool = False, + separate_global_qkv: bool = False, + global_dim_ratio: int = 1, + attn_linear_init_mode: str = "0", + ffn_linear_init_mode: str = "0", + conv_init_mode: str = "0", + up_linear_init_mode: str = "0", + norm_init_mode: str = "0", + ): + super(CuboidTransformerDecoder, self).__init__() + self.attn_linear_init_mode = attn_linear_init_mode + self.ffn_linear_init_mode = ffn_linear_init_mode + self.conv_init_mode = conv_init_mode + self.up_linear_init_mode = up_linear_init_mode + self.norm_init_mode = norm_init_mode + assert len(depth) == len(mem_shapes) + self.target_temporal_length = target_temporal_length + self.num_blocks = len(mem_shapes) + self.cross_start = cross_start + self.mem_shapes = mem_shapes + self.depth = depth + self.upsample_type = upsample_type + self.hierarchical_pos_embed = hierarchical_pos_embed + self.checkpoint_level = checkpoint_level + self.use_self_global = use_self_global + self.self_update_global = self_update_global + self.use_cross_global = use_cross_global + self.use_global_vector_ffn = use_global_vector_ffn + self.use_first_self_attn = use_first_self_attn + if block_self_attn_patterns is not None: + if isinstance(block_self_attn_patterns, (tuple, list)): + assert len(block_self_attn_patterns) == self.num_blocks + else: + block_self_attn_patterns = [ + block_self_attn_patterns for _ in range(self.num_blocks) + ] + block_self_cuboid_size = [] + block_self_cuboid_strategy = [] + block_self_shift_size = [] + for idx, key in enumerate(block_self_attn_patterns): + func = cuboid_utils.CuboidSelfAttentionPatterns.get(key) + cuboid_size, strategy, shift_size = func(mem_shapes[idx]) + block_self_cuboid_size.append(cuboid_size) + block_self_cuboid_strategy.append(strategy) + block_self_shift_size.append(shift_size) + else: + if not isinstance(block_self_cuboid_size[0][0], (list, tuple)): + block_self_cuboid_size = [ + block_self_cuboid_size for _ in range(self.num_blocks) + ] + else: + assert ( + len(block_self_cuboid_size) == self.num_blocks + ), f"Incorrect input format! Received block_self_cuboid_size={block_self_cuboid_size}" + if not isinstance(block_self_cuboid_strategy[0][0], (list, tuple)): + block_self_cuboid_strategy = [ + block_self_cuboid_strategy for _ in range(self.num_blocks) + ] + else: + assert ( + len(block_self_cuboid_strategy) == self.num_blocks + ), f"Incorrect input format! Received block_self_cuboid_strategy={block_self_cuboid_strategy}" + if not isinstance(block_self_shift_size[0][0], (list, tuple)): + block_self_shift_size = [ + block_self_shift_size for _ in range(self.num_blocks) + ] + else: + assert ( + len(block_self_shift_size) == self.num_blocks + ), f"Incorrect input format! Received block_self_shift_size={block_self_shift_size}" + self_blocks = [] + for i in range(self.num_blocks): + if not self.use_first_self_attn and i == self.num_blocks - 1: + ele_depth = depth[i] - 1 + else: + ele_depth = depth[i] + stack_cuboid_blocks = [ + cuboid_encoder.StackCuboidSelfAttentionBlock( + dim=self.mem_shapes[i][-1], + num_heads=num_heads, + block_cuboid_size=block_self_cuboid_size[i], + block_strategy=block_self_cuboid_strategy[i], + block_shift_size=block_self_shift_size[i], + attn_drop=attn_drop, + proj_drop=proj_drop, + ffn_drop=ffn_drop, + activation=ffn_activation, + gated_ffn=gated_ffn, + norm_layer=norm_layer, + use_inter_ffn=use_inter_ffn, + padding_type=padding_type, + use_global_vector=use_self_global, + use_global_vector_ffn=use_global_vector_ffn, + use_global_self_attn=use_global_self_attn, + separate_global_qkv=separate_global_qkv, + global_dim_ratio=global_dim_ratio, + checkpoint_level=checkpoint_level, + use_relative_pos=use_relative_pos, + use_final_proj=self_attn_use_final_proj, + attn_linear_init_mode=attn_linear_init_mode, + ffn_linear_init_mode=ffn_linear_init_mode, + norm_init_mode=norm_init_mode, + ) + for _ in range(ele_depth) + ] + self_blocks.append(nn.LayerList(sublayers=stack_cuboid_blocks)) + self.self_blocks = nn.LayerList(sublayers=self_blocks) + if block_cross_attn_patterns is not None: + if isinstance(block_cross_attn_patterns, (tuple, list)): + assert len(block_cross_attn_patterns) == self.num_blocks + else: + block_cross_attn_patterns = [ + block_cross_attn_patterns for _ in range(self.num_blocks) + ] + block_cross_cuboid_hw = [] + block_cross_cuboid_strategy = [] + block_cross_shift_hw = [] + block_cross_n_temporal = [] + for idx, key in enumerate(block_cross_attn_patterns): + if key == "last_frame_dst": + cuboid_hw = None + shift_hw = None + strategy = None + n_temporal = None + else: + func = cuboid_utils.CuboidCrossAttentionPatterns.get(key) + cuboid_hw, shift_hw, strategy, n_temporal = func(mem_shapes[idx]) + block_cross_cuboid_hw.append(cuboid_hw) + block_cross_cuboid_strategy.append(strategy) + block_cross_shift_hw.append(shift_hw) + block_cross_n_temporal.append(n_temporal) + else: + if not isinstance(block_cross_cuboid_hw[0][0], (list, tuple)): + block_cross_cuboid_hw = [ + block_cross_cuboid_hw for _ in range(self.num_blocks) + ] + else: + assert ( + len(block_cross_cuboid_hw) == self.num_blocks + ), f"Incorrect input format! Received block_cross_cuboid_hw={block_cross_cuboid_hw}" + if not isinstance(block_cross_cuboid_strategy[0][0], (list, tuple)): + block_cross_cuboid_strategy = [ + block_cross_cuboid_strategy for _ in range(self.num_blocks) + ] + else: + assert ( + len(block_cross_cuboid_strategy) == self.num_blocks + ), f"Incorrect input format! Received block_cross_cuboid_strategy={block_cross_cuboid_strategy}" + if not isinstance(block_cross_shift_hw[0][0], (list, tuple)): + block_cross_shift_hw = [ + block_cross_shift_hw for _ in range(self.num_blocks) + ] + else: + assert ( + len(block_cross_shift_hw) == self.num_blocks + ), f"Incorrect input format! Received block_cross_shift_hw={block_cross_shift_hw}" + if not isinstance(block_cross_n_temporal[0], (list, tuple)): + block_cross_n_temporal = [ + block_cross_n_temporal for _ in range(self.num_blocks) + ] + else: + assert ( + len(block_cross_n_temporal) == self.num_blocks + ), f"Incorrect input format! Received block_cross_n_temporal={block_cross_n_temporal}" + self.cross_blocks = nn.LayerList() + for i in range(self.cross_start, self.num_blocks): + cross_block = nn.LayerList( + sublayers=[ + StackCuboidCrossAttentionBlock( + dim=self.mem_shapes[i][-1], + num_heads=num_heads, + block_cuboid_hw=block_cross_cuboid_hw[i], + block_strategy=block_cross_cuboid_strategy[i], + block_shift_hw=block_cross_shift_hw[i], + block_n_temporal=block_cross_n_temporal[i], + cross_last_n_frames=cross_last_n_frames, + attn_drop=attn_drop, + proj_drop=proj_drop, + ffn_drop=ffn_drop, + gated_ffn=gated_ffn, + norm_layer=norm_layer, + use_inter_ffn=use_inter_ffn, + activation=ffn_activation, + max_temporal_relative=max_temporal_relative, + padding_type=padding_type, + use_global_vector=use_cross_global, + separate_global_qkv=separate_global_qkv, + global_dim_ratio=global_dim_ratio, + checkpoint_level=checkpoint_level, + use_relative_pos=use_relative_pos, + attn_linear_init_mode=attn_linear_init_mode, + ffn_linear_init_mode=ffn_linear_init_mode, + norm_init_mode=norm_init_mode, + ) + for _ in range(depth[i]) + ] + ) + self.cross_blocks.append(cross_block) + if self.num_blocks > 1: + if self.upsample_type == "upsample": + self.upsample_layers = nn.LayerList( + sublayers=[ + Upsample3DLayer( + dim=self.mem_shapes[i + 1][-1], + out_dim=self.mem_shapes[i][-1], + target_size=(target_temporal_length,) + + self.mem_shapes[i][1:3], + kernel_size=upsample_kernel_size, + temporal_upsample=False, + conv_init_mode=conv_init_mode, + ) + for i in range(self.num_blocks - 1) + ] + ) + else: + raise NotImplementedError(f"{self.upsample_type} is invalid.") + if self.hierarchical_pos_embed: + self.hierarchical_pos_embed_l = nn.LayerList( + sublayers=[ + PosEmbed( + embed_dim=self.mem_shapes[i][-1], + typ=pos_embed_type, + maxT=target_temporal_length, + maxH=self.mem_shapes[i][1], + maxW=self.mem_shapes[i][2], + ) + for i in range(self.num_blocks - 1) + ] + ) + self.reset_parameters() + + def reset_parameters(self): + for ms in self.self_blocks: + for m in ms: + m.reset_parameters() + for ms in self.cross_blocks: + for m in ms: + m.reset_parameters() + if self.num_blocks > 1: + for m in self.upsample_layers: + m.reset_parameters() + if self.hierarchical_pos_embed: + for m in self.hierarchical_pos_embed_l: + m.reset_parameters() + + def forward(self, x, mem_l, mem_global_vector_l=None): + """ + Args: + x : Shape (B, T_top, H_top, W_top, C). + mem_l : A list of memory tensors. + """ + + B, T_top, H_top, W_top, C = x.shape + assert T_top == self.target_temporal_length + assert (H_top, W_top) == (self.mem_shapes[-1][1], self.mem_shapes[-1][2]) + for i in range(self.num_blocks - 1, -1, -1): + mem_global_vector = ( + None if mem_global_vector_l is None else mem_global_vector_l[i] + ) + if not self.use_first_self_attn and i == self.num_blocks - 1: + if i >= self.cross_start: + x = self.cross_blocks[i - self.cross_start][0]( + x, mem_l[i], mem_global_vector + ) + for idx in range(self.depth[i] - 1): + if self.use_self_global: + if self.self_update_global: + x, mem_global_vector = self.self_blocks[i][idx]( + x, mem_global_vector + ) + else: + x, _ = self.self_blocks[i][idx](x, mem_global_vector) + else: + x = self.self_blocks[i][idx](x) + if i >= self.cross_start: + x = self.cross_blocks[i - self.cross_start][idx + 1]( + x, mem_l[i], mem_global_vector + ) + else: + for idx in range(self.depth[i]): + if self.use_self_global: + if self.self_update_global: + x, mem_global_vector = self.self_blocks[i][idx]( + x, mem_global_vector + ) + else: + x, _ = self.self_blocks[i][idx](x, mem_global_vector) + else: + x = self.self_blocks[i][idx](x) + if i >= self.cross_start: + x = self.cross_blocks[i - self.cross_start][idx]( + x, mem_l[i], mem_global_vector + ) + if i > 0: + x = self.upsample_layers[i - 1](x) + if self.hierarchical_pos_embed: + x = self.hierarchical_pos_embed_l[i - 1](x) + return x diff --git a/examples/smc_reac/ppsci/arch/cuboid_transformer_encoder.py b/examples/smc_reac/ppsci/arch/cuboid_transformer_encoder.py new file mode 100644 index 0000000000..79b2e6fd1d --- /dev/null +++ b/examples/smc_reac/ppsci/arch/cuboid_transformer_encoder.py @@ -0,0 +1,1515 @@ +from collections import OrderedDict +from functools import lru_cache +from typing import Tuple + +import numpy as np +import paddle +import paddle.nn.functional as F +from paddle import nn +from paddle.distributed import fleet + +import ppsci.arch.cuboid_transformer_utils as cuboid_utils +from ppsci.arch import activation as act_mod +from ppsci.utils import initializer + +NEGATIVE_SLOPE = 0.1 + + +class PatchMerging3D(nn.Layer): + """Patch Merging Layer + + Args: + dim (int): Number of input channels. + out_dim (int, optional): The dim of output. Defaults to None. + downsample (tuple, optional): Downsample factor. Defaults to (1, 2, 2). + norm_layer (str, optional): The normalization layer. Defaults to "layer_norm". + padding_type (str, optional): The type of padding. Defaults to "nearest". + linear_init_mode (str, optional): The mode of linear init. Defaults to "0". + norm_init_mode (str, optional): The mode of normalization init. Defaults to "0". + """ + + def __init__( + self, + dim: int, + out_dim: int = None, + downsample: Tuple[int, ...] = (1, 2, 2), + norm_layer: str = "layer_norm", + padding_type: str = "nearest", + linear_init_mode: str = "0", + norm_init_mode: str = "0", + ): + super().__init__() + self.linear_init_mode = linear_init_mode + self.norm_init_mode = norm_init_mode + self.dim = dim + if out_dim is None: + out_dim = max(downsample) * dim + self.out_dim = out_dim + self.downsample = downsample + self.padding_type = padding_type + self.reduction = nn.Linear( + in_features=downsample[0] * downsample[1] * downsample[2] * dim, + out_features=out_dim, + bias_attr=False, + ) + self.norm = cuboid_utils.get_norm_layer( + norm_layer, in_channels=downsample[0] * downsample[1] * downsample[2] * dim + ) + self.reset_parameters() + + def reset_parameters(self): + for m in self.children(): + cuboid_utils.apply_initialization( + m, linear_mode=self.linear_init_mode, norm_mode=self.norm_init_mode + ) + + def get_out_shape(self, data_shape): + T, H, W, C_in = data_shape + pad_t = (self.downsample[0] - T % self.downsample[0]) % self.downsample[0] + pad_h = (self.downsample[1] - H % self.downsample[1]) % self.downsample[1] + pad_w = (self.downsample[2] - W % self.downsample[2]) % self.downsample[2] + return ( + (T + pad_t) // self.downsample[0], + (H + pad_h) // self.downsample[1], + (W + pad_w) // self.downsample[2], + self.out_dim, + ) + + def forward(self, x): + """ + + Args: + x : (B, T, H, W, C) + + Returns: + out : Shape (B, T // downsample[0], H // downsample[1], W // downsample[2], out_dim) + """ + + B, T, H, W, C = x.shape + pad_t = (self.downsample[0] - T % self.downsample[0]) % self.downsample[0] + pad_h = (self.downsample[1] - H % self.downsample[1]) % self.downsample[1] + pad_w = (self.downsample[2] - W % self.downsample[2]) % self.downsample[2] + if pad_h or pad_h or pad_w: + T += pad_t + H += pad_h + W += pad_w + x = cuboid_utils.generalize_padding( + x, pad_t, pad_h, pad_w, padding_type=self.padding_type + ) + x = ( + x.reshape( + ( + B, + T // self.downsample[0], + self.downsample[0], + H // self.downsample[1], + self.downsample[1], + W // self.downsample[2], + self.downsample[2], + C, + ) + ) + .transpose(perm=[0, 1, 3, 5, 2, 4, 6, 7]) + .reshape( + [ + B, + T // self.downsample[0], + H // self.downsample[1], + W // self.downsample[2], + self.downsample[0] * self.downsample[1] * self.downsample[2] * C, + ] + ) + ) + x = self.norm(x) + x = self.reduction(x) + return x + + +class PositionwiseFFN(nn.Layer): + """The Position-wise FFN layer used in Transformer-like architectures + + If pre_norm is True: + norm(data) -> fc1 -> act -> act_dropout -> fc2 -> dropout -> res(+data) + Else: + data -> fc1 -> act -> act_dropout -> fc2 -> dropout -> norm(res(+data)) + Also, if we use gated projection. We will use + fc1_1 * act(fc1_2(data)) to map the data + + Args: + units (int, optional): The units. Defaults to 512. + hidden_size (int, optional): The size of hidden layer. Defaults to 2048. + activation_dropout (float, optional): The dropout of activate. Defaults to 0.0. + dropout (float, optional): The drop ratio used in DropPat. Defaults to 0.1. + gated_proj (bool, optional): Whether to use gate projection. Defaults to False. + activation (str, optional): The activate. Defaults to "relu". + normalization (str, optional): The normalization. Defaults to "layer_norm". + layer_norm_eps (float, optional): The epsilon of layer normalization. Defaults to 1e-05. + pre_norm (bool): Pre-layer normalization as proposed in the paper: + "[ACL2018] The Best of Both Worlds: Combining Recent Advances in Neural Machine Translation" This will stabilize the training of Transformers. + You may also refer to "[Arxiv2020] Understanding the Difficulty of Training Transformers". Defaults to False. + linear_init_mode (str, optional): The mode of linear initialization. Defaults to "0". + norm_init_mode (str, optional): The mode of normalization initialization. Defaults to "0". + """ + + def __init__( + self, + units: int = 512, + hidden_size: int = 2048, + activation_dropout: float = 0.0, + dropout: float = 0.1, + gated_proj: bool = False, + activation: str = "relu", + normalization: str = "layer_norm", + layer_norm_eps: float = 1e-05, + pre_norm: bool = False, + linear_init_mode: str = "0", + norm_init_mode: str = "0", + ): + super().__init__() + self.linear_init_mode = linear_init_mode + self.norm_init_mode = norm_init_mode + self._pre_norm = pre_norm + self._gated_proj = gated_proj + self._kwargs = OrderedDict( + [ + ("units", units), + ("hidden_size", hidden_size), + ("activation_dropout", activation_dropout), + ("activation", activation), + ("dropout", dropout), + ("normalization", normalization), + ("layer_norm_eps", layer_norm_eps), + ("gated_proj", gated_proj), + ("pre_norm", pre_norm), + ] + ) + self.dropout_layer = nn.Dropout(p=dropout) + self.activation_dropout_layer = nn.Dropout(p=activation_dropout) + self.ffn_1 = nn.Linear( + in_features=units, out_features=hidden_size, bias_attr=True + ) + if self._gated_proj: + self.ffn_1_gate = nn.Linear( + in_features=units, out_features=hidden_size, bias_attr=True + ) + if activation == "leaky_relu": + self.activation = nn.LeakyReLU(NEGATIVE_SLOPE) + else: + self.activation = act_mod.get_activation(activation) + self.ffn_2 = nn.Linear( + in_features=hidden_size, out_features=units, bias_attr=True + ) + self.layer_norm = cuboid_utils.get_norm_layer( + normalization=normalization, in_channels=units, epsilon=layer_norm_eps + ) + self.reset_parameters() + + def reset_parameters(self): + cuboid_utils.apply_initialization(self.ffn_1, linear_mode=self.linear_init_mode) + if self._gated_proj: + cuboid_utils.apply_initialization( + self.ffn_1_gate, linear_mode=self.linear_init_mode + ) + cuboid_utils.apply_initialization(self.ffn_2, linear_mode=self.linear_init_mode) + cuboid_utils.apply_initialization( + self.layer_norm, norm_mode=self.norm_init_mode + ) + + def forward(self, data): + """ + Args: + x : Shape (B, seq_length, C_in) + + Returns: + out : Shape (B, seq_length, C_out) + """ + + residual = data + if self._pre_norm: + data = self.layer_norm(data) + if self._gated_proj: + out = self.activation(self.ffn_1_gate(data)) * self.ffn_1(data) + else: + out = self.activation(self.ffn_1(data)) + out = self.activation_dropout_layer(out) + out = self.ffn_2(out) + out = self.dropout_layer(out) + out = out + residual + if not self._pre_norm: + out = self.layer_norm(out) + return out + + +def update_cuboid_size_shift_size(data_shape, cuboid_size, shift_size, strategy): + """Update the cuboid_size and shift_size + + Args: + data_shape (Tuple[int,...]): The shape of the data. + cuboid_size (Tuple[int,...]): Size of the cuboid. + shift_size (Tuple[int,...]): Size of the shift. + strategy (str): The strategy of attention. + + Returns: + new_cuboid_size (Tuple[int,...]): Size of the cuboid. + new_shift_size (Tuple[int,...]): Size of the shift. + """ + + new_cuboid_size = list(cuboid_size) + new_shift_size = list(shift_size) + for i in range(len(data_shape)): + if strategy[i] == "d": + new_shift_size[i] = 0 + if data_shape[i] <= cuboid_size[i]: + new_cuboid_size[i] = data_shape[i] + new_shift_size[i] = 0 + return tuple(new_cuboid_size), tuple(new_shift_size) + + +def cuboid_reorder(data, cuboid_size, strategy): + """Reorder the tensor into (B, num_cuboids, bT * bH * bW, C) + We assume that the tensor shapes are divisible to the cuboid sizes. + + Args: + data (paddle.Tensor): The input data. + cuboid_size (Tuple[int,...]): The size of the cuboid. + strategy (Tuple[int,...]): The cuboid strategy. + + Returns: + reordered_data (paddle.Tensor): Shape will be (B, num_cuboids, bT * bH * bW, C). + num_cuboids = T / bT * H / bH * W / bW + """ + + B, T, H, W, C = data.shape + num_cuboids = T // cuboid_size[0] * H // cuboid_size[1] * W // cuboid_size[2] + cuboid_volume = cuboid_size[0] * cuboid_size[1] * cuboid_size[2] + intermediate_shape = [] + nblock_axis = [] + block_axis = [] + for i, (block_size, total_size, ele_strategy) in enumerate( + zip(cuboid_size, (T, H, W), strategy) + ): + if ele_strategy == "l": + intermediate_shape.extend([total_size // block_size, block_size]) + nblock_axis.append(2 * i + 1) + block_axis.append(2 * i + 2) + elif ele_strategy == "d": + intermediate_shape.extend([block_size, total_size // block_size]) + nblock_axis.append(2 * i + 2) + block_axis.append(2 * i + 1) + else: + raise NotImplementedError(f"{ele_strategy} is invalid.") + data = data.reshape(list((B,) + tuple(intermediate_shape) + (C,))) + reordered_data = data.transpose( + perm=(0,) + tuple(nblock_axis) + tuple(block_axis) + (7,) + ) + reordered_data = reordered_data.reshape((B, num_cuboids, cuboid_volume, C)) + return reordered_data + + +@lru_cache() +def compute_cuboid_self_attention_mask( + data_shape, cuboid_size, shift_size, strategy, padding_type, device +): + """Compute the shift window attention mask + + Args: + data_shape (Tuple[int,....]): Should be (T, H, W). + cuboid_size (Tuple[int,....]): Size of the cuboid. + shift_size (Tuple[int,....]): The shift size. + strategy (str): The decomposition strategy. + padding_type (str): Type of the padding. + device (str): The device. + + Returns: + attn_mask (paddle.Tensor): Mask with shape (num_cuboid, cuboid_vol, cuboid_vol). + The padded values will always be masked. The other masks will ensure that the shifted windows + will only attend to those in the shifted windows. + """ + T, H, W = data_shape + pad_t = (cuboid_size[0] - T % cuboid_size[0]) % cuboid_size[0] + pad_h = (cuboid_size[1] - H % cuboid_size[1]) % cuboid_size[1] + pad_w = (cuboid_size[2] - W % cuboid_size[2]) % cuboid_size[2] + data_mask = None + if pad_t > 0 or pad_h > 0 or pad_w > 0: + if padding_type == "ignore": + data_mask = paddle.ones(shape=(1, T, H, W, 1), dtype="bool") + data_mask = F.pad( + data_mask, [0, 0, 0, pad_w, 0, pad_h, 0, pad_t], data_format="NDHWC" + ) + else: + data_mask = paddle.ones( + shape=(1, T + pad_t, H + pad_h, W + pad_w, 1), dtype="bool" + ) + if any(i > 0 for i in shift_size): + if padding_type == "ignore": + data_mask = paddle.roll( + x=data_mask, + shifts=(-shift_size[0], -shift_size[1], -shift_size[2]), + axis=(1, 2, 3), + ) + if padding_type == "ignore": + data_mask = cuboid_reorder(data_mask, cuboid_size, strategy=strategy) + data_mask = data_mask.squeeze(axis=-1).squeeze(axis=0) + shift_mask = np.zeros(shape=(1, T + pad_t, H + pad_h, W + pad_w, 1)) + cnt = 0 + for t in ( + slice(-cuboid_size[0]), + slice(-cuboid_size[0], -shift_size[0]), + slice(-shift_size[0], None), + ): + for h in ( + slice(-cuboid_size[1]), + slice(-cuboid_size[1], -shift_size[1]), + slice(-shift_size[1], None), + ): + for w in ( + slice(-cuboid_size[2]), + slice(-cuboid_size[2], -shift_size[2]), + slice(-shift_size[2], None), + ): + shift_mask[:, t, h, w, :] = cnt + cnt += 1 + shift_mask = paddle.to_tensor(shift_mask) + shift_mask = cuboid_reorder(shift_mask, cuboid_size, strategy=strategy) + shift_mask = shift_mask.squeeze(axis=-1).squeeze(axis=0) + attn_mask = shift_mask.unsqueeze(axis=1) - shift_mask.unsqueeze(axis=2) == 0 + if padding_type == "ignore": + attn_mask = ( + data_mask.unsqueeze(axis=1) * data_mask.unsqueeze(axis=2) * attn_mask + ) + return attn_mask + + +def masked_softmax(att_score, mask, axis: int = -1): + """Ignore the masked elements when calculating the softmax. + The mask can be broadcastable. + + Args: + att_score (paddle.Tensor): Shape (..., length, ...) + mask (paddle.Tensor): Shape (..., length, ...) + 1 --> The element is not masked + 0 --> The element is masked + axis (int): The axis to calculate the softmax. att_score.shape[axis] must be the same as mask.shape[axis] + + Returns: + att_weights (paddle.Tensor): Shape (..., length, ...). + """ + + if mask is not None: + if att_score.dtype == paddle.float16: + att_score = att_score.masked_fill(paddle.logical_not(mask), -1e4) + else: + att_score = att_score.masked_fill(paddle.logical_not(mask), -1e18) + att_weights = nn.functional.softmax(x=att_score, axis=axis) * mask + else: + att_weights = nn.functional.softmax(x=att_score, axis=axis) + return att_weights + + +def cuboid_reorder_reverse(data, cuboid_size, strategy, orig_data_shape): + """Reverse the reordered cuboid back to the original space + + Args: + data (paddle.Tensor): The input data. + cuboid_size (Tuple[int,...]): The size of cuboid. + strategy (str): The strategy of reordering. + orig_data_shape (Tuple[int,...]): The original shape of the data. + + Returns: + data (paddle.Tensor): The recovered data + """ + + B, num_cuboids, cuboid_volume, C = data.shape + T, H, W = orig_data_shape + permutation_axis = [0] + for i, (block_size, total_size, ele_strategy) in enumerate( + zip(cuboid_size, (T, H, W), strategy) + ): + if ele_strategy == "l": + permutation_axis.append(i + 1) + permutation_axis.append(i + 4) + elif ele_strategy == "d": + permutation_axis.append(i + 4) + permutation_axis.append(i + 1) + else: + raise NotImplementedError((f"{ele_strategy} is invalid.")) + permutation_axis.append(7) + data = data.reshape( + [ + B, + T // cuboid_size[0], + H // cuboid_size[1], + W // cuboid_size[2], + cuboid_size[0], + cuboid_size[1], + cuboid_size[2], + C, + ] + ) + data = data.transpose(perm=permutation_axis) + data = data.reshape((B, T, H, W, C)) + return data + + +class CuboidSelfAttentionLayer(nn.Layer): + """Implements the cuboid self attention. + + The idea of Cuboid Self Attention is to divide the input tensor (T, H, W) into several non-overlapping cuboids. + We apply self-attention inside each cuboid and all cuboid-level self attentions are executed in parallel. + + We adopt two mechanisms for decomposing the input tensor into cuboids: + + (1) local: + We group the tensors within a local window, e.g., X[t:(t+b_t), h:(h+b_h), w:(w+b_w)]. We can also apply the + shifted window strategy proposed in "[ICCV2021] Swin Transformer: Hierarchical Vision Transformer using Shifted Windows". + (2) dilated: + Inspired by the success of dilated convolution "[ICLR2016] Multi-Scale Context Aggregation by Dilated Convolutions", + we split the tensor with dilation factors that are tied to the size of the cuboid. For example, for a cuboid that has width `b_w`, + we sample the elements starting from 0 as 0, w / b_w, 2 * w / b_w, ..., (b_w - 1) * w / b_w. + + The cuboid attention can be viewed as a generalization of the attention mechanism proposed in Video Swin Transformer, https://arxiv.org/abs/2106.13230. + The computational complexity of CuboidAttention can be simply calculated as O(T H W * b_t b_h b_w). To cover multiple correlation patterns, + we are able to combine multiple CuboidAttention layers with different configurations such as cuboid size, shift size, and local / global decomposing strategy. + + In addition, it is straight-forward to extend the cuboid attention to other types of spatiotemporal data that are not described + as regular tensors. We need to define alternative approaches to partition the data into "cuboids". + + In addition, inspired by "[NeurIPS2021] Do Transformers Really Perform Badly for Graph Representation?", + "[NeurIPS2020] Big Bird: Transformers for Longer Sequences", "[EMNLP2021] Longformer: The Long-Document Transformer", we keep + $K$ global vectors to record the global status of the spatiotemporal system. These global vectors will attend to the whole tensor and + the vectors inside each individual cuboids will also attend to the global vectors so that they can peep into the global status of the system. + + Args: + dim (int): The dimension of the input tensor. + num_heads (int): The number of heads. + cuboid_size (tuple, optional): The size of cuboid. Defaults to (2, 7, 7). + shift_size (tuple, optional): The size of shift. Defaults to (0, 0, 0). + strategy (tuple, optional): The strategy. Defaults to ("l", "l", "l"). + padding_type (str, optional): The type of padding. Defaults to "ignore". + qkv_bias (bool, optional): Whether to enable bias in calculating qkv attention. Defaults to False. + qk_scale (float, optional): Whether to enable scale factor when calculating the attention. Defaults to None. + attn_drop (float, optional): The attention dropout. Defaults to 0.0. + proj_drop (float, optional): The projection dropout. Defaults to 0.0. + use_final_proj (bool, optional): Whether to use the final projection. Defaults to True. + norm_layer (str, optional): The normalization layer. Defaults to "layer_norm". + use_global_vector (bool, optional): Whether to use the global vector or not. Defaults to False. + use_global_self_attn (bool, optional): Whether to use self attention among global vectors. Defaults to False. + separate_global_qkv (bool, optional): Whether to use different network to calc q_global, k_global, v_global. Defaults to False. + global_dim_ratio (int, optional): The dim (channels) of global vectors is `global_dim_ratio*dim`. Defaults to 1. + checkpoint_level (bool, optional): Whether to enable gradient checkpointing. Defaults to True. + use_relative_pos (bool, optional): Whether to use relative pos. Defaults to True. + attn_linear_init_mode (str, optional): The mode of attention linear initialization. Defaults to "0". + ffn_linear_init_mode (str, optional): The mode of FFN linear initialization. Defaults to "0". + norm_init_mode (str, optional): The mode of normalization initialization. Defaults to "0". + """ + + def __init__( + self, + dim: int, + num_heads: int, + cuboid_size: Tuple[int, ...] = (2, 7, 7), + shift_size: Tuple[int, ...] = (0, 0, 0), + strategy: Tuple[str, ...] = ("l", "l", "l"), + padding_type: str = "ignore", + qkv_bias: bool = False, + qk_scale: float = None, + attn_drop: float = 0.0, + proj_drop: float = 0.0, + use_final_proj: bool = True, + norm_layer: str = "layer_norm", + use_global_vector: bool = False, + use_global_self_attn: bool = False, + separate_global_qkv: bool = False, + global_dim_ratio: int = 1, + checkpoint_level: bool = True, + use_relative_pos: bool = True, + attn_linear_init_mode: str = "0", + ffn_linear_init_mode: str = "0", + norm_init_mode: str = "0", + ): + super(CuboidSelfAttentionLayer, self).__init__() + self.attn_linear_init_mode = attn_linear_init_mode + self.ffn_linear_init_mode = ffn_linear_init_mode + self.norm_init_mode = norm_init_mode + assert dim % num_heads == 0 + self.num_heads = num_heads + self.dim = dim + self.cuboid_size = cuboid_size + self.shift_size = shift_size + self.strategy = strategy + self.padding_type = padding_type + self.use_final_proj = use_final_proj + self.use_relative_pos = use_relative_pos + self.use_global_vector = use_global_vector + self.use_global_self_attn = use_global_self_attn + self.separate_global_qkv = separate_global_qkv + if global_dim_ratio != 1: + assert ( + separate_global_qkv is True + ), "Setting global_dim_ratio != 1 requires separate_global_qkv == True." + self.global_dim_ratio = global_dim_ratio + assert self.padding_type in ["ignore", "zeros", "nearest"] + head_dim = dim // num_heads + self.scale = qk_scale or head_dim**-0.5 + if use_relative_pos: + init_data = paddle.zeros( + ( + (2 * cuboid_size[0] - 1) + * (2 * cuboid_size[1] - 1) + * (2 * cuboid_size[2] - 1), + num_heads, + ) + ) + self.relative_position_bias_table = paddle.create_parameter( + shape=init_data.shape, + dtype=init_data.dtype, + default_initializer=nn.initializer.Constant(0.0), + ) + self.relative_position_bias_table.stop_gradient = not True + self.relative_position_bias_table = initializer.trunc_normal_( + self.relative_position_bias_table, std=0.02 + ) + + coords_t = paddle.arange(end=self.cuboid_size[0]) + coords_h = paddle.arange(end=self.cuboid_size[1]) + coords_w = paddle.arange(end=self.cuboid_size[2]) + coords = paddle.stack(x=paddle.meshgrid(coords_t, coords_h, coords_w)) + coords_flatten = paddle.flatten(x=coords, start_axis=1) + relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] + relative_coords = relative_coords.transpose(perm=[1, 2, 0]) + relative_coords[:, :, 0] += self.cuboid_size[0] - 1 + relative_coords[:, :, 1] += self.cuboid_size[1] - 1 + relative_coords[:, :, 2] += self.cuboid_size[2] - 1 + relative_coords[:, :, 0] *= (2 * self.cuboid_size[1] - 1) * ( + 2 * self.cuboid_size[2] - 1 + ) + relative_coords[:, :, 1] *= 2 * self.cuboid_size[2] - 1 + relative_position_index = relative_coords.sum(axis=-1) + self.register_buffer( + name="relative_position_index", tensor=relative_position_index + ) + self.qkv = nn.Linear(in_features=dim, out_features=dim * 3, bias_attr=qkv_bias) + self.attn_drop = nn.Dropout(p=attn_drop) + if self.use_global_vector: + if self.separate_global_qkv: + self.l2g_q_net = nn.Linear( + in_features=dim, out_features=dim, bias_attr=qkv_bias + ) + self.l2g_global_kv_net = nn.Linear( + in_features=global_dim_ratio * dim, + out_features=dim * 2, + bias_attr=qkv_bias, + ) + self.g2l_global_q_net = nn.Linear( + in_features=global_dim_ratio * dim, + out_features=dim, + bias_attr=qkv_bias, + ) + self.g2l_k_net = nn.Linear( + in_features=dim, out_features=dim, bias_attr=qkv_bias + ) + self.g2l_v_net = nn.Linear( + in_features=dim, + out_features=global_dim_ratio * dim, + bias_attr=qkv_bias, + ) + if self.use_global_self_attn: + self.g2g_global_qkv_net = nn.Linear( + in_features=global_dim_ratio * dim, + out_features=global_dim_ratio * dim * 3, + bias_attr=qkv_bias, + ) + else: + self.global_qkv = nn.Linear( + in_features=dim, out_features=dim * 3, bias_attr=qkv_bias + ) + self.global_attn_drop = nn.Dropout(p=attn_drop) + if use_final_proj: + self.proj = nn.Linear(in_features=dim, out_features=dim) + self.proj_drop = nn.Dropout(p=proj_drop) + if self.use_global_vector: + self.global_proj = nn.Linear( + in_features=global_dim_ratio * dim, + out_features=global_dim_ratio * dim, + ) + self.norm = cuboid_utils.get_norm_layer(norm_layer, in_channels=dim) + if self.use_global_vector: + self.global_vec_norm = cuboid_utils.get_norm_layer( + norm_layer, in_channels=global_dim_ratio * dim + ) + self.checkpoint_level = checkpoint_level + self.reset_parameters() + + def reset_parameters(self): + cuboid_utils.apply_initialization( + self.qkv, linear_mode=self.attn_linear_init_mode + ) + if self.use_final_proj: + cuboid_utils.apply_initialization( + self.proj, linear_mode=self.ffn_linear_init_mode + ) + cuboid_utils.apply_initialization(self.norm, norm_mode=self.norm_init_mode) + if self.use_global_vector: + if self.separate_global_qkv: + cuboid_utils.apply_initialization( + self.l2g_q_net, linear_mode=self.attn_linear_init_mode + ) + cuboid_utils.apply_initialization( + self.l2g_global_kv_net, linear_mode=self.attn_linear_init_mode + ) + cuboid_utils.apply_initialization( + self.g2l_global_q_net, linear_mode=self.attn_linear_init_mode + ) + cuboid_utils.apply_initialization( + self.g2l_k_net, linear_mode=self.attn_linear_init_mode + ) + cuboid_utils.apply_initialization( + self.g2l_v_net, linear_mode=self.attn_linear_init_mode + ) + if self.use_global_self_attn: + cuboid_utils.apply_initialization( + self.g2g_global_qkv_net, linear_mode=self.attn_linear_init_mode + ) + else: + cuboid_utils.apply_initialization( + self.global_qkv, linear_mode=self.attn_linear_init_mode + ) + cuboid_utils.apply_initialization( + self.global_vec_norm, norm_mode=self.norm_init_mode + ) + + def forward(self, x, global_vectors=None): + x = self.norm(x) + + B, T, H, W, C_in = x.shape + assert C_in == self.dim + if self.use_global_vector: + _, num_global, _ = global_vectors.shape + global_vectors = self.global_vec_norm(global_vectors) + cuboid_size, shift_size = update_cuboid_size_shift_size( + (T, H, W), self.cuboid_size, self.shift_size, self.strategy + ) + + pad_t = (cuboid_size[0] - T % cuboid_size[0]) % cuboid_size[0] + pad_h = (cuboid_size[1] - H % cuboid_size[1]) % cuboid_size[1] + pad_w = (cuboid_size[2] - W % cuboid_size[2]) % cuboid_size[2] + x = cuboid_utils.generalize_padding(x, pad_t, pad_h, pad_w, self.padding_type) + + if any(i > 0 for i in shift_size): + shifted_x = paddle.roll( + x=x, + shifts=(-shift_size[0], -shift_size[1], -shift_size[2]), + axis=(1, 2, 3), + ) + else: + shifted_x = x + + reordered_x = cuboid_reorder( + shifted_x, cuboid_size=cuboid_size, strategy=self.strategy + ) + + _, num_cuboids, cuboid_volume, _ = reordered_x.shape + attn_mask = compute_cuboid_self_attention_mask( + (T, H, W), + cuboid_size, + shift_size=shift_size, + strategy=self.strategy, + padding_type=self.padding_type, + device=x.place, + ) + head_C = C_in // self.num_heads + qkv = ( + self.qkv(reordered_x) + .reshape([B, num_cuboids, cuboid_volume, 3, self.num_heads, head_C]) + .transpose(perm=[3, 0, 4, 1, 2, 5]) + ) + + q, k, v = qkv[0], qkv[1], qkv[2] + q = q * self.scale + perm_0 = list(range(k.ndim)) + perm_0[-2] = -1 + perm_0[-1] = -2 + attn_score = q @ k.transpose(perm=perm_0) + + if self.use_relative_pos: + relative_position_bias = self.relative_position_bias_table[ + self.relative_position_index[:cuboid_volume, :cuboid_volume].reshape( + [-1] + ) + ].reshape([cuboid_volume, cuboid_volume, -1]) + relative_position_bias = relative_position_bias.transpose( + perm=[2, 0, 1] + ).unsqueeze(axis=1) + attn_score = attn_score + relative_position_bias + + if self.use_global_vector: + global_head_C = self.global_dim_ratio * head_C + if self.separate_global_qkv: + l2g_q = ( + self.l2g_q_net(reordered_x) + .reshape([B, num_cuboids, cuboid_volume, self.num_heads, head_C]) + .transpose(perm=[0, 3, 1, 2, 4]) + ) + l2g_q = l2g_q * self.scale + l2g_global_kv = ( + self.l2g_global_kv_net(global_vectors) + .reshape([B, 1, num_global, 2, self.num_heads, head_C]) + .transpose(perm=[3, 0, 4, 1, 2, 5]) + ) + l2g_global_k, l2g_global_v = l2g_global_kv[0], l2g_global_kv[1] + g2l_global_q = ( + self.g2l_global_q_net(global_vectors) + .reshape([B, num_global, self.num_heads, head_C]) + .transpose(perm=[0, 2, 1, 3]) + ) + g2l_global_q = g2l_global_q * self.scale + g2l_k = ( + self.g2l_k_net(reordered_x) + .reshape([B, num_cuboids, cuboid_volume, self.num_heads, head_C]) + .transpose(perm=[0, 3, 1, 2, 4]) + ) + g2l_v = ( + self.g2l_v_net(reordered_x) + .reshape( + [B, num_cuboids, cuboid_volume, self.num_heads, global_head_C] + ) + .transpose(perm=[0, 3, 1, 2, 4]) + ) + if self.use_global_self_attn: + g2g_global_qkv = ( + self.g2g_global_qkv_net(global_vectors) + .reshape([B, 1, num_global, 3, self.num_heads, global_head_C]) + .transpose(perm=[3, 0, 4, 1, 2, 5]) + ) + g2g_global_q, g2g_global_k, g2g_global_v = ( + g2g_global_qkv[0], + g2g_global_qkv[1], + g2g_global_qkv[2], + ) + g2g_global_q = g2g_global_q.squeeze(axis=2) * self.scale + else: + q_global, k_global, v_global = ( + self.global_qkv(global_vectors) + .reshape([B, 1, num_global, 3, self.num_heads, head_C]) + .transpose(perm=[3, 0, 4, 1, 2, 5]) + ) + q_global = q_global.squeeze(axis=2) * self.scale + l2g_q, g2l_k, g2l_v = q, k, v + g2l_global_q, l2g_global_k, l2g_global_v = ( + q_global, + k_global, + v_global, + ) + if self.use_global_self_attn: + g2g_global_q, g2g_global_k, g2g_global_v = ( + q_global, + k_global, + v_global, + ) + + perm_1 = list(range(l2g_global_k.ndim)) + perm_1[-2] = -1 + perm_1[-1] = -2 + l2g_attn_score = l2g_q @ l2g_global_k.transpose(perm=perm_1) + attn_score_l2l_l2g = paddle.concat(x=(attn_score, l2g_attn_score), axis=-1) + + if attn_mask.ndim == 5: + attn_mask_l2l_l2g = F.pad( + attn_mask, [0, num_global], "constant", 1, data_format="NDHWC" + ) + elif attn_mask.ndim == 3: + attn_mask = attn_mask.astype("float32") + attn_mask_l2l_l2g = F.pad( + attn_mask, [0, num_global], "constant", 1, data_format="NCL" + ) + attn_mask_l2l_l2g = attn_mask_l2l_l2g.astype("bool") + else: + attn_mask_l2l_l2g = F.pad(attn_mask, [0, num_global], "constant", 1) + + v_l_g = paddle.concat( + x=( + v, + l2g_global_v.expand( + shape=[B, self.num_heads, num_cuboids, num_global, head_C] + ), + ), + axis=3, + ) + attn_score_l2l_l2g = masked_softmax( + attn_score_l2l_l2g, mask=attn_mask_l2l_l2g + ) + attn_score_l2l_l2g = self.attn_drop(attn_score_l2l_l2g) + reordered_x = ( + (attn_score_l2l_l2g @ v_l_g) + .transpose(perm=[0, 2, 3, 1, 4]) + .reshape([B, num_cuboids, cuboid_volume, self.dim]) + ) + if self.padding_type == "ignore": + g2l_attn_mask = paddle.ones(shape=(1, T, H, W, 1)) + if pad_t > 0 or pad_h > 0 or pad_w > 0: + g2l_attn_mask = F.pad( + g2l_attn_mask, + [0, 0, 0, pad_w, 0, pad_h, 0, pad_t], + data_format="NDHWC", + ) + if any(i > 0 for i in shift_size): + g2l_attn_mask = paddle.roll( + x=g2l_attn_mask, + shifts=(-shift_size[0], -shift_size[1], -shift_size[2]), + axis=(1, 2, 3), + ) + g2l_attn_mask = g2l_attn_mask.reshape((-1,)) + else: + g2l_attn_mask = None + temp = g2l_k.reshape( + [B, self.num_heads, num_cuboids * cuboid_volume, head_C] + ) + perm_2 = list(range(temp.ndim)) + perm_2[-2] = -1 + perm_2[-1] = -2 + g2l_attn_score = g2l_global_q @ temp.transpose(perm=perm_2) + if self.use_global_self_attn: + temp = g2g_global_k.squeeze(axis=2) + perm_3 = list(range(temp.ndim)) + perm_3[-2] = -1 + perm_3[-1] = -2 + g2g_attn_score = g2g_global_q @ temp.transpose(perm=perm_3) + g2all_attn_score = paddle.concat( + x=(g2l_attn_score, g2g_attn_score), axis=-1 + ) + if g2l_attn_mask is not None: + g2all_attn_mask = F.pad( + g2l_attn_mask, + [0, num_global], + "constant", + 1, + data_format="NDHWC", + ) + else: + g2all_attn_mask = None + new_v = paddle.concat( + x=( + g2l_v.reshape( + [ + B, + self.num_heads, + num_cuboids * cuboid_volume, + global_head_C, + ] + ), + g2g_global_v.reshape( + [B, self.num_heads, num_global, global_head_C] + ), + ), + axis=2, + ) + else: + g2all_attn_score = g2l_attn_score + g2all_attn_mask = g2l_attn_mask + new_v = g2l_v.reshape( + [B, self.num_heads, num_cuboids * cuboid_volume, global_head_C] + ) + g2all_attn_score = masked_softmax(g2all_attn_score, mask=g2all_attn_mask) + g2all_attn_score = self.global_attn_drop(g2all_attn_score) + new_global_vector = ( + (g2all_attn_score @ new_v) + .transpose(perm=[0, 2, 1, 3]) + .reshape([B, num_global, self.global_dim_ratio * self.dim]) + ) + else: + attn_score = masked_softmax(attn_score, mask=attn_mask) + attn_score = self.attn_drop(attn_score) + reordered_x = ( + (attn_score @ v) + .transpose(perm=[0, 2, 3, 1, 4]) + .reshape([B, num_cuboids, cuboid_volume, self.dim]) + ) + + if self.use_final_proj: + reordered_x = paddle.cast(reordered_x, dtype="float32") + reordered_x = self.proj_drop(self.proj(reordered_x)) + if self.use_global_vector: + new_global_vector = self.proj_drop(self.global_proj(new_global_vector)) + shifted_x = cuboid_reorder_reverse( + reordered_x, + cuboid_size=cuboid_size, + strategy=self.strategy, + orig_data_shape=(T + pad_t, H + pad_h, W + pad_w), + ) + if any(i > 0 for i in shift_size): + x = paddle.roll( + x=shifted_x, + shifts=(shift_size[0], shift_size[1], shift_size[2]), + axis=(1, 2, 3), + ) + else: + x = shifted_x + x = cuboid_utils.generalize_unpadding( + x, pad_t=pad_t, pad_h=pad_h, pad_w=pad_w, padding_type=self.padding_type + ) + if self.use_global_vector: + return x, new_global_vector + else: + return x + + +class StackCuboidSelfAttentionBlock(nn.Layer): + """ + - "use_inter_ffn" is True + x --> attn1 -----+-------> ffn1 ---+---> attn2 --> ... --> ffn_k --> out + | ^ | ^ + | | | | + |-------------| |-------------| + - "use_inter_ffn" is False + x --> attn1 -----+------> attn2 --> ... attnk --+----> ffnk ---+---> out + | ^ | ^ ^ | ^ + | | | | | | | + |-------------| |------------| ----------| |-----------| + If we have enabled global memory vectors, each attention will be a + + Args: + dim (int): The dimension of the input tensor. + num_heads (int): The number of heads. + block_cuboid_size (list, optional): The size of block cuboid . Defaults to [(4, 4, 4), (4, 4, 4)]. + block_shift_size (list, optional): The shift size of block. Defaults to [(0, 0, 0), (2, 2, 2)]. + block_strategy (list, optional): The strategy of block. Defaults to [("d", "d", "d"), ("l", "l", "l")]. + padding_type (str, optional): The type of padding. Defaults to "ignore". + qkv_bias (bool, optional): Whether to enable bias in calculating qkv attention. Defaults to False. + qk_scale (float, optional): Whether to enable scale factor when calculating the attention. Defaults to None. + attn_drop (float, optional): The attention dropout. Defaults to 0.0. + proj_drop (float, optional): The projection dropout. Defaults to 0.0. + use_final_proj (bool, optional): Whether to use the final projection. Defaults to True. + norm_layer (str, optional): The normalization layer. Defaults to "layer_norm". + use_global_vector (bool, optional): Whether to use the global vector or not. Defaults to False. + use_global_self_attn (bool, optional): Whether to use self attention among global vectors. Defaults to False. + separate_global_qkv (bool, optional): Whether to use different network to calc q_global, k_global, v_global. + Defaults to False. + global_dim_ratio (int, optional): The dim (channels) of global vectors is `global_dim_ratio*dim`. + Defaults to 1. + checkpoint_level (bool, optional): Whether to enable gradient checkpointing. Defaults to True. + use_relative_pos (bool, optional): Whether to use relative pos. Defaults to True. + use_relative_pos (bool, optional): Whether to use relative pos. Defaults to True. + attn_linear_init_mode (str, optional): The mode of attention linear initialization. Defaults to "0". + ffn_linear_init_mode (str, optional): The mode of FFN linear initialization. Defaults to "0". + norm_init_mode (str, optional): The mode of normalization initialization. Defaults to "0". + """ + + def __init__( + self, + dim: int, + num_heads: int, + block_cuboid_size: Tuple[Tuple[int, ...], ...] = [(4, 4, 4), (4, 4, 4)], + block_shift_size: Tuple[Tuple[int, ...], ...] = [(0, 0, 0), (2, 2, 2)], + block_strategy: Tuple[Tuple[str, ...], ...] = [ + ("d", "d", "d"), + ("l", "l", "l"), + ], + padding_type: str = "ignore", + qkv_bias: bool = False, + qk_scale: float = None, + attn_drop: float = 0.0, + proj_drop: float = 0.0, + ffn_drop: float = 0.0, + activation: str = "leaky", + gated_ffn: bool = False, + norm_layer: str = "layer_norm", + use_inter_ffn: bool = False, + use_global_vector: bool = False, + use_global_vector_ffn: bool = True, + use_global_self_attn: bool = False, + separate_global_qkv: bool = False, + global_dim_ratio: int = 1, + checkpoint_level: bool = True, + use_relative_pos: bool = True, + use_final_proj: bool = True, + attn_linear_init_mode: str = "0", + ffn_linear_init_mode: str = "0", + norm_init_mode: str = "0", + ): + super(StackCuboidSelfAttentionBlock, self).__init__() + self.attn_linear_init_mode = attn_linear_init_mode + self.ffn_linear_init_mode = ffn_linear_init_mode + self.norm_init_mode = norm_init_mode + if ( + len(block_cuboid_size[0]) <= 0 + or len(block_shift_size) <= 0 + or len(block_strategy) <= 0 + ): + raise ValueError( + "Format of the block cuboid size is not correct. block_cuboid_size={block_cuboid_size}" + ) + if len(block_cuboid_size) != len(block_shift_size) and len( + block_cuboid_size + ) != len(block_strategy): + raise ValueError( + "The lengths of block_cuboid_size, block_shift_size, and block_strategy must be equal." + ) + + self.num_attn = len(block_cuboid_size) + self.checkpoint_level = checkpoint_level + self.use_inter_ffn = use_inter_ffn + self.use_global_vector = use_global_vector + self.use_global_vector_ffn = use_global_vector_ffn + self.use_global_self_attn = use_global_self_attn + self.global_dim_ratio = global_dim_ratio + if self.use_inter_ffn: + self.ffn_l = nn.LayerList( + sublayers=[ + PositionwiseFFN( + units=dim, + hidden_size=4 * dim, + activation_dropout=ffn_drop, + dropout=ffn_drop, + gated_proj=gated_ffn, + activation=activation, + normalization=norm_layer, + pre_norm=True, + linear_init_mode=ffn_linear_init_mode, + norm_init_mode=norm_init_mode, + ) + for _ in range(self.num_attn) + ] + ) + if self.use_global_vector_ffn and self.use_global_vector: + self.global_ffn_l = nn.LayerList( + sublayers=[ + PositionwiseFFN( + units=global_dim_ratio * dim, + hidden_size=global_dim_ratio * 4 * dim, + activation_dropout=ffn_drop, + dropout=ffn_drop, + gated_proj=gated_ffn, + activation=activation, + normalization=norm_layer, + pre_norm=True, + linear_init_mode=ffn_linear_init_mode, + norm_init_mode=norm_init_mode, + ) + for _ in range(self.num_attn) + ] + ) + else: + self.ffn_l = nn.LayerList( + sublayers=[ + PositionwiseFFN( + units=dim, + hidden_size=4 * dim, + activation_dropout=ffn_drop, + dropout=ffn_drop, + gated_proj=gated_ffn, + activation=activation, + normalization=norm_layer, + pre_norm=True, + linear_init_mode=ffn_linear_init_mode, + norm_init_mode=norm_init_mode, + ) + ] + ) + if self.use_global_vector_ffn and self.use_global_vector: + self.global_ffn_l = nn.LayerList( + sublayers=[ + PositionwiseFFN( + units=global_dim_ratio * dim, + hidden_size=global_dim_ratio * 4 * dim, + activation_dropout=ffn_drop, + dropout=ffn_drop, + gated_proj=gated_ffn, + activation=activation, + normalization=norm_layer, + pre_norm=True, + linear_init_mode=ffn_linear_init_mode, + norm_init_mode=norm_init_mode, + ) + ] + ) + self.attn_l = nn.LayerList( + sublayers=[ + CuboidSelfAttentionLayer( + dim=dim, + num_heads=num_heads, + cuboid_size=ele_cuboid_size, + shift_size=ele_shift_size, + strategy=ele_strategy, + padding_type=padding_type, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + attn_drop=attn_drop, + proj_drop=proj_drop, + norm_layer=norm_layer, + use_global_vector=use_global_vector, + use_global_self_attn=use_global_self_attn, + separate_global_qkv=separate_global_qkv, + global_dim_ratio=global_dim_ratio, + checkpoint_level=checkpoint_level, + use_relative_pos=use_relative_pos, + use_final_proj=use_final_proj, + attn_linear_init_mode=attn_linear_init_mode, + ffn_linear_init_mode=ffn_linear_init_mode, + norm_init_mode=norm_init_mode, + ) + for ele_cuboid_size, ele_shift_size, ele_strategy in zip( + block_cuboid_size, block_shift_size, block_strategy + ) + ] + ) + + def reset_parameters(self): + for m in self.ffn_l: + m.reset_parameters() + if self.use_global_vector_ffn and self.use_global_vector: + for m in self.global_ffn_l: + m.reset_parameters() + for m in self.attn_l: + m.reset_parameters() + + def forward(self, x, global_vectors=None): + if self.use_inter_ffn: + if self.use_global_vector: + for idx, (attn, ffn) in enumerate(zip(self.attn_l, self.ffn_l)): + if self.checkpoint_level >= 2 and self.training: + x_out, global_vectors_out = fleet.utils.recompute( + attn, x, global_vectors + ) + else: + x_out, global_vectors_out = attn(x, global_vectors) + x = x + x_out + global_vectors = global_vectors + global_vectors_out + if self.checkpoint_level >= 1 and self.training: + x = fleet.utils.recompute(ffn, x) + if self.use_global_vector_ffn: + global_vectors = fleet.utils.recompute( + self.global_ffn_l[idx], global_vectors + ) + else: + x = ffn(x) + if self.use_global_vector_ffn: + global_vectors = self.global_ffn_l[idx](global_vectors) + return x, global_vectors + else: + for idx, (attn, ffn) in enumerate(zip(self.attn_l, self.ffn_l)): + if self.checkpoint_level >= 2 and self.training: + x = x + fleet.utils.recompute(attn, x) + else: + x = x + attn(x) + if self.checkpoint_level >= 1 and self.training: + x = fleet.utils.recompute(ffn, x) + else: + x = ffn(x) + return x + elif self.use_global_vector: + for idx, attn in enumerate(self.attn_l): + if self.checkpoint_level >= 2 and self.training: + x_out, global_vectors_out = fleet.utils.recompute( + attn, x, global_vectors + ) + else: + x_out, global_vectors_out = attn(x, global_vectors) + x = x + x_out + global_vectors = global_vectors + global_vectors_out + if self.checkpoint_level >= 1 and self.training: + x = fleet.utils.recompute(self.ffn_l[0], x) + if self.use_global_vector_ffn: + global_vectors = fleet.utils.recompute( + self.global_ffn_l[0], global_vectors + ) + else: + x = self.ffn_l[0](x) + if self.use_global_vector_ffn: + global_vectors = self.global_ffn_l[0](global_vectors) + return x, global_vectors + else: + for idx, attn in enumerate(self.attn_l): + if self.checkpoint_level >= 2 and self.training: + out = fleet.utils.recompute(attn, x) + else: + out = attn(x) + x = x + out + if self.checkpoint_level >= 1 and self.training: + x = fleet.utils.recompute(self.ffn_l[0], x) + else: + x = self.ffn_l[0](x) + return x + + +class CuboidTransformerEncoder(nn.Layer): + """Encoder of the CuboidTransformer + + x --> attn_block --> patch_merge --> attn_block --> patch_merge --> ... --> out + + Args: + input_shape (Tuple[int,...]): The shape of the input. Contains T, H, W, C + base_units (int, optional): The number of units. Defaults to 128. + block_units (int, optional): The number of block units. Defaults to None. + scale_alpha (float, optional): We scale up the channels based on the formula: + - round_to(base_units * max(downsample_scale) ** units_alpha, 4). Defaults to 1.0. + depth (list, optional): The number of layers for each block. Defaults to [4, 4, 4]. + downsample (int, optional): The downsample ratio. Defaults to 2. + downsample_type (str, optional): The type of downsample. Defaults to "patch_merge". + block_attn_patterns (str, optional): Attention pattern for the cuboid attention for each block. Defaults to None. + block_cuboid_size (list, optional): A list of cuboid size parameters. Defaults to [(4, 4, 4), (4, 4, 4)]. + block_strategy (list, optional): A list of cuboid strategies. Defaults to [("l", "l", "l"), ("d", "d", "d")]. + block_shift_size (list, optional): A list of shift sizes. Defaults to [(0, 0, 0), (0, 0, 0)]. + num_heads (int, optional): The number of heads. Defaults to 4. + attn_drop (float, optional): The ratio of attention dropout. Defaults to 0.0. + proj_drop (float, optional): The ratio of projection dropout. Defaults to 0.0. + ffn_drop (float, optional): The ratio of FFN dropout. Defaults to 0.0. + ffn_activation (str, optional): The FFN activation. Defaults to "leaky". + gated_ffn (bool, optional): Whether to use gate FFN. Defaults to False. + norm_layer (str, optional): The normalization layer. Defaults to "layer_norm". + use_inter_ffn (bool, optional): Whether to use inter FFN. Defaults to True. + padding_type (str, optional): The type of padding. Defaults to "ignore". + checkpoint_level (bool, optional): Whether to enable gradient checkpointing. Defaults to True. + use_relative_pos (bool, optional): Whether to use relative pos. Defaults to True. + self_attn_use_final_proj (bool, optional): Whether to use self attention for final projection. Defaults to True. + use_global_vector (bool, optional): Whether to use the global vector or not. Defaults to False. + use_global_vector_ffn (bool, optional): Whether to use FFN global vectors. Defaults to False. + use_global_self_attn (bool, optional): Whether to use global self attention. Defaults to False. + separate_global_qkv (bool, optional): Whether to use different network to calc q_global, k_global, v_global. + Defaults to False. + global_dim_ratio (int, optional): The dim (channels) of global vectors is `global_dim_ratio*dim`. + Defaults to 1. + attn_linear_init_mode (str, optional): The mode of attention linear initialization. Defaults to "0". + ffn_linear_init_mode (str, optional): The mode of FFN linear initialization. Defaults to "0". + conv_init_mode (str, optional): The mode of conv initialization. Defaults to "0". + down_linear_init_mode (str, optional): The mode of downsample linear initialization. Defaults to "0". + norm_init_mode (str, optional): The mode of normalization. Defaults to "0". + """ + + def __init__( + self, + input_shape: Tuple[int, ...], + base_units: int = 128, + block_units: int = None, + scale_alpha: float = 1.0, + depth: Tuple[int, ...] = [4, 4, 4], + downsample: int = 2, + downsample_type: str = "patch_merge", + block_attn_patterns: str = None, + block_cuboid_size: Tuple[Tuple[int, ...], ...] = [(4, 4, 4), (4, 4, 4)], + block_strategy: Tuple[Tuple[str, ...], ...] = [ + ("l", "l", "l"), + ("d", "d", "d"), + ], + block_shift_size: Tuple[Tuple[int, ...], ...] = [(0, 0, 0), (0, 0, 0)], + num_heads: int = 4, + attn_drop: float = 0.0, + proj_drop: float = 0.0, + ffn_drop: float = 0.0, + ffn_activation: str = "leaky", + gated_ffn: bool = False, + norm_layer: str = "layer_norm", + use_inter_ffn: bool = True, + padding_type: str = "ignore", + checkpoint_level: bool = True, + use_relative_pos: bool = True, + self_attn_use_final_proj: bool = True, + use_global_vector: bool = False, + use_global_vector_ffn: bool = True, + use_global_self_attn: bool = False, + separate_global_qkv: bool = False, + global_dim_ratio: int = 1, + attn_linear_init_mode: str = "0", + ffn_linear_init_mode: str = "0", + conv_init_mode: str = "0", + down_linear_init_mode: str = "0", + norm_init_mode: str = "0", + ): + super(CuboidTransformerEncoder, self).__init__() + self.attn_linear_init_mode = attn_linear_init_mode + self.ffn_linear_init_mode = ffn_linear_init_mode + self.conv_init_mode = conv_init_mode + self.down_linear_init_mode = down_linear_init_mode + self.norm_init_mode = norm_init_mode + self.input_shape = input_shape + self.depth = depth + self.num_blocks = len(depth) + self.base_units = base_units + self.scale_alpha = scale_alpha + if not isinstance(downsample, (tuple, list)): + downsample = 1, downsample, downsample + self.downsample = downsample + self.downsample_type = downsample_type + self.num_heads = num_heads + self.use_global_vector = use_global_vector + self.checkpoint_level = checkpoint_level + if block_units is None: + block_units = [ + cuboid_utils.round_to( + base_units * int((max(downsample) ** scale_alpha) ** i), 4 + ) + for i in range(self.num_blocks) + ] + else: + assert len(block_units) == self.num_blocks and block_units[0] == base_units + self.block_units = block_units + if self.num_blocks > 1: + if downsample_type == "patch_merge": + self.down_layers = nn.LayerList( + sublayers=[ + PatchMerging3D( + dim=self.block_units[i], + downsample=downsample, + padding_type=padding_type, + out_dim=self.block_units[i + 1], + linear_init_mode=down_linear_init_mode, + norm_init_mode=norm_init_mode, + ) + for i in range(self.num_blocks - 1) + ] + ) + else: + raise NotImplementedError(f"{downsample_type} is invalid.") + if self.use_global_vector: + self.down_layer_global_proj = nn.LayerList( + sublayers=[ + nn.Linear( + in_features=global_dim_ratio * self.block_units[i], + out_features=global_dim_ratio * self.block_units[i + 1], + ) + for i in range(self.num_blocks - 1) + ] + ) + if block_attn_patterns is not None: + mem_shapes = self.get_mem_shapes() + if isinstance(block_attn_patterns, (tuple, list)): + assert len(block_attn_patterns) == self.num_blocks + else: + block_attn_patterns = [ + block_attn_patterns for _ in range(self.num_blocks) + ] + block_cuboid_size = [] + block_strategy = [] + block_shift_size = [] + for idx, key in enumerate(block_attn_patterns): + func = cuboid_utils.CuboidSelfAttentionPatterns.get(key) + cuboid_size, strategy, shift_size = func(mem_shapes[idx]) + block_cuboid_size.append(cuboid_size) + block_strategy.append(strategy) + block_shift_size.append(shift_size) + else: + if not isinstance(block_cuboid_size[0][0], (list, tuple)): + block_cuboid_size = [block_cuboid_size for _ in range(self.num_blocks)] + else: + assert ( + len(block_cuboid_size) == self.num_blocks + ), f"Incorrect input format! Received block_cuboid_size={block_cuboid_size}" + if not isinstance(block_strategy[0][0], (list, tuple)): + block_strategy = [block_strategy for _ in range(self.num_blocks)] + else: + assert ( + len(block_strategy) == self.num_blocks + ), f"Incorrect input format! Received block_strategy={block_strategy}" + if not isinstance(block_shift_size[0][0], (list, tuple)): + block_shift_size = [block_shift_size for _ in range(self.num_blocks)] + else: + assert ( + len(block_shift_size) == self.num_blocks + ), f"Incorrect input format! Received block_shift_size={block_shift_size}" + self.block_cuboid_size = block_cuboid_size + self.block_strategy = block_strategy + self.block_shift_size = block_shift_size + self.blocks = nn.LayerList( + sublayers=[ + nn.Sequential( + *[ + StackCuboidSelfAttentionBlock( + dim=self.block_units[i], + num_heads=num_heads, + block_cuboid_size=block_cuboid_size[i], + block_strategy=block_strategy[i], + block_shift_size=block_shift_size[i], + attn_drop=attn_drop, + proj_drop=proj_drop, + ffn_drop=ffn_drop, + activation=ffn_activation, + gated_ffn=gated_ffn, + norm_layer=norm_layer, + use_inter_ffn=use_inter_ffn, + padding_type=padding_type, + use_global_vector=use_global_vector, + use_global_vector_ffn=use_global_vector_ffn, + use_global_self_attn=use_global_self_attn, + separate_global_qkv=separate_global_qkv, + global_dim_ratio=global_dim_ratio, + checkpoint_level=checkpoint_level, + use_relative_pos=use_relative_pos, + use_final_proj=self_attn_use_final_proj, + attn_linear_init_mode=attn_linear_init_mode, + ffn_linear_init_mode=ffn_linear_init_mode, + norm_init_mode=norm_init_mode, + ) + for _ in range(depth[i]) + ] + ) + for i in range(self.num_blocks) + ] + ) + self.reset_parameters() + + def reset_parameters(self): + if self.num_blocks > 1: + for m in self.down_layers: + m.reset_parameters() + if self.use_global_vector: + cuboid_utils.apply_initialization( + self.down_layer_global_proj, linear_mode=self.down_linear_init_mode + ) + for ms in self.blocks: + for m in ms: + m.reset_parameters() + + def get_mem_shapes(self): + """Get the shape of the output memory based on the input shape. This can be used for constructing the decoder. + + Returns: + mem_shapes : A list of shapes of the output memory + """ + + if self.num_blocks == 1: + return [self.input_shape] + else: + mem_shapes = [self.input_shape] + curr_shape = self.input_shape + for down_layer in self.down_layers: + curr_shape = down_layer.get_out_shape(curr_shape) + mem_shapes.append(curr_shape) + return mem_shapes + + def forward(self, x, global_vectors=None): + """ + Args: + x : Shape (B, T, H, W, C) + + Returns: + out (List[paddle.Tensor,..]): A list of tensors from the bottom layer to the top layer of the encoder. For + example, it can have shape + - (B, T, H, W, C1) + - (B, T, H // 2, W // 2, 2 * C1) + - (B, T, H // 4, W // 4, 4 * C1) + ... + global_mem_out (List,Optional): The output of the global vector. + """ + + B, T, H, W, C_in = x.shape + assert (T, H, W, C_in) == self.input_shape + + if self.use_global_vector: + out = [] + global_mem_out = [] + for i in range(self.num_blocks): + for l in self.blocks[i]: + x, global_vectors = l(x, global_vectors) + out.append(x) + global_mem_out.append(global_vectors) + if self.num_blocks > 1 and i < self.num_blocks - 1: + x = self.down_layers[i](x) + global_vectors = self.down_layer_global_proj[i](global_vectors) + return out, global_mem_out + else: + out = [] + for i in range(self.num_blocks): + x = self.blocks[i](x) + out.append(x) + if self.num_blocks > 1 and i < self.num_blocks - 1: + x = self.down_layers[i](x) + return out diff --git a/examples/smc_reac/ppsci/arch/cuboid_transformer_utils.py b/examples/smc_reac/ppsci/arch/cuboid_transformer_utils.py new file mode 100644 index 0000000000..3f7f366bc0 --- /dev/null +++ b/examples/smc_reac/ppsci/arch/cuboid_transformer_utils.py @@ -0,0 +1,347 @@ +import functools +from typing import Tuple + +import paddle +import paddle.nn.functional as F +from paddle import nn + +from ppsci.utils import initializer + + +def round_to(dat, c): + return dat + (dat - dat % c) % c + + +class RMSNorm(nn.Layer): + """Root Mean Square Layer Normalization proposed in "[NeurIPS2019] Root Mean Square Layer Normalization" + + Args: + d (Optional[int]): The model size. + p (float, optional): The partial RMSNorm, valid value [0, 1]. Defaults to -1.0. + eps (float, optional): The epsilon value. Defaults to 1e-08. + bias (bool, optional): Whether use bias term for RMSNorm, + because RMSNorm doesn't enforce re-centering invariance.Defaults to False. + """ + + def __init__( + self, + d: Tuple[int, ...], + p: float = -1.0, + eps: float = 1e-08, + bias: bool = False, + ): + super().__init__() + self.eps = eps + self.d = d + self.p = p + self.bias = bias + init_data = paddle.ones(d) + self.scale = paddle.create_parameter( + shape=init_data.shape, + dtype=init_data.dtype, + default_initializer=nn.initializer.Constant(1.0), + ) + self.scale.stop_gradient = False + self.add_parameter(name="scale", parameter=self.scale) + if self.bias: + init_data = paddle.zeros(d) + self.offset = paddle.create_parameter( + shape=init_data.shape, + dtype=init_data.dtype, + default_initializer=nn.initializer.Constant(0.0), + ) + self.offset.stop_gradient = False + self.add_parameter(name="offset", parameter=self.offset) + + def forward(self, x): + if self.p < 0.0 or self.p > 1.0: + norm_x = x.norm(p=2, axis=-1, keepdim=True) + d_x = self.d + else: + partial_size = int(self.d * self.p) + partial_x, _ = paddle.split( + x=x, num_or_sections=[partial_size, self.d - partial_size], axis=-1 + ) + norm_x = partial_x.norm(p=2, axis=-1, keepdim=True) + d_x = partial_size + rms_x = norm_x * d_x ** (-1.0 / 2) + x_normed = x / (rms_x + self.eps) + if self.bias: + return self.scale * x_normed + self.offset + return self.scale * x_normed + + +def get_norm_layer( + normalization: str = "layer_norm", + axis: int = -1, + epsilon: float = 1e-05, + in_channels: int = 0, + **kwargs, +): + """Get the normalization layer based on the provided type + + Args: + normalization (str): The type of the layer normalization from ['layer_norm']. + axis (float): The axis to normalize the. + epsilon (float): The epsilon of the normalization layer. + in_channels (int): Input channel. + + Returns: + norm_layer (norm): The layer normalization layer. + """ + + if isinstance(normalization, str): + if normalization == "layer_norm": + assert in_channels > 0 + assert axis == -1 + norm_layer = nn.LayerNorm( + normalized_shape=in_channels, epsilon=epsilon, **kwargs + ) + elif normalization == "rms_norm": + assert axis == -1 + norm_layer = RMSNorm(d=in_channels, eps=epsilon, **kwargs) + else: + raise NotImplementedError(f"normalization={normalization} is not supported") + return norm_layer + elif normalization is None: + return nn.Identity() + else: + raise NotImplementedError("The type of normalization must be str") + + +def generalize_padding(x, pad_t, pad_h, pad_w, padding_type, t_pad_left=False): + if pad_t == 0 and pad_h == 0 and pad_w == 0: + return x + assert padding_type in ["zeros", "ignore", "nearest"] + B, T, H, W, C = x.shape + if padding_type == "nearest": + return nn.functional.interpolate( + x=x.transpose(perm=[0, 4, 1, 2, 3]), size=(T + pad_t, H + pad_h, W + pad_w) + ).transpose(perm=[0, 2, 3, 4, 1]) + elif t_pad_left: + return F.pad(x, [0, 0, 0, pad_w, 0, pad_h, pad_t, 0], data_format="NDHWC") + else: + data_pad = F.pad( + x, [0, 0, pad_t, 0, pad_h, 0, pad_w, 0, 0, 0], data_format="NDHWC" + ) + data_pad = paddle.concat( + [data_pad[:, pad_t:, ...], data_pad[:, :pad_t, ...]], axis=1 + ) + return data_pad + + +def generalize_unpadding(x, pad_t, pad_h, pad_w, padding_type): + assert padding_type in ["zeros", "ignore", "nearest"] + B, T, H, W, C = x.shape + if pad_t == 0 and pad_h == 0 and pad_w == 0: + return x + if padding_type == "nearest": + return nn.functional.interpolate( + x=x.transpose(perm=[0, 4, 1, 2, 3]), size=(T - pad_t, H - pad_h, W - pad_w) + ).transpose(perm=[0, 2, 3, 4, 1]) + else: + return x[:, : T - pad_t, : H - pad_h, : W - pad_w, :] + + +def apply_initialization( + m: nn.Layer, + linear_mode: str = "0", + conv_mode: str = "0", + norm_mode: str = "0", + embed_mode: str = "0", +): + if isinstance(m, nn.Linear): + if linear_mode in ("0",): + m.weight = initializer.kaiming_normal_(m.weight, nonlinearity="linear") + elif linear_mode in ("1",): + m.weight = initializer.kaiming_normal_( + m.weight, a=0.1, mode="fan_out", nonlinearity="leaky_relu" + ) + else: + raise NotImplementedError(f"{linear_mode} is invalid.") + if hasattr(m, "bias") and m.bias is not None: + m.bias = initializer.zeros_(m.bias) + elif isinstance( + m, + ( + nn.Conv2D, + nn.Conv3D, + nn.Conv2DTranspose, + nn.Conv3DTranspose, + ), + ): + if conv_mode in ("0",): + m.weight = initializer.kaiming_normal_( + m.weight, a=0.1, mode="fan_out", nonlinearity="leaky_relu" + ) + else: + raise NotImplementedError(f"{conv_mode} is invalid.") + if hasattr(m, "bias") and m.bias is not None: + m.bias = initializer.zeros_(m.bias) + elif isinstance(m, nn.LayerNorm): + if norm_mode in ("0",): + m.weight = initializer.zeros_(m.weight) + m.bias = initializer.zeros_(m.bias) + else: + raise NotImplementedError(f"{norm_mode} is invalid.") + elif isinstance(m, nn.GroupNorm): + if norm_mode in ("0",): + m.weight = initializer.ones_(m.weight) + m.bias = initializer.zeros_(m.bias) + else: + raise NotImplementedError(f"{norm_mode} is invalid.") + elif isinstance(m, nn.Embedding): + if embed_mode in ("0",): + m.weight.data = initializer.trunc_normal_(m.weight.data, std=0.02) + else: + raise NotImplementedError(f"{embed_mode} is invalid.") + + else: + pass + + +class CuboidSelfAttentionPatterns: + def __init__(self): + super().__init__() + self.patterns = {} + self.patterns = { + "full": self.full_attention, + "axial": self.axial, + "divided_st": self.divided_space_time, + } + for p in [1, 2, 4, 8, 10]: + for m in [1, 2, 4, 8, 16, 32]: + key = f"video_swin_{p}x{m}" + self.patterns[key] = functools.partial(self.video_swin, P=p, M=m) + + for m in [1, 2, 4, 8, 16, 32]: + key = f"spatial_lg_{m}" + self.patterns[key] = functools.partial(self.spatial_lg_v1, M=m) + + for k in [2, 4, 8]: + key = f"axial_space_dilate_{k}" + self.patterns[key] = functools.partial(self.axial_space_dilate_K, K=k) + + def get(self, pattern_name): + return self.patterns[pattern_name] + + def full_attention(self, input_shape): + T, H, W, _ = input_shape + cuboid_size = [(T, H, W)] + strategy = [("l", "l", "l")] + shift_size = [(0, 0, 0)] + return cuboid_size, strategy, shift_size + + def axial(self, input_shape): + """Axial attention proposed in https://arxiv.org/abs/1912.12180 + + Args: + input_shape (Tuple[int,...]): The shape of the input tensor, T H W. + + Returns: + cuboid_size (Tuple[int,...]): The size of cuboid. + strategy (Tuple[str,...]): The strategy of the attention. + shift_size (Tuple[int,...]): The shift size of the attention. + """ + + T, H, W, _ = input_shape + cuboid_size = [(T, 1, 1), (1, H, 1), (1, 1, W)] + strategy = [("l", "l", "l"), ("l", "l", "l"), ("l", "l", "l")] + shift_size = [(0, 0, 0), (0, 0, 0), (0, 0, 0)] + return cuboid_size, strategy, shift_size + + def divided_space_time(self, input_shape): + T, H, W, _ = input_shape + cuboid_size = [(T, 1, 1), (1, H, W)] + strategy = [("l", "l", "l"), ("l", "l", "l")] + shift_size = [(0, 0, 0), (0, 0, 0)] + return cuboid_size, strategy, shift_size + + def video_swin(self, input_shape, P=2, M=4): + """Adopt the strategy in Video SwinTransformer https://arxiv.org/pdf/2106.13230.pdf""" + T, H, W, _ = input_shape + P = min(P, T) + M = min(M, H, W) + cuboid_size = [(P, M, M), (P, M, M)] + strategy = [("l", "l", "l"), ("l", "l", "l")] + shift_size = [(0, 0, 0), (P // 2, M // 2, M // 2)] + return cuboid_size, strategy, shift_size + + def spatial_lg_v1(self, input_shape, M=4): + T, H, W, _ = input_shape + if H <= M and W <= M: + cuboid_size = [(T, 1, 1), (1, H, W)] + strategy = [("l", "l", "l"), ("l", "l", "l")] + shift_size = [(0, 0, 0), (0, 0, 0)] + else: + cuboid_size = [(T, 1, 1), (1, M, M), (1, M, M)] + strategy = [("l", "l", "l"), ("l", "l", "l"), ("d", "d", "d")] + shift_size = [(0, 0, 0), (0, 0, 0), (0, 0, 0)] + return cuboid_size, strategy, shift_size + + def axial_space_dilate_K(self, input_shape, K=2): + T, H, W, _ = input_shape + K = min(K, H, W) + cuboid_size = [ + (T, 1, 1), + (1, H // K, 1), + (1, H // K, 1), + (1, 1, W // K), + (1, 1, W // K), + ] + strategy = [ + ("l", "l", "l"), + ("d", "d", "d"), + ("l", "l", "l"), + ("d", "d", "d"), + ("l", "l", "l"), + ] + shift_size = [(0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0)] + return cuboid_size, strategy, shift_size + + +class CuboidCrossAttentionPatterns: + def __init__(self): + super().__init__() + self.patterns = {} + for k in [1, 2, 4, 8]: + key1 = f"cross_{k}x{k}" + key2 = f"cross_{k}x{k}_lg" + key3 = f"cross_{k}x{k}_heter" + self.patterns[key1] = functools.partial(self.cross_KxK, K=k) + self.patterns[key2] = functools.partial(self.cross_KxK_lg, K=k) + self.patterns[key3] = functools.partial(self.cross_KxK_heter, K=k) + + def get(self, pattern_name): + return self.patterns[pattern_name] + + def cross_KxK(self, mem_shape, K): + T_mem, H, W, _ = mem_shape + K = min(K, H, W) + cuboid_hw = [(K, K)] + shift_hw = [(0, 0)] + strategy = [("l", "l", "l")] + n_temporal = [1] + return cuboid_hw, shift_hw, strategy, n_temporal + + def cross_KxK_lg(self, mem_shape, K): + T_mem, H, W, _ = mem_shape + K = min(K, H, W) + cuboid_hw = [(K, K), (K, K)] + shift_hw = [(0, 0), (0, 0)] + strategy = [("l", "l", "l"), ("d", "d", "d")] + n_temporal = [1, 1] + return cuboid_hw, shift_hw, strategy, n_temporal + + def cross_KxK_heter(self, mem_shape, K): + T_mem, H, W, _ = mem_shape + K = min(K, H, W) + cuboid_hw = [(K, K), (K, K), (K, K)] + shift_hw = [(0, 0), (0, 0), (K // 2, K // 2)] + strategy = [("l", "l", "l"), ("d", "d", "d"), ("l", "l", "l")] + n_temporal = [1, 1, 1] + return cuboid_hw, shift_hw, strategy, n_temporal + + +CuboidSelfAttentionPatterns = CuboidSelfAttentionPatterns() +CuboidCrossAttentionPatterns = CuboidCrossAttentionPatterns() diff --git a/examples/smc_reac/ppsci/arch/cvit.py b/examples/smc_reac/ppsci/arch/cvit.py new file mode 100644 index 0000000000..d39abd3118 --- /dev/null +++ b/examples/smc_reac/ppsci/arch/cvit.py @@ -0,0 +1,1095 @@ +# 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. + +import importlib + +try: + import einops +except ModuleNotFoundError: + pass +from typing import Callable +from typing import Optional +from typing import Sequence +from typing import Tuple + +import paddle +from paddle import nn +from paddle.nn import functional as F + +from ppsci.arch import base +from ppsci.utils import initializer + + +# Positional embedding from masked autoencoder https://arxiv.org/abs/2111.06377 +def get_1d_sincos_pos_embed_from_grid(embed_dim: int, pos: paddle.Tensor): + if embed_dim % 2 != 0: + raise ValueError(f"embedding dimension({embed_dim}) must be divisible by 2") + + omega = paddle.arange(embed_dim // 2, dtype=paddle.float32) + omega /= embed_dim / 2.0 + omega = 1.0 / 10000**omega # (D/2,) + + pos = pos.reshape([-1]) # (M,) + out = paddle.einsum("m,d->md", pos, omega) # (M, D/2), outer product + + emb_sin = paddle.sin(out) # (M, D/2) + emb_cos = paddle.cos(out) # (M, D/2) + + emb = paddle.concat([emb_sin, emb_cos], axis=1) # (M, D) + return emb + + +def get_1d_sincos_pos_embed(embed_dim: int, length: int): + return paddle.unsqueeze( + get_1d_sincos_pos_embed_from_grid( + embed_dim, paddle.arange(length, dtype=paddle.float32) + ), + 0, + ) + + +def get_2d_sincos_pos_embed(embed_dim: int, grid_size: Tuple[int, int]): + def get_2d_sincos_pos_embed_from_grid(embed_dim, grid): + if embed_dim % 2 != 0: + raise ValueError(f"embedding dimension({embed_dim}) must be divisible by 2") + + # use half of dimensions to encode grid_h + emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0]) # (H*W, D/2) + emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1]) # (H*W, D/2) + emb = paddle.concat([emb_h, emb_w], axis=1) # (H*W, D) + return emb + + grid_h = paddle.arange(grid_size[0], dtype=paddle.float32) + grid_w = paddle.arange(grid_size[1], dtype=paddle.float32) + grid = paddle.meshgrid(grid_w, grid_h, indexing="ij") # here w goes first + grid = paddle.stack(grid, axis=0) + + grid = grid.reshape([2, 1, grid_size[0], grid_size[1]]) + pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid) + + return paddle.unsqueeze(pos_embed, 0) + + +class MlpBlock(nn.Layer): + def __init__(self, in_dim: int, dim: int = 256, out_dim: int = 256): + super().__init__() + self.in_dim = in_dim + self.dim = dim + self.out_dim = out_dim + self.linear1 = nn.Linear(self.in_dim, self.dim) + self.act = nn.GELU(True) + self.linear2 = nn.Linear(self.dim, self.out_dim) + + self._init_weights() + + def forward(self, inputs): + x = self.linear1(inputs) + x = self.act(x) + x = self.linear2(x) + return x + + def _init_weights(self) -> None: + initializer.xavier_uniform_(self.linear1.weight) + initializer.constant_(self.linear1.bias, 0) + initializer.xavier_uniform_(self.linear2.weight) + initializer.constant_(self.linear2.bias, 0) + + +class SelfAttnBlock(nn.Layer): + def __init__( + self, num_heads: int, emb_dim: int, mlp_ratio: int, layer_norm_eps: float = 1e-5 + ): + super().__init__() + self.num_heads = num_heads + self.emb_dim = emb_dim + self.mlp_ratio = mlp_ratio + self.layer_norm1 = nn.LayerNorm(emb_dim, layer_norm_eps) + self.attn_layer = MultiHeadDotProductAttention( + self.emb_dim, + num_heads=self.num_heads, + qkv_features=self.emb_dim, + ) + self.layer_norm2 = nn.LayerNorm(emb_dim, layer_norm_eps) + self.mlp = MlpBlock(self.emb_dim, self.emb_dim * self.mlp_ratio, self.emb_dim) + + def forward(self, inputs): + # inputs: # [B, L/ps, self.emb_dim] + x = self.layer_norm1(inputs) + x = self.attn_layer(x, x) + x = x + inputs + y = self.layer_norm2(x) + y = self.mlp(y) + return x + y + + +class Mlp(nn.Layer): + def __init__( + self, + num_layers: int, + hidden_dim: int, + out_dim: int, + layer_norm_eps: float = 1e-5, + ): + super().__init__() + self.num_layers = num_layers + self.hidden_dim = hidden_dim + self.out_dim = out_dim + self.layer_norm_eps = layer_norm_eps + self.linears = nn.LayerList( + [ + nn.Linear( + self.hidden_dim, + self.hidden_dim, + ) + for _ in range(self.num_layers) + ] + ) + self.gelu = nn.GELU(True) + self.norms = nn.LayerList( + [ + nn.LayerNorm(self.hidden_dim, self.layer_norm_eps) + for _ in range(self.num_layers) + ] + ) + + self.linear_out = nn.Linear(self.hidden_dim, self.out_dim) + + self._init_weights() + + def forward(self, inputs): + x = inputs + for i in range(self.num_layers): + y = self.linears[i](x) + y = self.gelu(y) + x = x + y + x = self.norms[i](x) + + x = self.linear_out(x) + return x + + def _init_weights(self) -> None: + for linear in self.linears: + initializer.xavier_uniform_(linear.weight) + initializer.constant_(linear.bias, 0) + + +class PatchEmbed1D(nn.Layer): + def __init__( + self, + in_dim: int, + patch_size: Sequence[int] = (4,), + emb_dim: int = 768, + use_norm: bool = False, + layer_norm_eps: float = 1e-5, + ): + super().__init__() + self.patch_size = patch_size + self.emb_dim = emb_dim + self.use_norm = use_norm + self.layer_norm_eps = layer_norm_eps + self.conv = nn.Conv1D( + in_dim, + self.emb_dim, + self.patch_size[0], + self.patch_size[0], + data_format="NLC", + ) + self.norm = ( + nn.LayerNorm(self.emb_dim, self.layer_norm_eps) + if self.use_norm + else nn.Identity() + ) + self._init_weights() + + def forward(self, x): + x = self.conv(x) # [B, L, C] --> [B, L/ps, self.emb_dim] + if self.use_norm: + x = self.norm(x) + return x + + def _init_weights(self) -> None: + initializer.xavier_uniform_(self.conv.weight) + initializer.constant_(self.conv.bias, 0) + + +class PatchEmbed(nn.Layer): + def __init__( + self, + in_dim: int, + spatial_dims: Sequence[int], + patch_size: Tuple[int, ...] = (1, 16, 16), + emb_dim: int = 768, + use_norm: bool = False, + layer_norm_eps: float = 1e-5, + ): + super().__init__() + self.patch_size = patch_size + self.emb_dim = emb_dim + self.use_norm = use_norm + self.layer_norm_eps = layer_norm_eps + self.conv = nn.Conv3D( + in_dim, + self.emb_dim, + (self.patch_size[0], self.patch_size[1], self.patch_size[2]), + (self.patch_size[0], self.patch_size[1], self.patch_size[2]), + data_format="NDHWC", + ) + self.norm = ( + nn.LayerNorm(self.emb_dim, self.layer_norm_eps) + if self.use_norm + else nn.Identity() + ) + t, h, w = spatial_dims + self.num_patches = [ + t // self.patch_size[0], + h // self.patch_size[1], + w // self.patch_size[2], + ] + self._init_weights() + + def forward(self, x): + b, t, h, w, c = x.shape + + x = self.conv(x) # [B, L, C] --> [B, L/ps, self.emb_dim] + x = x.reshape( + [ + b, + self.num_patches[0], + self.num_patches[1] * self.num_patches[2], + self.emb_dim, + ] + ) + if self.use_norm: + x = self.norm(x) + return x + + def _init_weights(self) -> None: + initializer.xavier_uniform_(self.conv.weight) + initializer.constant_(self.conv.bias, 0) + + +class CrossAttnBlock(nn.Layer): + def __init__( + self, + num_heads: int, + emb_dim: int, + mlp_ratio: int, + layer_norm_eps: float = 1e-5, + out_features: int = None, + qkv_features: int = None, + ): + super().__init__() + self.num_heads = num_heads + self.emb_dim = emb_dim + self.mlp_ratio = mlp_ratio + self.layer_norm_eps = layer_norm_eps + self.head_dim = self.emb_dim // self.num_heads + + self.layer_norm_q = nn.LayerNorm(self.emb_dim, epsilon=self.layer_norm_eps) + self.layer_norm_kv = nn.LayerNorm(self.emb_dim, epsilon=self.layer_norm_eps) + + self.attn_layer = MultiHeadDotProductAttention( + self.emb_dim, + num_heads=num_heads, + qkv_features=qkv_features, + out_features=out_features, + ) + self.layer_norm_y = nn.LayerNorm(self.emb_dim, epsilon=self.layer_norm_eps) + self.mlp = MlpBlock(self.emb_dim, self.emb_dim * self.mlp_ratio, self.emb_dim) + + def forward(self, q_inputs, kv_inputs): + # [B, L/ps, self.dec_emb_dim] + q = self.layer_norm_q(q_inputs) + kv = self.layer_norm_kv(kv_inputs) + x = self.attn_layer(q, kv) + x = x + q_inputs + y = self.layer_norm_y(x) + y = self.mlp(y) + return x + y + + +class Encoder1D(nn.Layer): + def __init__( + self, + in_dim: int, + spatial_dims: int, + patch_size: int = (4,), + emb_dim: int = 256, + depth: int = 3, + num_heads: int = 8, + mlp_ratio: int = 1, + layer_norm_eps: float = 1e-5, + ): + super().__init__() + self.in_dim = in_dim + self.spatial_dims = spatial_dims + self.patch_size = patch_size + self.emb_dim = emb_dim + self.depth = depth + self.num_heads = num_heads + self.mlp_ratio = mlp_ratio + self.layer_norm_eps = layer_norm_eps + self.patch_embedding = PatchEmbed1D(in_dim, self.patch_size, self.emb_dim) + + self.self_attn_blocks = nn.LayerList( + [ + SelfAttnBlock( + self.num_heads, + self.emb_dim, + self.mlp_ratio, + self.layer_norm_eps, + ) + for _ in range(self.depth) + ] + ) + pos_emb = get_1d_sincos_pos_embed( + self.emb_dim, self.spatial_dims // self.patch_size[0] + ) + self.pos_emb = self.create_parameter( + pos_emb.shape, default_initializer=nn.initializer.Assign(pos_emb) + ) + + def forward(self, x): + x = self.patch_embedding(x) + x = x + self.pos_emb + + for _, block in enumerate(self.self_attn_blocks): + x = block(x) + + return x + + +class TimeAggregation(nn.Layer): + def __init__( + self, + emb_dim: int, + depth: int, + num_heads: int = 8, + num_latents: int = 64, + mlp_ratio: int = 1, + layer_norm_eps: float = 1e-5, + ): + super().__init__() + self.emb_dim = emb_dim + self.depth = depth + self.num_heads = num_heads + self.num_latents = num_latents + self.mlp_ratio = mlp_ratio + self.layer_norm_eps = layer_norm_eps + self.latents = self.create_parameter( + [self.num_latents, self.emb_dim], + default_initializer=nn.initializer.Normal(std=1e-2), + ) + self.cross_attn_blocks = nn.LayerList( + [ + CrossAttnBlock( + self.num_heads, self.emb_dim, self.mlp_ratio, self.layer_norm_eps + ) + for _ in range(self.depth) + ] + ) + + def forward(self, x): # (B, T, S, D) --> (B, T', S, D) + latents = einops.repeat( + self.latents, "t d -> b s t d", b=x.shape[0], s=x.shape[2] + ) # (B, T', S, D) + x = einops.rearrange(x, "b t s d -> b s t d") # (B, S, T, D) + + # Transformer + for i, block in enumerate(self.cross_attn_blocks): + latents = block(latents, x) + + latents = einops.rearrange(latents, "b s t d -> b t s d") # (B, T', S, D) + return latents + + +class Encoder(nn.Layer): + def __init__( + self, + in_dim: int, + spatial_dims: Sequence[int], + patch_size: int = (1, 16, 16), + emb_dim: int = 256, + depth: int = 3, + num_heads: int = 8, + mlp_ratio: int = 1, + layer_norm_eps: float = 1e-5, + ): + super().__init__() + self.in_dim = in_dim + self.spatial_dims = spatial_dims + self.patch_size = patch_size + self.emb_dim = emb_dim + self.depth = depth + self.num_heads = num_heads + self.mlp_ratio = mlp_ratio + self.layer_norm_eps = layer_norm_eps + self.patch_embedding = PatchEmbed( + in_dim, spatial_dims, self.patch_size, self.emb_dim + ) + + self.time_aggreator = TimeAggregation( + self.emb_dim, + 2, + self.num_heads, + 1, + self.mlp_ratio, + self.layer_norm_eps, + ) + self.norm = nn.LayerNorm(self.emb_dim, epsilon=self.layer_norm_eps) + + self.self_attn_blocks = nn.LayerList( + [ + SelfAttnBlock( + self.num_heads, + self.emb_dim, + self.mlp_ratio, + self.layer_norm_eps, + ) + for _ in range(self.depth) + ] + ) + t, h, w = spatial_dims + + time_emb = get_1d_sincos_pos_embed(self.emb_dim, t // self.patch_size[0]) + self.time_emb = self.create_parameter( + time_emb.shape, default_initializer=nn.initializer.Assign(time_emb) + ) + + pos_emb = get_2d_sincos_pos_embed( + self.emb_dim, (h // self.patch_size[1], w // self.patch_size[2]) + ) + self.pos_emb = self.create_parameter( + pos_emb.shape, default_initializer=nn.initializer.Assign(pos_emb) + ) + + def forward(self, x): + # patchify + x = self.patch_embedding(x) + + # add positional embedding + x = x + self.time_emb.unsqueeze(2) + self.pos_emb.unsqueeze(1) + + # aggregate along time dimension + x = self.time_aggreator(x) + x = self.norm(x) + x = einops.rearrange(x, "b t s d -> b (t s) d") + + for _, block in enumerate(self.self_attn_blocks): + x = block(x) + + return x + + +def dot_product_attention_weights( + query: paddle.Tensor, + key: paddle.Tensor, + bias: Optional[paddle.Tensor] = None, +): + """Computes dot-product attention weights given query and key. + + Used by :func:`dot_product_attention`, which is what you'll most likely use. + But if you want access to the attention weights for introspection, then + you can directly call this function and call einsum yourself. + + Args: + query: queries for calculating attention with shape of [batch..., q_length, + num_heads, qk_depth_per_head]. + key: keys for calculating attention with shape of [batch..., kv_length, + num_heads, qk_depth_per_head]. + bias: bias for the attention weights. This should be broadcastable to the + shape [batch..., num_heads, q_length, kv_length]. This can be used for + incorporating causal masks, padding masks, proximity bias, etc. + + Returns: + Output of shape [batch..., num_heads, q_length, kv_length]. + """ + dtype = query.dtype + + if paddle.in_dynamic_mode(): + assert query.ndim == key.ndim, "q, k must have same rank." + assert query.shape[:-3] == key.shape[:-3], "q, k batch dims must match." + assert query.shape[-2] == key.shape[-2], "q, k num_heads must match." + assert query.shape[-1] == key.shape[-1], "q, k depths must match." + + # calculate attention matrix + depth = query.shape[-1] + query = query / (depth**0.5) + # attn weight shape is (batch..., num_heads, q_length, kv_length) + attn_weights = paddle.einsum("...qhd,...khd->...hqk", query, key) + + # apply attention bias: masking, dropout, proximity bias, etc. + if bias is not None: + attn_weights = attn_weights + bias + + # normalize the attention weights + attn_weights = F.softmax(attn_weights).astype(dtype) + + # apply attention dropout + return attn_weights + + +def dot_product_attention( + query: paddle.Tensor, + key: paddle.Tensor, + value: paddle.Tensor, + bias: Optional[paddle.Tensor] = None, +) -> paddle.Tensor: + """Computes dot-product attention given query, key, and value. + + This is the core function for applying attention based on + https://arxiv.org/abs/1706.03762. It calculates the attention weights given + query and key and combines the values using the attention weights. + + Note: query, key, value needn't have any batch dimensions. + + Args: + query: queries for calculating attention with shape of [batch..., q_length, + num_heads, qk_depth_per_head]. + key: keys for calculating attention with shape of [batch..., kv_length, + num_heads, qk_depth_per_head]. + value: values to be used in attention with shape of [batch..., kv_length, + num_heads, v_depth_per_head]. + bias: bias for the attention weights. This should be broadcastable to the + shape [batch..., num_heads, q_length, kv_length]. This can be used for + incorporating causal masks, padding masks, proximity bias, etc. + + Returns: + paddle.Tensor: Output of shape [batch..., q_length, num_heads, v_depth_per_head]. + """ + if paddle.in_dynamic_mode(): + assert key.ndim == query.ndim == value.ndim, "q, k, v must have same rank." + assert ( + query.shape[:-3] == key.shape[:-3] == value.shape[:-3] + ), "q, k, v batch dims must match." + assert ( + query.shape[-2] == key.shape[-2] == value.shape[-2] + ), "q, k, v num_heads must match." + assert key.shape[-3] == value.shape[-3], "k, v lengths must match." + + # compute attention weights + attn_weights = dot_product_attention_weights( + query, + key, + bias, + ) + + # return weighted sum over values for each query position + return paddle.einsum("...hqk,...khd->...qhd", attn_weights, value) + + +class MultiHeadDotProductAttention(nn.Layer): + """Multi-head dot-product attention. + + Args: + in_dim: Number of input dimensions. + num_heads: Number of attention heads. Features (i.e. inputs_q.shape[-1]) + should be divisible by the number of heads. + qkv_features: dimension of the key, query, and value. + out_features: dimension of the last projection + use_bias: bool: whether pointwise QKVO dense transforms use bias. + attention_fn: dot_product_attention or compatible function. Accepts query, + key, value, and returns output of shape [bs, dim1, dim2, ..., dimN,, + num_heads, value_channels]` + normalize_qk: should QK normalization be applied (arxiv.org/abs/2302.05442). + """ + + def __init__( + self, + in_dim, + num_heads: int, + qkv_features: Optional[int] = None, + out_features: Optional[int] = None, + use_bias: bool = True, + attention_fn: Callable[..., paddle.Tensor] = dot_product_attention, + normalize_qk: bool = False, + ): + super().__init__() + self.num_heads = num_heads + self.qkv_features = qkv_features or in_dim + self.out_features = out_features or in_dim + self.use_bias = use_bias + self.attention_fn = attention_fn + self.normalize_qk = normalize_qk + assert self.qkv_features % self.num_heads == 0, ( + f"Memory dimension ({self.qkv_features}) must be divisible by number of" + f" heads ({self.num_heads})." + ) + self.head_dim = self.qkv_features // self.num_heads + + self.linear_q = nn.Linear( + in_dim, + self.qkv_features, + bias_attr=use_bias, + ) + self.linear_k = nn.Linear( + in_dim, + self.qkv_features, + bias_attr=use_bias, + ) + self.linear_v = nn.Linear( + in_dim, + self.qkv_features, + bias_attr=use_bias, + ) + self.query_ln = ( + nn.LayerNorm(self.qkv_features) if normalize_qk else nn.Identity() + ) + self.key_ln = nn.LayerNorm(self.qkv_features) if normalize_qk else nn.Identity() + self.linear_out = nn.Linear( + self.qkv_features, + self.out_features, + bias_attr=use_bias, + ) + + def forward( + self, + inputs_q: paddle.Tensor, + inputs_kv: Optional[paddle.Tensor] = None, + ): + # project inputs_q to multi-headed q/k/v + # dimensions are then [batch..., length, n_heads, n_features_per_head] + q_attn_shape = inputs_q.shape + q_attn_shape = q_attn_shape[:-1] + [self.num_heads, self.head_dim] + + kv_attn_shape = inputs_kv.shape + kv_attn_shape = kv_attn_shape[:-1] + [self.num_heads, self.head_dim] + query, key, value = ( + self.linear_q(inputs_q).reshape(q_attn_shape), + self.linear_k(inputs_kv).reshape(kv_attn_shape), + self.linear_v(inputs_kv).reshape(kv_attn_shape), + ) + + if self.normalize_qk: + # Normalizing query and key projections stabilizes training with higher + # LR. See ViT-22B paper http://arxiv.org/abs/2302.05442 for analysis. + query = self.query_ln(query) + key = self.key_ln(key) + + # apply attention + x = self.attention_fn( + query, + key, + value, + ) + # back to the original inputs dimensions + x = x.reshape(x.shape[:-2] + [x.shape[-2] * x.shape[-1]]) + out = self.linear_out(x) + return out + + +class CVit1D(base.Arch): + """ + 1D Convolutional Vision Transformer (CVit1D) class. + + [Bridging Operator Learning and Conditioned Neural Fields: A Unifying Perspective](https://arxiv.org/abs/2405.13998) + + Args: + input_keys (Sequence[str]): Keys identifying the input tensors. + output_keys (Sequence[str]): Keys identifying the output tensors. + spatial_dims (int): The spatial dimensions of the input data. + in_dim (int): The dimensionality of the input data. + coords_dim (int): The dimensionality of the positional encoding. + patch_size (Sequence[int], optional): Size of the patches. Defaults to (4,). + grid_size (Sequence[int], optional): Size of the grid. Defaults to (200,). + latent_dim (int, optional): Dimensionality of the latent space. Defaults to 256. + emb_dim (int, optional): Dimensionality of the embedding space. Defaults to 256. + depth (int, optional): Number of transformer encoder layers. Defaults to 3. + num_heads (int, optional): Number of attention heads. Defaults to 8. + dec_emb_dim (int, optional): Dimensionality of the decoder embedding space. Defaults to 256. + dec_num_heads (int, optional): Number of decoder attention heads. Defaults to 8. + dec_depth (int, optional): Number of decoder transformer layers. Defaults to 1. + num_mlp_layers (int, optional): Number of layers in the MLP. Defaults to 1. + mlp_ratio (int, optional): Ratio for determining the size of the MLP's hidden layer. Defaults to 1. + out_dim (int, optional): Dimensionality of the output data. Defaults to 1. + layer_norm_eps (float, optional): Epsilon for layer normalization. Defaults to 1e-5. + embedding_type (str, optional): Type of embedding to use ("grid" or other options). Defaults to "grid". + + Examples: + >>> import ppsci + >>> b, l, c = 2, 32, 1 + >>> l_query = 42 + >>> c_in = 1 + >>> c_out = 1 + >>> model = ppsci.arch.CVit1D( + ... input_keys=["u", "y"], + ... output_keys=["s"], + ... in_dim=c_in, + ... coords_dim=1, + ... spatial_dims=l, + ... patch_size=[4], + ... grid_size=[l], + ... latent_dim=32, + ... emb_dim=32, + ... depth=3, + ... num_heads=8, + ... dec_emb_dim=32, + ... dec_num_heads=8, + ... dec_depth=1, + ... num_mlp_layers=1, + ... mlp_ratio=1, + ... out_dim=c_out, + ... layer_norm_eps=1e-5, + ... embedding_type="grid", + ... ) + >>> x = paddle.randn([b, l, c_in]) + >>> coords = paddle.randn([l_query, 1]) + >>> out = model({"u": x, "y": coords})["s"] + >>> print(out.shape) # output shape should be [b, l_query, c_out] + [2, 42, 1] + """ + + def __init__( + self, + input_keys: Sequence[str], + output_keys: Sequence[str], + spatial_dims: int, + in_dim: int, + coords_dim: int, + patch_size: Sequence[int] = (4,), + grid_size: Sequence[int] = (200,), + latent_dim: int = 256, + emb_dim: int = 256, + depth: int = 3, + num_heads: int = 8, + dec_emb_dim: int = 256, + dec_num_heads: int = 8, + dec_depth: int = 1, + num_mlp_layers: int = 1, + mlp_ratio: int = 1, + out_dim: int = 1, + layer_norm_eps: float = 1e-5, + embedding_type: str = "grid", + ): + if not importlib.util.find_spec("einops"): + raise ModuleNotFoundError( + "Please install `einops` by running 'pip install einops'." + ) + super().__init__() + self.input_keys = input_keys + self.output_keys = output_keys + self.spatial_dims = spatial_dims + self.in_dim = in_dim + self.coords_dim = coords_dim + self.patch_size = patch_size + self.grid_size = grid_size + self.latent_dim = latent_dim + self.emb_dim = emb_dim + self.depth = depth + self.num_heads = num_heads + self.dec_emb_dim = dec_emb_dim + self.dec_num_heads = dec_num_heads + self.dec_depth = dec_depth + self.num_mlp_layers = num_mlp_layers + self.mlp_ratio = mlp_ratio + self.out_dim = out_dim + self.layer_norm_eps = layer_norm_eps + self.embedding_type = embedding_type + + if self.embedding_type == "grid": + # Create grid and latents + n_x = self.grid_size[0] + self.grid = paddle.linspace(0, 1, n_x) + self.latents = self.create_parameter( + [n_x, self.latent_dim], + default_initializer=nn.initializer.Normal(std=1e-2), + ) + self.fc = nn.Linear(self.latent_dim, self.dec_emb_dim) + self.norm = nn.LayerNorm(self.dec_emb_dim, self.layer_norm_eps) + elif self.embedding_type == "mlp": + self.mlp = MlpBlock(self.latent_dim, self.dec_emb_dim, self.dec_emb_dim) + self.norm = nn.LayerNorm(self.dec_emb_dim, self.layer_norm_eps) + + self.encoder = Encoder1D( + self.in_dim, + self.spatial_dims, + self.patch_size, + self.emb_dim, + self.depth, + self.num_heads, + self.mlp_ratio, + self.layer_norm_eps, + ) + self.enc_norm = nn.LayerNorm(self.emb_dim, self.layer_norm_eps) + self.fc1 = nn.Linear(self.emb_dim, self.dec_emb_dim) + self.cross_attn_blocks = nn.LayerList( + [ + CrossAttnBlock( + self.dec_num_heads, + self.dec_emb_dim, + self.mlp_ratio, + self.layer_norm_eps, + self.dec_emb_dim, + self.dec_emb_dim, + ) + for _ in range(self.dec_depth) + ] + ) + self.block_norm = nn.LayerNorm(self.dec_emb_dim, self.layer_norm_eps) + self.final_mlp = Mlp( + self.num_mlp_layers, + self.dec_emb_dim, + self.out_dim, + layer_norm_eps=self.layer_norm_eps, + ) + + def forward_tensor(self, x, coords): + b, h, c = x.shape + + # process query coordinates + if self.embedding_type == "grid": + d2 = (coords - self.grid.unsqueeze(0)) ** 2 + w = paddle.exp(-1e5 * d2) / paddle.exp(-1e5 * d2).sum(axis=1, keepdim=True) + coords = paddle.einsum("ic,pi->pc", self.latents, w) + coords = self.fc(coords) + coords = self.norm(coords) + elif self.embedding_type == "mlp": + coords = self.mlp(coords) + coords = self.norm(coords) + + coords = einops.repeat(coords, "n d -> b n d", b=b) + + # process input function(encoder) + x = self.encoder(x) + x = self.enc_norm(x) + x = self.fc1(x) + + # decoder + for i, block in enumerate(self.cross_attn_blocks): + coords = block(coords, x) + + # mlp + x = self.block_norm(coords) + x = self.final_mlp(x) + + return x + + def forward(self, x_dict): + if self._input_transform is not None: + x = self._input_transform(x_dict) + + x, coords = x_dict[self.input_keys[0]], x_dict[self.input_keys[1]] + if coords.ndim >= 3: + coords = coords[0] # [b, n, c] -> [n, c] + + y = self.forward_tensor(x, coords) + + y_dict = {self.output_keys[0]: y} + if self._output_transform is not None: + y_dict = self._output_transform(x_dict, y_dict) + + return y_dict + + +class CVit(base.Arch): + """ + CVit architecture. + + [Bridging Operator Learning and Conditioned Neural Fields: A Unifying Perspective](https://arxiv.org/abs/2405.13998) + + Args: + input_keys (Sequence[str]): Input keys. + output_keys (Sequence[str]): Output keys. + in_dim (int): Dimensionality of the input data. + coords_dim (int): Dimensionality of the coordinates. + spatial_dims (Sequence[int]): Spatial dimensions. + patch_size (Sequence[int], optional): Size of the patches. Defaults to (1, 16, 16). + grid_size (Sequence[int], optional): Size of the grid. Defaults to (128, 128). + latent_dim (int, optional): Dimensionality of the latent space. Defaults to 256. + emb_dim (int, optional): Dimensionality of the embedding space. Defaults to 256. + depth (int, optional): Number of transformer encoder layers. Defaults to 3. + num_heads (int, optional): Number of attention heads. Defaults to 8. + dec_emb_dim (int, optional): Dimensionality of the decoder embedding space. Defaults to 256. + dec_num_heads (int, optional): Number of decoder attention heads. Defaults to 8. + dec_depth (int, optional): Number of decoder transformer layers. Defaults to 1. + num_mlp_layers (int, optional): Number of MLP layers. Defaults to 1. + mlp_ratio (int, optional): Ratio of hidden units. Defaults to 1. + out_dim (int, optional): Dimensionality of the output. Defaults to 1. + layer_norm_eps (float, optional): Epsilon value for layer normalization. Defaults to 1e-5. + embedding_type (str, optional): Type of embedding. Defaults to "grid". + + Examples: + >>> import ppsci + >>> b, t, h, w, c_in = 2, 4, 8, 8, 3 + >>> c_out = 3 + >>> h_query, w_query = 32, 32 + >>> model = ppsci.arch.CVit( + ... input_keys=["u", "y"], + ... output_keys=["s"], + ... in_dim=c_in, + ... coords_dim=2, + ... spatial_dims=[t, h, w], + ... patch_size=(1, 4, 4), + ... grid_size=(h, w), + ... latent_dim=32, + ... emb_dim=32, + ... depth=3, + ... num_heads=8, + ... dec_emb_dim=32, + ... dec_num_heads=8, + ... dec_depth=1, + ... num_mlp_layers=1, + ... mlp_ratio=1, + ... out_dim=c_out, + ... layer_norm_eps=1e-5, + ... embedding_type="grid", + ... ) + >>> x = paddle.randn([b, t, h, w, c_in]) + >>> coords = paddle.randn([h_query * w_query, 2]) + >>> out = model({"u": x, "y": coords})["s"] + >>> print(out.shape) # output shape should be [b, h_query * w_query, c_out] + [2, 1024, 3] + """ + + def __init__( + self, + input_keys: Sequence[str], + output_keys: Sequence[str], + in_dim: int, + coords_dim: int, + spatial_dims: Sequence[int], + patch_size: Sequence[int] = (1, 16, 16), + grid_size: Sequence[int] = (128, 128), + latent_dim: int = 256, + emb_dim: int = 256, + depth: int = 3, + num_heads: int = 8, + dec_emb_dim: int = 256, + dec_num_heads: int = 8, + dec_depth: int = 1, + num_mlp_layers: int = 1, + mlp_ratio: int = 1, + out_dim: int = 1, + layer_norm_eps: float = 1e-5, + embedding_type: str = "grid", + ): + super().__init__() + self.input_keys = input_keys + self.output_keys = output_keys + self.spatial_dims = spatial_dims + self.in_dim = in_dim + self.coords_dim = coords_dim + self.patch_size = patch_size + self.grid_size = grid_size + self.latent_dim = latent_dim + self.emb_dim = emb_dim + self.depth = depth + self.num_heads = num_heads + self.dec_emb_dim = dec_emb_dim + self.dec_num_heads = dec_num_heads + self.dec_depth = dec_depth + self.num_mlp_layers = num_mlp_layers + self.mlp_ratio = mlp_ratio + self.out_dim = out_dim + self.layer_norm_eps = layer_norm_eps + self.embedding_type = embedding_type + + if self.embedding_type == "grid": + # Create grid and latents + n_x, n_y = self.grid_size[0], self.grid_size[1] + + x = paddle.linspace(0, 1, n_x) + y = paddle.linspace(0, 1, n_y) + xx, yy = paddle.meshgrid(x, y, indexing="ij") + + self.grid = paddle.hstack([xx.flatten()[:, None], yy.flatten()[:, None]]) + self.latents = self.create_parameter( + [n_x * n_y, self.latent_dim], + default_initializer=nn.initializer.Normal(std=1e-2), + ) + self.fc = nn.Linear(self.latent_dim, self.dec_emb_dim) + self.norm = nn.LayerNorm(self.dec_emb_dim, self.layer_norm_eps) + elif self.embedding_type == "mlp": + self.mlp = MlpBlock(self.latent_dim, self.dec_emb_dim, self.dec_emb_dim) + self.norm = nn.LayerNorm(self.dec_emb_dim, self.layer_norm_eps) + + self.encoder = Encoder( + self.in_dim, + self.spatial_dims, + self.patch_size, + self.emb_dim, + self.depth, + self.num_heads, + self.mlp_ratio, + self.layer_norm_eps, + ) + self.enc_norm = nn.LayerNorm(self.emb_dim, self.layer_norm_eps) + self.fc1 = nn.Linear(self.emb_dim, self.dec_emb_dim) + self.cross_attn_blocks = nn.LayerList( + [ + CrossAttnBlock( + self.dec_num_heads, + self.dec_emb_dim, + self.mlp_ratio, + self.layer_norm_eps, + self.dec_emb_dim, + self.dec_emb_dim, + ) + for _ in range(self.dec_depth) + ] + ) + self.block_norm = nn.LayerNorm(self.dec_emb_dim, self.layer_norm_eps) + self.final_mlp = Mlp( + self.num_mlp_layers, + self.dec_emb_dim, + self.out_dim, + layer_norm_eps=self.layer_norm_eps, + ) + + def forward_tensor(self, x, coords): + b, t, h, w, c = x.shape + + # process query coordinates + if self.embedding_type == "grid": + d2 = ((coords.unsqueeze(1) - self.grid.unsqueeze(0)) ** 2).sum(axis=2) + w = paddle.exp(-1e5 * d2) / paddle.exp(-1e5 * d2).sum(axis=1, keepdim=True) + coords = paddle.einsum("ic,pi->pc", self.latents, w) + coords = self.fc(coords) + coords = self.norm(coords) + elif self.embedding_type == "mlp": + coords = self.mlp(coords) + coords = self.norm(coords) + + coords = einops.repeat(coords, "n d -> b n d", b=b) + + # process input function(encoder) + x = self.encoder(x) + x = self.enc_norm(x) + x = self.fc1(x) + + # decoder + for i, block in enumerate(self.cross_attn_blocks): + coords = block(coords, x) + + # mlp + x = self.block_norm(coords) + x = self.final_mlp(x) + + return x + + def forward(self, x_dict): + if self._input_transform is not None: + x = self._input_transform(x_dict) + + x, coords = x_dict[self.input_keys[0]], x_dict[self.input_keys[1]] + if coords.ndim >= 3: + coords = coords[0] # [b, n, c] -> [n, c] + + y = self.forward_tensor(x, coords) + + y_dict = {self.output_keys[0]: y} + if self._output_transform is not None: + y_dict = self._output_transform(x_dict, y_dict) + + return y_dict diff --git a/examples/smc_reac/ppsci/arch/deeponet.py b/examples/smc_reac/ppsci/arch/deeponet.py new file mode 100644 index 0000000000..16a8807d81 --- /dev/null +++ b/examples/smc_reac/ppsci/arch/deeponet.py @@ -0,0 +1,154 @@ +# 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. + +from __future__ import annotations + +from typing import Tuple +from typing import Union + +import paddle +import paddle.nn as nn + +from ppsci.arch import activation as act_mod +from ppsci.arch import base +from ppsci.arch import mlp + + +class DeepONet(base.Arch): + """Deep operator network. + + [Lu et al. Learning nonlinear operators via DeepONet based on the universal approximation theorem of operators. Nat Mach Intell, 2021.](https://doi.org/10.1038/s42256-021-00302-5) + + Args: + u_key (str): Name of function data for input function u(x). + y_key (str): Name of location data for input function G(u). + G_key (str): Output name of predicted G(u)(y). + num_loc (int): Number of sampled u(x), i.e. `m` in paper. + num_features (int): Number of features extracted from u(x), same for y. + branch_num_layers (int): Number of hidden layers of branch net. + trunk_num_layers (int): Number of hidden layers of trunk net. + branch_hidden_size (Union[int, Tuple[int, ...]]): Number of hidden size of branch net. + An integer for all layers, or list of integer specify each layer's size. + trunk_hidden_size (Union[int, Tuple[int, ...]]): Number of hidden size of trunk net. + An integer for all layers, or list of integer specify each layer's size. + branch_skip_connection (bool, optional): Whether to use skip connection for branch net. Defaults to False. + trunk_skip_connection (bool, optional): Whether to use skip connection for trunk net. Defaults to False. + branch_activation (str, optional): Name of activation function. Defaults to "tanh". + trunk_activation (str, optional): Name of activation function. Defaults to "tanh". + branch_weight_norm (bool, optional): Whether to apply weight norm on parameter(s) for branch net. Defaults to False. + trunk_weight_norm (bool, optional): Whether to apply weight norm on parameter(s) for trunk net. Defaults to False. + use_bias (bool, optional): Whether to add bias on predicted G(u)(y). Defaults to True. + + Examples: + >>> import paddle + >>> import ppsci + >>> model = ppsci.arch.DeepONet( + ... "u", "y", "G", + ... 100, 40, + ... 1, 1, + ... 40, 40, + ... branch_activation="relu", trunk_activation="relu", + ... use_bias=True, + ... ) + >>> input_dict = {"u": paddle.rand([200, 100]), + ... "y": paddle.rand([200, 1])} + >>> output_dict = model(input_dict) + >>> print(output_dict["G"].shape) + [200, 1] + """ + + def __init__( + self, + u_key: str, + y_key: str, + G_key: str, + num_loc: int, + num_features: int, + branch_num_layers: int, + trunk_num_layers: int, + branch_hidden_size: Union[int, Tuple[int, ...]], + trunk_hidden_size: Union[int, Tuple[int, ...]], + branch_skip_connection: bool = False, + trunk_skip_connection: bool = False, + branch_activation: str = "tanh", + trunk_activation: str = "tanh", + branch_weight_norm: bool = False, + trunk_weight_norm: bool = False, + use_bias: bool = True, + ): + super().__init__() + self.u_key = u_key + self.y_key = y_key + self.input_keys = (u_key, y_key) + self.output_keys = (G_key,) + + self.branch_net = mlp.MLP( + (self.u_key,), + ("b",), + branch_num_layers, + branch_hidden_size, + branch_activation, + branch_skip_connection, + branch_weight_norm, + input_dim=num_loc, + output_dim=num_features, + ) + + self.trunk_net = mlp.MLP( + (self.y_key,), + ("t",), + trunk_num_layers, + trunk_hidden_size, + trunk_activation, + trunk_skip_connection, + trunk_weight_norm, + input_dim=1, + output_dim=num_features, + ) + self.trunk_act = act_mod.get_activation(trunk_activation) + + self.use_bias = use_bias + if use_bias: + # register bias to parameter for updating in optimizer and storage + self.b = self.create_parameter( + shape=(1,), + attr=nn.initializer.Constant(0.0), + ) + + def forward(self, x): + if self._input_transform is not None: + x = self._input_transform(x) + + # Branch net to encode the input function + u_features = self.branch_net(x)[self.branch_net.output_keys[0]] + + # Trunk net to encode the domain of the output function + y_features = self.trunk_net(x) + y_features = self.trunk_act(y_features[self.trunk_net.output_keys[0]]) + + # Dot product + G_u = paddle.einsum("bi,bi->b", u_features, y_features) # [batch_size, ] + G_u = paddle.reshape(G_u, [-1, 1]) # reshape [batch_size, ] to [batch_size, 1] + + # Add bias + if self.use_bias: + G_u += self.b + + result_dict = { + self.output_keys[0]: G_u, + } + if self._output_transform is not None: + result_dict = self._output_transform(x, result_dict) + + return result_dict diff --git a/examples/smc_reac/ppsci/arch/dgmr.py b/examples/smc_reac/ppsci/arch/dgmr.py new file mode 100644 index 0000000000..dd189bb2de --- /dev/null +++ b/examples/smc_reac/ppsci/arch/dgmr.py @@ -0,0 +1,1151 @@ +from typing import List +from typing import Tuple + +import paddle +import paddle.nn as nn + +from ppsci.arch import base + +try: + import einops +except ModuleNotFoundError: + pass + + +class DGMR(base.Arch): + """Deep Generative Model of Radar. + Nowcasting GAN is an attempt to recreate DeepMind's Skillful Nowcasting GAN from https://arxiv.org/abs/2104.00954. + but slightly modified for multiple satellite channels + + Args: + input_keys (Tuple[str, ...]): Name of input keys, such as ("input",). + output_keys (Tuple[str, ...]): Name of output keys, such as ("output",). + forecast_steps (int, optional): Number of steps to predict in the future + input_channels (int, optional): Number of input channels per image + gen_lr (float, optional): Learning rate for the generator + disc_lr (float, optional): Learning rate for the discriminators, shared for both temporal and spatial discriminator + conv_type (str, optional): Type of 2d convolution to use, see satflow/models/utils.py for options + beta1 (float, optional): Beta1 for Adam optimizer + beta2 (float, optional): Beta2 for Adam optimizer + num_samples (int, optional): Number of samples of the latent space to sample for training/validation + grid_lambda (float, optional): Lambda for the grid regularization loss + output_shape (int, optional): Shape of the output predictions, generally should be same as the input shape + generation_steps (int, optional): Number of generation steps to use in forward pass, in paper is 6 and the best is chosen for the loss + this results in huge amounts of GPU memory though, so less might work better for training. + context_channels (int, optional): Number of output channels for the lowest block of conditioning stack + latent_channels (int, optional): Number of channels that the latent space should be reshaped to, + input dimension into ConvGRU, also affects the number of channels for other linked inputs/outputs + + Examples: + >>> import ppsci + >>> import paddle + >>> model = ppsci.arch.DGMR(("input", ), ("output", )) + >>> input_dict = {"input": paddle.randn((1, 4, 1, 256, 256))} + >>> output_dict = model(input_dict) # doctest: +SKIP + >>> print(output_dict["output"].shape) # doctest: +SKIP + [1, 18, 1, 256, 256] + """ + + def __init__( + self, + input_keys: Tuple[str, ...], + output_keys: Tuple[str, ...], + forecast_steps: int = 18, + input_channels: int = 1, + output_shape: int = 256, + gen_lr: float = 5e-05, + disc_lr: float = 0.0002, + conv_type: str = "standard", + num_samples: int = 6, + grid_lambda: float = 20.0, + beta1: float = 0.0, + beta2: float = 0.999, + latent_channels: int = 768, + context_channels: int = 384, + generation_steps: int = 6, + ): + super().__init__() + self.input_keys = input_keys + self.output_keys = output_keys + self.gen_lr = gen_lr + self.disc_lr = disc_lr + self.beta1 = beta1 + self.beta2 = beta2 + self.grid_lambda = grid_lambda + self.num_samples = num_samples + self.latent_channels = latent_channels + self.context_channels = context_channels + self.input_channels = input_channels + self.generation_steps = generation_steps + self.conditioning_stack = ContextConditioningStack( + input_channels=input_channels, + conv_type=conv_type, + output_channels=self.context_channels, + ) + self.latent_stack = LatentConditioningStack( + shape=(8 * self.input_channels, output_shape // 32, output_shape // 32), + output_channels=self.latent_channels, + ) + self.sampler = Sampler( + forecast_steps=forecast_steps, + latent_channels=self.latent_channels, + context_channels=self.context_channels, + ) + self.generator = Generator( + self.conditioning_stack, self.latent_stack, self.sampler + ) + self.discriminator = Discriminator(input_channels) + self.global_iteration = 0 + self.automatic_optimization = False + + def split_to_dict( + self, data_tensors: Tuple[paddle.Tensor, ...], keys: Tuple[str, ...] + ): + return {key: data_tensors[i] for i, key in enumerate(keys)} + + def forward(self, x): + if self._input_transform is not None: + x = self._input_transform(x) + x_tensor = self.concat_to_tensor(x, self.input_keys) + y = [self.generator(x_tensor)] + y = self.split_to_dict(y, self.output_keys) + + if self._output_transform is not None: + y = self._output_transform(x, y) + return y + + +class Sampler(nn.Layer): + """ + Sampler from the Skillful Nowcasting, see https://arxiv.org/pdf/2104.00954.pdf + + The sampler takes the output from the Latent and Context conditioning stacks and + creates one stack of ConvGRU layers per future timestep. + + Args: + forecast_steps: Number of forecast steps + latent_channels: Number of input channels to the lowest ConvGRU layer + """ + + def __init__( + self, + forecast_steps: int = 18, + latent_channels: int = 768, + context_channels: int = 384, + output_channels: int = 1, + ): + super().__init__() + self.forecast_steps = forecast_steps + self.convGRU1 = ConvGRU( + input_channels=latent_channels + context_channels, + output_channels=context_channels, + kernel_size=3, + ) + self.gru_conv_1x1 = nn.utils.spectral_norm( + layer=nn.Conv2D( + in_channels=context_channels, + out_channels=latent_channels, + kernel_size=(1, 1), + ) + ) + self.g1 = GBlock( + input_channels=latent_channels, output_channels=latent_channels + ) + self.up_g1 = UpsampleGBlock( + input_channels=latent_channels, output_channels=latent_channels // 2 + ) + self.convGRU2 = ConvGRU( + input_channels=latent_channels // 2 + context_channels // 2, + output_channels=context_channels // 2, + kernel_size=3, + ) + self.gru_conv_1x1_2 = nn.utils.spectral_norm( + layer=nn.Conv2D( + in_channels=context_channels // 2, + out_channels=latent_channels // 2, + kernel_size=(1, 1), + ) + ) + self.g2 = GBlock( + input_channels=latent_channels // 2, output_channels=latent_channels // 2 + ) + self.up_g2 = UpsampleGBlock( + input_channels=latent_channels // 2, output_channels=latent_channels // 4 + ) + self.convGRU3 = ConvGRU( + input_channels=latent_channels // 4 + context_channels // 4, + output_channels=context_channels // 4, + kernel_size=3, + ) + self.gru_conv_1x1_3 = nn.utils.spectral_norm( + layer=nn.Conv2D( + in_channels=context_channels // 4, + out_channels=latent_channels // 4, + kernel_size=(1, 1), + ) + ) + self.g3 = GBlock( + input_channels=latent_channels // 4, output_channels=latent_channels // 4 + ) + self.up_g3 = UpsampleGBlock( + input_channels=latent_channels // 4, output_channels=latent_channels // 8 + ) + self.convGRU4 = ConvGRU( + input_channels=latent_channels // 8 + context_channels // 8, + output_channels=context_channels // 8, + kernel_size=3, + ) + self.gru_conv_1x1_4 = nn.utils.spectral_norm( + layer=nn.Conv2D( + in_channels=context_channels // 8, + out_channels=latent_channels // 8, + kernel_size=(1, 1), + ) + ) + self.g4 = GBlock( + input_channels=latent_channels // 8, output_channels=latent_channels // 8 + ) + self.up_g4 = UpsampleGBlock( + input_channels=latent_channels // 8, output_channels=latent_channels // 16 + ) + self.bn = nn.BatchNorm2D(num_features=latent_channels // 16) + self.relu = nn.ReLU() + self.conv_1x1 = nn.utils.spectral_norm( + layer=nn.Conv2D( + in_channels=latent_channels // 16, + out_channels=4 * output_channels, + kernel_size=(1, 1), + ) + ) + self.depth2space = nn.PixelShuffle(upscale_factor=2) + + def forward( + self, conditioning_states: List[paddle.Tensor], latent_dim: paddle.Tensor + ) -> paddle.Tensor: + """ + Perform the sampling from Skillful Nowcasting with GANs + + Args: + conditioning_states: Outputs from the `ContextConditioningStack` with the 4 input states, ordered from largest to smallest spatially + latent_dim: Output from `LatentConditioningStack` for input into the ConvGRUs + Returns: + forecast_steps-length output of images for future timesteps + + """ + init_states = conditioning_states + latent_dim = einops.repeat( + latent_dim, "b c h w -> (repeat b) c h w", repeat=init_states[0].shape[0] + ) + hidden_states = [latent_dim] * self.forecast_steps + + hidden_states = self.convGRU1(hidden_states, init_states[3]) + hidden_states = [self.gru_conv_1x1(h) for h in hidden_states] + hidden_states = [self.g1(h) for h in hidden_states] + hidden_states = [self.up_g1(h) for h in hidden_states] + hidden_states = self.convGRU2(hidden_states, init_states[2]) + hidden_states = [self.gru_conv_1x1_2(h) for h in hidden_states] + hidden_states = [self.g2(h) for h in hidden_states] + hidden_states = [self.up_g2(h) for h in hidden_states] + hidden_states = self.convGRU3(hidden_states, init_states[1]) + hidden_states = [self.gru_conv_1x1_3(h) for h in hidden_states] + hidden_states = [self.g3(h) for h in hidden_states] + hidden_states = [self.up_g3(h) for h in hidden_states] + hidden_states = self.convGRU4(hidden_states, init_states[0]) + hidden_states = [self.gru_conv_1x1_4(h) for h in hidden_states] + hidden_states = [self.g4(h) for h in hidden_states] + hidden_states = [self.up_g4(h) for h in hidden_states] + hidden_states = [nn.functional.relu(x=self.bn(h)) for h in hidden_states] + hidden_states = [self.conv_1x1(h) for h in hidden_states] + hidden_states = [self.depth2space(h) for h in hidden_states] + forecasts = paddle.stack(x=hidden_states, axis=1) + return forecasts + + +class Generator(nn.Layer): + """ + Wraps the three parts of the generator for simpler calling + + Args: + conditioning_stack: A layer representing the conditioning stack. + latent_stack: A layer representing the latent stack. + sampler: A layer representing the sampler. + """ + + def __init__( + self, + conditioning_stack: nn.Layer, + latent_stack: nn.Layer, + sampler: nn.Layer, + ): + super().__init__() + self.conditioning_stack = conditioning_stack + self.latent_stack = latent_stack + self.sampler = sampler + + def forward(self, x): + conditioning_states = self.conditioning_stack(x) + latent_dim = self.latent_stack(x) + x = self.sampler(conditioning_states, latent_dim) + return x + + +class Discriminator(nn.Layer): + def __init__( + self, + input_channels: int = 12, + num_spatial_frames: int = 8, + conv_type: str = "standard", + ): + super().__init__() + self.spatial_discriminator = SpatialDiscriminator( + input_channels=input_channels, + num_timesteps=num_spatial_frames, + conv_type=conv_type, + ) + self.temporal_discriminator = TemporalDiscriminator( + input_channels=input_channels, conv_type=conv_type + ) + + def forward(self, x: paddle.Tensor) -> paddle.Tensor: + spatial_loss = self.spatial_discriminator(x) + temporal_loss = self.temporal_discriminator(x) + return paddle.concat(x=[spatial_loss, temporal_loss], axis=1) + + +class TemporalDiscriminator(nn.Layer): + """ + Temporal Discriminator from the Skillful Nowcasting, see https://arxiv.org/pdf/2104.00954.pdf + + Args: + input_channels: Number of channels per timestep + crop_size: Size of the crop, in the paper half the width of the input images + num_layers: Number of intermediate DBlock layers to use + conv_type: Type of 2d convolutions to use, see satflow/models/utils.py for options + """ + + def __init__( + self, + input_channels: int = 12, + num_layers: int = 3, + conv_type: str = "standard", + ): + super().__init__() + self.downsample = nn.AvgPool3D( + kernel_size=(1, 2, 2), stride=(1, 2, 2), exclusive=False + ) + self.space2depth = nn.PixelUnshuffle(downscale_factor=2) + internal_chn = 48 + self.d1 = DBlock( + input_channels=4 * input_channels, + output_channels=internal_chn * input_channels, + conv_type="3d", + first_relu=False, + ) + self.d2 = DBlock( + input_channels=internal_chn * input_channels, + output_channels=2 * internal_chn * input_channels, + conv_type="3d", + ) + self.intermediate_dblocks = nn.LayerList() + for _ in range(num_layers): + internal_chn *= 2 + self.intermediate_dblocks.append( + DBlock( + input_channels=internal_chn * input_channels, + output_channels=2 * internal_chn * input_channels, + conv_type=conv_type, + ) + ) + self.d_last = DBlock( + input_channels=2 * internal_chn * input_channels, + output_channels=2 * internal_chn * input_channels, + keep_same_output=True, + conv_type=conv_type, + ) + self.fc = nn.utils.spectral_norm( + layer=nn.Linear( + in_features=2 * internal_chn * input_channels, out_features=1 + ) + ) + self.relu = nn.ReLU() + self.bn = nn.BatchNorm1D(num_features=2 * internal_chn * input_channels) + + def forward(self, x: paddle.Tensor) -> paddle.Tensor: + x = self.downsample(x) + if len(x.shape) == 4: + x = self.space2depth(x) + elif len(x.shape) == 5: + B, T = x.shape[0], x.shape[1] + x_reshaped = paddle.reshape(x, [-1] + list(x.shape[2:])) + x = self.space2depth(x_reshaped) + x = paddle.reshape(x, [B, T] + list(x.shape[1:])) + x = paddle.transpose(x=x, perm=(0, 2, 1, 3, 4)) + x = self.d1(x) + x = self.d2(x) + x = paddle.transpose(x=x, perm=(0, 2, 1, 3, 4)) + representations = [] + for idx in range(x.shape[1]): + rep = x[:, idx, :, :, :] + for d in self.intermediate_dblocks: + rep = d(rep) + rep = self.d_last(rep) + rep = paddle.sum(x=nn.functional.relu(x=rep), axis=[2, 3]) + rep = self.bn(rep) + rep = self.fc(rep) + representations.append(rep) + x = paddle.stack(x=representations, axis=1) + x = paddle.sum(x=x, keepdim=True, axis=1) + return x + + +class SpatialDiscriminator(nn.Layer): + """ + Spatial discriminator from Skillful Nowcasting, see https://arxiv.org/pdf/2104.00954.pdf + + Args: + input_channels: Number of input channels per timestep + num_timesteps: Number of timesteps to use, in the paper 8/18 timesteps were chosen + num_layers: Number of intermediate DBlock layers to use + conv_type: Type of 2d convolutions to use, see satflow/models/utils.py for options + """ + + def __init__( + self, + input_channels: int = 12, + num_timesteps: int = 8, + num_layers: int = 4, + conv_type: str = "standard", + ): + super().__init__() + self.num_timesteps = num_timesteps + self.mean_pool = nn.AvgPool2D(kernel_size=2, exclusive=False) + self.space2depth = nn.PixelUnshuffle(downscale_factor=2) + internal_chn = 24 + self.d1 = DBlock( + input_channels=4 * input_channels, + output_channels=2 * internal_chn * input_channels, + first_relu=False, + conv_type=conv_type, + ) + self.intermediate_dblocks = nn.LayerList() + for _ in range(num_layers): + internal_chn *= 2 + self.intermediate_dblocks.append( + DBlock( + input_channels=internal_chn * input_channels, + output_channels=2 * internal_chn * input_channels, + conv_type=conv_type, + ) + ) + self.d6 = DBlock( + input_channels=2 * internal_chn * input_channels, + output_channels=2 * internal_chn * input_channels, + keep_same_output=True, + conv_type=conv_type, + ) + self.fc = nn.utils.spectral_norm( + layer=nn.Linear( + in_features=2 * internal_chn * input_channels, out_features=1 + ) + ) + self.relu = nn.ReLU() + self.bn = nn.BatchNorm1D(num_features=2 * internal_chn * input_channels) + + def forward(self, x: paddle.Tensor) -> paddle.Tensor: + idxs = paddle.randint(low=0, high=x.shape[1], shape=(self.num_timesteps,)) + representations = [] + for idx in idxs: + rep = self.mean_pool(x[:, idx, :, :, :]) + if len(rep.shape) == 4: + rep = self.space2depth(rep) + elif len(rep.shape) == 5: + B, T = rep.shape[0], rep.shape[1] + rep_reshaped = paddle.reshape(rep, [-1] + list(rep.shape[2:])) + rep = self.space2depth(rep_reshaped) + rep = paddle.reshape(rep, [B, T] + list(rep.shape[1:])) + rep = self.d1(rep) + for d in self.intermediate_dblocks: + rep = d(rep) + rep = self.d6(rep) + rep = paddle.sum(x=nn.functional.relu(x=rep), axis=[2, 3]) + rep = self.bn(rep) + rep = self.fc(rep) + """ + Pseudocode from DeepMind + # Sum-pool the representations and feed to spectrally normalized lin. layer. + y = tf.reduce_sum(tf.nn.relu(y), axis=[1, 2]) + y = layers.BatchNorm(calc_sigma=False)(y) + output_layer = layers.Linear(output_size=1) + output = output_layer(y) + + # Take the sum across the t samples. Note: we apply the ReLU to + # (1 - score_real) and (1 + score_generated) in the loss. + output = tf.reshape(output, [b, n, 1]) + output = tf.reduce_sum(output, keepdims=True, axis=1) + return output + """ + representations.append(rep) + x = paddle.stack(x=representations, axis=1) + x = paddle.sum(x=x, keepdim=True, axis=1) + return x + + +class GBlock(nn.Layer): + """Residual generator block without upsampling. G Block from Skillful Nowcasting, see https://arxiv.org/pdf/2104.00954.pdf + + Args: + input_channels: Number of input channels + output_channels: Number of output channels + conv_type: Type of convolution desired, see satflow/models/utils.py for options + """ + + def __init__( + self, + input_channels: int = 12, + output_channels: int = 12, + conv_type: str = "standard", + spectral_normalized_eps=0.0001, + ): + super().__init__() + self.output_channels = output_channels + self.bn1 = nn.BatchNorm2D(num_features=input_channels) + self.bn2 = nn.BatchNorm2D(num_features=input_channels) + self.relu = nn.ReLU() + conv2d = get_conv_layer(conv_type) + self.conv_1x1 = nn.utils.spectral_norm( + layer=conv2d( + in_channels=input_channels, out_channels=output_channels, kernel_size=1 + ), + eps=spectral_normalized_eps, + ) + self.first_conv_3x3 = nn.utils.spectral_norm( + layer=conv2d( + in_channels=input_channels, + out_channels=input_channels, + kernel_size=3, + padding=1, + ), + eps=spectral_normalized_eps, + ) + self.last_conv_3x3 = nn.utils.spectral_norm( + layer=conv2d( + in_channels=input_channels, + out_channels=output_channels, + kernel_size=3, + padding=1, + ), + eps=spectral_normalized_eps, + ) + + def forward(self, x: paddle.Tensor) -> paddle.Tensor: + if x.shape[1] != self.output_channels: + sc = self.conv_1x1(x) + else: + sc = x + x2 = self.bn1(x) + x2 = self.relu(x2) + x2 = self.first_conv_3x3(x2) + x2 = self.bn2(x2) + x2 = self.relu(x2) + x2 = self.last_conv_3x3(x2) + x = x2 + sc + return x + + +class UpsampleGBlock(nn.Layer): + """Residual generator block with upsampling + G Block from Skillful Nowcasting, see https://arxiv.org/pdf/2104.00954.pdf + + Args: + input_channels: Number of input channels + output_channels: Number of output channels + conv_type: Type of convolution desired, see satflow/models/utils.py for options + """ + + def __init__( + self, + input_channels: int = 12, + output_channels: int = 12, + conv_type: str = "standard", + spectral_normalized_eps=0.0001, + ): + super().__init__() + self.output_channels = output_channels + self.bn1 = nn.BatchNorm2D(num_features=input_channels) + self.bn2 = nn.BatchNorm2D(num_features=input_channels) + self.relu = nn.ReLU() + conv2d = get_conv_layer(conv_type) + self.conv_1x1 = nn.utils.spectral_norm( + layer=conv2d( + in_channels=input_channels, out_channels=output_channels, kernel_size=1 + ), + eps=spectral_normalized_eps, + ) + self.upsample = nn.Upsample(scale_factor=2, mode="nearest") + self.first_conv_3x3 = nn.utils.spectral_norm( + layer=conv2d( + in_channels=input_channels, + out_channels=input_channels, + kernel_size=3, + padding=1, + ), + eps=spectral_normalized_eps, + ) + self.last_conv_3x3 = nn.utils.spectral_norm( + layer=conv2d( + in_channels=input_channels, + out_channels=output_channels, + kernel_size=3, + padding=1, + ), + eps=spectral_normalized_eps, + ) + + def forward(self, x: paddle.Tensor) -> paddle.Tensor: + sc = self.upsample(x) + sc = self.conv_1x1(sc) + x2 = self.bn1(x) + x2 = self.relu(x2) + x2 = self.upsample(x2) + x2 = self.first_conv_3x3(x2) + x2 = self.bn2(x2) + x2 = self.relu(x2) + x2 = self.last_conv_3x3(x2) + x = x2 + sc + return x + + +class DBlock(nn.Layer): + """ + D and 3D Block from Skillful Nowcasting, see https://arxiv.org/pdf/2104.00954.pdf + + Args: + input_channels: Number of input channels + output_channels: Number of output channels + conv_type: Convolution type, see satflow/models/utils.py for options + first_relu: Whether to have an ReLU before the first 3x3 convolution + keep_same_output: Whether the output should have the same spatial dimensions as input, if False, downscales by 2 + """ + + def __init__( + self, + input_channels: int = 12, + output_channels: int = 12, + conv_type: str = "standard", + first_relu: bool = True, + keep_same_output: bool = False, + ): + super().__init__() + self.input_channels = input_channels + self.output_channels = output_channels + self.first_relu = first_relu + self.keep_same_output = keep_same_output + self.conv_type = conv_type + conv2d = get_conv_layer(conv_type) + if conv_type == "3d": + self.pooling = nn.AvgPool3D(kernel_size=2, stride=2, exclusive=False) + else: + self.pooling = nn.AvgPool2D(kernel_size=2, stride=2, exclusive=False) + self.conv_1x1 = nn.utils.spectral_norm( + layer=conv2d( + in_channels=input_channels, out_channels=output_channels, kernel_size=1 + ) + ) + self.first_conv_3x3 = nn.utils.spectral_norm( + layer=conv2d( + in_channels=input_channels, + out_channels=output_channels, + kernel_size=3, + padding=1, + ) + ) + self.last_conv_3x3 = nn.utils.spectral_norm( + layer=conv2d( + in_channels=output_channels, + out_channels=output_channels, + kernel_size=3, + padding=1, + stride=1, + ) + ) + self.relu = nn.ReLU() + + def forward(self, x: paddle.Tensor) -> paddle.Tensor: + if self.input_channels != self.output_channels: + x1 = self.conv_1x1(x) + if not self.keep_same_output: + x1 = self.pooling(x1) + else: + x1 = x + if self.first_relu: + x = self.relu(x) + x = self.first_conv_3x3(x) + x = self.relu(x) + x = self.last_conv_3x3(x) + if not self.keep_same_output: + x = self.pooling(x) + x = x1 + x + return x + + +class LBlock(nn.Layer): + """Residual block for the Latent Stack. + L-Block for increasing the number of channels in the input + from Skillful Nowcasting, see https://arxiv.org/pdf/2104.00954.pdf + + Args: + input_channels: Number of input channels + output_channels: Number of output channels + conv_type: Which type of convolution desired, see satflow/models/utils.py for options + """ + + def __init__( + self, + input_channels: int = 12, + output_channels: int = 12, + kernel_size: int = 3, + conv_type: str = "standard", + ): + super().__init__() + self.input_channels = input_channels + self.output_channels = output_channels + conv2d = get_conv_layer(conv_type) + self.conv_1x1 = conv2d( + in_channels=input_channels, + out_channels=output_channels - input_channels, + kernel_size=1, + ) + self.first_conv_3x3 = conv2d( + input_channels, + out_channels=output_channels, + kernel_size=kernel_size, + padding=1, + stride=1, + ) + self.relu = nn.ReLU() + self.last_conv_3x3 = conv2d( + in_channels=output_channels, + out_channels=output_channels, + kernel_size=kernel_size, + padding=1, + stride=1, + ) + + def forward(self, x) -> paddle.Tensor: + if self.input_channels < self.output_channels: + sc = self.conv_1x1(x) + sc = paddle.concat(x=[x, sc], axis=1) + else: + sc = x + x2 = self.relu(x) + x2 = self.first_conv_3x3(x2) + x2 = self.relu(x2) + x2 = self.last_conv_3x3(x2) + return x2 + sc + + +class ContextConditioningStack(nn.Layer): + """ + Conditioning Stack using the context images from Skillful Nowcasting, , see https://arxiv.org/pdf/2104.00954.pdf + + Args: + input_channels: Number of input channels per timestep + output_channels: Number of output channels for the lowest block + conv_type: Type of 2D convolution to use, see satflow/models/utils.py for options + """ + + def __init__( + self, + input_channels: int = 1, + output_channels: int = 768, + num_context_steps: int = 4, + conv_type: str = "standard", + ): + super().__init__() + conv2d = get_conv_layer(conv_type) + self.space2depth = nn.PixelUnshuffle(downscale_factor=2) + self.d1 = DBlock( + input_channels=4 * input_channels, + output_channels=output_channels // 4 * input_channels // num_context_steps, + conv_type=conv_type, + ) + self.d2 = DBlock( + input_channels=output_channels // 4 * input_channels // num_context_steps, + output_channels=output_channels // 2 * input_channels // num_context_steps, + conv_type=conv_type, + ) + self.d3 = DBlock( + input_channels=output_channels // 2 * input_channels // num_context_steps, + output_channels=output_channels * input_channels // num_context_steps, + conv_type=conv_type, + ) + self.d4 = DBlock( + input_channels=output_channels * input_channels // num_context_steps, + output_channels=output_channels * 2 * input_channels // num_context_steps, + conv_type=conv_type, + ) + self.conv1 = nn.utils.spectral_norm( + layer=conv2d( + in_channels=output_channels // 4 * input_channels, + out_channels=output_channels // 8 * input_channels, + kernel_size=3, + padding=1, + ) + ) + self.conv2 = nn.utils.spectral_norm( + layer=conv2d( + in_channels=output_channels // 2 * input_channels, + out_channels=output_channels // 4 * input_channels, + kernel_size=3, + padding=1, + ) + ) + self.conv3 = nn.utils.spectral_norm( + layer=conv2d( + in_channels=output_channels * input_channels, + out_channels=output_channels // 2 * input_channels, + kernel_size=3, + padding=1, + ) + ) + self.conv4 = nn.utils.spectral_norm( + layer=conv2d( + in_channels=output_channels * 2 * input_channels, + out_channels=output_channels * input_channels, + kernel_size=3, + padding=1, + ) + ) + self.relu = nn.ReLU() + + def forward( + self, x: paddle.Tensor + ) -> Tuple[paddle.Tensor, paddle.Tensor, paddle.Tensor, paddle.Tensor]: + if len(x.shape) == 4: + x = self.space2depth(x) + elif len(x.shape) == 5: + B, T = x.shape[0], x.shape[1] + x_reshaped = paddle.reshape(x, [-1] + list(x.shape[2:])) + x = self.space2depth(x_reshaped) + x = paddle.reshape(x, [B, T] + list(x.shape[1:])) + steps = x.shape[1] + scale_1 = [] + scale_2 = [] + scale_3 = [] + scale_4 = [] + for i in range(steps): + s1 = self.d1(x[:, i, :, :, :]) + s2 = self.d2(s1) + s3 = self.d3(s2) + s4 = self.d4(s3) + scale_1.append(s1) + scale_2.append(s2) + scale_3.append(s3) + scale_4.append(s4) + scale_1 = paddle.stack(x=scale_1, axis=1) + scale_2 = paddle.stack(x=scale_2, axis=1) + scale_3 = paddle.stack(x=scale_3, axis=1) + scale_4 = paddle.stack(x=scale_4, axis=1) + scale_1 = self._mixing_layer(scale_1, self.conv1) + scale_2 = self._mixing_layer(scale_2, self.conv2) + scale_3 = self._mixing_layer(scale_3, self.conv3) + scale_4 = self._mixing_layer(scale_4, self.conv4) + return scale_1, scale_2, scale_3, scale_4 + + def _mixing_layer(self, inputs, conv_block): + stacked_inputs = einops.rearrange(inputs, "b t c h w -> b (c t) h w") + return nn.functional.relu(x=conv_block(stacked_inputs)) + + +class LatentConditioningStack(nn.Layer): + """ + Latent conditioning stack from Skillful Nowcasting, see https://arxiv.org/pdf/2104.00954.pdf + + Args: + shape: Shape of the latent space, Should be (H/32,W/32,x) of the final image shape + output_channels: Number of output channels for the conditioning stack + use_attention: Whether to have a self-attention block or not + """ + + def __init__( + self, + shape: (int, int, int) = (8, 8, 8), + output_channels: int = 768, + use_attention: bool = True, + ): + super().__init__() + self.shape = shape + self.use_attention = use_attention + self.distribution = paddle.distribution.Normal( + loc=paddle.to_tensor(data=[0.0], dtype="float32"), + scale=paddle.to_tensor(data=[2.0], dtype="float32"), + ) + self.conv_3x3 = nn.utils.spectral_norm( + layer=nn.Conv2D( + in_channels=shape[0], + out_channels=shape[0], + kernel_size=(3, 3), + padding=1, + ) + ) + self.l_block1 = LBlock( + input_channels=shape[0], output_channels=output_channels // 32 + ) + self.l_block2 = LBlock( + input_channels=output_channels // 32, output_channels=output_channels // 16 + ) + self.l_block3 = LBlock( + input_channels=output_channels // 16, output_channels=output_channels // 4 + ) + if self.use_attention: + self.att_block = AttentionLayer( + input_channels=output_channels // 4, + output_channels=output_channels // 4, + ) + self.l_block4 = LBlock( + input_channels=output_channels // 4, output_channels=output_channels + ) + + def forward(self, x: paddle.Tensor) -> paddle.Tensor: + """ + Args: + x: tensor on the correct device, to move over the latent distribution + Returns: z + """ + z = self.distribution.sample(self.shape) + z = paddle.transpose(x=z, perm=(3, 0, 1, 2)).astype(dtype=x.dtype) + z = self.conv_3x3(z) + z = self.l_block1(z) + z = self.l_block2(z) + z = self.l_block3(z) + z = self.att_block(z) + z = self.l_block4(z) + return z + + +def attention_einsum(q, k, v): + """Apply the attention operator to tensors of shape [h, w, c].""" + k = einops.rearrange(k, "h w c -> (h w) c") + v = einops.rearrange(v, "h w c -> (h w) c") + beta = nn.functional.softmax(x=paddle.einsum("hwc, Lc->hwL", q, k), axis=-1) + out = paddle.einsum("hwL, Lc->hwc", beta, v) + return out + + +class AttentionLayer(nn.Layer): + """Attention Module""" + + def __init__( + self, input_channels: int, output_channels: int, ratio_kq=8, ratio_v=8 + ): + super().__init__() + self.ratio_kq = ratio_kq + self.ratio_v = ratio_v + self.output_channels = output_channels + self.input_channels = input_channels + self.query = nn.Conv2D( + in_channels=input_channels, + out_channels=self.output_channels // self.ratio_kq, + kernel_size=(1, 1), + padding="valid", + bias_attr=False, + ) + self.key = nn.Conv2D( + in_channels=input_channels, + out_channels=self.output_channels // self.ratio_kq, + kernel_size=(1, 1), + padding="valid", + bias_attr=False, + ) + self.value = nn.Conv2D( + in_channels=input_channels, + out_channels=self.output_channels // self.ratio_v, + kernel_size=(1, 1), + padding="valid", + bias_attr=False, + ) + self.last_conv = nn.Conv2D( + in_channels=self.output_channels // 8, + out_channels=self.output_channels, + kernel_size=(1, 1), + padding="valid", + bias_attr=False, + ) + gamma = paddle.create_parameter( + shape=paddle.zeros(shape=[1]).shape, + dtype=paddle.zeros(shape=[1]).numpy().dtype, + default_initializer=nn.initializer.Assign(paddle.zeros(shape=[1])), + ) + gamma.stop_gradient = not True + self.gamma = gamma + + def forward(self, x: paddle.Tensor) -> paddle.Tensor: + query = self.query(x) + key = self.key(x) + value = self.value(x) + out = [] + for b in range(x.shape[0]): + out.append(attention_einsum(query[b], key[b], value[b])) + out = paddle.stack(x=out, axis=0) + out = self.gamma * self.last_conv(out) + return out + x + + +class AddCoords(nn.Layer): + def __init__(self, with_r=False): + super().__init__() + self.with_r = with_r + + def forward(self, input_tensor): + """ + Args: + input_tensor: shape(batch, channel, x_dim, y_dim) + """ + batch_size, _, x_dim, y_dim = input_tensor.shape + xx_channel = paddle.arange(end=x_dim).tile([1, y_dim, 1]) + x = paddle.arange(end=y_dim).tile([1, x_dim, 1]) + perm_0 = list(range(x.ndim)) + perm_0[1] = 2 + perm_0[2] = 1 + yy_channel = x.transpose(perm=perm_0) + xx_channel = xx_channel.astype(dtype="float32") / (x_dim - 1) + yy_channel = yy_channel.astype(dtype="float32") / (y_dim - 1) + xx_channel = xx_channel * 2 - 1 + yy_channel = yy_channel * 2 - 1 + x = xx_channel.tile([batch_size, 1, 1, 1]) + perm_1 = list(range(x.ndim)) + perm_1[2] = 3 + perm_1[3] = 2 + xx_channel = x.transpose(perm=perm_1) + x = yy_channel.tile([batch_size, 1, 1, 1]) + perm_2 = list(range(x.ndim)) + perm_2[2] = 3 + perm_2[3] = 2 + yy_channel = x.transpose(perm=perm_2) + ret = paddle.concat( + x=[ + input_tensor, + xx_channel.astype(dtype=input_tensor.dtype), + yy_channel.astype(dtype=input_tensor.dtype), + ], + axis=1, + ) + if self.with_r: + rr = paddle.sqrt( + x=paddle.pow(x=xx_channel.astype(dtype=input_tensor.dtype) - 0.5, y=2) + + paddle.pow(x=yy_channel.astype(dtype=input_tensor.dtype) - 0.5, y=2) + ) + ret = paddle.concat(x=[ret, rr], axis=1) + return ret + + +class CoordConv(nn.Layer): + def __init__(self, in_channels, out_channels, with_r=False): + super().__init__() + self.addcoords = AddCoords(with_r=with_r) + in_size = in_channels + 2 + if with_r: + in_size += 1 + self.conv = nn.Conv2D(in_size, out_channels) + + def forward(self, x): + ret = self.addcoords(x) + ret = self.conv(ret) + return ret + + +class ConvGRUCell(nn.Layer): + """A ConvGRU implementation. + + Args: + kernel_size: kernel size of the convolutions. Default: 3. + sn_eps: constant for spectral normalization. Default: 1e-4. + """ + + def __init__( + self, input_channels: int, output_channels: int, kernel_size=3, sn_eps=0.0001 + ): + super().__init__() + self._kernel_size = kernel_size + self._sn_eps = sn_eps + self.read_gate_conv = nn.utils.spectral_norm( + layer=nn.Conv2D( + in_channels=input_channels, + out_channels=output_channels, + kernel_size=(kernel_size, kernel_size), + padding=1, + ), + eps=sn_eps, + ) + self.update_gate_conv = nn.utils.spectral_norm( + layer=nn.Conv2D( + in_channels=input_channels, + out_channels=output_channels, + kernel_size=(kernel_size, kernel_size), + padding=1, + ), + eps=sn_eps, + ) + self.output_conv = nn.utils.spectral_norm( + layer=nn.Conv2D( + in_channels=input_channels, + out_channels=output_channels, + kernel_size=(kernel_size, kernel_size), + padding=1, + ), + eps=sn_eps, + ) + + def forward(self, x, prev_state): + """ + ConvGRU forward, returning the current+new state + + Args: + x: Input tensor + prev_state: Previous state + + Returns: + New tensor plus the new state + """ + xh = paddle.concat(x=[x, prev_state], axis=1) + read_gate = nn.functional.sigmoid(x=self.read_gate_conv(xh)) + update_gate = nn.functional.sigmoid(x=self.update_gate_conv(xh)) + gated_input = paddle.concat(x=[x, read_gate * prev_state], axis=1) + c = nn.functional.relu(x=self.output_conv(gated_input)) + out = update_gate * prev_state + (1.0 - update_gate) * c + new_state = out + return out, new_state + + +class ConvGRU(nn.Layer): + """ConvGRU Cell wrapper to replace tf.static_rnn in TF implementation""" + + def __init__( + self, + input_channels: int, + output_channels: int, + kernel_size: int = 3, + sn_eps=0.0001, + ): + super().__init__() + self.cell = ConvGRUCell(input_channels, output_channels, kernel_size, sn_eps) + + def forward(self, x: paddle.Tensor, hidden_state=None) -> paddle.Tensor: + outputs = [] + for step in range(len(x)): + output, hidden_state = self.cell(x[step], hidden_state) + outputs.append(output) + outputs = paddle.stack(x=outputs, axis=0) + return outputs + + +def get_conv_layer(conv_type: str = "standard") -> nn.Layer: + if conv_type == "standard": + conv_layer = nn.Conv2D + elif conv_type == "coord": + conv_layer = CoordConv + elif conv_type == "3d": + conv_layer = nn.Conv3D + else: + raise ValueError(f"{conv_type} is not a recognized Conv method") + return conv_layer diff --git a/examples/smc_reac/ppsci/arch/embedding_koopman.py b/examples/smc_reac/ppsci/arch/embedding_koopman.py new file mode 100644 index 0000000000..367b5cd3ca --- /dev/null +++ b/examples/smc_reac/ppsci/arch/embedding_koopman.py @@ -0,0 +1,544 @@ +# 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. + +""" +Code below is heavily based on [transformer-physx](https://github.com/zabaras/transformer-physx) +""" + +from __future__ import annotations + +from typing import Optional +from typing import Tuple + +import numpy as np +import paddle +from paddle import nn +from paddle.nn.initializer import Constant +from paddle.nn.initializer import Uniform + +from ppsci.arch import base + +zeros_ = Constant(value=0.0) +ones_ = Constant(value=1.0) + + +class LorenzEmbedding(base.Arch): + """Embedding Koopman model for the Lorenz ODE system. + + Args: + input_keys (Tuple[str, ...]): Input keys, such as ("states",). + output_keys (Tuple[str, ...]): Output keys, such as ("pred_states", "recover_states"). + mean (Optional[Tuple[float, ...]]): Mean of training dataset. Defaults to None. + std (Optional[Tuple[float, ...]]): Standard Deviation of training dataset. Defaults to None. + input_size (int, optional): Size of input data. Defaults to 3. + hidden_size (int, optional): Number of hidden size. Defaults to 500. + embed_size (int, optional): Number of embedding size. Defaults to 32. + drop (float, optional): Probability of dropout the units. Defaults to 0.0. + + Examples: + >>> import ppsci + >>> model = ppsci.arch.LorenzEmbedding( + ... input_keys=("x", "y"), + ... output_keys=("u", "v"), + ... input_size=3, + ... hidden_size=500, + ... embed_size=32, + ... drop=0.0, + ... mean=None, + ... std=None, + ... ) + >>> x_shape = [8, 3, 2] + >>> y_shape = [8, 3, 1] + >>> input_dict = {"x": paddle.rand(x_shape), + ... "y": paddle.rand(y_shape)} + >>> output_dict = model(input_dict) + >>> print(output_dict["u"].shape) + [8, 2, 3] + >>> print(output_dict["v"].shape) + [8, 3, 3] + """ + + def __init__( + self, + input_keys: Tuple[str, ...], + output_keys: Tuple[str, ...], + mean: Optional[Tuple[float, ...]] = None, + std: Optional[Tuple[float, ...]] = None, + input_size: int = 3, + hidden_size: int = 500, + embed_size: int = 32, + drop: float = 0.0, + ): + super().__init__() + self.input_keys = input_keys + self.output_keys = output_keys + self.input_size = input_size + self.hidden_size = hidden_size + self.embed_size = embed_size + + # build observable network + self.encoder_net = self.build_encoder(input_size, hidden_size, embed_size, drop) + # build koopman operator + self.k_diag, self.k_ut = self.build_koopman_operator(embed_size) + # build recovery network + self.decoder_net = self.build_decoder(input_size, hidden_size, embed_size) + + mean = [0.0, 0.0, 0.0] if mean is None else mean + std = [1.0, 1.0, 1.0] if std is None else std + self.register_buffer("mean", paddle.to_tensor(mean).reshape([1, 3])) + self.register_buffer("std", paddle.to_tensor(std).reshape([1, 3])) + + self.apply(self._init_weights) + + def _init_weights(self, m: nn.Layer): + if isinstance(m, nn.Linear): + k = 1 / m.weight.shape[0] + uniform = Uniform(-(k**0.5), k**0.5) + uniform(m.weight) + if m.bias is not None: + uniform(m.bias) + elif isinstance(m, nn.LayerNorm): + zeros_(m.bias) + ones_(m.weight) + + def build_encoder( + self, input_size: int, hidden_size: int, embed_size: int, drop: float = 0.0 + ): + net = nn.Sequential( + nn.Linear(input_size, hidden_size), + nn.ReLU(), + nn.Linear(hidden_size, embed_size), + nn.LayerNorm(embed_size), + nn.Dropout(drop), + ) + return net + + def build_decoder(self, input_size: int, hidden_size: int, embed_size: int): + net = nn.Sequential( + nn.Linear(embed_size, hidden_size), + nn.ReLU(), + nn.Linear(hidden_size, input_size), + ) + return net + + def build_koopman_operator(self, embed_size: int): + # Learned Koopman operator + data = paddle.linspace(1, 0, embed_size) + k_diag = paddle.create_parameter( + shape=data.shape, + dtype=paddle.get_default_dtype(), + default_initializer=nn.initializer.Assign(data), + ) + + data = 0.1 * paddle.rand([2 * embed_size - 3]) + k_ut = paddle.create_parameter( + shape=data.shape, + dtype=paddle.get_default_dtype(), + default_initializer=nn.initializer.Assign(data), + ) + return k_diag, k_ut + + def encoder(self, x: paddle.Tensor): + x = self._normalize(x) + g = self.encoder_net(x) + return g + + def decoder(self, g: paddle.Tensor): + out = self.decoder_net(g) + x = self._unnormalize(out) + return x + + def koopman_operation(self, embed_data: paddle.Tensor, k_matrix: paddle.Tensor): + # Apply Koopman operation + embed_pred_data = paddle.bmm( + k_matrix.expand( + [embed_data.shape[0], k_matrix.shape[0], k_matrix.shape[1]] + ), + embed_data.transpose([0, 2, 1]), + ).transpose([0, 2, 1]) + return embed_pred_data + + def _normalize(self, x: paddle.Tensor): + return (x - self.mean) / self.std + + def _unnormalize(self, x: paddle.Tensor): + return self.std * x + self.mean + + def get_koopman_matrix(self): + # # Koopman operator + k_ut_tensor = self.k_ut * 1 + k_ut_tensor = paddle.diag( + k_ut_tensor[0 : self.embed_size - 1], offset=1 + ) + paddle.diag(k_ut_tensor[self.embed_size - 1 :], offset=2) + k_matrix = k_ut_tensor + (-1) * k_ut_tensor.t() + k_matrix = k_matrix + paddle.diag(self.k_diag) + return k_matrix + + def forward_tensor(self, x): + k_matrix = self.get_koopman_matrix() + embed_data = self.encoder(x) + recover_data = self.decoder(embed_data) + + embed_pred_data = self.koopman_operation(embed_data, k_matrix) + pred_data = self.decoder(embed_pred_data) + + return (pred_data[:, :-1, :], recover_data, k_matrix) + + @staticmethod + def split_to_dict(data_tensors: Tuple[paddle.Tensor, ...], keys: Tuple[str, ...]): + return {key: data_tensors[i] for i, key in enumerate(keys)} + + def forward(self, x): + if self._input_transform is not None: + x = self._input_transform(x) + + x_tensor = self.concat_to_tensor(x, self.input_keys, axis=-1) + y = self.forward_tensor(x_tensor) + y = self.split_to_dict(y, self.output_keys) + + if self._output_transform is not None: + y = self._output_transform(x, y) + return y + + +class RosslerEmbedding(LorenzEmbedding): + """Embedding Koopman model for the Rossler ODE system. + + Args: + input_keys (Tuple[str, ...]): Input keys, such as ("states",). + output_keys (Tuple[str, ...]): Output keys, such as ("pred_states", "recover_states"). + mean (Optional[Tuple[float, ...]]): Mean of training dataset. Defaults to None. + std (Optional[Tuple[float, ...]]): Standard Deviation of training dataset. Defaults to None. + input_size (int, optional): Size of input data. Defaults to 3. + hidden_size (int, optional): Number of hidden size. Defaults to 500. + embed_size (int, optional): Number of embedding size. Defaults to 32. + drop (float, optional): Probability of dropout the units. Defaults to 0.0. + + Examples: + >>> import ppsci + >>> model = ppsci.arch.RosslerEmbedding( + ... input_keys=("x", "y"), + ... output_keys=("u", "v"), + ... input_size=3, + ... hidden_size=500, + ... embed_size=32, + ... drop=0.0, + ... mean=None, + ... std=None, + ... ) + >>> x_shape = [8, 3, 2] + >>> y_shape = [8, 3, 1] + >>> input_dict = {"x": paddle.rand(x_shape), + ... "y": paddle.rand(y_shape)} + >>> output_dict = model(input_dict) + >>> print(output_dict["u"].shape) + [8, 2, 3] + >>> print(output_dict["v"].shape) + [8, 3, 3] + """ + + def __init__( + self, + input_keys: Tuple[str, ...], + output_keys: Tuple[str, ...], + mean: Optional[Tuple[float, ...]] = None, + std: Optional[Tuple[float, ...]] = None, + input_size: int = 3, + hidden_size: int = 500, + embed_size: int = 32, + drop: float = 0.0, + ): + super().__init__( + input_keys, + output_keys, + mean, + std, + input_size, + hidden_size, + embed_size, + drop, + ) + + +class CylinderEmbedding(base.Arch): + """Embedding Koopman model for the Cylinder system. + + Args: + input_keys (Tuple[str, ...]): Input keys, such as ("states", "visc"). + output_keys (Tuple[str, ...]): Output keys, such as ("pred_states", "recover_states"). + mean (Optional[Tuple[float, ...]]): Mean of training dataset. Defaults to None. + std (Optional[Tuple[float, ...]]): Standard Deviation of training dataset. Defaults to None. + embed_size (int, optional): Number of embedding size. Defaults to 128. + encoder_channels (Optional[Tuple[int, ...]]): Number of channels in encoder network. Defaults to None. + decoder_channels (Optional[Tuple[int, ...]]): Number of channels in decoder network. Defaults to None. + drop (float, optional): Probability of dropout the units. Defaults to 0.0. + + Examples: + >>> import paddle + >>> import ppsci + >>> model = ppsci.arch.CylinderEmbedding(("states", "visc"), ("pred_states", "recover_states")) + >>> states_shape = [32, 10, 3, 64, 128] + >>> visc_shape = [32, 1] + >>> input_dict = {"states" : paddle.rand(states_shape), + ... "visc" : paddle.rand(visc_shape)} + >>> out_dict = model(input_dict) + >>> print(out_dict["pred_states"].shape) + [32, 9, 3, 64, 128] + >>> print(out_dict["recover_states"].shape) + [32, 10, 3, 64, 128] + """ + + def __init__( + self, + input_keys: Tuple[str, ...], + output_keys: Tuple[str, ...], + mean: Optional[Tuple[float, ...]] = None, + std: Optional[Tuple[float, ...]] = None, + embed_size: int = 128, + encoder_channels: Optional[Tuple[int, ...]] = None, + decoder_channels: Optional[Tuple[int, ...]] = None, + drop: float = 0.0, + ): + super().__init__() + self.input_keys = input_keys + self.output_keys = output_keys + self.embed_size = embed_size + + X, Y = np.meshgrid(np.linspace(-2, 14, 128), np.linspace(-4, 4, 64)) + self.mask = paddle.to_tensor(np.sqrt(X**2 + Y**2)).unsqueeze(0).unsqueeze(0) + + encoder_channels = ( + [4, 16, 32, 64, 128] if encoder_channels is None else encoder_channels + ) + decoder_channels = ( + [embed_size // 32, 128, 64, 32, 16] + if decoder_channels is None + else decoder_channels + ) + self.encoder_net = self.build_encoder(embed_size, encoder_channels, drop) + self.k_diag_net, self.k_ut_net, self.k_lt_net = self.build_koopman_operator( + embed_size + ) + self.decoder_net = self.build_decoder(decoder_channels) + + xidx = [] + yidx = [] + for i in range(1, 5): + yidx.append(np.arange(i, embed_size)) + xidx.append(np.arange(0, embed_size - i)) + self.xidx = paddle.to_tensor(np.concatenate(xidx), dtype="int64") + self.yidx = paddle.to_tensor(np.concatenate(yidx), dtype="int64") + + mean = [0.0, 0.0, 0.0, 0.0] if mean is None else mean + std = [1.0, 1.0, 1.0, 1.0] if std is None else std + self.register_buffer("mean", paddle.to_tensor(mean).reshape([1, 4, 1, 1])) + self.register_buffer("std", paddle.to_tensor(std).reshape([1, 4, 1, 1])) + + self.apply(self._init_weights) + + def _init_weights(self, m): + if isinstance(m, nn.Linear): + k = 1 / m.weight.shape[0] + uniform = Uniform(-(k**0.5), k**0.5) + uniform(m.weight) + if m.bias is not None: + uniform(m.bias) + elif isinstance(m, nn.LayerNorm): + zeros_(m.bias) + ones_(m.weight) + elif isinstance(m, nn.Conv2D): + k = 1 / (m.weight.shape[1] * m.weight.shape[2] * m.weight.shape[3]) + uniform = Uniform(-(k**0.5), k**0.5) + uniform(m.weight) + if m.bias is not None: + uniform(m.bias) + + def _build_conv_relu_list( + self, in_channels: Tuple[int, ...], out_channels: Tuple[int, ...] + ): + net_list = [ + nn.Conv2D( + in_channels, + out_channels, + kernel_size=(3, 3), + stride=2, + padding=1, + padding_mode="replicate", + ), + nn.ReLU(), + ] + return net_list + + def build_encoder( + self, embed_size: int, channels: Tuple[int, ...], drop: float = 0.0 + ): + net = [] + for i in range(1, len(channels)): + net.extend(self._build_conv_relu_list(channels[i - 1], channels[i])) + net.append( + nn.Conv2D( + channels[-1], + embed_size // 32, + kernel_size=(3, 3), + padding=1, + padding_mode="replicate", + ) + ) + net.append( + nn.LayerNorm( + (4, 4, 8), + ) + ) + net.append(nn.Dropout(drop)) + net = nn.Sequential(*net) + return net + + def _build_upsample_conv_relu( + self, in_channels: Tuple[int, ...], out_channels: Tuple[int, ...] + ): + net_list = [ + nn.Upsample(scale_factor=2, mode="bilinear", align_corners=True), + nn.Conv2D( + in_channels, + out_channels, + kernel_size=(3, 3), + stride=1, + padding=1, + padding_mode="replicate", + ), + nn.ReLU(), + ] + return net_list + + def build_decoder(self, channels: Tuple[int, ...]): + net = [] + for i in range(1, len(channels)): + net.extend(self._build_upsample_conv_relu(channels[i - 1], channels[i])) + net.append( + nn.Conv2D( + channels[-1], + 3, + kernel_size=(3, 3), + stride=1, + padding=1, + padding_mode="replicate", + ), + ) + net = nn.Sequential(*net) + return net + + def build_koopman_operator(self, embed_size: int): + # Learned Koopman operator parameters + k_diag_net = nn.Sequential( + nn.Linear(1, 50), nn.ReLU(), nn.Linear(50, embed_size) + ) + + k_ut_net = nn.Sequential( + nn.Linear(1, 50), nn.ReLU(), nn.Linear(50, 4 * embed_size - 10) + ) + k_lt_net = nn.Sequential( + nn.Linear(1, 50), nn.ReLU(), nn.Linear(50, 4 * embed_size - 10) + ) + return k_diag_net, k_ut_net, k_lt_net + + def encoder(self, x: paddle.Tensor, viscosity: paddle.Tensor): + B, T, C, H, W = x.shape + x = x.reshape((B * T, C, H, W)) + viscosity = viscosity.repeat_interleave(T, axis=1).reshape((B * T, 1)) + x = paddle.concat( + [x, viscosity.unsqueeze(-1).unsqueeze(-1) * paddle.ones_like(x[:, :1])], + axis=1, + ) + x = self._normalize(x) + g = self.encoder_net(x) + g = g.reshape([B, T, -1]) + return g + + def decoder(self, g: paddle.Tensor): + B, T, _ = g.shape + x = self.decoder_net(g.reshape([-1, self.embed_size // 32, 4, 8])) + x = self._unnormalize(x) + mask0 = ( + self.mask.repeat_interleave(x.shape[1], axis=1).repeat_interleave( + x.shape[0], axis=0 + ) + < 1 + ) + x[mask0] = 0 + _, C, H, W = x.shape + x = x.reshape([B, T, C, H, W]) + return x + + def get_koopman_matrix(self, g: paddle.Tensor, visc: paddle.Tensor): + # # Koopman operator + kMatrix = paddle.zeros([g.shape[0], self.embed_size, self.embed_size]) + kMatrix.stop_gradient = False + # Populate the off diagonal terms + kMatrixUT_data = self.k_ut_net(100 * visc) + kMatrixLT_data = self.k_lt_net(100 * visc) + + kMatrix = kMatrix.transpose([1, 2, 0]) + kMatrixUT_data_t = kMatrixUT_data.transpose([1, 0]) + kMatrixLT_data_t = kMatrixLT_data.transpose([1, 0]) + kMatrix[self.xidx, self.yidx] = kMatrixUT_data_t + kMatrix[self.yidx, self.xidx] = kMatrixLT_data_t + + # Populate the diagonal + ind = np.diag_indices(kMatrix.shape[1]) + ind = paddle.to_tensor(ind, dtype="int64") + + kMatrixDiag = self.k_diag_net(100 * visc) + kMatrixDiag_t = kMatrixDiag.transpose([1, 0]) + kMatrix[ind[0], ind[1]] = kMatrixDiag_t + return kMatrix.transpose([2, 0, 1]) + + def koopman_operation(self, embed_data: paddle.Tensor, k_matrix: paddle.Tensor): + embed_pred_data = paddle.bmm( + k_matrix, embed_data.transpose([0, 2, 1]) + ).transpose([0, 2, 1]) + return embed_pred_data + + def _normalize(self, x: paddle.Tensor): + x = (x - self.mean) / self.std + return x + + def _unnormalize(self, x: paddle.Tensor): + return self.std[:, :3] * x + self.mean[:, :3] + + def forward_tensor(self, states, visc): + # states.shape=(B, T, C, H, W) + embed_data = self.encoder(states, visc) + recover_data = self.decoder(embed_data) + + k_matrix = self.get_koopman_matrix(embed_data, visc) + embed_pred_data = self.koopman_operation(embed_data, k_matrix) + pred_data = self.decoder(embed_pred_data) + + return (pred_data[:, :-1], recover_data, k_matrix) + + @staticmethod + def split_to_dict(data_tensors: Tuple[paddle.Tensor, ...], keys: Tuple[str, ...]): + return {key: data_tensors[i] for i, key in enumerate(keys)} + + def forward(self, x): + + if self._input_transform is not None: + x = self._input_transform(x) + + y = self.forward_tensor(**x) + y = self.split_to_dict(y, self.output_keys) + + if self._output_transform is not None: + y = self._output_transform(x, y) + return y diff --git a/examples/smc_reac/ppsci/arch/epnn.py b/examples/smc_reac/ppsci/arch/epnn.py new file mode 100644 index 0000000000..0f6a9ffed6 --- /dev/null +++ b/examples/smc_reac/ppsci/arch/epnn.py @@ -0,0 +1,126 @@ +# 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. + +# Elasto-Plastic Neural Network (EPNN) + +# DEVELOPED AT: +# COMPUTATIONAL GEOMECHANICS LABORATORY +# DEPARTMENT OF CIVIL ENGINEERING +# UNIVERSITY OF CALGARY, AB, CANADA +# DIRECTOR: Prof. Richard Wan + +# DEVELOPED BY: +# MAHDAD EGHBALIAN + +# MIT License + +# Copyright (c) 2022 Mahdad Eghbalian + +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: + +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. + +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +from typing import Tuple + +import paddle.nn as nn + +from ppsci.arch import activation as act_mod +from ppsci.arch import base + + +class Epnn(base.Arch): + """Builds a feedforward network with arbitrary layers. + + Args: + input_keys (Tuple[str, ...]): Name of input keys, such as ("x", "y", "z"). + output_keys (Tuple[str, ...]): Name of output keys, such as ("u", "v", "w"). + node_sizes (Tuple[int, ...]): The tuple of node size. + activations (Tuple[str, ...]): Name of activation functions. + drop_p (float): The parameter p of nn.Dropout. + + Examples: + >>> import ppsci + >>> ann_node_sizes_state = [1, 20] + >>> model = ppsci.arch.Epnn( + ... ("x",), + ... ("y",), + ... node_sizes=ann_node_sizes_state, + ... activations=("leaky_relu",), + ... drop_p=0.0, + ... ) + """ + + def __init__( + self, + input_keys: Tuple[str, ...], + output_keys: Tuple[str, ...], + node_sizes: Tuple[int, ...], + activations: Tuple[str, ...], + drop_p: float, + ): + super().__init__() + self.active_func = [ + act_mod.get_activation(act_name) for act_name in activations + ] + self.node_sizes = node_sizes + self.drop_p = drop_p + self.layers = [] + self.layers.append( + nn.Linear(in_features=node_sizes[0], out_features=node_sizes[1]) + ) + layer_sizes = zip(node_sizes[1:-2], node_sizes[2:-1]) + self.layers.extend( + [nn.Linear(in_features=h1, out_features=h2) for h1, h2 in layer_sizes] + ) + self.layers.append( + nn.Linear( + in_features=node_sizes[-2], out_features=node_sizes[-1], bias_attr=False + ) + ) + + self.layers = nn.LayerList(self.layers) + self.dropout = nn.Dropout(p=drop_p) + self.input_keys = input_keys + self.output_keys = output_keys + + def forward(self, x): + if self._input_transform is not None: + x = self._input_transform(x) + + y = x[self.input_keys[0]] + for ilayer in range(len(self.layers)): + y = self.layers[ilayer](y) + if ilayer != len(self.layers) - 1: + y = self.active_func[ilayer + 1](y) + if ilayer != len(self.layers) - 1: + y = self.dropout(y) + y = self.split_to_dict(y, self.output_keys, axis=-1) + + if self._output_transform is not None: + y = self._output_transform(x, y) + return y diff --git a/examples/smc_reac/ppsci/arch/extformer_moe_cuboid.py b/examples/smc_reac/ppsci/arch/extformer_moe_cuboid.py new file mode 100644 index 0000000000..7c57d23948 --- /dev/null +++ b/examples/smc_reac/ppsci/arch/extformer_moe_cuboid.py @@ -0,0 +1,996 @@ +from typing import Sequence +from typing import Tuple +from typing import Union + +import paddle +from paddle import nn + +import ppsci.arch.extformer_moe_cuboid_decoder as cuboid_decoder +import ppsci.arch.extformer_moe_cuboid_encoder as cuboid_encoder +import ppsci.arch.extformer_moe_cuboid_utils as cuboid_utils +from ppsci.arch import activation as act_mod +from ppsci.arch import base +from ppsci.arch import extformer_moe_utils +from ppsci.arch.extformer_moe_cuboid_encoder import NEGATIVE_SLOPE +from ppsci.utils import initializer + +"""A space-time Transformer with Cuboid Attention""" + + +class InitialEncoder(nn.Layer): + def __init__( + self, + dim, + out_dim, + downsample_scale: Union[int, Sequence[int]], + num_conv_layers: int = 2, + activation: str = "leaky", + padding_type: str = "nearest", + conv_init_mode: str = "0", + linear_init_mode: str = "0", + norm_init_mode: str = "0", + moe_config: dict = None, + ): + super(InitialEncoder, self).__init__() + self.num_conv_layers = num_conv_layers + self.conv_init_mode = conv_init_mode + self.linear_init_mode = linear_init_mode + self.norm_init_mode = norm_init_mode + conv_block = [] + for i in range(num_conv_layers): + if i == 0: + conv_block.append( + nn.Conv2D( + kernel_size=(3, 3), + padding=(1, 1), + in_channels=dim, + out_channels=out_dim, + ) + ) + conv_block.append(nn.GroupNorm(num_groups=16, num_channels=out_dim)) + conv_block.append( + act_mod.get_activation(activation) + if activation != "leaky_relu" + else nn.LeakyReLU(NEGATIVE_SLOPE) + ) + else: + conv_block.append( + nn.Conv2D( + kernel_size=(3, 3), + padding=(1, 1), + in_channels=out_dim, + out_channels=out_dim, + ) + ) + conv_block.append(nn.GroupNorm(num_groups=16, num_channels=out_dim)) + conv_block.append( + act_mod.get_activation(activation) + if activation != "leaky_relu" + else nn.LeakyReLU(NEGATIVE_SLOPE) + ) + self.conv_block = nn.Sequential(*conv_block) + if isinstance(downsample_scale, int): + patch_merge_downsample = (1, downsample_scale, downsample_scale) + elif len(downsample_scale) == 2: + patch_merge_downsample = (1, *downsample_scale) + elif len(downsample_scale) == 3: + patch_merge_downsample = tuple(downsample_scale) + else: + raise NotImplementedError( + f"downsample_scale {downsample_scale} format not supported!" + ) + self.patch_merge = cuboid_encoder.PatchMerging3D( + dim=out_dim, + out_dim=out_dim, + padding_type=padding_type, + downsample=patch_merge_downsample, + linear_init_mode=linear_init_mode, + norm_init_mode=norm_init_mode, + ) + self.reset_parameters() + + def reset_parameters(self): + for m in self.children(): + cuboid_utils.apply_initialization( + m, + conv_mode=self.conv_init_mode, + linear_mode=self.linear_init_mode, + norm_mode=self.norm_init_mode, + ) + + def forward(self, x): + """x --> [K x Conv2D] --> PatchMerge + + Args: + x: (B, T, H, W, C) + + Returns: + out: (B, T, H_new, W_new, C_out) + """ + + B, T, H, W, C = x.shape + + if self.num_conv_layers > 0: + x = x.reshape([B * T, H, W, C]).transpose(perm=[0, 3, 1, 2]) + x = self.conv_block(x).transpose(perm=[0, 2, 3, 1]) + x = self.patch_merge(x.reshape([B, T, H, W, -1])) + else: + x = self.patch_merge(x) + return x + + +class FinalDecoder(nn.Layer): + def __init__( + self, + target_thw: Tuple[int, ...], + dim: int, + num_conv_layers: int = 2, + activation: str = "leaky", + conv_init_mode: str = "0", + linear_init_mode: str = "0", + norm_init_mode: str = "0", + moe_config: dict = None, + ): + super(FinalDecoder, self).__init__() + self.target_thw = target_thw + self.dim = dim + self.num_conv_layers = num_conv_layers + self.conv_init_mode = conv_init_mode + self.linear_init_mode = linear_init_mode + self.norm_init_mode = norm_init_mode + conv_block = [] + for i in range(num_conv_layers): + conv_block.append( + nn.Conv2D( + kernel_size=(3, 3), + padding=(1, 1), + in_channels=dim, + out_channels=dim, + ) + ) + conv_block.append(nn.GroupNorm(num_groups=16, num_channels=dim)) + conv_block.append( + act_mod.get_activation(activation) + if activation != "leaky_relu" + else nn.LeakyReLU(NEGATIVE_SLOPE) + ) + self.conv_block = nn.Sequential(*conv_block) + self.upsample = cuboid_decoder.Upsample3DLayer( + dim=dim, + out_dim=dim, + target_size=target_thw, + kernel_size=3, + conv_init_mode=conv_init_mode, + ) + self.reset_parameters() + + def reset_parameters(self): + for m in self.children(): + cuboid_utils.apply_initialization( + m, + conv_mode=self.conv_init_mode, + linear_mode=self.linear_init_mode, + norm_mode=self.norm_init_mode, + ) + + def forward(self, x): + """x --> Upsample --> [K x Conv2D] + + Args: + x: (B, T, H, W, C) + + Returns: + out: (B, T, H_new, W_new, C) + """ + + x = self.upsample(x) + if self.num_conv_layers > 0: + B, T, H, W, C = x.shape + x = x.reshape([B * T, H, W, C]).transpose(perm=[0, 3, 1, 2]) + x = ( + self.conv_block(x) + .transpose(perm=[0, 2, 3, 1]) + .reshape([B, T, H, W, -1]) + ) + return x + + +class InitialStackPatchMergingEncoder(nn.Layer): + def __init__( + self, + num_merge: int, + in_dim: int, + out_dim_list: Tuple[int, ...], + downsample_scale_list: Tuple[float, ...], + num_conv_per_merge_list: Tuple[int, ...] = None, + activation: str = "leaky", + padding_type: str = "nearest", + conv_init_mode: str = "0", + linear_init_mode: str = "0", + norm_init_mode: str = "0", + moe_config: dict = None, + ): + super(InitialStackPatchMergingEncoder, self).__init__() + self.conv_init_mode = conv_init_mode + self.linear_init_mode = linear_init_mode + self.norm_init_mode = norm_init_mode + self.num_merge = num_merge + self.in_dim = in_dim + self.out_dim_list = out_dim_list[:num_merge] + self.downsample_scale_list = downsample_scale_list[:num_merge] + self.num_conv_per_merge_list = num_conv_per_merge_list + self.num_group_list = [max(1, out_dim // 4) for out_dim in self.out_dim_list] + self.conv_block_list = nn.LayerList() + self.patch_merge_list = nn.LayerList() + for i in range(num_merge): + if i == 0: + in_dim = in_dim + else: + in_dim = self.out_dim_list[i - 1] + out_dim = self.out_dim_list[i] + downsample_scale = self.downsample_scale_list[i] + conv_block = [] + for j in range(self.num_conv_per_merge_list[i]): + if j == 0: + conv_in_dim = in_dim + else: + conv_in_dim = out_dim + conv_block.append( + nn.Conv2D( + kernel_size=(3, 3), + padding=(1, 1), + in_channels=conv_in_dim, + out_channels=out_dim, + ) + ) + conv_block.append( + nn.GroupNorm( + num_groups=self.num_group_list[i], num_channels=out_dim + ) + ) + conv_block.append( + act_mod.get_activation(activation) + if activation != "leaky_relu" + else nn.LeakyReLU(NEGATIVE_SLOPE) + ) + conv_block = nn.Sequential(*conv_block) + self.conv_block_list.append(conv_block) + patch_merge = cuboid_encoder.PatchMerging3D( + dim=out_dim, + out_dim=out_dim, + padding_type=padding_type, + downsample=(1, downsample_scale, downsample_scale), + linear_init_mode=linear_init_mode, + norm_init_mode=norm_init_mode, + ) + self.patch_merge_list.append(patch_merge) + self.reset_parameters() + + def reset_parameters(self): + for m in self.children(): + cuboid_utils.apply_initialization( + m, + conv_mode=self.conv_init_mode, + linear_mode=self.linear_init_mode, + norm_mode=self.norm_init_mode, + ) + + def get_out_shape_list(self, input_shape): + out_shape_list = [] + for patch_merge in self.patch_merge_list: + input_shape = patch_merge.get_out_shape(input_shape) + out_shape_list.append(input_shape) + return out_shape_list + + def forward(self, x): + """x --> [K x Conv2D] --> PatchMerge --> ... --> [K x Conv2D] --> PatchMerge + + Args: + x: (B, T, H, W, C) + + Returns: + out: (B, T, H_new, W_new, C_out) + """ + + for i, (conv_block, patch_merge) in enumerate( + zip(self.conv_block_list, self.patch_merge_list) + ): + B, T, H, W, C = x.shape + if self.num_conv_per_merge_list[i] > 0: + x = x.reshape([B * T, H, W, C]).transpose(perm=[0, 3, 1, 2]) + x = conv_block(x).transpose(perm=[0, 2, 3, 1]).reshape([B, T, H, W, -1]) + x = patch_merge(x) + return x + + +class FinalStackUpsamplingDecoder(nn.Layer): + def __init__( + self, + target_shape_list: Tuple[Tuple[int, ...]], + in_dim: int, + num_conv_per_up_list: Tuple[int, ...] = None, + activation: str = "leaky", + conv_init_mode: str = "0", + linear_init_mode: str = "0", + norm_init_mode: str = "0", + moe_config: dict = None, + ): + super(FinalStackUpsamplingDecoder, self).__init__() + self.conv_init_mode = conv_init_mode + self.linear_init_mode = linear_init_mode + self.norm_init_mode = norm_init_mode + self.target_shape_list = target_shape_list + self.out_dim_list = [ + target_shape[-1] for target_shape in self.target_shape_list + ] + self.num_upsample = len(target_shape_list) + self.in_dim = in_dim + self.num_conv_per_up_list = num_conv_per_up_list + self.num_group_list = [max(1, out_dim // 4) for out_dim in self.out_dim_list] + self.conv_block_list = nn.LayerList() + self.upsample_list = nn.LayerList() + for i in range(self.num_upsample): + if i == 0: + in_dim = in_dim + else: + in_dim = self.out_dim_list[i - 1] + out_dim = self.out_dim_list[i] + upsample = cuboid_decoder.Upsample3DLayer( + dim=in_dim, + out_dim=in_dim, + target_size=target_shape_list[i][:-1], + kernel_size=3, + conv_init_mode=conv_init_mode, + ) + self.upsample_list.append(upsample) + conv_block = [] + for j in range(num_conv_per_up_list[i]): + if j == 0: + conv_in_dim = in_dim + else: + conv_in_dim = out_dim + conv_block.append( + nn.Conv2D( + kernel_size=(3, 3), + padding=(1, 1), + in_channels=conv_in_dim, + out_channels=out_dim, + ) + ) + conv_block.append( + nn.GroupNorm( + num_groups=self.num_group_list[i], num_channels=out_dim + ) + ) + conv_block.append( + act_mod.get_activation(activation) + if activation != "leaky_relu" + else nn.LeakyReLU(NEGATIVE_SLOPE) + ) + conv_block = nn.Sequential(*conv_block) + self.conv_block_list.append(conv_block) + self.reset_parameters() + + def reset_parameters(self): + for m in self.children(): + cuboid_utils.apply_initialization( + m, + conv_mode=self.conv_init_mode, + linear_mode=self.linear_init_mode, + norm_mode=self.norm_init_mode, + ) + + @staticmethod + def get_init_params(enc_input_shape, enc_out_shape_list, large_channel=False): + dec_target_shape_list = list(enc_out_shape_list[:-1])[::-1] + [ + tuple(enc_input_shape) + ] + if large_channel: + dec_target_shape_list_large_channel = [] + for i, enc_out_shape in enumerate(enc_out_shape_list[::-1]): + dec_target_shape_large_channel = list(dec_target_shape_list[i]) + dec_target_shape_large_channel[-1] = enc_out_shape[-1] + dec_target_shape_list_large_channel.append( + tuple(dec_target_shape_large_channel) + ) + dec_target_shape_list = dec_target_shape_list_large_channel + dec_in_dim = enc_out_shape_list[-1][-1] + return dec_target_shape_list, dec_in_dim + + def forward(self, x): + """x --> Upsample --> [K x Conv2D] --> ... --> Upsample --> [K x Conv2D] + + Args: + x: Shape (B, T, H, W, C) + + Returns: + out: Shape (B, T, H_new, W_new, C) + """ + for i, (conv_block, upsample) in enumerate( + zip(self.conv_block_list, self.upsample_list) + ): + x = upsample(x) + if self.num_conv_per_up_list[i] > 0: + B, T, H, W, C = x.shape + x = x.reshape([B * T, H, W, C]).transpose(perm=[0, 3, 1, 2]) + x = conv_block(x).transpose(perm=[0, 2, 3, 1]).reshape([B, T, H, W, -1]) + return x + + +class ExtFormerMoECuboid(base.Arch): + """Cuboid Transformer for spatiotemporal forecasting + + We adopt the Non-autoregressive encoder-decoder architecture. + The decoder takes the multi-scale memory output from the encoder. + + The initial downsampling / upsampling layers will be + Downsampling: [K x Conv2D --> PatchMerge] + Upsampling: [Nearest Interpolation-based Upsample --> K x Conv2D] + + x --> downsample (optional) ---> (+pos_embed) ---> enc --> mem_l initial_z (+pos_embed) ---> FC + | | + |------------| + | + | + y <--- upsample (optional) <--- dec <---------- + + Args: + input_keys (Tuple[str, ...]): Name of input keys, such as ("input",). + output_keys (Tuple[str, ...]): Name of output keys, such as ("output",). + input_shape (Tuple[int, ...]): The shape of the input data. + target_shape (Tuple[int, ...]): The shape of the target data. + base_units (int, optional): The base units. Defaults to 128. + block_units (int, optional): The block units. Defaults to None. + scale_alpha (float, optional): We scale up the channels based on the formula: + - round_to(base_units * max(downsample_scale) ** units_alpha, 4). Defaults to 1.0. + num_heads (int, optional): The number of heads. Defaults to 4. + attn_drop (float, optional): The attention dropout. Defaults to 0.0. + proj_drop (float, optional): The projection dropout. Defaults to 0.0. + ffn_drop (float, optional): The ffn dropout. Defaults to 0.0. + downsample (int, optional): The rate of downsample. Defaults to 2. + downsample_type (str, optional): The type of downsample. Defaults to "patch_merge". + upsample_type (str, optional): The rate of upsample. Defaults to "upsample". + upsample_kernel_size (int, optional): The kernel size of upsample. Defaults to 3. + enc_depth (list, optional): The depth of encoder. Defaults to [4, 4, 4]. + enc_attn_patterns (str, optional): The pattern of encoder attention. Defaults to None. + enc_cuboid_size (list, optional): The cuboid size of encoder. Defaults to [(4, 4, 4), (4, 4, 4)]. + enc_cuboid_strategy (list, optional): The cuboid strategy of encoder. Defaults to [("l", "l", "l"), ("d", "d", "d")]. + enc_shift_size (list, optional): The shift size of encoder. Defaults to [(0, 0, 0), (0, 0, 0)]. + enc_use_inter_ffn (bool, optional): Whether to use intermediate FFN for encoder. Defaults to True. + dec_depth (list, optional): The depth of decoder. Defaults to [2, 2]. + dec_cross_start (int, optional): The cross start of decoder. Defaults to 0. + dec_self_attn_patterns (str, optional): The partterns of decoder. Defaults to None. + dec_self_cuboid_size (list, optional): The cuboid size of decoder. Defaults to [(4, 4, 4), (4, 4, 4)]. + dec_self_cuboid_strategy (list, optional): The strategy of decoder. Defaults to [("l", "l", "l"), ("d", "d", "d")]. + dec_self_shift_size (list, optional): The shift size of decoder. Defaults to [(1, 1, 1), (0, 0, 0)]. + dec_cross_attn_patterns (_type_, optional): The cross attention patterns of decoder. Defaults to None. + dec_cross_cuboid_hw (list, optional): The cuboid_hw of decoder. Defaults to [(4, 4), (4, 4)]. + dec_cross_cuboid_strategy (list, optional): The cuboid strategy of decoder. Defaults to [("l", "l", "l"), ("d", "l", "l")]. + dec_cross_shift_hw (list, optional): The shift_hw of decoder. Defaults to [(0, 0), (0, 0)]. + dec_cross_n_temporal (list, optional): The cross_n_temporal of decoder. Defaults to [1, 2]. + dec_cross_last_n_frames (int, optional): The cross_last_n_frames of decoder. Defaults to None. + dec_use_inter_ffn (bool, optional): Whether to use intermediate FFN for decoder. Defaults to True. + dec_hierarchical_pos_embed (bool, optional): Whether to use hierarchical pos_embed for decoder. Defaults to False. + num_global_vectors (int, optional): The num of global vectors. Defaults to 4. + use_dec_self_global (bool, optional): Whether to use global vector for decoder. Defaults to True. + dec_self_update_global (bool, optional): Whether to update global vector for decoder. Defaults to True. + use_dec_cross_global (bool, optional): Whether to use cross global vector for decoder. Defaults to True. + use_global_vector_ffn (bool, optional): Whether to use global vector FFN. Defaults to True. + use_global_self_attn (bool, optional): Whether to use global attentions. Defaults to False. + separate_global_qkv (bool, optional): Whether to separate global qkv. Defaults to False. + global_dim_ratio (int, optional): The ratio of global dim. Defaults to 1. + self_pattern (str, optional): The pattern. Defaults to "axial". + cross_self_pattern (str, optional): The self cross pattern. Defaults to "axial". + cross_pattern (str, optional): The cross pattern. Defaults to "cross_1x1". + z_init_method (str, optional): How the initial input to the decoder is initialized. Defaults to "nearest_interp". + initial_downsample_type (str, optional): The downsample type of initial. Defaults to "conv". + initial_downsample_activation (str, optional): The downsample activation of initial. Defaults to "leaky". + initial_downsample_scale (int, optional): The downsample scale of initial. Defaults to 1. + initial_downsample_conv_layers (int, optional): The conv layer of downsample of initial. Defaults to 2. + final_upsample_conv_layers (int, optional): The conv layer of final upsample. Defaults to 2. + initial_downsample_stack_conv_num_layers (int, optional): The num of stack conv layer of initial downsample. Defaults to 1. + initial_downsample_stack_conv_dim_list (list, optional): The dim list of stack conv of initial downsample. Defaults to None. + initial_downsample_stack_conv_downscale_list (list, optional): The downscale list of stack conv of initial downsample. Defaults to [1]. + initial_downsample_stack_conv_num_conv_list (list, optional): The num of stack conv list of initial downsample. Defaults to [2]. + ffn_activation (str, optional): The activation of FFN. Defaults to "leaky". + gated_ffn (bool, optional): Whether to use gate FFN. Defaults to False. + norm_layer (str, optional): The type of normilize. Defaults to "layer_norm". + padding_type (str, optional): The type of padding. Defaults to "ignore". + pos_embed_type (str, optional): The type of pos embedding. Defaults to "t+hw". + checkpoint_level (bool, optional): Whether to use checkpoint. Defaults to True. + use_relative_pos (bool, optional): Whether to use relative pose. Defaults to True. + self_attn_use_final_proj (bool, optional): Whether to use final projection. Defaults to True. + dec_use_first_self_attn (bool, optional): Whether to use first self attention for decoder. Defaults to False. + attn_linear_init_mode (str, optional): The mode of attention linear init. Defaults to "0". + ffn_linear_init_mode (str, optional): The mode of FFN linear init. Defaults to "0". + conv_init_mode (str, optional): The mode of conv init. Defaults to "0". + down_up_linear_init_mode (str, optional): The mode of downsample and upsample linear init. Defaults to "0". + norm_init_mode (str, optional): The mode of normalization init. Defaults to "0". + """ + + def __init__( + self, + input_keys: Tuple[str, ...], + output_keys: Tuple[str, ...], + input_shape: Tuple[int, ...], + target_shape: Tuple[int, ...], + base_units: int = 128, + block_units: int = None, + scale_alpha: float = 1.0, + num_heads: int = 4, + attn_drop: float = 0.0, + proj_drop: float = 0.0, + ffn_drop: float = 0.0, + downsample: int = 2, + downsample_type: str = "patch_merge", + upsample_type: str = "upsample", + upsample_kernel_size: int = 3, + enc_depth: Tuple[int, ...] = [4, 4, 4], + enc_attn_patterns: str = None, + enc_cuboid_size: Tuple[Tuple[int, ...], ...] = [(4, 4, 4), (4, 4, 4)], + enc_cuboid_strategy: Tuple[Tuple[str, ...], ...] = [ + ("l", "l", "l"), + ("d", "d", "d"), + ], + enc_shift_size: Tuple[Tuple[int, ...], ...] = [(0, 0, 0), (0, 0, 0)], + enc_use_inter_ffn: bool = True, + dec_depth: Tuple[int, ...] = [2, 2], + dec_cross_start: int = 0, + dec_self_attn_patterns: str = None, + dec_self_cuboid_size: Tuple[Tuple[int, ...], ...] = [(4, 4, 4), (4, 4, 4)], + dec_self_cuboid_strategy: Tuple[Tuple[str, ...], ...] = [ + ("l", "l", "l"), + ("d", "d", "d"), + ], + dec_self_shift_size: Tuple[Tuple[int, ...], ...] = [(1, 1, 1), (0, 0, 0)], + dec_cross_attn_patterns: str = None, + dec_cross_cuboid_hw: Tuple[Tuple[int, ...], ...] = [(4, 4), (4, 4)], + dec_cross_cuboid_strategy: Tuple[Tuple[str, ...], ...] = [ + ("l", "l", "l"), + ("d", "l", "l"), + ], + dec_cross_shift_hw: Tuple[Tuple[int, ...], ...] = [(0, 0), (0, 0)], + dec_cross_n_temporal: Tuple[int, ...] = [1, 2], + dec_cross_last_n_frames: int = None, + dec_use_inter_ffn: bool = True, + dec_hierarchical_pos_embed: bool = False, + num_global_vectors: int = 4, + use_dec_self_global: bool = True, + dec_self_update_global: bool = True, + use_dec_cross_global: bool = True, + use_global_vector_ffn: bool = True, + use_global_self_attn: bool = False, + separate_global_qkv: bool = False, + global_dim_ratio: int = 1, + self_pattern: str = "axial", + cross_self_pattern: str = "axial", + cross_pattern: str = "cross_1x1", + z_init_method: str = "nearest_interp", + initial_downsample_type: str = "conv", + initial_downsample_activation: str = "leaky", + initial_downsample_scale: int = 1, + initial_downsample_conv_layers: int = 2, + final_upsample_conv_layers: int = 2, + initial_downsample_stack_conv_num_layers: int = 1, + initial_downsample_stack_conv_dim_list: Tuple[int, ...] = None, + initial_downsample_stack_conv_downscale_list: Tuple[int, ...] = [1], + initial_downsample_stack_conv_num_conv_list: Tuple[int, ...] = [2], + ffn_activation: str = "leaky", + gated_ffn: bool = False, + norm_layer: str = "layer_norm", + padding_type: str = "ignore", + pos_embed_type: str = "t+hw", + checkpoint_level: bool = True, + use_relative_pos: bool = True, + self_attn_use_final_proj: bool = True, + dec_use_first_self_attn: bool = False, + attn_linear_init_mode: str = "0", + ffn_linear_init_mode: str = "0", + conv_init_mode: str = "0", + down_up_linear_init_mode: str = "0", + norm_init_mode: str = "0", + moe_config: dict = None, + rnc_config: dict = None, + ): + super().__init__() + self.input_keys = input_keys + self.output_keys = output_keys + self.attn_linear_init_mode = attn_linear_init_mode + self.ffn_linear_init_mode = ffn_linear_init_mode + self.conv_init_mode = conv_init_mode + self.down_up_linear_init_mode = down_up_linear_init_mode + self.norm_init_mode = norm_init_mode + assert len(enc_depth) == len(dec_depth) + self.base_units = base_units + self.num_global_vectors = num_global_vectors + self.moe_config = moe_config + self.rnc_config = rnc_config + self.checkpoint_level = checkpoint_level + + num_blocks = len(enc_depth) + if isinstance(self_pattern, str): + enc_attn_patterns = [self_pattern] * num_blocks + if isinstance(cross_self_pattern, str): + dec_self_attn_patterns = [cross_self_pattern] * num_blocks + if isinstance(cross_pattern, str): + dec_cross_attn_patterns = [cross_pattern] * num_blocks + if global_dim_ratio != 1: + assert ( + separate_global_qkv is True + ), "Setting global_dim_ratio != 1 requires separate_global_qkv == True." + self.global_dim_ratio = global_dim_ratio + self.z_init_method = z_init_method + assert self.z_init_method in ["zeros", "nearest_interp", "last", "mean"] + self.input_shape = input_shape + self.target_shape = target_shape + T_in, H_in, W_in, C_in = input_shape + T_out, H_out, W_out, C_out = target_shape + assert H_in == H_out and W_in == W_out + if self.num_global_vectors > 0: + init_data = paddle.zeros( + (self.num_global_vectors, global_dim_ratio * base_units) + ) + self.init_global_vectors = paddle.create_parameter( + shape=init_data.shape, + dtype=init_data.dtype, + default_initializer=nn.initializer.Constant(0.0), + ) + + self.init_global_vectors.stop_gradient = not True + new_input_shape = self.get_initial_encoder_final_decoder( + initial_downsample_scale=initial_downsample_scale, + initial_downsample_type=initial_downsample_type, + activation=initial_downsample_activation, + initial_downsample_conv_layers=initial_downsample_conv_layers, + final_upsample_conv_layers=final_upsample_conv_layers, + padding_type=padding_type, + initial_downsample_stack_conv_num_layers=initial_downsample_stack_conv_num_layers, + initial_downsample_stack_conv_dim_list=initial_downsample_stack_conv_dim_list, + initial_downsample_stack_conv_downscale_list=initial_downsample_stack_conv_downscale_list, + initial_downsample_stack_conv_num_conv_list=initial_downsample_stack_conv_num_conv_list, + ) + T_in, H_in, W_in, _ = new_input_shape + self.encoder = cuboid_encoder.CuboidTransformerEncoder( + input_shape=(T_in, H_in, W_in, base_units), + base_units=base_units, + block_units=block_units, + scale_alpha=scale_alpha, + depth=enc_depth, + downsample=downsample, + downsample_type=downsample_type, + block_attn_patterns=enc_attn_patterns, + block_cuboid_size=enc_cuboid_size, + block_strategy=enc_cuboid_strategy, + block_shift_size=enc_shift_size, + num_heads=num_heads, + attn_drop=attn_drop, + proj_drop=proj_drop, + ffn_drop=ffn_drop, + gated_ffn=gated_ffn, + ffn_activation=ffn_activation, + norm_layer=norm_layer, + use_inter_ffn=enc_use_inter_ffn, + padding_type=padding_type, + use_global_vector=num_global_vectors > 0, + use_global_vector_ffn=use_global_vector_ffn, + use_global_self_attn=use_global_self_attn, + separate_global_qkv=separate_global_qkv, + global_dim_ratio=global_dim_ratio, + checkpoint_level=checkpoint_level, + use_relative_pos=use_relative_pos, + self_attn_use_final_proj=self_attn_use_final_proj, + attn_linear_init_mode=attn_linear_init_mode, + ffn_linear_init_mode=ffn_linear_init_mode, + conv_init_mode=conv_init_mode, + down_linear_init_mode=down_up_linear_init_mode, + norm_init_mode=norm_init_mode, + moe_config=moe_config, + ) + self.enc_pos_embed = cuboid_decoder.PosEmbed( + embed_dim=base_units, typ=pos_embed_type, maxH=H_in, maxW=W_in, maxT=T_in + ) + mem_shapes = self.encoder.get_mem_shapes() + self.z_proj = nn.Linear( + in_features=mem_shapes[-1][-1], out_features=mem_shapes[-1][-1] + ) + self.dec_pos_embed = cuboid_decoder.PosEmbed( + embed_dim=mem_shapes[-1][-1], + typ=pos_embed_type, + maxT=T_out, + maxH=mem_shapes[-1][1], + maxW=mem_shapes[-1][2], + ) + self.decoder = cuboid_decoder.CuboidTransformerDecoder( + target_temporal_length=T_out, + mem_shapes=mem_shapes, + cross_start=dec_cross_start, + depth=dec_depth, + upsample_type=upsample_type, + block_self_attn_patterns=dec_self_attn_patterns, + block_self_cuboid_size=dec_self_cuboid_size, + block_self_shift_size=dec_self_shift_size, + block_self_cuboid_strategy=dec_self_cuboid_strategy, + block_cross_attn_patterns=dec_cross_attn_patterns, + block_cross_cuboid_hw=dec_cross_cuboid_hw, + block_cross_shift_hw=dec_cross_shift_hw, + block_cross_cuboid_strategy=dec_cross_cuboid_strategy, + block_cross_n_temporal=dec_cross_n_temporal, + cross_last_n_frames=dec_cross_last_n_frames, + num_heads=num_heads, + attn_drop=attn_drop, + proj_drop=proj_drop, + ffn_drop=ffn_drop, + upsample_kernel_size=upsample_kernel_size, + ffn_activation=ffn_activation, + gated_ffn=gated_ffn, + norm_layer=norm_layer, + use_inter_ffn=dec_use_inter_ffn, + max_temporal_relative=T_in + T_out, + padding_type=padding_type, + hierarchical_pos_embed=dec_hierarchical_pos_embed, + pos_embed_type=pos_embed_type, + use_self_global=num_global_vectors > 0 and use_dec_self_global, + self_update_global=dec_self_update_global, + use_cross_global=num_global_vectors > 0 and use_dec_cross_global, + use_global_vector_ffn=use_global_vector_ffn, + use_global_self_attn=use_global_self_attn, + separate_global_qkv=separate_global_qkv, + global_dim_ratio=global_dim_ratio, + checkpoint_level=checkpoint_level, + use_relative_pos=use_relative_pos, + self_attn_use_final_proj=self_attn_use_final_proj, + use_first_self_attn=dec_use_first_self_attn, + attn_linear_init_mode=attn_linear_init_mode, + ffn_linear_init_mode=ffn_linear_init_mode, + conv_init_mode=conv_init_mode, + up_linear_init_mode=down_up_linear_init_mode, + norm_init_mode=norm_init_mode, + moe_config=moe_config, + ) + + if rnc_config["use_rnc"]: + self.rnc_cri = extformer_moe_utils.RnCLoss(rnc_config) + + self.reset_parameters() + + def get_initial_encoder_final_decoder( + self, + initial_downsample_type, + activation, + initial_downsample_scale, + initial_downsample_conv_layers, + final_upsample_conv_layers, + padding_type, + initial_downsample_stack_conv_num_layers, + initial_downsample_stack_conv_dim_list, + initial_downsample_stack_conv_downscale_list, + initial_downsample_stack_conv_num_conv_list, + ): + T_in, H_in, W_in, C_in = self.input_shape + T_out, H_out, W_out, C_out = self.target_shape + self.initial_downsample_type = initial_downsample_type + if self.initial_downsample_type == "conv": + if isinstance(initial_downsample_scale, int): + initial_downsample_scale = ( + 1, + initial_downsample_scale, + initial_downsample_scale, + ) + elif len(initial_downsample_scale) == 2: + initial_downsample_scale = 1, *initial_downsample_scale + elif len(initial_downsample_scale) == 3: + initial_downsample_scale = tuple(initial_downsample_scale) + else: + raise NotImplementedError( + f"initial_downsample_scale {initial_downsample_scale} format not supported!" + ) + self.initial_encoder = InitialEncoder( + dim=C_in, + out_dim=self.base_units, + downsample_scale=initial_downsample_scale, + num_conv_layers=initial_downsample_conv_layers, + padding_type=padding_type, + activation=activation, + conv_init_mode=self.conv_init_mode, + linear_init_mode=self.down_up_linear_init_mode, + norm_init_mode=self.norm_init_mode, + ) + + self.final_decoder = FinalDecoder( + dim=self.base_units, + target_thw=(T_out, H_out, W_out), + num_conv_layers=final_upsample_conv_layers, + activation=activation, + conv_init_mode=self.conv_init_mode, + linear_init_mode=self.down_up_linear_init_mode, + norm_init_mode=self.norm_init_mode, + ) + new_input_shape = self.initial_encoder.patch_merge.get_out_shape( + self.input_shape + ) + self.dec_final_proj = nn.Linear( + in_features=self.base_units, out_features=C_out + ) + elif self.initial_downsample_type == "stack_conv": + if initial_downsample_stack_conv_dim_list is None: + initial_downsample_stack_conv_dim_list = [ + self.base_units + ] * initial_downsample_stack_conv_num_layers + self.initial_encoder = InitialStackPatchMergingEncoder( + num_merge=initial_downsample_stack_conv_num_layers, + in_dim=C_in, + out_dim_list=initial_downsample_stack_conv_dim_list, + downsample_scale_list=initial_downsample_stack_conv_downscale_list, + num_conv_per_merge_list=initial_downsample_stack_conv_num_conv_list, + padding_type=padding_type, + activation=activation, + conv_init_mode=self.conv_init_mode, + linear_init_mode=self.down_up_linear_init_mode, + norm_init_mode=self.norm_init_mode, + ) + initial_encoder_out_shape_list = self.initial_encoder.get_out_shape_list( + self.target_shape + ) + ( + dec_target_shape_list, + dec_in_dim, + ) = FinalStackUpsamplingDecoder.get_init_params( + enc_input_shape=self.target_shape, + enc_out_shape_list=initial_encoder_out_shape_list, + large_channel=True, + ) + self.final_decoder = FinalStackUpsamplingDecoder( + target_shape_list=dec_target_shape_list, + in_dim=dec_in_dim, + num_conv_per_up_list=initial_downsample_stack_conv_num_conv_list[::-1], + activation=activation, + conv_init_mode=self.conv_init_mode, + linear_init_mode=self.down_up_linear_init_mode, + norm_init_mode=self.norm_init_mode, + ) + self.dec_final_proj = nn.Linear( + in_features=dec_target_shape_list[-1][-1], out_features=C_out + ) + new_input_shape = self.initial_encoder.get_out_shape_list(self.input_shape)[ + -1 + ] + else: + raise NotImplementedError(f"{self.initial_downsample_type} is invalid.") + self.input_shape_after_initial_downsample = new_input_shape + T_in, H_in, W_in, _ = new_input_shape + return new_input_shape + + def reset_parameters(self): + if self.num_global_vectors > 0: + self.init_global_vectors = initializer.trunc_normal_( + self.init_global_vectors, std=0.02 + ) + if hasattr(self.initial_encoder, "reset_parameters"): + self.initial_encoder.reset_parameters() + else: + cuboid_utils.apply_initialization( + self.initial_encoder, + conv_mode=self.conv_init_mode, + linear_mode=self.down_up_linear_init_mode, + norm_mode=self.norm_init_mode, + ) + if hasattr(self.final_decoder, "reset_parameters"): + self.final_decoder.reset_parameters() + else: + cuboid_utils.apply_initialization( + self.final_decoder, + conv_mode=self.conv_init_mode, + linear_mode=self.down_up_linear_init_mode, + norm_mode=self.norm_init_mode, + ) + cuboid_utils.apply_initialization( + self.dec_final_proj, linear_mode=self.down_up_linear_init_mode + ) + self.encoder.reset_parameters() + self.enc_pos_embed.reset_parameters() + self.decoder.reset_parameters() + self.dec_pos_embed.reset_parameters() + cuboid_utils.apply_initialization(self.z_proj, linear_mode="0") + + def get_initial_z(self, final_mem, T_out): + B = final_mem.shape[0] + if self.z_init_method == "zeros": + z_shape = list((1, T_out)) + final_mem.shape[2:] + initial_z = paddle.zeros(shape=z_shape, dtype=final_mem.dtype) + initial_z = self.z_proj(self.dec_pos_embed(initial_z)).expand( + shape=[B, -1, -1, -1, -1] + ) + elif self.z_init_method == "nearest_interp": + initial_z = nn.functional.interpolate( + x=final_mem.transpose(perm=[0, 4, 1, 2, 3]), + size=(T_out, final_mem.shape[2], final_mem.shape[3]), + ).transpose(perm=[0, 2, 3, 4, 1]) + initial_z = self.z_proj(initial_z) + elif self.z_init_method == "last": + initial_z = paddle.broadcast_to( + x=final_mem[:, -1:, :, :, :], shape=(B, T_out) + final_mem.shape[2:] + ) + initial_z = self.z_proj(initial_z) + elif self.z_init_method == "mean": + initial_z = paddle.broadcast_to( + x=final_mem.mean(axis=1, keepdims=True), + shape=(B, T_out) + final_mem.shape[2:], + ) + initial_z = self.z_proj(initial_z) + else: + raise NotImplementedError + return initial_z + + def forward(self, x: "paddle.Tensor", verbose: bool = False) -> "paddle.Tensor": + """ + Args: + x (paddle.Tensor): Tensor with shape (B, T, H, W, C). + verbose (bool): if True, print intermediate shapes. + + Returns: + out (paddle.Tensor): The output Shape (B, T_out, H, W, C_out) + """ + + labels = x["sst_target"] + x = self.concat_to_tensor(x, self.input_keys) + flag_ndim = x.ndim + if flag_ndim == 6: + x = x.reshape([-1, *x.shape[2:]]) + B, _, _, _, _ = x.shape + + T_out = self.target_shape[0] + x = self.initial_encoder(x) + x = self.enc_pos_embed(x) + + if self.num_global_vectors > 0: + init_global_vectors = self.init_global_vectors.expand( + shape=[ + B, + self.num_global_vectors, + self.global_dim_ratio * self.base_units, + ] + ) + mem_l, mem_global_vector_l = self.encoder(x, init_global_vectors) + else: + mem_l = self.encoder(x) + + if verbose: + for i, mem in enumerate(mem_l): + print(f"mem[{i}].shape = {mem.shape}") + initial_z = self.get_initial_z(final_mem=mem_l[-1], T_out=T_out) + + if self.num_global_vectors > 0: + dec_out = self.decoder(initial_z, mem_l, mem_global_vector_l) + else: + dec_out = self.decoder(initial_z, mem_l) + + dec_out = self.final_decoder(dec_out) + out = self.dec_final_proj(dec_out) + + if flag_ndim == 6: + out = out.reshape([-1, *out.shape]) + + out_dict = {key: out for key in self.output_keys[:2]} + + # moe loss + if self.training: + aux_losses = extformer_moe_utils.aggregate_aux_losses(self) + if len(aux_losses) > 0: + aux_loss = paddle.concat(aux_losses).mean() + else: + aux_loss = None + else: + aux_loss = None + assert "aux_loss" in self.output_keys + out_dict["aux_loss"] = aux_loss + + # rnc + if self.training and self.rnc_config["use_rnc"]: + rank_loss = self.rnc_cri(dec_out, labels) + rank_loss = rank_loss.unsqueeze(0) + else: + rank_loss = None + assert "rank_loss" in self.output_keys + out_dict["rank_loss"] = rank_loss + + return out_dict diff --git a/examples/smc_reac/ppsci/arch/extformer_moe_cuboid_decoder.py b/examples/smc_reac/ppsci/arch/extformer_moe_cuboid_decoder.py new file mode 100644 index 0000000000..b16311a0e7 --- /dev/null +++ b/examples/smc_reac/ppsci/arch/extformer_moe_cuboid_decoder.py @@ -0,0 +1,1475 @@ +from functools import lru_cache +from typing import Tuple + +import numpy as np +import paddle +import paddle.nn.functional as F +from paddle import nn +from paddle.distributed import fleet + +import ppsci.arch.extformer_moe_cuboid_encoder as cuboid_encoder +import ppsci.arch.extformer_moe_cuboid_utils as cuboid_utils +import ppsci.arch.extformer_moe_utils as moe_utils +from ppsci.utils import initializer + + +class PosEmbed(nn.Layer): + """pose embedding + + Args: + embed_dim (int): The dimension of embedding. + maxT (int): The embedding max time. + maxH (int): The embedding max height. + maxW (int): The embedding max width. + typ (str): + The type of the positional embedding. + - t+h+w: + Embed the spatial position to embeddings + - t+hw: + Embed the spatial position to embeddings + """ + + def __init__( + self, + embed_dim, + maxT, + maxH, + maxW, + typ: str = "t+h+w", + moe_config: dict = None, + ): + super(PosEmbed, self).__init__() + self.typ = typ + assert self.typ in ["t+h+w", "t+hw"] + self.maxT = maxT + self.maxH = maxH + self.maxW = maxW + self.embed_dim = embed_dim + if self.typ == "t+h+w": + self.T_embed = nn.Embedding(num_embeddings=maxT, embedding_dim=embed_dim) + self.H_embed = nn.Embedding(num_embeddings=maxH, embedding_dim=embed_dim) + self.W_embed = nn.Embedding(num_embeddings=maxW, embedding_dim=embed_dim) + elif self.typ == "t+hw": + self.T_embed = nn.Embedding(num_embeddings=maxT, embedding_dim=embed_dim) + self.HW_embed = nn.Embedding( + num_embeddings=maxH * maxW, embedding_dim=embed_dim + ) + else: + raise NotImplementedError(f"{self.typ} is invalid.") + self.reset_parameters() + + def reset_parameters(self): + for m in self.children(): + cuboid_utils.apply_initialization(m, embed_mode="0") + + def forward(self, x): + """ + Args: + x : Shape (B, T, H, W, C) + + Returns: + out : the x + positional embeddings + """ + + _, T, H, W, _ = x.shape + t_idx = paddle.arange(end=T) + h_idx = paddle.arange(end=H) + w_idx = paddle.arange(end=W) + if self.typ == "t+h+w": + return ( + x + + self.T_embed(t_idx).reshape([T, 1, 1, self.embed_dim]) + + self.H_embed(h_idx).reshape([1, H, 1, self.embed_dim]) + + self.W_embed(w_idx).reshape([1, 1, W, self.embed_dim]) + ) + elif self.typ == "t+hw": + spatial_idx = h_idx.unsqueeze(axis=-1) * self.maxW + w_idx + return ( + x + + self.T_embed(t_idx).reshape([T, 1, 1, self.embed_dim]) + + self.HW_embed(spatial_idx) + ) + else: + raise NotImplementedError(f"{self.typ} is invalid.") + + +@lru_cache() +def compute_cuboid_cross_attention_mask( + T_x, T_mem, H, W, n_temporal, cuboid_hw, shift_hw, strategy, padding_type, device +): + pad_t_mem = (n_temporal - T_mem % n_temporal) % n_temporal + pad_t_x = (n_temporal - T_x % n_temporal) % n_temporal + pad_h = (cuboid_hw[0] - H % cuboid_hw[0]) % cuboid_hw[0] + pad_w = (cuboid_hw[1] - W % cuboid_hw[1]) % cuboid_hw[1] + mem_cuboid_size = ((T_mem + pad_t_mem) // n_temporal,) + cuboid_hw + x_cuboid_size = ((T_x + pad_t_x) // n_temporal,) + cuboid_hw + if pad_t_mem > 0 or pad_h > 0 or pad_w > 0: + if padding_type == "ignore": + mem_mask = paddle.ones(shape=(1, T_mem, H, W, 1), dtype="bool") + mem_mask = F.pad( + mem_mask, [0, 0, 0, pad_w, 0, pad_h, pad_t_mem, 0], data_format="NDHWC" + ) + else: + mem_mask = paddle.ones( + shape=(1, T_mem + pad_t_mem, H + pad_h, W + pad_w, 1), dtype="bool" + ) + if pad_t_x > 0 or pad_h > 0 or pad_w > 0: + if padding_type == "ignore": + x_mask = paddle.ones(shape=(1, T_x, H, W, 1), dtype="bool") + x_mask = F.pad( + x_mask, [0, 0, 0, pad_w, 0, pad_h, 0, pad_t_x], data_format="NDHWC" + ) + else: + x_mask = paddle.ones( + shape=(1, T_x + pad_t_x, H + pad_h, W + pad_w, 1), dtype="bool" + ) + if any(i > 0 for i in shift_hw): + if padding_type == "ignore": + x_mask = paddle.roll( + x=x_mask, shifts=(-shift_hw[0], -shift_hw[1]), axis=(2, 3) + ) + mem_mask = paddle.roll( + x=mem_mask, shifts=(-shift_hw[0], -shift_hw[1]), axis=(2, 3) + ) + x_mask = cuboid_encoder.cuboid_reorder(x_mask, x_cuboid_size, strategy=strategy) + x_mask = x_mask.squeeze(axis=-1).squeeze(axis=0) + num_cuboids, x_cuboid_volume = x_mask.shape + mem_mask = cuboid_encoder.cuboid_reorder( + mem_mask, mem_cuboid_size, strategy=strategy + ) + mem_mask = mem_mask.squeeze(axis=-1).squeeze(axis=0) + _, mem_cuboid_volume = mem_mask.shape + shift_mask = np.zeros(shape=(1, n_temporal, H + pad_h, W + pad_w, 1)) + cnt = 0 + for h in ( + slice(-cuboid_hw[0]), + slice(-cuboid_hw[0], -shift_hw[0]), + slice(-shift_hw[0], None), + ): + for w in ( + slice(-cuboid_hw[1]), + slice(-cuboid_hw[1], -shift_hw[1]), + slice(-shift_hw[1], None), + ): + shift_mask[:, :, h, w, :] = cnt + cnt += 1 + shift_mask = paddle.to_tensor(shift_mask) + shift_mask = cuboid_encoder.cuboid_reorder( + shift_mask, (1,) + cuboid_hw, strategy=strategy + ) + shift_mask = shift_mask.squeeze(axis=-1).squeeze(axis=0) + shift_mask = shift_mask.unsqueeze(axis=1) - shift_mask.unsqueeze(axis=2) == 0 + bh_bw = cuboid_hw[0] * cuboid_hw[1] + attn_mask = ( + shift_mask.reshape((num_cuboids, 1, bh_bw, 1, bh_bw)) + * x_mask.reshape((num_cuboids, -1, bh_bw, 1, 1)) + * mem_mask.reshape([num_cuboids, 1, 1, -1, bh_bw]) + ) + attn_mask = attn_mask.reshape([num_cuboids, x_cuboid_volume, mem_cuboid_volume]) + return attn_mask + + +class CuboidCrossAttentionLayer(nn.Layer): + """Implements the cuboid cross attention. + + The idea of Cuboid Cross Attention is to extend the idea of cuboid self attention to work for the + encoder-decoder-type cross attention. + + Assume that there is a memory tensor with shape (T1, H, W, C) and another query tensor with shape (T2, H, W, C), + + Here, we decompose the query tensor and the memory tensor into the same number of cuboids and attend the cuboid in + the query tensor with the corresponding cuboid in the memory tensor. + + For the height and width axes, we reuse the grid decomposition techniques described in the cuboid self-attention. + For the temporal axis, the layer supports the "n_temporal" parameter, that controls the number of cuboids we can + get after cutting the tensors. For example, if the temporal dilation is 2, both the query and + memory will be decomposed into 2 cuboids along the temporal axis. Like in the Cuboid Self-attention, + we support "local" and "dilated" decomposition strategy. + + The complexity of the layer is O((T2 / n_t * Bh * Bw) * (T1 / n_t * Bh * Bw) * n_t (H / Bh) (W / Bw)) = O(T2 * T1 / n_t H W Bh Bw) + + Args: + dim (int): The dimension of input tensor. + num_heads (int): The number of head. + n_temporal (int, optional): The num of temporal. Defaults to 1. + cuboid_hw (tuple, optional): The height and width of cuboid. Defaults to (7, 7). + shift_hw (tuple, optional): The height and width of shift. Defaults to (0, 0). + strategy (tuple, optional): The strategy. Defaults to ("d", "l", "l"). + padding_type (str, optional): The type of padding. Defaults to "ignore". + cross_last_n_frames (int, optional): The cross_last_n_frames of decoder. Defaults to None. + qkv_bias (bool, optional): Whether to enable bias in calculating qkv attention. Defaults to False. + qk_scale (float, optional): Whether to enable scale factor when calculating the attention. Defaults to None. + attn_drop (float, optional): The attention dropout. Defaults to 0.0. + proj_drop (float, optional): The projrction dropout. Defaults to 0.0. + max_temporal_relative (int, optional): The max temporal. Defaults to 50. + norm_layer (str, optional): The normalization layer. Defaults to "layer_norm". + use_global_vector (bool, optional): Whether to use the global vector or not. Defaults to True. + separate_global_qkv (bool, optional): Whether to use different network to calc q_global, k_global, v_global. Defaults to False. + global_dim_ratio (int, optional): The dim (channels) of global vectors is `global_dim_ratio*dim`. Defaults to 1. + checkpoint_level (int, optional): Whether to enable gradient checkpointing. Defaults to 1. + use_relative_pos (bool, optional): Whether to use relative pos. Defaults to True. + attn_linear_init_mode (str, optional): The mode of attention linear initialization. Defaults to "0". + ffn_linear_init_mode (str, optional): The mode of FFN linear initialization. Defaults to "0". + norm_init_mode (str, optional): The mode of normalization initialization. Defaults to "0". + """ + + def __init__( + self, + dim: int, + num_heads: int, + n_temporal: int = 1, + cuboid_hw: Tuple[int, ...] = (7, 7), + shift_hw: Tuple[int, ...] = (0, 0), + strategy: Tuple[str, ...] = ("d", "l", "l"), + padding_type: str = "ignore", + cross_last_n_frames: int = None, + qkv_bias: bool = False, + qk_scale: float = None, + attn_drop: float = 0.0, + proj_drop: float = 0.0, + max_temporal_relative: int = 50, + norm_layer: str = "layer_norm", + use_global_vector: bool = True, + separate_global_qkv: bool = False, + global_dim_ratio: int = 1, + checkpoint_level: int = 1, + use_relative_pos: bool = True, + attn_linear_init_mode: str = "0", + ffn_linear_init_mode: str = "0", + norm_init_mode: str = "0", + moe_config: dict = None, + ): + super(CuboidCrossAttentionLayer, self).__init__() + self.attn_linear_init_mode = attn_linear_init_mode + self.ffn_linear_init_mode = ffn_linear_init_mode + self.norm_init_mode = norm_init_mode + self.dim = dim + self.num_heads = num_heads + self.n_temporal = n_temporal + assert n_temporal > 0 + head_dim = dim // num_heads + self.scale = qk_scale or head_dim**-0.5 + shift_hw = list(shift_hw) + if strategy[1] == "d": + shift_hw[0] = 0 + if strategy[2] == "d": + shift_hw[1] = 0 + self.cuboid_hw = cuboid_hw + self.shift_hw = tuple(shift_hw) + self.strategy = strategy + self.padding_type = padding_type + self.max_temporal_relative = max_temporal_relative + self.cross_last_n_frames = cross_last_n_frames + self.use_relative_pos = use_relative_pos + self.use_global_vector = use_global_vector + self.separate_global_qkv = separate_global_qkv + if global_dim_ratio != 1 and separate_global_qkv is False: + raise ValueError( + "Setting global_dim_ratio != 1 requires separate_global_qkv == True." + ) + self.global_dim_ratio = global_dim_ratio + if self.padding_type not in ["ignore", "zeros", "nearest"]: + raise ValueError('padding_type should be ["ignore", "zeros", "nearest"]') + if use_relative_pos: + init_data = paddle.zeros( + ( + (2 * max_temporal_relative - 1) + * (2 * cuboid_hw[0] - 1) + * (2 * cuboid_hw[1] - 1), + num_heads, + ) + ) + self.relative_position_bias_table = paddle.create_parameter( + shape=init_data.shape, + dtype=init_data.dtype, + default_initializer=nn.initializer.Constant(0.0), + ) + self.relative_position_bias_table.stop_gradient = not True + self.relative_position_bias_table = initializer.trunc_normal_( + self.relative_position_bias_table, std=0.02 + ) + + coords_t = paddle.arange(end=max_temporal_relative) + coords_h = paddle.arange(end=self.cuboid_hw[0]) + coords_w = paddle.arange(end=self.cuboid_hw[1]) + coords = paddle.stack(x=paddle.meshgrid(coords_t, coords_h, coords_w)) + coords_flatten = paddle.flatten(x=coords, start_axis=1) + relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] + relative_coords = relative_coords.transpose(perm=[1, 2, 0]) + relative_coords[:, :, 0] += max_temporal_relative - 1 + relative_coords[:, :, 1] += self.cuboid_hw[0] - 1 + relative_coords[:, :, 2] += self.cuboid_hw[1] - 1 + relative_position_index = ( + relative_coords[:, :, 0] + * (2 * self.cuboid_hw[0] - 1) + * (2 * self.cuboid_hw[1] - 1) + + relative_coords[:, :, 1] * (2 * self.cuboid_hw[1] - 1) + + relative_coords[:, :, 2] + ) + self.register_buffer( + name="relative_position_index", tensor=relative_position_index + ) + self.q_proj = nn.Linear(in_features=dim, out_features=dim, bias_attr=qkv_bias) + self.kv_proj = nn.Linear( + in_features=dim, out_features=dim * 2, bias_attr=qkv_bias + ) + self.attn_drop = nn.Dropout(p=attn_drop) + self.proj = nn.Linear(in_features=dim, out_features=dim) + self.proj_drop = nn.Dropout(p=proj_drop) + if self.use_global_vector: + if self.separate_global_qkv: + self.l2g_q_net = nn.Linear( + in_features=dim, out_features=dim, bias_attr=qkv_bias + ) + self.l2g_global_kv_net = nn.Linear( + in_features=global_dim_ratio * dim, + out_features=dim * 2, + bias_attr=qkv_bias, + ) + self.norm = cuboid_utils.get_norm_layer(norm_layer, in_channels=dim) + self._checkpoint_level = checkpoint_level + self.reset_parameters() + + def reset_parameters(self): + cuboid_utils.apply_initialization( + self.q_proj, linear_mode=self.attn_linear_init_mode + ) + cuboid_utils.apply_initialization( + self.kv_proj, linear_mode=self.attn_linear_init_mode + ) + cuboid_utils.apply_initialization( + self.proj, linear_mode=self.ffn_linear_init_mode + ) + cuboid_utils.apply_initialization(self.norm, norm_mode=self.norm_init_mode) + if self.use_global_vector: + if self.separate_global_qkv: + cuboid_utils.apply_initialization( + self.l2g_q_net, linear_mode=self.attn_linear_init_mode + ) + cuboid_utils.apply_initialization( + self.l2g_global_kv_net, linear_mode=self.attn_linear_init_mode + ) + + def forward(self, x, mem, mem_global_vectors=None): + """Calculate the forward + + Along the temporal axis, we pad the mem tensor from the left and the x tensor from the right so that the + relative position encoding can be calculated correctly. For example: + + mem: 0, 1, 2, 3, 4 + x: 0, 1, 2, 3, 4, 5 + + n_temporal = 1 + mem: 0, 1, 2, 3, 4 x: 0, 1, 2, 3, 4, 5 + + n_temporal = 2 + mem: pad, 1, 3 x: 0, 2, 4 + mem: 0, 2, 4 x: 1, 3, 5 + + n_temporal = 3 + mem: pad, 2 dec: 0, 3 + mem: 0, 3 dec: 1, 4 + mem: 1, 4 dec: 2, 5 + + Args: + x (paddle.Tensor): The input of the layer. It will have shape (B, T, H, W, C) + mem (paddle.Tensor): The memory. It will have shape (B, T_mem, H, W, C) + mem_global_vectors (paddle.Tensor): The global vectors from the memory. It will have shape (B, N, C) + + Returns: + out (paddle.Tensor): Output tensor should have shape (B, T, H, W, C_out) + """ + + if self.cross_last_n_frames is not None: + cross_last_n_frames = int(min(self.cross_last_n_frames, mem.shape[1])) + mem = mem[:, -cross_last_n_frames:, ...] + if self.use_global_vector: + _, num_global, _ = mem_global_vectors.shape + x = self.norm(x) + B, T_x, H, W, C_in = x.shape + B_mem, T_mem, H_mem, W_mem, C_mem = mem.shape + assert T_x < self.max_temporal_relative and T_mem < self.max_temporal_relative + cuboid_hw = self.cuboid_hw + n_temporal = self.n_temporal + shift_hw = self.shift_hw + assert ( + B_mem == B and H == H_mem and W == W_mem and C_in == C_mem + ), f"Shape of memory and the input tensor does not match. x.shape={x.shape}, mem.shape={mem.shape}" + pad_t_mem = (n_temporal - T_mem % n_temporal) % n_temporal + pad_t_x = (n_temporal - T_x % n_temporal) % n_temporal + pad_h = (cuboid_hw[0] - H % cuboid_hw[0]) % cuboid_hw[0] + pad_w = (cuboid_hw[1] - W % cuboid_hw[1]) % cuboid_hw[1] + mem = cuboid_utils.generalize_padding( + mem, pad_t_mem, pad_h, pad_w, self.padding_type, t_pad_left=True + ) + + x = cuboid_utils.generalize_padding( + x, pad_t_x, pad_h, pad_w, self.padding_type, t_pad_left=False + ) + + if any(i > 0 for i in shift_hw): + shifted_x = paddle.roll( + x=x, shifts=(-shift_hw[0], -shift_hw[1]), axis=(2, 3) + ) + shifted_mem = paddle.roll( + x=mem, shifts=(-shift_hw[0], -shift_hw[1]), axis=(2, 3) + ) + else: + shifted_x = x + shifted_mem = mem + mem_cuboid_size = (mem.shape[1] // n_temporal,) + cuboid_hw + x_cuboid_size = (x.shape[1] // n_temporal,) + cuboid_hw + reordered_mem = cuboid_encoder.cuboid_reorder( + shifted_mem, cuboid_size=mem_cuboid_size, strategy=self.strategy + ) + reordered_x = cuboid_encoder.cuboid_reorder( + shifted_x, cuboid_size=x_cuboid_size, strategy=self.strategy + ) + _, num_cuboids_mem, mem_cuboid_volume, _ = reordered_mem.shape + _, num_cuboids, x_cuboid_volume, _ = reordered_x.shape + assert ( + num_cuboids_mem == num_cuboids + ), f"Number of cuboids do not match. num_cuboids={num_cuboids}, num_cuboids_mem={num_cuboids_mem}" + attn_mask = compute_cuboid_cross_attention_mask( + T_x, + T_mem, + H, + W, + n_temporal, + cuboid_hw, + shift_hw, + strategy=self.strategy, + padding_type=self.padding_type, + device=x.place, + ) + head_C = C_in // self.num_heads + kv = ( + self.kv_proj(reordered_mem) + .reshape([B, num_cuboids, mem_cuboid_volume, 2, self.num_heads, head_C]) + .transpose(perm=[3, 0, 4, 1, 2, 5]) + ) + k, v = kv[0], kv[1] + q = ( + self.q_proj(reordered_x) + .reshape([B, num_cuboids, x_cuboid_volume, self.num_heads, head_C]) + .transpose(perm=[0, 3, 1, 2, 4]) + ) + q = q * self.scale + perm_4 = list(range(k.ndim)) + perm_4[-2] = -1 + perm_4[-1] = -2 + attn_score = q @ k.transpose(perm=perm_4) + if self.use_relative_pos: + relative_position_bias = self.relative_position_bias_table[ + self.relative_position_index[ + :x_cuboid_volume, :mem_cuboid_volume + ].reshape([-1]) + ].reshape([x_cuboid_volume, mem_cuboid_volume, -1]) + relative_position_bias = relative_position_bias.transpose( + perm=[2, 0, 1] + ).unsqueeze(axis=1) + attn_score = attn_score + relative_position_bias + if self.use_global_vector: + if self.separate_global_qkv: + l2g_q = ( + self.l2g_q_net(reordered_x) + .reshape([B, num_cuboids, x_cuboid_volume, self.num_heads, head_C]) + .transpose(perm=[0, 3, 1, 2, 4]) + ) + l2g_q = l2g_q * self.scale + l2g_global_kv = ( + self.l2g_global_kv_net(mem_global_vectors) + .reshape([B, 1, num_global, 2, self.num_heads, head_C]) + .transpose(perm=[3, 0, 4, 1, 2, 5]) + ) + l2g_global_k, l2g_global_v = l2g_global_kv[0], l2g_global_kv[1] + else: + kv_global = ( + self.kv_proj(mem_global_vectors) + .reshape([B, 1, num_global, 2, self.num_heads, head_C]) + .transpose(perm=[3, 0, 4, 1, 2, 5]) + ) + l2g_global_k, l2g_global_v = kv_global[0], kv_global[1] + l2g_q = q + perm_5 = list(range(l2g_global_k.ndim)) + perm_5[-2] = -1 + perm_5[-1] = -2 + l2g_attn_score = l2g_q @ l2g_global_k.transpose(perm=perm_5) + attn_score_l2l_l2g = paddle.concat(x=(attn_score, l2g_attn_score), axis=-1) + if attn_mask.ndim == 5: + attn_mask_l2l_l2g = F.pad( + attn_mask, [0, num_global], "constant", 1, data_format="NDHWC" + ) + else: + attn_mask_l2l_l2g = F.pad(attn_mask, [0, num_global], "constant", 1) + v_l_g = paddle.concat( + x=( + v, + l2g_global_v.expand( + shape=[B, self.num_heads, num_cuboids, num_global, head_C] + ), + ), + axis=3, + ) + attn_score_l2l_l2g = cuboid_encoder.masked_softmax( + attn_score_l2l_l2g, mask=attn_mask_l2l_l2g + ) + attn_score_l2l_l2g = self.attn_drop(attn_score_l2l_l2g) + reordered_x = ( + (attn_score_l2l_l2g @ v_l_g) + .transpose(perm=[0, 2, 3, 1, 4]) + .reshape(B, num_cuboids, x_cuboid_volume, self.dim) + ) + else: + attn_score = cuboid_encoder.masked_softmax(attn_score, mask=attn_mask) + attn_score = self.attn_drop(attn_score) + reordered_x = ( + (attn_score @ v) + .transpose(perm=[0, 2, 3, 1, 4]) + .reshape([B, num_cuboids, x_cuboid_volume, self.dim]) + ) + reordered_x = paddle.cast(reordered_x, dtype="float32") + reordered_x = self.proj_drop(self.proj(reordered_x)) + shifted_x = cuboid_encoder.cuboid_reorder_reverse( + reordered_x, + cuboid_size=x_cuboid_size, + strategy=self.strategy, + orig_data_shape=(x.shape[1], x.shape[2], x.shape[3]), + ) + if any(i > 0 for i in shift_hw): + x = paddle.roll(x=shifted_x, shifts=(shift_hw[0], shift_hw[1]), axis=(2, 3)) + else: + x = shifted_x + x = cuboid_utils.generalize_unpadding( + x, pad_t=pad_t_x, pad_h=pad_h, pad_w=pad_w, padding_type=self.padding_type + ) + return x + + +class StackCuboidCrossAttentionBlock(nn.Layer): + """A stack of cuboid cross attention layers. + + The advantage of cuboid attention is that we can combine cuboid attention building blocks with different + hyper-parameters to mimic a broad range of space-time correlation patterns. + + - "use_inter_ffn" is True + x, mem --> attn1 -----+-------> ffn1 ---+---> attn2 --> ... --> ffn_k --> out + | ^ | ^ + | | | | + |-------------|----|-------------| + - "use_inter_ffn" is False + x, mem --> attn1 -----+------> attn2 --> ... attnk --+----> ffnk ---+---> out, mem + | ^ | ^ ^ | ^ + | | | | | | | + |-------------|----|------------|-- ----------|--|-----------| + + Args: + dim (int): The dimension of the input. + num_heads (int): The number of head. + block_cuboid_hw (list, optional): The height and width of block cuboid.Defaults to [(4, 4), (4, 4)]. + block_shift_hw (list, optional): The height and width of shift cuboid . Defaults to [(0, 0), (2, 2)]. + block_n_temporal (list, optional): The length of block temporal. Defaults to [1, 2]. + block_strategy (list, optional): The strategy of block. Defaults to [("d", "d", "d"), ("l", "l", "l")]. + padding_type (str, optional): The type of paddling. Defaults to "ignore". + cross_last_n_frames (int, optional): The num of cross_last_n_frames. Defaults to None. + qkv_bias (bool, optional): Whether to enable bias in calculating qkv attention. Defaults to False. + qk_scale (float, optional): Whether to enable scale factor when calculating the attention. Defaults to None. + attn_drop (float, optional): The attention dropout. Defaults to 0.0. + proj_drop (float, optional): The projection dropout. Defaults to 0.0. + ffn_drop (float, optional): The ratio of FFN dropout. Defaults to 0.0. + activation (str, optional): The activation. Defaults to "leaky". + gated_ffn (bool, optional): Whether to use gate FFN. Defaults to False. + norm_layer (str, optional): The normalization layer. Defaults to "layer_norm". + use_inter_ffn (bool, optional): Whether to use inter FFN. Defaults to True. + max_temporal_relative (int, optional): The max temporal. Defaults to 50. + checkpoint_level (int, optional): Whether to enable gradient checkpointing. Defaults to 1. + use_relative_pos (bool, optional): Whether to use relative pos. Defaults to True. + use_global_vector (bool, optional): Whether to use the global vector or not. Defaults to False. + separate_global_qkv (bool, optional): Whether to use different network to calc q_global, k_global, v_global. Defaults to False. + global_dim_ratio (int, optional): The dim (channels) of global vectors is `global_dim_ratio*dim`. Defaults to 1. + attn_linear_init_mode (str, optional): The mode of attention linear initialization. Defaults to "0". + ffn_linear_init_mode (str, optional): The mode of FFN linear initialization. Defaults to "0". + norm_init_mode (str, optional): The mode of normalization. Defaults to "0". + """ + + def __init__( + self, + dim: int, + num_heads: int, + block_cuboid_hw: Tuple[Tuple[int, ...], ...] = [(4, 4), (4, 4)], + block_shift_hw: Tuple[Tuple[int, ...], ...] = [(0, 0), (2, 2)], + block_n_temporal: Tuple[int, ...] = [1, 2], + block_strategy: Tuple[Tuple[str, ...], ...] = [ + ("d", "d", "d"), + ("l", "l", "l"), + ], + padding_type: str = "ignore", + cross_last_n_frames: int = None, + qkv_bias: bool = False, + qk_scale: float = None, + attn_drop: float = 0.0, + proj_drop: float = 0.0, + ffn_drop: float = 0.0, + activation: str = "leaky", + gated_ffn: bool = False, + norm_layer: str = "layer_norm", + use_inter_ffn: bool = True, + max_temporal_relative: int = 50, + checkpoint_level: int = 1, + use_relative_pos: bool = True, + use_global_vector: bool = False, + separate_global_qkv: bool = False, + global_dim_ratio: int = 1, + attn_linear_init_mode: str = "0", + ffn_linear_init_mode: str = "0", + norm_init_mode: str = "0", + moe_config: dict = None, + expert_shape: tuple = None, + ): + super(StackCuboidCrossAttentionBlock, self).__init__() + self.attn_linear_init_mode = attn_linear_init_mode + self.ffn_linear_init_mode = ffn_linear_init_mode + self.norm_init_mode = norm_init_mode + if ( + len(block_cuboid_hw[0]) <= 0 + or len(block_shift_hw) <= 0 + or len(block_strategy) <= 0 + ): + raise ValueError( + "Incorrect format.The lengths of block_cuboid_hw[0], block_shift_hw, and block_strategy must be greater than zero." + ) + if len(block_cuboid_hw) != len(block_shift_hw) and len(block_shift_hw) == len( + block_strategy + ): + raise ValueError( + "The lengths of block_cuboid_size, block_shift_size, and block_strategy must be equal." + ) + + self.num_attn = len(block_cuboid_hw) + self.checkpoint_level = checkpoint_level + self.use_inter_ffn = use_inter_ffn + self.use_global_vector = use_global_vector + if self.use_inter_ffn: + if moe_config["use_ffn_moe"]: + self.ffn_l = nn.LayerList( + sublayers=[ + cuboid_encoder.MixtureFFN( + units=dim, + hidden_size=4 * dim, + activation_dropout=ffn_drop, + dropout=ffn_drop, + gated_proj=gated_ffn, + activation=activation, + normalization=norm_layer, + pre_norm=True, + linear_init_mode=ffn_linear_init_mode, + norm_init_mode=norm_init_mode, + expert_shape=expert_shape, + moe_config=moe_config, + ) + for _ in range(self.num_attn) + ] + ) + else: + self.ffn_l = nn.LayerList( + sublayers=[ + cuboid_encoder.PositionwiseFFN( + units=dim, + hidden_size=4 * dim, + activation_dropout=ffn_drop, + dropout=ffn_drop, + gated_proj=gated_ffn, + activation=activation, + normalization=norm_layer, + pre_norm=True, + linear_init_mode=ffn_linear_init_mode, + norm_init_mode=norm_init_mode, + expert_shape=expert_shape, + moe_config=moe_config, + ) + for _ in range(self.num_attn) + ] + ) + else: + if moe_config["use_ffn_moe"]: + self.ffn_l = nn.LayerList( + sublayers=[ + cuboid_encoder.MixtureFFN( + units=dim, + hidden_size=4 * dim, + activation_dropout=ffn_drop, + dropout=ffn_drop, + gated_proj=gated_ffn, + activation=activation, + normalization=norm_layer, + pre_norm=True, + linear_init_mode=ffn_linear_init_mode, + norm_init_mode=norm_init_mode, + expert_shape=expert_shape, + moe_config=moe_config, + ) + ] + ) + else: + self.ffn_l = nn.LayerList( + sublayers=[ + cuboid_encoder.PositionwiseFFN( + units=dim, + hidden_size=4 * dim, + activation_dropout=ffn_drop, + dropout=ffn_drop, + gated_proj=gated_ffn, + activation=activation, + normalization=norm_layer, + pre_norm=True, + linear_init_mode=ffn_linear_init_mode, + norm_init_mode=norm_init_mode, + expert_shape=expert_shape, + moe_config=moe_config, + ) + ] + ) + + if moe_config["use_attn_moe"]: + self.attn_l = nn.LayerList( + sublayers=[ + MixtureCrossAttention( + dim=dim, + num_heads=num_heads, + cuboid_hw=ele_cuboid_hw, + shift_hw=ele_shift_hw, + strategy=ele_strategy, + n_temporal=ele_n_temporal, + cross_last_n_frames=cross_last_n_frames, + padding_type=padding_type, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + attn_drop=attn_drop, + proj_drop=proj_drop, + norm_layer=norm_layer, + max_temporal_relative=max_temporal_relative, + use_global_vector=use_global_vector, + separate_global_qkv=separate_global_qkv, + global_dim_ratio=global_dim_ratio, + checkpoint_level=checkpoint_level, + use_relative_pos=use_relative_pos, + attn_linear_init_mode=attn_linear_init_mode, + ffn_linear_init_mode=ffn_linear_init_mode, + norm_init_mode=norm_init_mode, + expert_shape=expert_shape, + moe_config=moe_config, + ) + for ele_cuboid_hw, ele_shift_hw, ele_strategy, ele_n_temporal in zip( + block_cuboid_hw, + block_shift_hw, + block_strategy, + block_n_temporal, + ) + ] + ) + else: + self.attn_l = nn.LayerList( + sublayers=[ + CuboidCrossAttentionLayer( + dim=dim, + num_heads=num_heads, + cuboid_hw=ele_cuboid_hw, + shift_hw=ele_shift_hw, + strategy=ele_strategy, + n_temporal=ele_n_temporal, + cross_last_n_frames=cross_last_n_frames, + padding_type=padding_type, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + attn_drop=attn_drop, + proj_drop=proj_drop, + norm_layer=norm_layer, + max_temporal_relative=max_temporal_relative, + use_global_vector=use_global_vector, + separate_global_qkv=separate_global_qkv, + global_dim_ratio=global_dim_ratio, + checkpoint_level=checkpoint_level, + use_relative_pos=use_relative_pos, + attn_linear_init_mode=attn_linear_init_mode, + ffn_linear_init_mode=ffn_linear_init_mode, + norm_init_mode=norm_init_mode, + ) + for ele_cuboid_hw, ele_shift_hw, ele_strategy, ele_n_temporal in zip( + block_cuboid_hw, + block_shift_hw, + block_strategy, + block_n_temporal, + ) + ] + ) + + def reset_parameters(self): + for m in self.ffn_l: + m.reset_parameters() + for m in self.attn_l: + m.reset_parameters() + + def forward(self, x, mem, mem_global_vector=None): + """ + Args: + x (paddle.Tensor): Shape (B, T_x, H, W, C) + mem (paddle.Tensor): Shape (B, T_mem, H, W, C) + mem_global_vector (paddle.Tensor): Shape (B, N_global, C) + + Returns: + out (paddle.Tensor): (B, T_x, H, W, C_out) + """ + + if self.use_inter_ffn: + for attn, ffn in zip(self.attn_l, self.ffn_l): + if self.checkpoint_level >= 2 and self.training: + x = x + fleet.utils.recompute(attn, x, mem, mem_global_vector) + else: + x = x + attn(x, mem, mem_global_vector) + if self.checkpoint_level >= 1 and self.training: + x = fleet.utils.recompute(ffn, x) + else: + x = ffn(x) + return x + else: + for attn in self.attn_l: + if self.checkpoint_level >= 2 and self.training: + x = x + fleet.utils.recompute(attn, x, mem, mem_global_vector) + else: + x = x + attn(x, mem, mem_global_vector) + if self.checkpoint_level >= 1 and self.training: + x = fleet.utils.recompute(self.ffn_l[0], x) + else: + x = self.ffn_l[0](x) + return x + + +class Upsample3DLayer(nn.Layer): + """Upsampling based on nn.UpSampling and Conv3x3. + + If the temporal dimension remains the same: + x --> interpolation-2d (nearest) --> conv3x3(dim, out_dim) + Else: + x --> interpolation-3d (nearest) --> conv3x3x3(dim, out_dim) + + Args: + dim (int): The dimension of the input tensor. + out_dim (int): The dimension of the output tensor. + target_size (Tuple[int,...]): The size of output tensor. + temporal_upsample (bool, optional): Whether the temporal axis will go through upsampling. Defaults to False. + kernel_size (int, optional): The kernel size of the Conv2D layer. Defaults to 3. + layout (str, optional): The layout of the inputs. Defaults to "THWC". + conv_init_mode (str, optional): The mode of conv initialization. Defaults to "0". + """ + + def __init__( + self, + dim: int, + out_dim: int, + target_size: Tuple[int, ...], + temporal_upsample: bool = False, + kernel_size: int = 3, + layout: str = "THWC", + conv_init_mode: str = "0", + moe_config: dict = None, + ): + super(Upsample3DLayer, self).__init__() + self.conv_init_mode = conv_init_mode + self.target_size = target_size + self.out_dim = out_dim + self.temporal_upsample = temporal_upsample + if temporal_upsample: + self.up = nn.Upsample(size=target_size, mode="nearest") + else: + self.up = nn.Upsample(size=(target_size[1], target_size[2]), mode="nearest") + self.conv = nn.Conv2D( + in_channels=dim, + out_channels=out_dim, + kernel_size=(kernel_size, kernel_size), + padding=(kernel_size // 2, kernel_size // 2), + ) + assert layout in ["THWC", "CTHW"] + self.layout = layout + self.reset_parameters() + + def reset_parameters(self): + for m in self.children(): + cuboid_utils.apply_initialization(m, conv_mode=self.conv_init_mode) + + def forward(self, x): + """ + + Args: + x : (B, T, H, W, C) or (B, C, T, H, W) + + Returns: + out : (B, T, H_new, W_out, C_out) or (B, C, T, H_out, W_out) + """ + + if self.layout == "THWC": + B, T, H, W, C = x.shape + if self.temporal_upsample: + x = x.transpose(perm=[0, 4, 1, 2, 3]) + return self.conv(self.up(x)).transpose(perm=[0, 2, 3, 4, 1]) + else: + assert self.target_size[0] == T + x = x.reshape([B * T, H, W, C]).transpose(perm=[0, 3, 1, 2]) + x = self.up(x) + return ( + self.conv(x) + .transpose(perm=[0, 2, 3, 1]) + .reshape(list((B,) + self.target_size + (self.out_dim,))) + ) + elif self.layout == "CTHW": + B, C, T, H, W = x.shape + if self.temporal_upsample: + return self.conv(self.up(x)) + else: + assert self.output_size[0] == T + x = x.transpose(perm=[0, 2, 1, 3, 4]) + x = x.reshape([B * T, C, H, W]) + return ( + self.conv(self.up(x)) + .reshape( + [ + B, + self.target_size[0], + self.out_dim, + self.target_size[1], + self.target_size[2], + ] + ) + .transpose(perm=[0, 2, 1, 3, 4]) + ) + + +class CuboidTransformerDecoder(nn.Layer): + """Decoder of the CuboidTransformer. + + For each block, we first apply the StackCuboidSelfAttention and then apply the StackCuboidCrossAttention + + Repeat the following structure K times + + x --> StackCuboidSelfAttention --> | + |----> StackCuboidCrossAttention (If used) --> out + mem --> | + + Args: + target_temporal_length (int): The temporal length of the target. + mem_shapes (Tuple[int,...]): The mem shapes of the decoder. + cross_start (int, optional): The block to start cross attention. Defaults to 0. + depth (list, optional): The number of layers for each block. Defaults to [2, 2]. + upsample_type (str, optional): The type of upsample. Defaults to "upsample". + upsample_kernel_size (int, optional): The kernel size of upsample. Defaults to 3. + block_self_attn_patterns (str, optional): The patterns of block attention. Defaults to None. + block_self_cuboid_size (list, optional): The size of cuboid block. Defaults to [(4, 4, 4), (4, 4, 4)]. + block_self_cuboid_strategy (list, optional): The strategy of cuboid. Defaults to [("l", "l", "l"), ("d", "d", "d")]. + block_self_shift_size (list, optional): The size of shift. Defaults to [(1, 1, 1), (0, 0, 0)]. + block_cross_attn_patterns (str, optional): The patterns of cross attentions. Defaults to None. + block_cross_cuboid_hw (list, optional): The height and width of cross cuboid. Defaults to [(4, 4), (4, 4)]. + block_cross_cuboid_strategy (list, optional): The strategy of cross cuboid. Defaults to [("l", "l", "l"), ("d", "l", "l")]. + block_cross_shift_hw (list, optional): The height and width of cross shift. Defaults to [(0, 0), (0, 0)]. + block_cross_n_temporal (list, optional): The cross temporal of block. Defaults to [1, 2]. + cross_last_n_frames (int, optional): The num of cross last frames. Defaults to None. + num_heads (int, optional): The num of head. Defaults to 4. + attn_drop (float, optional): The ratio of attention dropout. Defaults to 0.0. + proj_drop (float, optional): The ratio of projection dropout. Defaults to 0.0. + ffn_drop (float, optional): The ratio of FFN dropout. Defaults to 0.0. + ffn_activation (str, optional): The activation layer of FFN. Defaults to "leaky". + gated_ffn (bool, optional): Whether to use gate FFN. Defaults to False. + norm_layer (str, optional): The normalization layer. Defaults to "layer_norm". + use_inter_ffn (bool, optional): Whether to use inter FFN. Defaults to False. + hierarchical_pos_embed (bool, optional): Whether to use hierarchical pos_embed. Defaults to False. + pos_embed_type (str, optional): The type of pos embedding. Defaults to "t+hw". + max_temporal_relative (int, optional): The max number of teemporal relative. Defaults to 50. + padding_type (str, optional): The type of padding. Defaults to "ignore". + checkpoint_level (bool, optional): Whether to enable gradient checkpointing. Defaults to True. + use_relative_pos (bool, optional): Whether to use relative pos. Defaults to True. + self_attn_use_final_proj (bool, optional): Whether to use self attention for final projection. Defaults to True. + use_first_self_attn (bool, optional): Whether to use first self attention. Defaults to False. + use_self_global (bool, optional): Whether to use self global vector. Defaults to False. + self_update_global (bool, optional): Whether to update global vector. Defaults to True. + use_cross_global (bool, optional): Whether to use cross global vector. Defaults to False. + use_global_vector_ffn (bool, optional): Whether to use FFN global vectors. Defaults to True. + use_global_self_attn (bool, optional): Whether to use global self attention. Defaults to False. + separate_global_qkv (bool, optional): Whether to use different network to calc q_global, k_global, v_global. Defaults to False. + global_dim_ratio (int, optional): The dim (channels) of global vectors is `global_dim_ratio*dim`. Defaults to 1. + attn_linear_init_mode (str, optional): The mode of attention linear initialization. Defaults to "0". + ffn_linear_init_mode (str, optional): The mode of FFN linear initialization. Defaults to "0". + conv_init_mode (str, optional): The mode of conv initialization. Defaults to "0". + up_linear_init_mode (str, optional): The mode of up linear initialization. Defaults to "0". + norm_init_mode (str, optional): The mode of normalization initialization. Defaults to "0". + """ + + def __init__( + self, + target_temporal_length: int, + mem_shapes: Tuple[int, ...], + cross_start: int = 0, + depth: Tuple[int, ...] = [2, 2], + upsample_type: str = "upsample", + upsample_kernel_size: int = 3, + block_self_attn_patterns: str = None, + block_self_cuboid_size: Tuple[Tuple[int, ...], ...] = [(4, 4, 4), (4, 4, 4)], + block_self_cuboid_strategy: Tuple[Tuple[str, ...], ...] = [ + ("l", "l", "l"), + ("d", "d", "d"), + ], + block_self_shift_size: Tuple[Tuple[int, ...], ...] = [(1, 1, 1), (0, 0, 0)], + block_cross_attn_patterns: str = None, + block_cross_cuboid_hw: Tuple[Tuple[int, ...], ...] = [(4, 4), (4, 4)], + block_cross_cuboid_strategy: Tuple[Tuple[str, ...], ...] = [ + ("l", "l", "l"), + ("d", "l", "l"), + ], + block_cross_shift_hw: Tuple[Tuple[int, ...], ...] = [(0, 0), (0, 0)], + block_cross_n_temporal: Tuple[int, ...] = [1, 2], + cross_last_n_frames: int = None, + num_heads: int = 4, + attn_drop: float = 0.0, + proj_drop: float = 0.0, + ffn_drop: float = 0.0, + ffn_activation: str = "leaky", + gated_ffn: bool = False, + norm_layer: str = "layer_norm", + use_inter_ffn: bool = False, + hierarchical_pos_embed: bool = False, + pos_embed_type: str = "t+hw", + max_temporal_relative: int = 50, + padding_type: str = "ignore", + checkpoint_level: bool = True, + use_relative_pos: bool = True, + self_attn_use_final_proj: bool = True, + use_first_self_attn: bool = False, + use_self_global: bool = False, + self_update_global: bool = True, + use_cross_global: bool = False, + use_global_vector_ffn: bool = True, + use_global_self_attn: bool = False, + separate_global_qkv: bool = False, + global_dim_ratio: int = 1, + attn_linear_init_mode: str = "0", + ffn_linear_init_mode: str = "0", + conv_init_mode: str = "0", + up_linear_init_mode: str = "0", + norm_init_mode: str = "0", + moe_config: dict = None, + ): + super(CuboidTransformerDecoder, self).__init__() + self.attn_linear_init_mode = attn_linear_init_mode + self.ffn_linear_init_mode = ffn_linear_init_mode + self.conv_init_mode = conv_init_mode + self.up_linear_init_mode = up_linear_init_mode + self.norm_init_mode = norm_init_mode + assert len(depth) == len(mem_shapes) + self.target_temporal_length = target_temporal_length + self.num_blocks = len(mem_shapes) + self.cross_start = cross_start + self.mem_shapes = mem_shapes + self.depth = depth + self.upsample_type = upsample_type + self.hierarchical_pos_embed = hierarchical_pos_embed + self.checkpoint_level = checkpoint_level + self.use_self_global = use_self_global + self.self_update_global = self_update_global + self.use_cross_global = use_cross_global + self.use_global_vector_ffn = use_global_vector_ffn + self.use_first_self_attn = use_first_self_attn + if block_self_attn_patterns is not None: + if isinstance(block_self_attn_patterns, (tuple, list)): + assert len(block_self_attn_patterns) == self.num_blocks + else: + block_self_attn_patterns = [ + block_self_attn_patterns for _ in range(self.num_blocks) + ] + block_self_cuboid_size = [] + block_self_cuboid_strategy = [] + block_self_shift_size = [] + for idx, key in enumerate(block_self_attn_patterns): + func = cuboid_utils.CuboidSelfAttentionPatterns.get(key) + cuboid_size, strategy, shift_size = func(mem_shapes[idx]) + block_self_cuboid_size.append(cuboid_size) + block_self_cuboid_strategy.append(strategy) + block_self_shift_size.append(shift_size) + else: + if not isinstance(block_self_cuboid_size[0][0], (list, tuple)): + block_self_cuboid_size = [ + block_self_cuboid_size for _ in range(self.num_blocks) + ] + else: + assert ( + len(block_self_cuboid_size) == self.num_blocks + ), f"Incorrect input format! Received block_self_cuboid_size={block_self_cuboid_size}" + if not isinstance(block_self_cuboid_strategy[0][0], (list, tuple)): + block_self_cuboid_strategy = [ + block_self_cuboid_strategy for _ in range(self.num_blocks) + ] + else: + assert ( + len(block_self_cuboid_strategy) == self.num_blocks + ), f"Incorrect input format! Received block_self_cuboid_strategy={block_self_cuboid_strategy}" + if not isinstance(block_self_shift_size[0][0], (list, tuple)): + block_self_shift_size = [ + block_self_shift_size for _ in range(self.num_blocks) + ] + else: + assert ( + len(block_self_shift_size) == self.num_blocks + ), f"Incorrect input format! Received block_self_shift_size={block_self_shift_size}" + + expert_shape_list = [ + (target_temporal_length,) + mem_shape[1:] for mem_shape in mem_shapes + ] + self_blocks = [] + for i in range(self.num_blocks): + if not self.use_first_self_attn and i == self.num_blocks - 1: + ele_depth = depth[i] - 1 + else: + ele_depth = depth[i] + stack_cuboid_blocks = [ + cuboid_encoder.StackCuboidSelfAttentionBlock( + dim=self.mem_shapes[i][-1], + num_heads=num_heads, + block_cuboid_size=block_self_cuboid_size[i], + block_strategy=block_self_cuboid_strategy[i], + block_shift_size=block_self_shift_size[i], + attn_drop=attn_drop, + proj_drop=proj_drop, + ffn_drop=ffn_drop, + activation=ffn_activation, + gated_ffn=gated_ffn, + norm_layer=norm_layer, + use_inter_ffn=use_inter_ffn, + padding_type=padding_type, + use_global_vector=use_self_global, + use_global_vector_ffn=use_global_vector_ffn, + use_global_self_attn=use_global_self_attn, + separate_global_qkv=separate_global_qkv, + global_dim_ratio=global_dim_ratio, + checkpoint_level=checkpoint_level, + use_relative_pos=use_relative_pos, + use_final_proj=self_attn_use_final_proj, + attn_linear_init_mode=attn_linear_init_mode, + ffn_linear_init_mode=ffn_linear_init_mode, + norm_init_mode=norm_init_mode, + expert_shape=expert_shape_list[i], + moe_config=moe_config, + ) + for _ in range(ele_depth) + ] + self_blocks.append(nn.LayerList(sublayers=stack_cuboid_blocks)) + self.self_blocks = nn.LayerList(sublayers=self_blocks) + + if block_cross_attn_patterns is not None: + if isinstance(block_cross_attn_patterns, (tuple, list)): + assert len(block_cross_attn_patterns) == self.num_blocks + else: + block_cross_attn_patterns = [ + block_cross_attn_patterns for _ in range(self.num_blocks) + ] + block_cross_cuboid_hw = [] + block_cross_cuboid_strategy = [] + block_cross_shift_hw = [] + block_cross_n_temporal = [] + for idx, key in enumerate(block_cross_attn_patterns): + if key == "last_frame_dst": + cuboid_hw = None + shift_hw = None + strategy = None + n_temporal = None + else: + func = cuboid_utils.CuboidCrossAttentionPatterns.get(key) + cuboid_hw, shift_hw, strategy, n_temporal = func(mem_shapes[idx]) + block_cross_cuboid_hw.append(cuboid_hw) + block_cross_cuboid_strategy.append(strategy) + block_cross_shift_hw.append(shift_hw) + block_cross_n_temporal.append(n_temporal) + else: + if not isinstance(block_cross_cuboid_hw[0][0], (list, tuple)): + block_cross_cuboid_hw = [ + block_cross_cuboid_hw for _ in range(self.num_blocks) + ] + else: + assert ( + len(block_cross_cuboid_hw) == self.num_blocks + ), f"Incorrect input format! Received block_cross_cuboid_hw={block_cross_cuboid_hw}" + if not isinstance(block_cross_cuboid_strategy[0][0], (list, tuple)): + block_cross_cuboid_strategy = [ + block_cross_cuboid_strategy for _ in range(self.num_blocks) + ] + else: + assert ( + len(block_cross_cuboid_strategy) == self.num_blocks + ), f"Incorrect input format! Received block_cross_cuboid_strategy={block_cross_cuboid_strategy}" + if not isinstance(block_cross_shift_hw[0][0], (list, tuple)): + block_cross_shift_hw = [ + block_cross_shift_hw for _ in range(self.num_blocks) + ] + else: + assert ( + len(block_cross_shift_hw) == self.num_blocks + ), f"Incorrect input format! Received block_cross_shift_hw={block_cross_shift_hw}" + if not isinstance(block_cross_n_temporal[0], (list, tuple)): + block_cross_n_temporal = [ + block_cross_n_temporal for _ in range(self.num_blocks) + ] + else: + assert ( + len(block_cross_n_temporal) == self.num_blocks + ), f"Incorrect input format! Received block_cross_n_temporal={block_cross_n_temporal}" + self.cross_blocks = nn.LayerList() + assert self.cross_start == 0 + for i in range(self.cross_start, self.num_blocks): + cross_block = nn.LayerList( + sublayers=[ + StackCuboidCrossAttentionBlock( + dim=self.mem_shapes[i][-1], + num_heads=num_heads, + block_cuboid_hw=block_cross_cuboid_hw[i], + block_strategy=block_cross_cuboid_strategy[i], + block_shift_hw=block_cross_shift_hw[i], + block_n_temporal=block_cross_n_temporal[i], + cross_last_n_frames=cross_last_n_frames, + attn_drop=attn_drop, + proj_drop=proj_drop, + ffn_drop=ffn_drop, + gated_ffn=gated_ffn, + norm_layer=norm_layer, + use_inter_ffn=use_inter_ffn, + activation=ffn_activation, + max_temporal_relative=max_temporal_relative, + padding_type=padding_type, + use_global_vector=use_cross_global, + separate_global_qkv=separate_global_qkv, + global_dim_ratio=global_dim_ratio, + checkpoint_level=checkpoint_level, + use_relative_pos=use_relative_pos, + attn_linear_init_mode=attn_linear_init_mode, + ffn_linear_init_mode=ffn_linear_init_mode, + norm_init_mode=norm_init_mode, + expert_shape=expert_shape_list[i], + moe_config=moe_config, + ) + for _ in range(depth[i]) + ] + ) + self.cross_blocks.append(cross_block) + if self.num_blocks > 1: + if self.upsample_type == "upsample": + self.upsample_layers = nn.LayerList( + sublayers=[ + Upsample3DLayer( + dim=self.mem_shapes[i + 1][-1], + out_dim=self.mem_shapes[i][-1], + target_size=(target_temporal_length,) + + self.mem_shapes[i][1:3], + kernel_size=upsample_kernel_size, + temporal_upsample=False, + conv_init_mode=conv_init_mode, + ) + for i in range(self.num_blocks - 1) + ] + ) + else: + raise NotImplementedError(f"{self.upsample_type} is invalid.") + if self.hierarchical_pos_embed: + self.hierarchical_pos_embed_l = nn.LayerList( + sublayers=[ + PosEmbed( + embed_dim=self.mem_shapes[i][-1], + typ=pos_embed_type, + maxT=target_temporal_length, + maxH=self.mem_shapes[i][1], + maxW=self.mem_shapes[i][2], + ) + for i in range(self.num_blocks - 1) + ] + ) + self.reset_parameters() + + def reset_parameters(self): + for ms in self.self_blocks: + for m in ms: + m.reset_parameters() + for ms in self.cross_blocks: + for m in ms: + m.reset_parameters() + if self.num_blocks > 1: + for m in self.upsample_layers: + m.reset_parameters() + if self.hierarchical_pos_embed: + for m in self.hierarchical_pos_embed_l: + m.reset_parameters() + + def forward(self, x, mem_l, mem_global_vector_l=None): + """ + Args: + x : Shape (B, T_top, H_top, W_top, C). + mem_l : A list of memory tensors. + """ + + B, T_top, H_top, W_top, C = x.shape + assert T_top == self.target_temporal_length + assert (H_top, W_top) == (self.mem_shapes[-1][1], self.mem_shapes[-1][2]) + for i in range(self.num_blocks - 1, -1, -1): + mem_global_vector = ( + None if mem_global_vector_l is None else mem_global_vector_l[i] + ) + if not self.use_first_self_attn and i == self.num_blocks - 1: + if i >= self.cross_start: + x = self.cross_blocks[i - self.cross_start][0]( + x, mem_l[i], mem_global_vector + ) + for idx in range(self.depth[i] - 1): + if self.use_self_global: + if self.self_update_global: + x, mem_global_vector = self.self_blocks[i][idx]( + x, mem_global_vector + ) + else: + x, _ = self.self_blocks[i][idx](x, mem_global_vector) + else: + x = self.self_blocks[i][idx](x) + if i >= self.cross_start: + x = self.cross_blocks[i - self.cross_start][idx + 1]( + x, mem_l[i], mem_global_vector + ) + else: + for idx in range(self.depth[i]): + if self.use_self_global: + if self.self_update_global: + x, mem_global_vector = self.self_blocks[i][idx]( + x, mem_global_vector + ) + else: + x, _ = self.self_blocks[i][idx](x, mem_global_vector) + else: + x = self.self_blocks[i][idx](x) + if i >= self.cross_start: + x = self.cross_blocks[i - self.cross_start][idx]( + x, mem_l[i], mem_global_vector + ) + if i > 0: + x = self.upsample_layers[i - 1](x) + if self.hierarchical_pos_embed: + x = self.hierarchical_pos_embed_l[i - 1](x) + return x + + +class MixtureCrossAttention(nn.Layer): + def __init__( + self, + dim, + num_heads, + cuboid_hw, + shift_hw, + strategy, + n_temporal, + cross_last_n_frames, + padding_type, + qkv_bias, + qk_scale, + attn_drop, + proj_drop, + norm_layer, + max_temporal_relative, + use_global_vector, + separate_global_qkv, + global_dim_ratio, + checkpoint_level, + use_relative_pos, + attn_linear_init_mode, + ffn_linear_init_mode, + norm_init_mode, + expert_shape, + moe_config, + ): + super().__init__() + + self.in_dim = dim + self.out_dim = dim + self.expert_shape = expert_shape # T, H, W, C + self.num_experts = moe_config["num_experts"] + self.out_planes = moe_config["out_planes"] + self.moe_config = moe_config + assert expert_shape is not None and moe_config["use_attn_moe"] + assert not use_global_vector + + if moe_config["gate_style"] == "linear": + self.gate = moe_utils.LinearGatingNet(moe_config, expert_shape, dim) + elif moe_config["gate_style"] == "spatial-latent": + self.gate = moe_utils.SpatialLatentGatingNet(moe_config, expert_shape, dim) + elif moe_config["gate_style"] == "cuboid-latent": + self.gate = moe_utils.CuboidLatentGatingNet(moe_config, expert_shape, dim) + elif moe_config["gate_style"] == "spatial-latent-linear": + self.gate = moe_utils.SpatialLatentLinearGatingNet( + moe_config, expert_shape, dim + ) + elif moe_config["gate_style"] == "cuboid-latent-linear": + self.gate = moe_utils.CuboidLatentLinearGatingNet( + moe_config, expert_shape, dim + ) + else: + raise NotImplementedError + + self.experts = nn.LayerList( + [ + CuboidCrossAttentionLayer( + dim=dim, + num_heads=num_heads, + cuboid_hw=cuboid_hw, + shift_hw=shift_hw, + strategy=strategy, + n_temporal=n_temporal, + cross_last_n_frames=cross_last_n_frames, + padding_type=padding_type, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + attn_drop=attn_drop, + proj_drop=proj_drop, + norm_layer=norm_layer, + max_temporal_relative=max_temporal_relative, + use_global_vector=use_global_vector, + separate_global_qkv=separate_global_qkv, + global_dim_ratio=global_dim_ratio, + checkpoint_level=checkpoint_level, + use_relative_pos=use_relative_pos, + attn_linear_init_mode=attn_linear_init_mode, + ffn_linear_init_mode=ffn_linear_init_mode, + norm_init_mode=norm_init_mode, + ) + for _ in range(self.num_experts) + ] + ) + + def forward(self, x, mem, mem_global_vectors=None): + + B, T_x, H, W, C = x.shape + _, T_m, _, _, _ = mem.shape + E = self.num_experts + assert C == self.in_dim and list(self.expert_shape)[:-1] == x.shape[1:-1] + ( + dense_routing_weights, + sparse_routing_weights, + sparse_routing_inds, + self.aux_loss, + ) = self.gate( + x + ) # dense: B, T_x, H, W, E + + dispatcher = moe_utils.DenseDispatcher( + E, + sparse_routing_weights.reshape([B * T_x * H * W, -1]), + sparse_routing_inds.reshape([B * T_x * H * W, -1]), + ) + expert_outputs = paddle.stack( + [self.experts[i](x, mem, mem_global_vectors) for i in range(E)], axis=-2 + ).reshape([B * T_x * H * W, E, C]) + y = dispatcher.combine(expert_outputs).reshape([B, T_x, H, W, C]) + + return y + + def reset_parameters(self): + + for i in range(len(self.experts)): + self.experts[i].reset_parameters() diff --git a/examples/smc_reac/ppsci/arch/extformer_moe_cuboid_encoder.py b/examples/smc_reac/ppsci/arch/extformer_moe_cuboid_encoder.py new file mode 100644 index 0000000000..c26b3837a5 --- /dev/null +++ b/examples/smc_reac/ppsci/arch/extformer_moe_cuboid_encoder.py @@ -0,0 +1,1992 @@ +from collections import OrderedDict +from functools import lru_cache +from typing import Tuple + +import numpy as np +import paddle +import paddle.nn.functional as F +from paddle import nn +from paddle.distributed import fleet + +import ppsci.arch.extformer_moe_cuboid_utils as cuboid_utils +import ppsci.arch.extformer_moe_utils as moe_utils +from ppsci.arch import activation as act_mod +from ppsci.utils import initializer + +NEGATIVE_SLOPE = 0.1 + + +class PatchMerging3D(nn.Layer): + """Patch Merging Layer + + Args: + dim (int): Number of input channels. + out_dim (int, optional): The dim of output. Defaults to None. + downsample (tuple, optional): Downsample factor. Defaults to (1, 2, 2). + norm_layer (str, optional): The normalization layer. Defaults to "layer_norm". + padding_type (str, optional): The type of padding. Defaults to "nearest". + linear_init_mode (str, optional): The mode of linear init. Defaults to "0". + norm_init_mode (str, optional): The mode of normalization init. Defaults to "0". + """ + + def __init__( + self, + dim: int, + out_dim: int = None, + downsample: Tuple[int, ...] = (1, 2, 2), + norm_layer: str = "layer_norm", + padding_type: str = "nearest", + linear_init_mode: str = "0", + norm_init_mode: str = "0", + moe_config: dict = None, + ): + super().__init__() + self.linear_init_mode = linear_init_mode + self.norm_init_mode = norm_init_mode + self.dim = dim + if out_dim is None: + out_dim = max(downsample) * dim + self.out_dim = out_dim + self.downsample = downsample + self.padding_type = padding_type + self.reduction = nn.Linear( + in_features=downsample[0] * downsample[1] * downsample[2] * dim, + out_features=out_dim, + bias_attr=False, + ) + self.norm = cuboid_utils.get_norm_layer( + norm_layer, in_channels=downsample[0] * downsample[1] * downsample[2] * dim + ) + self.reset_parameters() + + def reset_parameters(self): + for m in self.children(): + cuboid_utils.apply_initialization( + m, linear_mode=self.linear_init_mode, norm_mode=self.norm_init_mode + ) + + def get_out_shape(self, data_shape): + T, H, W, C_in = data_shape + pad_t = (self.downsample[0] - T % self.downsample[0]) % self.downsample[0] + pad_h = (self.downsample[1] - H % self.downsample[1]) % self.downsample[1] + pad_w = (self.downsample[2] - W % self.downsample[2]) % self.downsample[2] + return ( + (T + pad_t) // self.downsample[0], + (H + pad_h) // self.downsample[1], + (W + pad_w) // self.downsample[2], + self.out_dim, + ) + + def forward(self, x): + """ + + Args: + x : (B, T, H, W, C) + + Returns: + out : Shape (B, T // downsample[0], H // downsample[1], W // downsample[2], out_dim) + """ + + B, T, H, W, C = x.shape + pad_t = (self.downsample[0] - T % self.downsample[0]) % self.downsample[0] + pad_h = (self.downsample[1] - H % self.downsample[1]) % self.downsample[1] + pad_w = (self.downsample[2] - W % self.downsample[2]) % self.downsample[2] + if pad_h or pad_h or pad_w: + T += pad_t + H += pad_h + W += pad_w + x = cuboid_utils.generalize_padding( + x, pad_t, pad_h, pad_w, padding_type=self.padding_type + ) + x = ( + x.reshape( + ( + B, + T // self.downsample[0], + self.downsample[0], + H // self.downsample[1], + self.downsample[1], + W // self.downsample[2], + self.downsample[2], + C, + ) + ) + .transpose(perm=[0, 1, 3, 5, 2, 4, 6, 7]) + .reshape( + [ + B, + T // self.downsample[0], + H // self.downsample[1], + W // self.downsample[2], + self.downsample[0] * self.downsample[1] * self.downsample[2] * C, + ] + ) + ) + x = self.norm(x) + x = self.reduction(x) + return x + + +class PositionwiseFFN(nn.Layer): + """The Position-wise FFN layer used in Transformer-like architectures + + If pre_norm is True: + norm(data) -> fc1 -> act -> act_dropout -> fc2 -> dropout -> res(+data) + Else: + data -> fc1 -> act -> act_dropout -> fc2 -> dropout -> norm(res(+data)) + Also, if we use gated projection. We will use + fc1_1 * act(fc1_2(data)) to map the data + + Args: + units (int, optional): The units. Defaults to 512. + hidden_size (int, optional): The size of hidden layer. Defaults to 2048. + activation_dropout (float, optional): The dropout of activate. Defaults to 0.0. + dropout (float, optional): The drop ratio used in DropPat. Defaults to 0.1. + gated_proj (bool, optional): Whether to use gate projection. Defaults to False. + activation (str, optional): The activate. Defaults to "relu". + normalization (str, optional): The normalization. Defaults to "layer_norm". + layer_norm_eps (float, optional): The epsilon of layer normalization. Defaults to 1e-05. + pre_norm (bool): Pre-layer normalization as proposed in the paper: + "[ACL2018] The Best of Both Worlds: Combining Recent Advances in Neural Machine Translation" This will stabilize the training of Transformers. + You may also refer to "[Arxiv2020] Understanding the Difficulty of Training Transformers". Defaults to False. + linear_init_mode (str, optional): The mode of linear initialization. Defaults to "0". + norm_init_mode (str, optional): The mode of normalization initialization. Defaults to "0". + """ + + def __init__( + self, + units: int = 512, + hidden_size: int = 2048, + activation_dropout: float = 0.0, + dropout: float = 0.1, + gated_proj: bool = False, + activation: str = "relu", + normalization: str = "layer_norm", + layer_norm_eps: float = 1e-05, + pre_norm: bool = False, + linear_init_mode: str = "0", + norm_init_mode: str = "0", + moe_config: dict = None, + expert_shape: tuple = None, + ): + super().__init__() + self.linear_init_mode = linear_init_mode + self.norm_init_mode = norm_init_mode + self._pre_norm = pre_norm + self._gated_proj = gated_proj + self._kwargs = OrderedDict( + [ + ("units", units), + ("hidden_size", hidden_size), + ("activation_dropout", activation_dropout), + ("activation", activation), + ("dropout", dropout), + ("normalization", normalization), + ("layer_norm_eps", layer_norm_eps), + ("gated_proj", gated_proj), + ("pre_norm", pre_norm), + ] + ) + self.dropout_layer = nn.Dropout(p=dropout) + self.activation_dropout_layer = nn.Dropout(p=activation_dropout) + + if moe_config["use_linear_moe"]: + self.ffn_1 = MixtureLinear( + in_dim=units, + out_dim=hidden_size, + bias_attr=True, + expert_shape=expert_shape[:-1] + (hidden_size,), + moe_config=moe_config, + ) + else: + self.ffn_1 = nn.Linear( + in_features=units, out_features=hidden_size, bias_attr=True + ) + if self._gated_proj: + self.ffn_1_gate = nn.Linear( + in_features=units, out_features=hidden_size, bias_attr=True + ) + if activation == "leaky_relu": + self.activation = nn.LeakyReLU(NEGATIVE_SLOPE) + else: + self.activation = act_mod.get_activation(activation) + + if moe_config["use_linear_moe"]: + self.ffn_2 = MixtureLinear( + in_dim=hidden_size, + out_dim=units, + bias_attr=True, + expert_shape=expert_shape, + moe_config=moe_config, + ) + else: + self.ffn_2 = nn.Linear( + in_features=hidden_size, out_features=units, bias_attr=True + ) + self.layer_norm = cuboid_utils.get_norm_layer( + normalization=normalization, in_channels=units, epsilon=layer_norm_eps + ) + self.reset_parameters() + + def reset_parameters(self): + cuboid_utils.apply_initialization(self.ffn_1, linear_mode=self.linear_init_mode) + if self._gated_proj: + cuboid_utils.apply_initialization( + self.ffn_1_gate, linear_mode=self.linear_init_mode + ) + cuboid_utils.apply_initialization(self.ffn_2, linear_mode=self.linear_init_mode) + cuboid_utils.apply_initialization( + self.layer_norm, norm_mode=self.norm_init_mode + ) + + def forward(self, data): + """ + Args: + x : Shape (B, seq_length, C_in) + + Returns: + out : Shape (B, seq_length, C_out) + """ + + residual = data + if self._pre_norm: + data = self.layer_norm(data) + if self._gated_proj: + out = self.activation(self.ffn_1_gate(data)) * self.ffn_1(data) + else: + out = self.activation(self.ffn_1(data)) + out = self.activation_dropout_layer(out) + out = self.ffn_2(out) + out = self.dropout_layer(out) + out = out + residual + if not self._pre_norm: + out = self.layer_norm(out) + return out + + +def update_cuboid_size_shift_size(data_shape, cuboid_size, shift_size, strategy): + """Update the cuboid_size and shift_size + + Args: + data_shape (Tuple[int,...]): The shape of the data. + cuboid_size (Tuple[int,...]): Size of the cuboid. + shift_size (Tuple[int,...]): Size of the shift. + strategy (str): The strategy of attention. + + Returns: + new_cuboid_size (Tuple[int,...]): Size of the cuboid. + new_shift_size (Tuple[int,...]): Size of the shift. + """ + + new_cuboid_size = list(cuboid_size) + new_shift_size = list(shift_size) + for i in range(len(data_shape)): + if strategy[i] == "d": + new_shift_size[i] = 0 + if data_shape[i] <= cuboid_size[i]: + new_cuboid_size[i] = data_shape[i] + new_shift_size[i] = 0 + return tuple(new_cuboid_size), tuple(new_shift_size) + + +def cuboid_reorder(data, cuboid_size, strategy): + """Reorder the tensor into (B, num_cuboids, bT * bH * bW, C) + We assume that the tensor shapes are divisible to the cuboid sizes. + + Args: + data (paddle.Tensor): The input data. + cuboid_size (Tuple[int,...]): The size of the cuboid. + strategy (Tuple[int,...]): The cuboid strategy. + + Returns: + reordered_data (paddle.Tensor): Shape will be (B, num_cuboids, bT * bH * bW, C). + num_cuboids = T / bT * H / bH * W / bW + """ + + B, T, H, W, C = data.shape + num_cuboids = T // cuboid_size[0] * H // cuboid_size[1] * W // cuboid_size[2] + cuboid_volume = cuboid_size[0] * cuboid_size[1] * cuboid_size[2] + intermediate_shape = [] + nblock_axis = [] + block_axis = [] + for i, (block_size, total_size, ele_strategy) in enumerate( + zip(cuboid_size, (T, H, W), strategy) + ): + if ele_strategy == "l": + intermediate_shape.extend([total_size // block_size, block_size]) + nblock_axis.append(2 * i + 1) + block_axis.append(2 * i + 2) + elif ele_strategy == "d": + intermediate_shape.extend([block_size, total_size // block_size]) + nblock_axis.append(2 * i + 2) + block_axis.append(2 * i + 1) + else: + raise NotImplementedError(f"{ele_strategy} is invalid.") + data = data.reshape(list((B,) + tuple(intermediate_shape) + (C,))) + reordered_data = data.transpose( + perm=(0,) + tuple(nblock_axis) + tuple(block_axis) + (7,) + ) + reordered_data = reordered_data.reshape((B, num_cuboids, cuboid_volume, C)) + return reordered_data + + +@lru_cache() +def compute_cuboid_self_attention_mask( + data_shape, cuboid_size, shift_size, strategy, padding_type, device +): + """Compute the shift window attention mask + + Args: + data_shape (Tuple[int,....]): Should be (T, H, W). + cuboid_size (Tuple[int,....]): Size of the cuboid. + shift_size (Tuple[int,....]): The shift size. + strategy (str): The decomposition strategy. + padding_type (str): Type of the padding. + device (str): The device. + + Returns: + attn_mask (paddle.Tensor): Mask with shape (num_cuboid, cuboid_vol, cuboid_vol). + The padded values will always be masked. The other masks will ensure that the shifted windows + will only attend to those in the shifted windows. + """ + T, H, W = data_shape + pad_t = (cuboid_size[0] - T % cuboid_size[0]) % cuboid_size[0] + pad_h = (cuboid_size[1] - H % cuboid_size[1]) % cuboid_size[1] + pad_w = (cuboid_size[2] - W % cuboid_size[2]) % cuboid_size[2] + data_mask = None + if pad_t > 0 or pad_h > 0 or pad_w > 0: + if padding_type == "ignore": + data_mask = paddle.ones(shape=(1, T, H, W, 1), dtype="bool") + data_mask = F.pad( + data_mask, [0, 0, 0, pad_w, 0, pad_h, 0, pad_t], data_format="NDHWC" + ) + else: + data_mask = paddle.ones( + shape=(1, T + pad_t, H + pad_h, W + pad_w, 1), dtype="bool" + ) + if any(i > 0 for i in shift_size): + if padding_type == "ignore": + data_mask = paddle.roll( + x=data_mask, + shifts=(-shift_size[0], -shift_size[1], -shift_size[2]), + axis=(1, 2, 3), + ) + if padding_type == "ignore": + data_mask = cuboid_reorder(data_mask, cuboid_size, strategy=strategy) + data_mask = data_mask.squeeze(axis=-1).squeeze(axis=0) + shift_mask = np.zeros(shape=(1, T + pad_t, H + pad_h, W + pad_w, 1)) + cnt = 0 + for t in ( + slice(-cuboid_size[0]), + slice(-cuboid_size[0], -shift_size[0]), + slice(-shift_size[0], None), + ): + for h in ( + slice(-cuboid_size[1]), + slice(-cuboid_size[1], -shift_size[1]), + slice(-shift_size[1], None), + ): + for w in ( + slice(-cuboid_size[2]), + slice(-cuboid_size[2], -shift_size[2]), + slice(-shift_size[2], None), + ): + shift_mask[:, t, h, w, :] = cnt + cnt += 1 + shift_mask = paddle.to_tensor(shift_mask) + shift_mask = cuboid_reorder(shift_mask, cuboid_size, strategy=strategy) + shift_mask = shift_mask.squeeze(axis=-1).squeeze(axis=0) + attn_mask = shift_mask.unsqueeze(axis=1) - shift_mask.unsqueeze(axis=2) == 0 + if padding_type == "ignore": + attn_mask = ( + data_mask.unsqueeze(axis=1) * data_mask.unsqueeze(axis=2) * attn_mask + ) + return attn_mask + + +def masked_softmax(att_score, mask, axis: int = -1): + """Ignore the masked elements when calculating the softmax. + The mask can be broadcastable. + + Args: + att_score (paddle.Tensor): Shape (..., length, ...) + mask (paddle.Tensor): Shape (..., length, ...) + 1 --> The element is not masked + 0 --> The element is masked + axis (int): The axis to calculate the softmax. att_score.shape[axis] must be the same as mask.shape[axis] + + Returns: + att_weights (paddle.Tensor): Shape (..., length, ...). + """ + + if mask is not None: + if att_score.dtype == paddle.float16: + att_score = att_score.masked_fill(paddle.logical_not(mask), -1e4) + else: + att_score = att_score.masked_fill(paddle.logical_not(mask), -1e18) + att_weights = nn.functional.softmax(x=att_score, axis=axis) * mask.astype( + att_score.dtype + ) + else: + att_weights = nn.functional.softmax(x=att_score, axis=axis) + return att_weights + + +def cuboid_reorder_reverse(data, cuboid_size, strategy, orig_data_shape): + """Reverse the reordered cuboid back to the original space + + Args: + data (paddle.Tensor): The input data. + cuboid_size (Tuple[int,...]): The size of cuboid. + strategy (str): The strategy of reordering. + orig_data_shape (Tuple[int,...]): The original shape of the data. + + Returns: + data (paddle.Tensor): The recovered data + """ + + B, num_cuboids, cuboid_volume, C = data.shape + T, H, W = orig_data_shape + permutation_axis = [0] + for i, (block_size, total_size, ele_strategy) in enumerate( + zip(cuboid_size, (T, H, W), strategy) + ): + if ele_strategy == "l": + permutation_axis.append(i + 1) + permutation_axis.append(i + 4) + elif ele_strategy == "d": + permutation_axis.append(i + 4) + permutation_axis.append(i + 1) + else: + raise NotImplementedError((f"{ele_strategy} is invalid.")) + permutation_axis.append(7) + data = data.reshape( + [ + B, + T // cuboid_size[0], + H // cuboid_size[1], + W // cuboid_size[2], + cuboid_size[0], + cuboid_size[1], + cuboid_size[2], + C, + ] + ) + data = data.transpose(perm=permutation_axis) + data = data.reshape((B, T, H, W, C)) + return data + + +class CuboidSelfAttentionLayer(nn.Layer): + """Implements the cuboid self attention. + + The idea of Cuboid Self Attention is to divide the input tensor (T, H, W) into several non-overlapping cuboids. + We apply self-attention inside each cuboid and all cuboid-level self attentions are executed in parallel. + + We adopt two mechanisms for decomposing the input tensor into cuboids: + + (1) local: + We group the tensors within a local window, e.g., X[t:(t+b_t), h:(h+b_h), w:(w+b_w)]. We can also apply the + shifted window strategy proposed in "[ICCV2021] Swin Transformer: Hierarchical Vision Transformer using Shifted Windows". + (2) dilated: + Inspired by the success of dilated convolution "[ICLR2016] Multi-Scale Context Aggregation by Dilated Convolutions", + we split the tensor with dilation factors that are tied to the size of the cuboid. For example, for a cuboid that has width `b_w`, + we sample the elements starting from 0 as 0, w / b_w, 2 * w / b_w, ..., (b_w - 1) * w / b_w. + + The cuboid attention can be viewed as a generalization of the attention mechanism proposed in Video Swin Transformer, https://arxiv.org/abs/2106.13230. + The computational complexity of CuboidAttention can be simply calculated as O(T H W * b_t b_h b_w). To cover multiple correlation patterns, + we are able to combine multiple CuboidAttention layers with different configurations such as cuboid size, shift size, and local / global decomposing strategy. + + In addition, it is straight-forward to extend the cuboid attention to other types of spatiotemporal data that are not described + as regular tensors. We need to define alternative approaches to partition the data into "cuboids". + + In addition, inspired by "[NeurIPS2021] Do Transformers Really Perform Badly for Graph Representation?", + "[NeurIPS2020] Big Bird: Transformers for Longer Sequences", "[EMNLP2021] Longformer: The Long-Document Transformer", we keep + $K$ global vectors to record the global status of the spatiotemporal system. These global vectors will attend to the whole tensor and + the vectors inside each individual cuboids will also attend to the global vectors so that they can peep into the global status of the system. + + Args: + dim (int): The dimension of the input tensor. + num_heads (int): The number of heads. + cuboid_size (tuple, optional): The size of cuboid. Defaults to (2, 7, 7). + shift_size (tuple, optional): The size of shift. Defaults to (0, 0, 0). + strategy (tuple, optional): The strategy. Defaults to ("l", "l", "l"). + padding_type (str, optional): The type of padding. Defaults to "ignore". + qkv_bias (bool, optional): Whether to enable bias in calculating qkv attention. Defaults to False. + qk_scale (float, optional): Whether to enable scale factor when calculating the attention. Defaults to None. + attn_drop (float, optional): The attention dropout. Defaults to 0.0. + proj_drop (float, optional): The projection dropout. Defaults to 0.0. + use_final_proj (bool, optional): Whether to use the final projection. Defaults to True. + norm_layer (str, optional): The normalization layer. Defaults to "layer_norm". + use_global_vector (bool, optional): Whether to use the global vector or not. Defaults to False. + use_global_self_attn (bool, optional): Whether to use self attention among global vectors. Defaults to False. + separate_global_qkv (bool, optional): Whether to use different network to calc q_global, k_global, v_global. Defaults to False. + global_dim_ratio (int, optional): The dim (channels) of global vectors is `global_dim_ratio*dim`. Defaults to 1. + checkpoint_level (bool, optional): Whether to enable gradient checkpointing. Defaults to True. + use_relative_pos (bool, optional): Whether to use relative pos. Defaults to True. + attn_linear_init_mode (str, optional): The mode of attention linear initialization. Defaults to "0". + ffn_linear_init_mode (str, optional): The mode of FFN linear initialization. Defaults to "0". + norm_init_mode (str, optional): The mode of normalization initialization. Defaults to "0". + """ + + def __init__( + self, + dim: int, + num_heads: int, + cuboid_size: Tuple[int, ...] = (2, 7, 7), + shift_size: Tuple[int, ...] = (0, 0, 0), + strategy: Tuple[str, ...] = ("l", "l", "l"), + padding_type: str = "ignore", + qkv_bias: bool = False, + qk_scale: float = None, + attn_drop: float = 0.0, + proj_drop: float = 0.0, + use_final_proj: bool = True, + norm_layer: str = "layer_norm", + use_global_vector: bool = False, + use_global_self_attn: bool = False, + separate_global_qkv: bool = False, + global_dim_ratio: int = 1, + checkpoint_level: bool = True, + use_relative_pos: bool = True, + attn_linear_init_mode: str = "0", + ffn_linear_init_mode: str = "0", + norm_init_mode: str = "0", + moe_config: dict = None, + ): + super(CuboidSelfAttentionLayer, self).__init__() + self.attn_linear_init_mode = attn_linear_init_mode + self.ffn_linear_init_mode = ffn_linear_init_mode + self.norm_init_mode = norm_init_mode + assert dim % num_heads == 0 + self.num_heads = num_heads + self.dim = dim + self.cuboid_size = cuboid_size + self.shift_size = shift_size + self.strategy = strategy + self.padding_type = padding_type + self.use_final_proj = use_final_proj + self.use_relative_pos = use_relative_pos + self.use_global_vector = use_global_vector + self.use_global_self_attn = use_global_self_attn + self.separate_global_qkv = separate_global_qkv + if global_dim_ratio != 1: + assert ( + separate_global_qkv is True + ), "Setting global_dim_ratio != 1 requires separate_global_qkv == True." + self.global_dim_ratio = global_dim_ratio + assert self.padding_type in ["ignore", "zeros", "nearest"] + head_dim = dim // num_heads + self.scale = qk_scale or head_dim**-0.5 + if use_relative_pos: + init_data = paddle.zeros( + ( + (2 * cuboid_size[0] - 1) + * (2 * cuboid_size[1] - 1) + * (2 * cuboid_size[2] - 1), + num_heads, + ) + ) + self.relative_position_bias_table = paddle.create_parameter( + shape=init_data.shape, + dtype=init_data.dtype, + default_initializer=nn.initializer.Constant(0.0), + ) + self.relative_position_bias_table.stop_gradient = not True + self.relative_position_bias_table = initializer.trunc_normal_( + self.relative_position_bias_table, std=0.02 + ) + + coords_t = paddle.arange(end=self.cuboid_size[0]) + coords_h = paddle.arange(end=self.cuboid_size[1]) + coords_w = paddle.arange(end=self.cuboid_size[2]) + coords = paddle.stack(x=paddle.meshgrid(coords_t, coords_h, coords_w)) + coords_flatten = paddle.flatten(x=coords, start_axis=1) + relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] + relative_coords = relative_coords.transpose(perm=[1, 2, 0]) + relative_coords[:, :, 0] += self.cuboid_size[0] - 1 + relative_coords[:, :, 1] += self.cuboid_size[1] - 1 + relative_coords[:, :, 2] += self.cuboid_size[2] - 1 + relative_coords[:, :, 0] *= (2 * self.cuboid_size[1] - 1) * ( + 2 * self.cuboid_size[2] - 1 + ) + relative_coords[:, :, 1] *= 2 * self.cuboid_size[2] - 1 + relative_position_index = relative_coords.sum(axis=-1) + self.register_buffer( + name="relative_position_index", tensor=relative_position_index + ) + self.qkv = nn.Linear(in_features=dim, out_features=dim * 3, bias_attr=qkv_bias) + self.attn_drop = nn.Dropout(p=attn_drop) + if self.use_global_vector: + if self.separate_global_qkv: + self.l2g_q_net = nn.Linear( + in_features=dim, out_features=dim, bias_attr=qkv_bias + ) + self.l2g_global_kv_net = nn.Linear( + in_features=global_dim_ratio * dim, + out_features=dim * 2, + bias_attr=qkv_bias, + ) + self.g2l_global_q_net = nn.Linear( + in_features=global_dim_ratio * dim, + out_features=dim, + bias_attr=qkv_bias, + ) + self.g2l_k_net = nn.Linear( + in_features=dim, out_features=dim, bias_attr=qkv_bias + ) + self.g2l_v_net = nn.Linear( + in_features=dim, + out_features=global_dim_ratio * dim, + bias_attr=qkv_bias, + ) + if self.use_global_self_attn: + self.g2g_global_qkv_net = nn.Linear( + in_features=global_dim_ratio * dim, + out_features=global_dim_ratio * dim * 3, + bias_attr=qkv_bias, + ) + else: + self.global_qkv = nn.Linear( + in_features=dim, out_features=dim * 3, bias_attr=qkv_bias + ) + self.global_attn_drop = nn.Dropout(p=attn_drop) + if use_final_proj: + self.proj = nn.Linear(in_features=dim, out_features=dim) + self.proj_drop = nn.Dropout(p=proj_drop) + if self.use_global_vector: + self.global_proj = nn.Linear( + in_features=global_dim_ratio * dim, + out_features=global_dim_ratio * dim, + ) + self.norm = cuboid_utils.get_norm_layer(norm_layer, in_channels=dim) + if self.use_global_vector: + self.global_vec_norm = cuboid_utils.get_norm_layer( + norm_layer, in_channels=global_dim_ratio * dim + ) + self.checkpoint_level = checkpoint_level + self.reset_parameters() + + def reset_parameters(self): + cuboid_utils.apply_initialization( + self.qkv, linear_mode=self.attn_linear_init_mode + ) + if self.use_final_proj: + cuboid_utils.apply_initialization( + self.proj, linear_mode=self.ffn_linear_init_mode + ) + cuboid_utils.apply_initialization(self.norm, norm_mode=self.norm_init_mode) + if self.use_global_vector: + if self.separate_global_qkv: + cuboid_utils.apply_initialization( + self.l2g_q_net, linear_mode=self.attn_linear_init_mode + ) + cuboid_utils.apply_initialization( + self.l2g_global_kv_net, linear_mode=self.attn_linear_init_mode + ) + cuboid_utils.apply_initialization( + self.g2l_global_q_net, linear_mode=self.attn_linear_init_mode + ) + cuboid_utils.apply_initialization( + self.g2l_k_net, linear_mode=self.attn_linear_init_mode + ) + cuboid_utils.apply_initialization( + self.g2l_v_net, linear_mode=self.attn_linear_init_mode + ) + if self.use_global_self_attn: + cuboid_utils.apply_initialization( + self.g2g_global_qkv_net, linear_mode=self.attn_linear_init_mode + ) + else: + cuboid_utils.apply_initialization( + self.global_qkv, linear_mode=self.attn_linear_init_mode + ) + cuboid_utils.apply_initialization( + self.global_vec_norm, norm_mode=self.norm_init_mode + ) + + def forward(self, x, global_vectors=None): + x = self.norm(x) + + B, T, H, W, C_in = x.shape + assert C_in == self.dim + if self.use_global_vector: + _, num_global, _ = global_vectors.shape + global_vectors = self.global_vec_norm(global_vectors) + cuboid_size, shift_size = update_cuboid_size_shift_size( + (T, H, W), self.cuboid_size, self.shift_size, self.strategy + ) + + pad_t = (cuboid_size[0] - T % cuboid_size[0]) % cuboid_size[0] + pad_h = (cuboid_size[1] - H % cuboid_size[1]) % cuboid_size[1] + pad_w = (cuboid_size[2] - W % cuboid_size[2]) % cuboid_size[2] + x = cuboid_utils.generalize_padding(x, pad_t, pad_h, pad_w, self.padding_type) + + if any(i > 0 for i in shift_size): + shifted_x = paddle.roll( + x=x, + shifts=(-shift_size[0], -shift_size[1], -shift_size[2]), + axis=(1, 2, 3), + ) + else: + shifted_x = x + + reordered_x = cuboid_reorder( + shifted_x, cuboid_size=cuboid_size, strategy=self.strategy + ) + + _, num_cuboids, cuboid_volume, _ = reordered_x.shape + attn_mask = compute_cuboid_self_attention_mask( + (T, H, W), + cuboid_size, + shift_size=shift_size, + strategy=self.strategy, + padding_type=self.padding_type, + device=x.place, + ) + head_C = C_in // self.num_heads + qkv = ( + self.qkv(reordered_x) + .reshape([B, num_cuboids, cuboid_volume, 3, self.num_heads, head_C]) + .transpose(perm=[3, 0, 4, 1, 2, 5]) + ) + + q, k, v = qkv[0], qkv[1], qkv[2] + q = q * self.scale + perm_0 = list(range(k.ndim)) + perm_0[-2] = -1 + perm_0[-1] = -2 + attn_score = q @ k.transpose(perm=perm_0) + + if self.use_relative_pos: + relative_position_bias = self.relative_position_bias_table[ + self.relative_position_index[:cuboid_volume, :cuboid_volume].reshape( + [-1] + ) + ].reshape([cuboid_volume, cuboid_volume, -1]) + relative_position_bias = relative_position_bias.transpose( + perm=[2, 0, 1] + ).unsqueeze(axis=1) + attn_score = attn_score + relative_position_bias + + if self.use_global_vector: + global_head_C = self.global_dim_ratio * head_C + if self.separate_global_qkv: + l2g_q = ( + self.l2g_q_net(reordered_x) + .reshape([B, num_cuboids, cuboid_volume, self.num_heads, head_C]) + .transpose(perm=[0, 3, 1, 2, 4]) + ) + l2g_q = l2g_q * self.scale + l2g_global_kv = ( + self.l2g_global_kv_net(global_vectors) + .reshape([B, 1, num_global, 2, self.num_heads, head_C]) + .transpose(perm=[3, 0, 4, 1, 2, 5]) + ) + l2g_global_k, l2g_global_v = l2g_global_kv[0], l2g_global_kv[1] + g2l_global_q = ( + self.g2l_global_q_net(global_vectors) + .reshape([B, num_global, self.num_heads, head_C]) + .transpose(perm=[0, 2, 1, 3]) + ) + g2l_global_q = g2l_global_q * self.scale + g2l_k = ( + self.g2l_k_net(reordered_x) + .reshape([B, num_cuboids, cuboid_volume, self.num_heads, head_C]) + .transpose(perm=[0, 3, 1, 2, 4]) + ) + g2l_v = ( + self.g2l_v_net(reordered_x) + .reshape( + [B, num_cuboids, cuboid_volume, self.num_heads, global_head_C] + ) + .transpose(perm=[0, 3, 1, 2, 4]) + ) + if self.use_global_self_attn: + g2g_global_qkv = ( + self.g2g_global_qkv_net(global_vectors) + .reshape([B, 1, num_global, 3, self.num_heads, global_head_C]) + .transpose(perm=[3, 0, 4, 1, 2, 5]) + ) + g2g_global_q, g2g_global_k, g2g_global_v = ( + g2g_global_qkv[0], + g2g_global_qkv[1], + g2g_global_qkv[2], + ) + g2g_global_q = g2g_global_q.squeeze(axis=2) * self.scale + else: + q_global, k_global, v_global = ( + self.global_qkv(global_vectors) + .reshape([B, 1, num_global, 3, self.num_heads, head_C]) + .transpose(perm=[3, 0, 4, 1, 2, 5]) + ) + q_global = q_global.squeeze(axis=2) * self.scale + l2g_q, g2l_k, g2l_v = q, k, v + g2l_global_q, l2g_global_k, l2g_global_v = ( + q_global, + k_global, + v_global, + ) + if self.use_global_self_attn: + g2g_global_q, g2g_global_k, g2g_global_v = ( + q_global, + k_global, + v_global, + ) + + perm_1 = list(range(l2g_global_k.ndim)) + perm_1[-2] = -1 + perm_1[-1] = -2 + l2g_attn_score = l2g_q @ l2g_global_k.transpose(perm=perm_1) + attn_score_l2l_l2g = paddle.concat(x=(attn_score, l2g_attn_score), axis=-1) + + if attn_mask.ndim == 5: + attn_mask_l2l_l2g = F.pad( + attn_mask, [0, num_global], "constant", 1, data_format="NDHWC" + ) + elif attn_mask.ndim == 3: + attn_mask = attn_mask.astype("float32") + attn_mask_l2l_l2g = F.pad( + attn_mask, [0, num_global], "constant", 1, data_format="NCL" + ) + attn_mask_l2l_l2g = attn_mask_l2l_l2g.astype("bool") + else: + attn_mask_l2l_l2g = F.pad(attn_mask, [0, num_global], "constant", 1) + + v_l_g = paddle.concat( + x=( + v, + l2g_global_v.expand( + shape=[B, self.num_heads, num_cuboids, num_global, head_C] + ), + ), + axis=3, + ) + attn_score_l2l_l2g = masked_softmax( + attn_score_l2l_l2g, mask=attn_mask_l2l_l2g + ) + attn_score_l2l_l2g = self.attn_drop(attn_score_l2l_l2g) + reordered_x = ( + (attn_score_l2l_l2g @ v_l_g) + .transpose(perm=[0, 2, 3, 1, 4]) + .reshape([B, num_cuboids, cuboid_volume, self.dim]) + ) + if self.padding_type == "ignore": + g2l_attn_mask = paddle.ones(shape=(1, T, H, W, 1)) + if pad_t > 0 or pad_h > 0 or pad_w > 0: + g2l_attn_mask = F.pad( + g2l_attn_mask, + [0, 0, 0, pad_w, 0, pad_h, 0, pad_t], + data_format="NDHWC", + ) + if any(i > 0 for i in shift_size): + g2l_attn_mask = paddle.roll( + x=g2l_attn_mask, + shifts=(-shift_size[0], -shift_size[1], -shift_size[2]), + axis=(1, 2, 3), + ) + g2l_attn_mask = g2l_attn_mask.reshape((-1,)) + else: + g2l_attn_mask = None + temp = g2l_k.reshape( + [B, self.num_heads, num_cuboids * cuboid_volume, head_C] + ) + perm_2 = list(range(temp.ndim)) + perm_2[-2] = -1 + perm_2[-1] = -2 + g2l_attn_score = g2l_global_q @ temp.transpose(perm=perm_2) + if self.use_global_self_attn: + temp = g2g_global_k.squeeze(axis=2) + perm_3 = list(range(temp.ndim)) + perm_3[-2] = -1 + perm_3[-1] = -2 + g2g_attn_score = g2g_global_q @ temp.transpose(perm=perm_3) + g2all_attn_score = paddle.concat( + x=(g2l_attn_score, g2g_attn_score), axis=-1 + ) + if g2l_attn_mask is not None: + g2all_attn_mask = F.pad( + g2l_attn_mask, + [0, num_global], + "constant", + 1, + data_format="NDHWC", + ) + else: + g2all_attn_mask = None + new_v = paddle.concat( + x=( + g2l_v.reshape( + [ + B, + self.num_heads, + num_cuboids * cuboid_volume, + global_head_C, + ] + ), + g2g_global_v.reshape( + [B, self.num_heads, num_global, global_head_C] + ), + ), + axis=2, + ) + else: + g2all_attn_score = g2l_attn_score + g2all_attn_mask = g2l_attn_mask + new_v = g2l_v.reshape( + [B, self.num_heads, num_cuboids * cuboid_volume, global_head_C] + ) + g2all_attn_score = masked_softmax(g2all_attn_score, mask=g2all_attn_mask) + g2all_attn_score = self.global_attn_drop(g2all_attn_score) + new_global_vector = ( + (g2all_attn_score @ new_v) + .transpose(perm=[0, 2, 1, 3]) + .reshape([B, num_global, self.global_dim_ratio * self.dim]) + ) + else: + attn_score = masked_softmax(attn_score, mask=attn_mask) + attn_score = self.attn_drop(attn_score) + reordered_x = ( + (attn_score @ v) + .transpose(perm=[0, 2, 3, 1, 4]) + .reshape([B, num_cuboids, cuboid_volume, self.dim]) + ) + + if self.use_final_proj: + reordered_x = paddle.cast(reordered_x, dtype="float32") + reordered_x = self.proj_drop(self.proj(reordered_x)) + if self.use_global_vector: + new_global_vector = self.proj_drop(self.global_proj(new_global_vector)) + shifted_x = cuboid_reorder_reverse( + reordered_x, + cuboid_size=cuboid_size, + strategy=self.strategy, + orig_data_shape=(T + pad_t, H + pad_h, W + pad_w), + ) + if any(i > 0 for i in shift_size): + x = paddle.roll( + x=shifted_x, + shifts=(shift_size[0], shift_size[1], shift_size[2]), + axis=(1, 2, 3), + ) + else: + x = shifted_x + x = cuboid_utils.generalize_unpadding( + x, pad_t=pad_t, pad_h=pad_h, pad_w=pad_w, padding_type=self.padding_type + ) + if self.use_global_vector: + return x, new_global_vector + else: + return x + + +class StackCuboidSelfAttentionBlock(nn.Layer): + """ + - "use_inter_ffn" is True + x --> attn1 -----+-------> ffn1 ---+---> attn2 --> ... --> ffn_k --> out + | ^ | ^ + | | | | + |-------------| |-------------| + - "use_inter_ffn" is False + x --> attn1 -----+------> attn2 --> ... attnk --+----> ffnk ---+---> out + | ^ | ^ ^ | ^ + | | | | | | | + |-------------| |------------| ----------| |-----------| + If we have enabled global memory vectors, each attention will be a + + Args: + dim (int): The dimension of the input tensor. + num_heads (int): The number of heads. + block_cuboid_size (list, optional): The size of block cuboid . Defaults to [(4, 4, 4), (4, 4, 4)]. + block_shift_size (list, optional): The shift size of block. Defaults to [(0, 0, 0), (2, 2, 2)]. + block_strategy (list, optional): The strategy of block. Defaults to [("d", "d", "d"), ("l", "l", "l")]. + padding_type (str, optional): The type of padding. Defaults to "ignore". + qkv_bias (bool, optional): Whether to enable bias in calculating qkv attention. Defaults to False. + qk_scale (float, optional): Whether to enable scale factor when calculating the attention. Defaults to None. + attn_drop (float, optional): The attention dropout. Defaults to 0.0. + proj_drop (float, optional): The projection dropout. Defaults to 0.0. + use_final_proj (bool, optional): Whether to use the final projection. Defaults to True. + norm_layer (str, optional): The normalization layer. Defaults to "layer_norm". + use_global_vector (bool, optional): Whether to use the global vector or not. Defaults to False. + use_global_self_attn (bool, optional): Whether to use self attention among global vectors. Defaults to False. + separate_global_qkv (bool, optional): Whether to use different network to calc q_global, k_global, v_global. + Defaults to False. + global_dim_ratio (int, optional): The dim (channels) of global vectors is `global_dim_ratio*dim`. + Defaults to 1. + checkpoint_level (bool, optional): Whether to enable gradient checkpointing. Defaults to True. + use_relative_pos (bool, optional): Whether to use relative pos. Defaults to True. + use_relative_pos (bool, optional): Whether to use relative pos. Defaults to True. + attn_linear_init_mode (str, optional): The mode of attention linear initialization. Defaults to "0". + ffn_linear_init_mode (str, optional): The mode of FFN linear initialization. Defaults to "0". + norm_init_mode (str, optional): The mode of normalization initialization. Defaults to "0". + """ + + def __init__( + self, + dim: int, + num_heads: int, + block_cuboid_size: Tuple[Tuple[int, ...], ...] = [(4, 4, 4), (4, 4, 4)], + block_shift_size: Tuple[Tuple[int, ...], ...] = [(0, 0, 0), (2, 2, 2)], + block_strategy: Tuple[Tuple[str, ...], ...] = [ + ("d", "d", "d"), + ("l", "l", "l"), + ], + padding_type: str = "ignore", + qkv_bias: bool = False, + qk_scale: float = None, + attn_drop: float = 0.0, + proj_drop: float = 0.0, + ffn_drop: float = 0.0, + activation: str = "leaky", + gated_ffn: bool = False, + norm_layer: str = "layer_norm", + use_inter_ffn: bool = False, + use_global_vector: bool = False, + use_global_vector_ffn: bool = True, + use_global_self_attn: bool = False, + separate_global_qkv: bool = False, + global_dim_ratio: int = 1, + checkpoint_level: bool = True, + use_relative_pos: bool = True, + use_final_proj: bool = True, + attn_linear_init_mode: str = "0", + ffn_linear_init_mode: str = "0", + norm_init_mode: str = "0", + moe_config: dict = None, + expert_shape: tuple = None, + ): + super(StackCuboidSelfAttentionBlock, self).__init__() + self.attn_linear_init_mode = attn_linear_init_mode + self.ffn_linear_init_mode = ffn_linear_init_mode + self.norm_init_mode = norm_init_mode + if ( + len(block_cuboid_size[0]) <= 0 + or len(block_shift_size) <= 0 + or len(block_strategy) <= 0 + ): + raise ValueError( + "Format of the block cuboid size is not correct. block_cuboid_size={block_cuboid_size}" + ) + if len(block_cuboid_size) != len(block_shift_size) and len( + block_cuboid_size + ) != len(block_strategy): + raise ValueError( + "The lengths of block_cuboid_size, block_shift_size, and block_strategy must be equal." + ) + + self.num_attn = len(block_cuboid_size) + self.checkpoint_level = checkpoint_level + self.use_inter_ffn = use_inter_ffn + self.use_global_vector = use_global_vector + self.use_global_vector_ffn = use_global_vector_ffn + self.use_global_self_attn = use_global_self_attn + self.global_dim_ratio = global_dim_ratio + if self.use_inter_ffn: + if moe_config["use_ffn_moe"]: + self.ffn_l = nn.LayerList( + sublayers=[ + MixtureFFN( + units=dim, + hidden_size=4 * dim, + activation_dropout=ffn_drop, + dropout=ffn_drop, + gated_proj=gated_ffn, + activation=activation, + normalization=norm_layer, + pre_norm=True, + linear_init_mode=ffn_linear_init_mode, + norm_init_mode=norm_init_mode, + expert_shape=expert_shape, + moe_config=moe_config, + ) + for _ in range(self.num_attn) + ] + ) + else: + self.ffn_l = nn.LayerList( + sublayers=[ + PositionwiseFFN( + units=dim, + hidden_size=4 * dim, + activation_dropout=ffn_drop, + dropout=ffn_drop, + gated_proj=gated_ffn, + activation=activation, + normalization=norm_layer, + pre_norm=True, + linear_init_mode=ffn_linear_init_mode, + norm_init_mode=norm_init_mode, + expert_shape=expert_shape, + moe_config=moe_config, + ) + for _ in range(self.num_attn) + ] + ) + if self.use_global_vector_ffn and self.use_global_vector: + if moe_config["use_ffn_moe"]: + self.global_ffn_l = nn.LayerList( + sublayers=[ + MixtureFFN( + units=global_dim_ratio * dim, + hidden_size=global_dim_ratio * 4 * dim, + activation_dropout=ffn_drop, + dropout=ffn_drop, + gated_proj=gated_ffn, + activation=activation, + normalization=norm_layer, + pre_norm=True, + linear_init_mode=ffn_linear_init_mode, + norm_init_mode=norm_init_mode, + expert_shape=expert_shape, + moe_config=moe_config, + ) + for _ in range(self.num_attn) + ] + ) + else: + self.global_ffn_l = nn.LayerList( + sublayers=[ + PositionwiseFFN( + units=global_dim_ratio * dim, + hidden_size=global_dim_ratio * 4 * dim, + activation_dropout=ffn_drop, + dropout=ffn_drop, + gated_proj=gated_ffn, + activation=activation, + normalization=norm_layer, + pre_norm=True, + linear_init_mode=ffn_linear_init_mode, + norm_init_mode=norm_init_mode, + expert_shape=expert_shape, + moe_config=moe_config, + ) + for _ in range(self.num_attn) + ] + ) + else: + if moe_config["use_ffn_moe"]: + self.ffn_l = nn.LayerList( + sublayers=[ + MixtureFFN( + units=dim, + hidden_size=4 * dim, + activation_dropout=ffn_drop, + dropout=ffn_drop, + gated_proj=gated_ffn, + activation=activation, + normalization=norm_layer, + pre_norm=True, + linear_init_mode=ffn_linear_init_mode, + norm_init_mode=norm_init_mode, + expert_shape=expert_shape, + moe_config=moe_config, + ) + ] + ) + else: + self.ffn_l = nn.LayerList( + sublayers=[ + PositionwiseFFN( + units=dim, + hidden_size=4 * dim, + activation_dropout=ffn_drop, + dropout=ffn_drop, + gated_proj=gated_ffn, + activation=activation, + normalization=norm_layer, + pre_norm=True, + linear_init_mode=ffn_linear_init_mode, + norm_init_mode=norm_init_mode, + expert_shape=expert_shape, + moe_config=moe_config, + ) + ] + ) + if self.use_global_vector_ffn and self.use_global_vector: + if moe_config["use_ffn_moe"]: + self.global_ffn_l = nn.LayerList( + sublayers=[ + MixtureFFN( + units=global_dim_ratio * dim, + hidden_size=global_dim_ratio * 4 * dim, + activation_dropout=ffn_drop, + dropout=ffn_drop, + gated_proj=gated_ffn, + activation=activation, + normalization=norm_layer, + pre_norm=True, + linear_init_mode=ffn_linear_init_mode, + norm_init_mode=norm_init_mode, + expert_shape=expert_shape, + moe_config=moe_config, + ) + ] + ) + else: + self.global_ffn_l = nn.LayerList( + sublayers=[ + PositionwiseFFN( + units=global_dim_ratio * dim, + hidden_size=global_dim_ratio * 4 * dim, + activation_dropout=ffn_drop, + dropout=ffn_drop, + gated_proj=gated_ffn, + activation=activation, + normalization=norm_layer, + pre_norm=True, + linear_init_mode=ffn_linear_init_mode, + norm_init_mode=norm_init_mode, + expert_shape=expert_shape, + moe_config=moe_config, + ) + ] + ) + + if moe_config["use_attn_moe"]: + self.attn_l = nn.LayerList( + sublayers=[ + MixtureSelfAttention( + dim=dim, + num_heads=num_heads, + cuboid_size=ele_cuboid_size, + shift_size=ele_shift_size, + strategy=ele_strategy, + padding_type=padding_type, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + attn_drop=attn_drop, + proj_drop=proj_drop, + norm_layer=norm_layer, + use_global_vector=use_global_vector, + use_global_self_attn=use_global_self_attn, + separate_global_qkv=separate_global_qkv, + global_dim_ratio=global_dim_ratio, + checkpoint_level=checkpoint_level, + use_relative_pos=use_relative_pos, + use_final_proj=use_final_proj, + attn_linear_init_mode=attn_linear_init_mode, + ffn_linear_init_mode=ffn_linear_init_mode, + norm_init_mode=norm_init_mode, + expert_shape=expert_shape, + moe_config=moe_config, + ) + for ele_cuboid_size, ele_shift_size, ele_strategy in zip( + block_cuboid_size, block_shift_size, block_strategy + ) + ] + ) + else: + self.attn_l = nn.LayerList( + sublayers=[ + CuboidSelfAttentionLayer( + dim=dim, + num_heads=num_heads, + cuboid_size=ele_cuboid_size, + shift_size=ele_shift_size, + strategy=ele_strategy, + padding_type=padding_type, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + attn_drop=attn_drop, + proj_drop=proj_drop, + norm_layer=norm_layer, + use_global_vector=use_global_vector, + use_global_self_attn=use_global_self_attn, + separate_global_qkv=separate_global_qkv, + global_dim_ratio=global_dim_ratio, + checkpoint_level=checkpoint_level, + use_relative_pos=use_relative_pos, + use_final_proj=use_final_proj, + attn_linear_init_mode=attn_linear_init_mode, + ffn_linear_init_mode=ffn_linear_init_mode, + norm_init_mode=norm_init_mode, + ) + for ele_cuboid_size, ele_shift_size, ele_strategy in zip( + block_cuboid_size, block_shift_size, block_strategy + ) + ] + ) + + def reset_parameters(self): + for m in self.ffn_l: + m.reset_parameters() + if self.use_global_vector_ffn and self.use_global_vector: + for m in self.global_ffn_l: + m.reset_parameters() + for m in self.attn_l: + m.reset_parameters() + + def forward(self, x, global_vectors=None): + if self.use_inter_ffn: + if self.use_global_vector: + for idx, (attn, ffn) in enumerate(zip(self.attn_l, self.ffn_l)): + if self.checkpoint_level >= 2 and self.training: + x_out, global_vectors_out = fleet.utils.recompute( + attn, x, global_vectors + ) + else: + x_out, global_vectors_out = attn(x, global_vectors) + x = x + x_out + global_vectors = global_vectors + global_vectors_out + if self.checkpoint_level >= 1 and self.training: + x = fleet.utils.recompute(ffn, x) + if self.use_global_vector_ffn: + global_vectors = fleet.utils.recompute( + self.global_ffn_l[idx], global_vectors + ) + else: + x = ffn(x) + if self.use_global_vector_ffn: + global_vectors = self.global_ffn_l[idx](global_vectors) + return x, global_vectors + else: + for idx, (attn, ffn) in enumerate(zip(self.attn_l, self.ffn_l)): + if self.checkpoint_level >= 2 and self.training: + x = x + fleet.utils.recompute(attn, x) + else: + x = x + attn(x) + if self.checkpoint_level >= 1 and self.training: + x = fleet.utils.recompute(ffn, x) + else: + x = ffn(x) + return x + elif self.use_global_vector: + for idx, attn in enumerate(self.attn_l): + if self.checkpoint_level >= 2 and self.training: + x_out, global_vectors_out = fleet.utils.recompute( + attn, x, global_vectors + ) + else: + x_out, global_vectors_out = attn(x, global_vectors) + x = x + x_out + global_vectors = global_vectors + global_vectors_out + if self.checkpoint_level >= 1 and self.training: + x = fleet.utils.recompute(self.ffn_l[0], x) + if self.use_global_vector_ffn: + global_vectors = fleet.utils.recompute( + self.global_ffn_l[0], global_vectors + ) + else: + x = self.ffn_l[0](x) + if self.use_global_vector_ffn: + global_vectors = self.global_ffn_l[0](global_vectors) + return x, global_vectors + else: + for idx, attn in enumerate(self.attn_l): + if self.checkpoint_level >= 2 and self.training: + out = fleet.utils.recompute(attn, x) + else: + out = attn(x) + x = x + out + if self.checkpoint_level >= 1 and self.training: + x = fleet.utils.recompute(self.ffn_l[0], x) + else: + x = self.ffn_l[0](x) + return x + + +class CuboidTransformerEncoder(nn.Layer): + """Encoder of the CuboidTransformer + + x --> attn_block --> patch_merge --> attn_block --> patch_merge --> ... --> out + + Args: + input_shape (Tuple[int,...]): The shape of the input. Contains T, H, W, C + base_units (int, optional): The number of units. Defaults to 128. + block_units (int, optional): The number of block units. Defaults to None. + scale_alpha (float, optional): We scale up the channels based on the formula: + - round_to(base_units * max(downsample_scale) ** units_alpha, 4). Defaults to 1.0. + depth (list, optional): The number of layers for each block. Defaults to [4, 4, 4]. + downsample (int, optional): The downsample ratio. Defaults to 2. + downsample_type (str, optional): The type of downsample. Defaults to "patch_merge". + block_attn_patterns (str, optional): Attention pattern for the cuboid attention for each block. Defaults to None. + block_cuboid_size (list, optional): A list of cuboid size parameters. Defaults to [(4, 4, 4), (4, 4, 4)]. + block_strategy (list, optional): A list of cuboid strategies. Defaults to [("l", "l", "l"), ("d", "d", "d")]. + block_shift_size (list, optional): A list of shift sizes. Defaults to [(0, 0, 0), (0, 0, 0)]. + num_heads (int, optional): The number of heads. Defaults to 4. + attn_drop (float, optional): The ratio of attention dropout. Defaults to 0.0. + proj_drop (float, optional): The ratio of projection dropout. Defaults to 0.0. + ffn_drop (float, optional): The ratio of FFN dropout. Defaults to 0.0. + ffn_activation (str, optional): The FFN activation. Defaults to "leaky". + gated_ffn (bool, optional): Whether to use gate FFN. Defaults to False. + norm_layer (str, optional): The normalization layer. Defaults to "layer_norm". + use_inter_ffn (bool, optional): Whether to use inter FFN. Defaults to True. + padding_type (str, optional): The type of padding. Defaults to "ignore". + checkpoint_level (bool, optional): Whether to enable gradient checkpointing. Defaults to True. + use_relative_pos (bool, optional): Whether to use relative pos. Defaults to True. + self_attn_use_final_proj (bool, optional): Whether to use self attention for final projection. Defaults to True. + use_global_vector (bool, optional): Whether to use the global vector or not. Defaults to False. + use_global_vector_ffn (bool, optional): Whether to use FFN global vectors. Defaults to False. + use_global_self_attn (bool, optional): Whether to use global self attention. Defaults to False. + separate_global_qkv (bool, optional): Whether to use different network to calc q_global, k_global, v_global. + Defaults to False. + global_dim_ratio (int, optional): The dim (channels) of global vectors is `global_dim_ratio*dim`. + Defaults to 1. + attn_linear_init_mode (str, optional): The mode of attention linear initialization. Defaults to "0". + ffn_linear_init_mode (str, optional): The mode of FFN linear initialization. Defaults to "0". + conv_init_mode (str, optional): The mode of conv initialization. Defaults to "0". + down_linear_init_mode (str, optional): The mode of downsample linear initialization. Defaults to "0". + norm_init_mode (str, optional): The mode of normalization. Defaults to "0". + """ + + def __init__( + self, + input_shape: Tuple[int, ...], + base_units: int = 128, + block_units: int = None, + scale_alpha: float = 1.0, + depth: Tuple[int, ...] = [4, 4, 4], + downsample: int = 2, + downsample_type: str = "patch_merge", + block_attn_patterns: str = None, + block_cuboid_size: Tuple[Tuple[int, ...], ...] = [(4, 4, 4), (4, 4, 4)], + block_strategy: Tuple[Tuple[str, ...], ...] = [ + ("l", "l", "l"), + ("d", "d", "d"), + ], + block_shift_size: Tuple[Tuple[int, ...], ...] = [(0, 0, 0), (0, 0, 0)], + num_heads: int = 4, + attn_drop: float = 0.0, + proj_drop: float = 0.0, + ffn_drop: float = 0.0, + ffn_activation: str = "leaky", + gated_ffn: bool = False, + norm_layer: str = "layer_norm", + use_inter_ffn: bool = True, + padding_type: str = "ignore", + checkpoint_level: bool = True, + use_relative_pos: bool = True, + self_attn_use_final_proj: bool = True, + use_global_vector: bool = False, + use_global_vector_ffn: bool = True, + use_global_self_attn: bool = False, + separate_global_qkv: bool = False, + global_dim_ratio: int = 1, + attn_linear_init_mode: str = "0", + ffn_linear_init_mode: str = "0", + conv_init_mode: str = "0", + down_linear_init_mode: str = "0", + norm_init_mode: str = "0", + moe_config: dict = None, + ): + super(CuboidTransformerEncoder, self).__init__() + self.attn_linear_init_mode = attn_linear_init_mode + self.ffn_linear_init_mode = ffn_linear_init_mode + self.conv_init_mode = conv_init_mode + self.down_linear_init_mode = down_linear_init_mode + self.norm_init_mode = norm_init_mode + self.input_shape = input_shape + self.depth = depth + self.num_blocks = len(depth) + self.base_units = base_units + self.scale_alpha = scale_alpha + if not isinstance(downsample, (tuple, list)): + downsample = 1, downsample, downsample + self.downsample = downsample + self.downsample_type = downsample_type + self.num_heads = num_heads + self.use_global_vector = use_global_vector + self.checkpoint_level = checkpoint_level + + if block_units is None: + block_units = [ + cuboid_utils.round_to( + base_units * int((max(downsample) ** scale_alpha) ** i), 4 + ) + for i in range(self.num_blocks) + ] + else: + assert len(block_units) == self.num_blocks and block_units[0] == base_units + self.block_units = block_units + if self.num_blocks > 1: + if downsample_type == "patch_merge": + self.down_layers = nn.LayerList( + sublayers=[ + PatchMerging3D( + dim=self.block_units[i], + downsample=downsample, + padding_type=padding_type, + out_dim=self.block_units[i + 1], + linear_init_mode=down_linear_init_mode, + norm_init_mode=norm_init_mode, + ) + for i in range(self.num_blocks - 1) + ] + ) + else: + raise NotImplementedError(f"{downsample_type} is invalid.") + if self.use_global_vector: + self.down_layer_global_proj = nn.LayerList( + sublayers=[ + nn.Linear( + in_features=global_dim_ratio * self.block_units[i], + out_features=global_dim_ratio * self.block_units[i + 1], + ) + for i in range(self.num_blocks - 1) + ] + ) + if block_attn_patterns is not None: + mem_shapes = self.get_mem_shapes() + if isinstance(block_attn_patterns, (tuple, list)): + assert len(block_attn_patterns) == self.num_blocks + else: + block_attn_patterns = [ + block_attn_patterns for _ in range(self.num_blocks) + ] + block_cuboid_size = [] + block_strategy = [] + block_shift_size = [] + for idx, key in enumerate(block_attn_patterns): + func = cuboid_utils.CuboidSelfAttentionPatterns.get(key) + cuboid_size, strategy, shift_size = func(mem_shapes[idx]) + block_cuboid_size.append(cuboid_size) + block_strategy.append(strategy) + block_shift_size.append(shift_size) + else: + if not isinstance(block_cuboid_size[0][0], (list, tuple)): + block_cuboid_size = [block_cuboid_size for _ in range(self.num_blocks)] + else: + assert ( + len(block_cuboid_size) == self.num_blocks + ), f"Incorrect input format! Received block_cuboid_size={block_cuboid_size}" + if not isinstance(block_strategy[0][0], (list, tuple)): + block_strategy = [block_strategy for _ in range(self.num_blocks)] + else: + assert ( + len(block_strategy) == self.num_blocks + ), f"Incorrect input format! Received block_strategy={block_strategy}" + if not isinstance(block_shift_size[0][0], (list, tuple)): + block_shift_size = [block_shift_size for _ in range(self.num_blocks)] + else: + assert ( + len(block_shift_size) == self.num_blocks + ), f"Incorrect input format! Received block_shift_size={block_shift_size}" + self.block_cuboid_size = block_cuboid_size + self.block_strategy = block_strategy + self.block_shift_size = block_shift_size + + expert_shape_list = self.get_mem_shapes() + self.blocks = nn.LayerList( + sublayers=[ + nn.Sequential( + *[ + StackCuboidSelfAttentionBlock( + dim=self.block_units[i], + num_heads=num_heads, + block_cuboid_size=block_cuboid_size[i], + block_strategy=block_strategy[i], + block_shift_size=block_shift_size[i], + attn_drop=attn_drop, + proj_drop=proj_drop, + ffn_drop=ffn_drop, + activation=ffn_activation, + gated_ffn=gated_ffn, + norm_layer=norm_layer, + use_inter_ffn=use_inter_ffn, + padding_type=padding_type, + use_global_vector=use_global_vector, + use_global_vector_ffn=use_global_vector_ffn, + use_global_self_attn=use_global_self_attn, + separate_global_qkv=separate_global_qkv, + global_dim_ratio=global_dim_ratio, + checkpoint_level=checkpoint_level, + use_relative_pos=use_relative_pos, + use_final_proj=self_attn_use_final_proj, + attn_linear_init_mode=attn_linear_init_mode, + ffn_linear_init_mode=ffn_linear_init_mode, + norm_init_mode=norm_init_mode, + expert_shape=expert_shape_list[i], + moe_config=moe_config, + ) + for _ in range(depth[i]) + ] + ) + for i in range(self.num_blocks) + ] + ) + self.reset_parameters() + + def reset_parameters(self): + if self.num_blocks > 1: + for m in self.down_layers: + m.reset_parameters() + if self.use_global_vector: + cuboid_utils.apply_initialization( + self.down_layer_global_proj, linear_mode=self.down_linear_init_mode + ) + for ms in self.blocks: + for m in ms: + m.reset_parameters() + + def get_mem_shapes(self): + """Get the shape of the output memory based on the input shape. This can be used for constructing the decoder. + + Returns: + mem_shapes : A list of shapes of the output memory + """ + + if self.num_blocks == 1: + return [self.input_shape] + else: + mem_shapes = [self.input_shape] + curr_shape = self.input_shape + for down_layer in self.down_layers: + curr_shape = down_layer.get_out_shape(curr_shape) + mem_shapes.append(curr_shape) + return mem_shapes + + def forward(self, x, global_vectors=None): + """ + Args: + x : Shape (B, T, H, W, C) + + Returns: + out (List[paddle.Tensor,..]): A list of tensors from the bottom layer to the top layer of the encoder. For + example, it can have shape + - (B, T, H, W, C1) + - (B, T, H // 2, W // 2, 2 * C1) + - (B, T, H // 4, W // 4, 4 * C1) + ... + global_mem_out (List,Optional): The output of the global vector. + """ + + B, T, H, W, C_in = x.shape + assert (T, H, W, C_in) == self.input_shape + + if self.use_global_vector: + out = [] + global_mem_out = [] + for i in range(self.num_blocks): + for l in self.blocks[i]: + x, global_vectors = l(x, global_vectors) + out.append(x) + global_mem_out.append(global_vectors) + if self.num_blocks > 1 and i < self.num_blocks - 1: + x = self.down_layers[i](x) + global_vectors = self.down_layer_global_proj[i](global_vectors) + return out, global_mem_out + else: + out = [] + for i in range(self.num_blocks): + x = self.blocks[i](x) + out.append(x) + if self.num_blocks > 1 and i < self.num_blocks - 1: + x = self.down_layers[i](x) + return out + + +class MixtureLinear(nn.Layer): + def __init__(self, in_dim, out_dim, expert_shape, moe_config, bias_attr=True): + super().__init__() + + self.in_dim = in_dim + self.out_dim = out_dim + self.bias = bias_attr + self.expert_shape = expert_shape # T, H, W, C_o + self.num_experts = moe_config["num_experts"] + self.out_planes = moe_config["out_planes"] + self.moe_config = moe_config + assert expert_shape is not None and moe_config["use_linear_moe"] + + if moe_config["gate_style"] == "linear": + self.gate = moe_utils.LinearGatingNet(moe_config, expert_shape, in_dim) + elif moe_config["gate_style"] == "spatial-latent": + self.gate = moe_utils.SpatialLatentGatingNet( + moe_config, expert_shape, in_dim + ) + elif moe_config["gate_style"] == "cuboid-latent": + self.gate = moe_utils.CuboidLatentGatingNet( + moe_config, expert_shape, in_dim + ) + elif moe_config["gate_style"] == "spatial-latent-linear": + self.gate = moe_utils.SpatialLatentLinearGatingNet( + moe_config, expert_shape, in_dim + ) + elif moe_config["gate_style"] == "cuboid-latent-linear": + self.gate = moe_utils.CuboidLatentLinearGatingNet( + moe_config, expert_shape, in_dim + ) + else: + raise NotImplementedError + + self.experts = nn.LayerList( + [ + nn.Linear(in_features=in_dim, out_features=out_dim, bias_attr=bias_attr) + for _ in range(self.num_experts) + ] + ) + + def forward(self, x): + + B, T, H, W, C = x.shape + E = self.num_experts + assert C == self.in_dim and list(self.expert_shape)[:-1] == x.shape[1:-1] + ( + dense_routing_weights, + sparse_routing_weights, + sparse_routing_inds, + self.aux_loss, + ) = self.gate( + x + ) # dense: B, T, H, W, E + + if self.moe_config["dispatch_style"] == "dense": + dispatcher = moe_utils.DenseDispatcher( + E, + sparse_routing_weights.reshape([B * T * H * W, -1]), + sparse_routing_inds.reshape([B * T * H * W, -1]), + ) + expert_outputs = paddle.stack( + [self.experts[i](x.reshape([B * T * H * W, -1])) for i in range(E)], + axis=-2, + ) + y = dispatcher.combine(expert_outputs).reshape([B, T, H, W, -1]) + elif self.moe_config["dispatch_style"] == "sparse": + dispatcher = moe_utils.SparseDispatcher( + E, + sparse_routing_weights.reshape([B * T * H * W, -1]), + sparse_routing_inds.reshape([B * T * H * W, -1]), + ) + expert_inputs = dispatcher.dispatch(x.reshape([B * T * H * W, -1])) + expert_outputs = [ + self.experts[i](expert_inputs[i]) + if expert_inputs[i].shape[0] > 0 + else paddle.zeros([0, self.out_dim]) + for i in range(E) + ] + y = dispatcher.combine(expert_outputs).reshape([B, T, H, W, -1]) + else: + raise NotImplementedError + + return y + + +class MixtureFFN(nn.Layer): + def __init__( + self, + units, + hidden_size, + activation_dropout, + dropout, + gated_proj, + activation, + normalization, + pre_norm, + linear_init_mode, + norm_init_mode, + expert_shape, + moe_config, + ): + super().__init__() + + self.in_dim = units + self.out_dim = units + self.expert_shape = expert_shape # T, H, W, C_o + self.num_experts = moe_config["num_experts"] + self.out_planes = moe_config["out_planes"] + self.moe_config = moe_config + assert expert_shape is not None and moe_config["use_ffn_moe"] + + if moe_config["gate_style"] == "linear": + self.gate = moe_utils.LinearGatingNet(moe_config, expert_shape, units) + elif moe_config["gate_style"] == "spatial-latent": + self.gate = moe_utils.SpatialLatentGatingNet( + moe_config, expert_shape, units + ) + elif moe_config["gate_style"] == "cuboid-latent": + self.gate = moe_utils.CuboidLatentGatingNet(moe_config, expert_shape, units) + elif moe_config["gate_style"] == "spatial-latent-linear": + self.gate = moe_utils.SpatialLatentLinearGatingNet( + moe_config, expert_shape, units + ) + elif moe_config["gate_style"] == "cuboid-latent-linear": + self.gate = moe_utils.CuboidLatentLinearGatingNet( + moe_config, expert_shape, units + ) + else: + raise NotImplementedError + + self.experts = nn.LayerList( + [ + PositionwiseFFN( + units=units, + hidden_size=hidden_size, + activation_dropout=activation_dropout, + dropout=dropout, + gated_proj=gated_proj, + activation=activation, + normalization=normalization, + pre_norm=pre_norm, + linear_init_mode=linear_init_mode, + norm_init_mode=norm_init_mode, + moe_config=moe_config, + expert_shape=expert_shape, + ) + for _ in range(self.num_experts) + ] + ) + + def forward(self, x): + + B, T, H, W, C = x.shape + E = self.num_experts + assert C == self.in_dim and list(self.expert_shape)[:-1] == x.shape[1:-1] + ( + dense_routing_weights, + sparse_routing_weights, + sparse_routing_inds, + self.aux_loss, + ) = self.gate( + x + ) # dense: B, T, H, W, E + + if self.moe_config["dispatch_style"] == "dense": + dispatcher = moe_utils.DenseDispatcher( + E, + sparse_routing_weights.reshape([B * T * H * W, -1]), + sparse_routing_inds.reshape([B * T * H * W, -1]), + ) + expert_outputs = paddle.stack( + [self.experts[i](x.reshape([B * T * H * W, -1])) for i in range(E)], + axis=-2, + ) + y = dispatcher.combine(expert_outputs).reshape([B, T, H, W, C]) + elif self.moe_config["dispatch_style"] == "sparse": + dispatcher = moe_utils.SparseDispatcher( + E, + sparse_routing_weights.reshape([B * T * H * W, -1]), + sparse_routing_inds.reshape([B * T * H * W, -1]), + ) + expert_inputs = dispatcher.dispatch(x.reshape([B * T * H * W, -1])) + expert_outputs = [ + self.experts[i](expert_inputs[i]) + if expert_inputs[i].shape[0] > 0 + else paddle.zeros([0, self.out_dim]) + for i in range(E) + ] + y = dispatcher.combine(expert_outputs).reshape([B, T, H, W, C]) + else: + raise NotImplementedError + + return y + + def reset_parameters(self): + + for i in range(len(self.experts)): + self.experts[i].reset_parameters() + + +class MixtureSelfAttention(nn.Layer): + def __init__( + self, + dim, + num_heads, + cuboid_size, + shift_size, + strategy, + padding_type, + qkv_bias, + qk_scale, + attn_drop, + proj_drop, + norm_layer, + use_global_vector, + use_global_self_attn, + separate_global_qkv, + global_dim_ratio, + checkpoint_level, + use_relative_pos, + use_final_proj, + attn_linear_init_mode, + ffn_linear_init_mode, + norm_init_mode, + expert_shape, + moe_config, + ): + super().__init__() + + self.in_dim = dim + self.out_dim = dim + self.expert_shape = expert_shape # T, H, W, C + self.num_experts = moe_config["num_experts"] + self.out_planes = moe_config["out_planes"] + self.moe_config = moe_config + assert expert_shape is not None and moe_config["use_attn_moe"] + assert not use_global_vector + + if moe_config["gate_style"] == "linear": + self.gate = moe_utils.LinearGatingNet(moe_config, expert_shape, dim) + elif moe_config["gate_style"] == "spatial-latent": + self.gate = moe_utils.SpatialLatentGatingNet(moe_config, expert_shape, dim) + elif moe_config["gate_style"] == "cuboid-latent": + self.gate = moe_utils.CuboidLatentGatingNet(moe_config, expert_shape, dim) + elif moe_config["gate_style"] == "spatial-latent-linear": + self.gate = moe_utils.SpatialLatentLinearGatingNet( + moe_config, expert_shape, dim + ) + elif moe_config["gate_style"] == "cuboid-latent-linear": + self.gate = moe_utils.CuboidLatentLinearGatingNet( + moe_config, expert_shape, dim + ) + else: + raise NotImplementedError + + self.experts = nn.LayerList( + [ + CuboidSelfAttentionLayer( + dim=dim, + num_heads=num_heads, + cuboid_size=cuboid_size, + shift_size=shift_size, + strategy=strategy, + padding_type=padding_type, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + attn_drop=attn_drop, + proj_drop=proj_drop, + norm_layer=norm_layer, + use_global_vector=use_global_vector, + use_global_self_attn=use_global_self_attn, + separate_global_qkv=separate_global_qkv, + global_dim_ratio=global_dim_ratio, + checkpoint_level=checkpoint_level, + use_relative_pos=use_relative_pos, + use_final_proj=use_final_proj, + attn_linear_init_mode=attn_linear_init_mode, + ffn_linear_init_mode=ffn_linear_init_mode, + norm_init_mode=norm_init_mode, + ) + for _ in range(self.num_experts) + ] + ) + + def forward(self, x, global_vectors=None): + + B, T, H, W, C = x.shape + E = self.num_experts + assert C == self.in_dim and list(self.expert_shape)[:-1] == x.shape[1:-1] + ( + dense_routing_weights, + sparse_routing_weights, + sparse_routing_inds, + self.aux_loss, + ) = self.gate( + x + ) # dense: B, T, H, W, E + + dispatcher = moe_utils.DenseDispatcher( + E, + sparse_routing_weights.reshape([B * T * H * W, -1]), + sparse_routing_inds.reshape([B * T * H * W, -1]), + ) + expert_outputs = paddle.stack( + [self.experts[i](x, global_vectors) for i in range(E)], axis=-2 + ).reshape([B * T * H * W, E, C]) + y = dispatcher.combine(expert_outputs).reshape([B, T, H, W, C]) + + return y + + def reset_parameters(self): + + for i in range(len(self.experts)): + self.experts[i].reset_parameters() diff --git a/examples/smc_reac/ppsci/arch/extformer_moe_cuboid_utils.py b/examples/smc_reac/ppsci/arch/extformer_moe_cuboid_utils.py new file mode 100644 index 0000000000..20531c82d6 --- /dev/null +++ b/examples/smc_reac/ppsci/arch/extformer_moe_cuboid_utils.py @@ -0,0 +1,350 @@ +import functools +from typing import Tuple + +import paddle +import paddle.nn.functional as F +from paddle import nn + +from ppsci.utils import initializer + + +def round_to(dat, c): + return dat + (dat - dat % c) % c + + +class RMSNorm(nn.Layer): + """Root Mean Square Layer Normalization proposed in "[NeurIPS2019] Root Mean Square Layer Normalization" + + Args: + d (Optional[int]): The model size. + p (float, optional): The partial RMSNorm, valid value [0, 1]. Defaults to -1.0. + eps (float, optional): The epsilon value. Defaults to 1e-08. + bias (bool, optional): Whether use bias term for RMSNorm, + because RMSNorm doesn't enforce re-centering invariance.Defaults to False. + """ + + def __init__( + self, + d: Tuple[int, ...], + p: float = -1.0, + eps: float = 1e-08, + bias: bool = False, + ): + super().__init__() + self.eps = eps + self.d = d + self.p = p + self.bias = bias + init_data = paddle.ones(d) + self.scale = paddle.create_parameter( + shape=init_data.shape, + dtype=init_data.dtype, + default_initializer=nn.initializer.Constant(1.0), + ) + self.scale.stop_gradient = False + self.add_parameter(name="scale", parameter=self.scale) + if self.bias: + init_data = paddle.zeros(d) + self.offset = paddle.create_parameter( + shape=init_data.shape, + dtype=init_data.dtype, + default_initializer=nn.initializer.Constant(0.0), + ) + self.offset.stop_gradient = False + self.add_parameter(name="offset", parameter=self.offset) + + def forward(self, x): + if self.p < 0.0 or self.p > 1.0: + norm_x = x.norm(p=2, axis=-1, keepdim=True) + d_x = self.d + else: + partial_size = int(self.d * self.p) + partial_x, _ = paddle.split( + x=x, num_or_sections=[partial_size, self.d - partial_size], axis=-1 + ) + norm_x = partial_x.norm(p=2, axis=-1, keepdim=True) + d_x = partial_size + rms_x = norm_x * d_x ** (-1.0 / 2) + x_normed = x / (rms_x + self.eps) + if self.bias: + return self.scale * x_normed + self.offset + return self.scale * x_normed + + +def get_norm_layer( + normalization: str = "layer_norm", + axis: int = -1, + epsilon: float = 1e-05, + in_channels: int = 0, + **kwargs, +): + """Get the normalization layer based on the provided type + + Args: + normalization (str): The type of the layer normalization from ['layer_norm']. + axis (float): The axis to normalize the. + epsilon (float): The epsilon of the normalization layer. + in_channels (int): Input channel. + + Returns: + norm_layer (norm): The layer normalization layer. + """ + + if isinstance(normalization, str): + if normalization == "layer_norm": + assert in_channels > 0 + assert axis == -1 + norm_layer = nn.LayerNorm( + normalized_shape=in_channels, epsilon=epsilon, **kwargs + ) + elif normalization == "rms_norm": + assert axis == -1 + norm_layer = RMSNorm(d=in_channels, eps=epsilon, **kwargs) + else: + raise NotImplementedError(f"normalization={normalization} is not supported") + return norm_layer + elif normalization is None: + return nn.Identity() + else: + raise NotImplementedError("The type of normalization must be str") + + +def generalize_padding(x, pad_t, pad_h, pad_w, padding_type, t_pad_left=False): + if pad_t == 0 and pad_h == 0 and pad_w == 0: + return x + assert padding_type in ["zeros", "ignore", "nearest"] + B, T, H, W, C = x.shape + if padding_type == "nearest": + return nn.functional.interpolate( + x=x.transpose(perm=[0, 4, 1, 2, 3]), size=(T + pad_t, H + pad_h, W + pad_w) + ).transpose(perm=[0, 2, 3, 4, 1]) + elif t_pad_left: + return F.pad(x, [0, 0, 0, pad_w, 0, pad_h, pad_t, 0], data_format="NDHWC") + else: + data_pad = F.pad( + x, [0, 0, pad_t, 0, pad_h, 0, pad_w, 0, 0, 0], data_format="NDHWC" + ) + data_pad = paddle.concat( + [data_pad[:, pad_t:, ...], data_pad[:, :pad_t, ...]], axis=1 + ) + return data_pad + + +def generalize_unpadding(x, pad_t, pad_h, pad_w, padding_type): + assert padding_type in ["zeros", "ignore", "nearest"] + B, T, H, W, C = x.shape + if pad_t == 0 and pad_h == 0 and pad_w == 0: + return x + if padding_type == "nearest": + return nn.functional.interpolate( + x=x.transpose(perm=[0, 4, 1, 2, 3]), size=(T - pad_t, H - pad_h, W - pad_w) + ).transpose(perm=[0, 2, 3, 4, 1]) + else: + return x[:, : T - pad_t, : H - pad_h, : W - pad_w, :] + + +def apply_initialization( + m: nn.Layer, + linear_mode: str = "0", + conv_mode: str = "0", + norm_mode: str = "0", + embed_mode: str = "0", +): + if isinstance(m, nn.Linear): + if linear_mode in ("0",): + m.weight = initializer.kaiming_normal_(m.weight, nonlinearity="linear") + elif linear_mode in ("1",): + m.weight = initializer.kaiming_normal_( + m.weight, a=0.1, mode="fan_out", nonlinearity="leaky_relu" + ) + else: + raise NotImplementedError(f"{linear_mode} is invalid.") + if hasattr(m, "bias") and m.bias is not None: + m.bias = initializer.zeros_(m.bias) + elif isinstance( + m, + ( + nn.Conv2D, + nn.Conv3D, + nn.Conv2DTranspose, + nn.Conv3DTranspose, + ), + ): + if conv_mode in ("0",): + m.weight = initializer.kaiming_normal_( + m.weight, a=0.1, mode="fan_out", nonlinearity="leaky_relu" + ) + else: + raise NotImplementedError(f"{conv_mode} is invalid.") + if hasattr(m, "bias") and m.bias is not None: + m.bias = initializer.zeros_(m.bias) + elif isinstance(m, nn.LayerNorm): + if norm_mode in ("0",): + m.weight = initializer.zeros_(m.weight) + m.bias = initializer.zeros_(m.bias) + else: + raise NotImplementedError(f"{norm_mode} is invalid.") + elif isinstance(m, nn.GroupNorm): + if norm_mode in ("0",): + m.weight = initializer.ones_(m.weight) + m.bias = initializer.zeros_(m.bias) + else: + raise NotImplementedError(f"{norm_mode} is invalid.") + elif isinstance(m, nn.Embedding): + if embed_mode in ("0",): + m.weight.data = initializer.trunc_normal_(m.weight.data, std=0.02) + else: + raise NotImplementedError(f"{embed_mode} is invalid.") + elif isinstance(m, nn.Layer) and hasattr(m, "experts"): + for lin in m.experts: + assert isinstance(lin, nn.Linear) + apply_initialization(lin, linear_mode=linear_mode) + else: + pass + + +class CuboidSelfAttentionPatterns: + def __init__(self): + super().__init__() + self.patterns = {} + self.patterns = { + "full": self.full_attention, + "axial": self.axial, + "divided_st": self.divided_space_time, + } + for p in [1, 2, 4, 8, 10]: + for m in [1, 2, 4, 8, 16, 32]: + key = f"video_swin_{p}x{m}" + self.patterns[key] = functools.partial(self.video_swin, P=p, M=m) + + for m in [1, 2, 4, 8, 16, 32]: + key = f"spatial_lg_{m}" + self.patterns[key] = functools.partial(self.spatial_lg_v1, M=m) + + for k in [2, 4, 8]: + key = f"axial_space_dilate_{k}" + self.patterns[key] = functools.partial(self.axial_space_dilate_K, K=k) + + def get(self, pattern_name): + return self.patterns[pattern_name] + + def full_attention(self, input_shape): + T, H, W, _ = input_shape + cuboid_size = [(T, H, W)] + strategy = [("l", "l", "l")] + shift_size = [(0, 0, 0)] + return cuboid_size, strategy, shift_size + + def axial(self, input_shape): + """Axial attention proposed in https://arxiv.org/abs/1912.12180 + + Args: + input_shape (Tuple[int,...]): The shape of the input tensor, T H W. + + Returns: + cuboid_size (Tuple[int,...]): The size of cuboid. + strategy (Tuple[str,...]): The strategy of the attention. + shift_size (Tuple[int,...]): The shift size of the attention. + """ + + T, H, W, _ = input_shape + cuboid_size = [(T, 1, 1), (1, H, 1), (1, 1, W)] + strategy = [("l", "l", "l"), ("l", "l", "l"), ("l", "l", "l")] + shift_size = [(0, 0, 0), (0, 0, 0), (0, 0, 0)] + return cuboid_size, strategy, shift_size + + def divided_space_time(self, input_shape): + T, H, W, _ = input_shape + cuboid_size = [(T, 1, 1), (1, H, W)] + strategy = [("l", "l", "l"), ("l", "l", "l")] + shift_size = [(0, 0, 0), (0, 0, 0)] + return cuboid_size, strategy, shift_size + + def video_swin(self, input_shape, P=2, M=4): + """Adopt the strategy in Video SwinTransformer https://arxiv.org/pdf/2106.13230.pdf""" + T, H, W, _ = input_shape + P = min(P, T) + M = min(M, H, W) + cuboid_size = [(P, M, M), (P, M, M)] + strategy = [("l", "l", "l"), ("l", "l", "l")] + shift_size = [(0, 0, 0), (P // 2, M // 2, M // 2)] + return cuboid_size, strategy, shift_size + + def spatial_lg_v1(self, input_shape, M=4): + T, H, W, _ = input_shape + if H <= M and W <= M: + cuboid_size = [(T, 1, 1), (1, H, W)] + strategy = [("l", "l", "l"), ("l", "l", "l")] + shift_size = [(0, 0, 0), (0, 0, 0)] + else: + cuboid_size = [(T, 1, 1), (1, M, M), (1, M, M)] + strategy = [("l", "l", "l"), ("l", "l", "l"), ("d", "d", "d")] + shift_size = [(0, 0, 0), (0, 0, 0), (0, 0, 0)] + return cuboid_size, strategy, shift_size + + def axial_space_dilate_K(self, input_shape, K=2): + T, H, W, _ = input_shape + K = min(K, H, W) + cuboid_size = [ + (T, 1, 1), + (1, H // K, 1), + (1, H // K, 1), + (1, 1, W // K), + (1, 1, W // K), + ] + strategy = [ + ("l", "l", "l"), + ("d", "d", "d"), + ("l", "l", "l"), + ("d", "d", "d"), + ("l", "l", "l"), + ] + shift_size = [(0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0)] + return cuboid_size, strategy, shift_size + + +class CuboidCrossAttentionPatterns: + def __init__(self): + super().__init__() + self.patterns = {} + for k in [1, 2, 4, 8]: + key1 = f"cross_{k}x{k}" + key2 = f"cross_{k}x{k}_lg" + key3 = f"cross_{k}x{k}_heter" + self.patterns[key1] = functools.partial(self.cross_KxK, K=k) + self.patterns[key2] = functools.partial(self.cross_KxK_lg, K=k) + self.patterns[key3] = functools.partial(self.cross_KxK_heter, K=k) + + def get(self, pattern_name): + return self.patterns[pattern_name] + + def cross_KxK(self, mem_shape, K): + T_mem, H, W, _ = mem_shape + K = min(K, H, W) + cuboid_hw = [(K, K)] + shift_hw = [(0, 0)] + strategy = [("l", "l", "l")] + n_temporal = [1] + return cuboid_hw, shift_hw, strategy, n_temporal + + def cross_KxK_lg(self, mem_shape, K): + T_mem, H, W, _ = mem_shape + K = min(K, H, W) + cuboid_hw = [(K, K), (K, K)] + shift_hw = [(0, 0), (0, 0)] + strategy = [("l", "l", "l"), ("d", "d", "d")] + n_temporal = [1, 1] + return cuboid_hw, shift_hw, strategy, n_temporal + + def cross_KxK_heter(self, mem_shape, K): + T_mem, H, W, _ = mem_shape + K = min(K, H, W) + cuboid_hw = [(K, K), (K, K), (K, K)] + shift_hw = [(0, 0), (0, 0), (K // 2, K // 2)] + strategy = [("l", "l", "l"), ("d", "d", "d"), ("l", "l", "l")] + n_temporal = [1, 1, 1] + return cuboid_hw, shift_hw, strategy, n_temporal + + +CuboidSelfAttentionPatterns = CuboidSelfAttentionPatterns() +CuboidCrossAttentionPatterns = CuboidCrossAttentionPatterns() diff --git a/examples/smc_reac/ppsci/arch/extformer_moe_utils.py b/examples/smc_reac/ppsci/arch/extformer_moe_utils.py new file mode 100644 index 0000000000..3332b356c8 --- /dev/null +++ b/examples/smc_reac/ppsci/arch/extformer_moe_utils.py @@ -0,0 +1,563 @@ +import math + +import paddle +from paddle import nn + +# MoE Gating + + +class GatingNet(nn.Layer): + def __init__(self, moe_config, input_shape, in_channels): + super().__init__() + + self.num_experts = moe_config["num_experts"] + self.out_planes = moe_config["out_planes"] + self.aux_loss_style = moe_config["aux_loss_style"] + assert self.out_planes > 1 and self.out_planes <= self.num_experts + assert len(input_shape) == 4 + self.input_shape = input_shape + + self.noise_lin = nn.Linear( + in_features=in_channels, out_features=self.num_experts, bias_attr=False + ) + self.noise_eps = 1e-2 + self.softplus = nn.Softplus() + self.softmax = nn.Softmax(axis=-1) + + self.importance_weight = moe_config["importance_weight"] + self.load_weight = moe_config["load_weight"] + + def cv_squared(self, x, eps=1e-25): + return x.var(axis=-1) / (x.mean(axis=-1) ** 2 + eps) + + def intra_cdf(self, value, loc=0.0, scale=1.0): + return 0.5 * (1 + paddle.erf((value - loc) / scale / math.sqrt(2))) + + def importance_loss_cell(self, routing_weights): + importance_loss = self.cv_squared(routing_weights.sum(axis=0)).mean() + return importance_loss + + def load_loss_cell( + self, clean_values, noisy_values, noise_stddev, noisy_top_values + ): + B, T, H, W, E = clean_values.shape + M = noisy_top_values.shape[-1] + clean_values = clean_values.transpose([1, 2, 3, 0, 4]) + noisy_values = noisy_values.transpose([1, 2, 3, 0, 4]) + noise_stddev = noise_stddev.transpose([1, 2, 3, 0, 4]) + top_values_flat = noisy_top_values.transpose([1, 2, 3, 0, 4]).reshape( + [T, H, W, B * M] + ) + + threshold_positions_if_in = paddle.arange(B) * M + self.out_planes + threshold_if_in = paddle.take_along_axis( + top_values_flat, + axis=-1, + indices=threshold_positions_if_in.unsqueeze(axis=[0, 1, 2]), + ).unsqueeze( + -1 + ) # T, H, W, B, 1 + is_in = noisy_values > threshold_if_in # T, H, W, B, E + threshold_positions_if_out = threshold_positions_if_in - 1 + threshold_if_out = paddle.take_along_axis( + top_values_flat, + axis=-1, + indices=threshold_positions_if_out.unsqueeze(axis=[0, 1, 2]), + ).unsqueeze(-1) + + prob_if_in = self.intra_cdf( + (clean_values - threshold_if_in) / noise_stddev + ) # T, H, W, B, E + prob_if_out = self.intra_cdf( + (clean_values - threshold_if_out) / noise_stddev + ) # T, H, W, B, E + prob = paddle.where(is_in, prob_if_in, prob_if_out) # T, H, W, B, E + + load_loss = self.cv_squared(prob.sum(axis=-2)).mean() + return load_loss + + def importance_loss_all(self, routing_weights): + importance_loss = self.cv_squared(routing_weights.sum(axis=0)) + return importance_loss + + def load_loss_all(self, clean_values, noisy_values, noise_stddev, noisy_top_values): + B, E = clean_values.shape + M = noisy_top_values.shape[-1] + top_values_flat = noisy_top_values.flatten() # B * M + + threshold_positions_if_in = paddle.arange(B) * M + self.out_planes # B + threshold_if_in = paddle.take_along_axis( + top_values_flat, axis=-1, indices=threshold_positions_if_in + ).unsqueeze( + -1 + ) # B, 1 + is_in = noisy_values > threshold_if_in # B, E + threshold_positions_if_out = threshold_positions_if_in - 1 # B + threshold_if_out = paddle.take_along_axis( + top_values_flat, axis=-1, indices=threshold_positions_if_out + ).unsqueeze( + -1 + ) # B, 1 + + prob_if_in = self.intra_cdf( + (clean_values - threshold_if_in) / noise_stddev + ) # B, E + prob_if_out = self.intra_cdf( + (clean_values - threshold_if_out) / noise_stddev + ) # B, E + prob = paddle.where(is_in, prob_if_in, prob_if_out) # B, E + + load_loss = self.cv_squared(prob.sum(axis=0)) + return load_loss + + def forward(self, x, t_map=None, eps=1e-25, dense_routing=False): + assert x.shape[1:-1] == list(self.input_shape)[:-1] + B, T, H, W, C = x.shape + E = self.num_experts + + raw_logits = self.gating(x, t_map) + if self.training: + noise = self.softplus(self.noise_lin(x)) + self.noise_eps + noisy_logits = raw_logits + paddle.randn(shape=raw_logits.shape) * noise + logits = noisy_logits + else: + logits = raw_logits + + assert logits.shape[-1] == self.num_experts + logits = self.softmax(logits) # [B, T, H, W, E] + top_logits, top_indices = logits.topk( + min(self.out_planes + 1, self.num_experts), axis=-1 + ) + top_k_logits = top_logits[:, :, :, :, : self.out_planes] + top_k_indices = top_indices[:, :, :, :, : self.out_planes] + top_k_gates = top_k_logits / ( + top_k_logits.sum(axis=-1, keepdim=True) + eps + ) # normalization + + if dense_routing: + # zeros = paddle.zeros_like(logits) + # zeros.stop_gradient = False + # print(zeros.shape) + # print(top_k_gates.shape, top_k_gates[0, 0, 0, 0]) + # routing_weights = paddle.put_along_axis(zeros, axis=-1, indices=top_k_indices, values=top_k_gates) + # print(routing_weights.shape, routing_weights.stop_gradient) + pass + else: + routing_weights = None + + if self.training: + if self.aux_loss_style == "cell": + # importance_loss = self.importance_loss(routing_weights) + importance_loss = self.importance_loss_cell(logits) + load_loss = self.load_loss_cell( + raw_logits, noisy_logits, noise, top_logits + ) + elif self.aux_loss_style == "all": + importance_loss = self.importance_loss_all( + logits.reshape([B * T * H * W, E]) + ) + load_loss = self.load_loss_all( + raw_logits.reshape([B * T * H * W, E]), + noisy_logits.reshape([B * T * H * W, E]), + noise.reshape([B * T * H * W, E]), + top_logits.reshape([B * T * H * W, -1]), + ) + else: + raise NotImplementedError + loss = ( + self.importance_weight * importance_loss + self.load_weight * load_loss + ) + else: + loss = None + + return routing_weights, top_k_gates, top_k_indices, loss + + +class LinearGatingNet(GatingNet): + def __init__(self, moe_config, input_shape, in_channels): + super().__init__(moe_config, input_shape, in_channels) + assert len(input_shape) == 4 + T, H, W, C = input_shape + + self.lin = nn.Linear( + in_features=in_channels, out_features=self.num_experts, bias_attr=False + ) + + def gating(self, x, t_map=None): + routing_weights = self.lin(x) # [B, T, H, W, E] + return routing_weights + + +class SpatialLatentGatingNet(GatingNet): + def __init__(self, moe_config, input_shape, in_channels): + super().__init__(moe_config, input_shape, in_channels) + assert len(input_shape) == 4 + T, H, W, C = input_shape + + gain = 1.0 + fan = self.out_planes / self.num_experts + bound = gain * math.sqrt(3.0 / fan) + self.routing_weights = paddle.create_parameter( + shape=[H, W, self.num_experts], + dtype="float32", + default_initializer=nn.initializer.Uniform(-bound, bound), + ) + + def gating(self, x, t_map=None): + # assert t_map is not None + routing_weights = self.routing_weights.unsqueeze(0).tile( + [x.shape[0], x.shape[1], 1, 1, 1] + ) # [B, T, H, W, E] + return routing_weights + + +class SpatialLatentLinearGatingNet(GatingNet): + def __init__(self, moe_config, input_shape, in_channels): + super().__init__(moe_config, input_shape, in_channels) + assert len(input_shape) == 4 + T, H, W, C = input_shape + + gain = 1.0 + fan = self.out_planes / self.num_experts + bound = gain * math.sqrt(3.0 / fan) + self.spatial_routing_weights = paddle.create_parameter( + shape=[H, W, self.num_experts], + dtype="float32", + default_initializer=nn.initializer.Uniform(-bound, bound), + ) + self.lin = nn.Linear( + in_features=in_channels, out_features=self.num_experts, bias_attr=False + ) + + self.combine_weight = paddle.create_parameter( + shape=[H, W, self.num_experts, 2], + dtype="float32", + default_initializer=nn.initializer.Uniform(-bound, bound), + ) + + def gating(self, x, t_map=None): + # assert t_map is not None + spatial_routing_weights = self.spatial_routing_weights.tile( + [x.shape[0], x.shape[1], 1, 1, 1] + ) # [B, T, H, W, E] + linear_routing_weights = self.lin(x) # [B, T, H, W, E] + routing_weights = paddle.stack( + [spatial_routing_weights, linear_routing_weights], axis=-1 + ) # [B, T, H, W, E, 2] + combine_weight = self.combine_weight.tile( + [x.shape[0], x.shape[1], 1, 1, 1, 1] + ) # [B, T, H, W, E, 2] + routing_weights = (routing_weights * combine_weight).sum(-1) # [B, T, H, W, E] + return routing_weights + + +class CuboidLatentGatingNet(GatingNet): + def __init__(self, moe_config, input_shape, in_channels): + super().__init__(moe_config, input_shape, in_channels) + assert len(input_shape) == 4 + T, H, W, C = input_shape + + gain = 1.0 + fan = self.out_planes / self.num_experts + bound = gain * math.sqrt(3.0 / fan) + self.routing_weights = paddle.create_parameter( + shape=[T, H, W, self.num_experts], + dtype="float32", + default_initializer=nn.initializer.Uniform(-bound, bound), + ) + + def gating(self, x, t_map=None): + # assert t_map is not None + routing_weights = self.routing_weights.unsqueeze(0).tile( + [x.shape[0], 1, 1, 1, 1] + ) # [B, T, H, W, E] + return routing_weights + + +class CuboidLatentLinearGatingNet(GatingNet): + def __init__(self, moe_config, input_shape, in_channels): + super().__init__(moe_config, input_shape, in_channels) + assert len(input_shape) == 4 + T, H, W, C = input_shape + + gain = 1.0 + fan = self.out_planes / self.num_experts + bound = gain * math.sqrt(3.0 / fan) + self.cuboid_routing_weights = paddle.create_parameter( + shape=[T, H, W, self.num_experts], + dtype="float32", + default_initializer=nn.initializer.Uniform(-bound, bound), + ) + + self.lin = nn.Linear( + in_features=in_channels, out_features=self.num_experts, bias_attr=False + ) + + self.combine_weight = paddle.create_parameter( + shape=[T, H, W, self.num_experts, 2], + dtype="float32", + default_initializer=nn.initializer.Uniform(-bound, bound), + ) + + def gating(self, x, t_map=None): + # assert t_map is not None + cuboid_routing_weights = self.cuboid_routing_weights.unsqueeze(0).tile( + [x.shape[0], 1, 1, 1, 1] + ) # [B, T, H, W, E] + linear_routing_weights = self.lin(x) # [B, T, H, W, E] + routing_weights = paddle.stack( + [cuboid_routing_weights, linear_routing_weights], axis=-1 + ) # [B, T, H, W, E, 2] + combine_weight = self.combine_weight.tile( + [x.shape[0], 1, 1, 1, 1, 1] + ) # [B, T, H, W, E, 2] + routing_weights = (routing_weights * combine_weight).sum(-1) # [B, T, H, W, E] + return routing_weights + + +def aggregate_aux_losses(net): + aux_losses = [] + for module in net.sublayers(): + if hasattr(module, "aux_loss"): + aux_losses.append(module.aux_loss.unsqueeze(0)) + return aux_losses + + +# MoE Routing + + +class SparseDispatcherScatter(object): + def __init__(self, num_experts, gates): + self._gates = gates + self._num_experts = num_experts + sorted_experts, index_sorted_experts = paddle.nonzero(gates).sort( + 0 + ), paddle.nonzero(gates).argsort(0) + _, self._expert_index = sorted_experts.split(1, axis=1) + self._batch_index = paddle.nonzero(gates)[index_sorted_experts[:, 1], 0] + self._part_sizes = (gates > 0).sum(0).tolist() + gates_exp = gates[self._batch_index.flatten()] + self._nonzero_gates = paddle.take_along_axis( + gates_exp, axis=1, indices=self._expert_index + ) + + def dispatch(self, inp): + inp_exp = inp[self._batch_index].squeeze(1) + return paddle.split(inp_exp, self._part_sizes, axis=0) + + def combine(self, expert_out, multiply_by_gates=True): + stitched = paddle.concat(expert_out, 0) + if multiply_by_gates: + stitched = stitched.multiply(self._nonzero_gates) + zeros = paddle.zeros([self._gates.shape[0], expert_out[-1].shape[1]]) + zeros.stop_gradient = False + # combine samples that have been processed by the same k experts + combined = zeros.index_add(0, self._batch_index, stitched.float()) + return combined + + +class SparseDispatcher(object): + def __init__(self, num_experts, top_k_gates, top_k_indices): + self.num_experts = num_experts + self.gates = top_k_gates # [B, K] + self.gate_inds = top_k_indices # [B, K] + E = num_experts + B, K = top_k_gates.shape + self.batch_index_per_expert = paddle.stack( + [ + (top_k_indices == expert_id).sum(-1).astype("bool") + for expert_id in range(E) + ], + axis=0, + ) # [E, B] + self.gates_per_expert = paddle.concat( + [top_k_gates[top_k_indices == expert_id] for expert_id in range(E)] + ) # B * K + self.batch_index_all = paddle.nonzero(self.batch_index_per_expert)[ + :, 1 + ] # B * K + self.expert_size = self.batch_index_per_expert.sum(-1) # [E] + + def dispatch(self, x): + B, C = x.shape + dispatched_x = [ + x[batch_index] for batch_index in self.batch_index_per_expert + ] # E * [B_e, C] + return dispatched_x + + def combine(self, expert_out): + # expert_out: E * [B_e, C] + assert len(expert_out) == self.num_experts + com_res = paddle.concat(expert_out, axis=0) # [B * K, C] + zeros = paddle.zeros([self.gates.shape[0], com_res.shape[1]]) + zeros.stop_gradient = False + combined_res = zeros.index_add( + axis=0, + index=self.batch_index_all, + value=com_res * self.gates_per_expert.unsqueeze(-1), + ) + return combined_res + + +class DenseDispatcher(object): + def __init__(self, num_experts, top_k_gates, top_k_indices): + self.num_experts = num_experts + self.gates = top_k_gates # [B, K] + self.gate_inds = top_k_indices # [B, K] + + def combine(self, expert_out): + # expert_out: [B, E, C] + B, E, C = expert_out.shape + assert E == self.num_experts + selected_out = paddle.take_along_axis( + expert_out, axis=1, indices=self.gate_inds.unsqueeze(-1) + ) # [B, K, C] + combined_res = (selected_out * self.gates.unsqueeze(-1)).sum(1) + return combined_res + + +# RNC + + +class LabelDifference(nn.Layer): + def __init__(self, distance_type="l1"): + super().__init__() + self.distance_type = distance_type + + def forward(self, labels): + # labels: [bs, label_dim] + # output: [bs, bs] + assert labels.ndim == 3 + if self.distance_type == "l1": + return paddle.abs(labels[:, :, None, :] - labels[:, None, :, :]).sum( + axis=-1 + ) + else: + raise ValueError(self.distance_type) + + +class FeatureSimilarity(nn.Layer): + def __init__(self, similarity_type="l2", temperature=2): + super().__init__() + self.similarity_type = similarity_type + self.t = temperature + + def forward(self, features): + # labels: [bs, feat_dim] + # output: [bs, bs] + assert features.ndim == 3 + if self.similarity_type == "l2": + logits = -(features[:, :, None, :] - features[:, None, :, :]).norm( + 2, axis=-1 + ) + logits /= self.t + logits_max = paddle.max(logits, axis=1, keepdim=True) + logits -= logits_max.detach() + return logits + elif self.similarity_type == "cosine": + cos_func = nn.CosineSimilarity(axis=-1) + logits = cos_func(features[:, :, None, :], features[:, None, :, :]) + logits /= self.t + return logits + else: + raise ValueError(self.similarity_type) + + +class RnCLoss(nn.Layer): + def __init__(self, rnc_config): + super().__init__() + + self.rank_mode = rnc_config["rank_imbalance_style"] + self.t = rnc_config["rank_imbalance_temp"] + self.label_diff_fn = LabelDifference(rnc_config["label_difference_style"]) + self.feature_sim_fn = FeatureSimilarity( + rnc_config["feature_similarity_style"], self.t + ) + self.rnc_weight = rnc_config["rank_reg_coeff"] + self.loss_cal_mode = rnc_config["loss_cal_style"] + self.softmax_cri = nn.Softmax(axis=-1) + + def cal_loss(self, features, labels): + + B = features.shape[0] + assert B > 1 + label_diffs = self.label_diff_fn(labels) + logits = self.feature_sim_fn(features) + exp_logits = logits.exp() + n = logits.shape[1] + + # remove diagonal + logits = logits.masked_select( + (1 - paddle.eye(n)).astype("bool").unsqueeze(0).tile([B, 1, 1]) + ).reshape([B, n, n - 1]) + exp_logits = exp_logits.masked_select( + (1 - paddle.eye(n)).astype("bool").unsqueeze(0).tile([B, 1, 1]) + ).reshape([B, n, n - 1]) + label_diffs = label_diffs.masked_select( + (1 - paddle.eye(n)).astype("bool").unsqueeze(0).tile([B, 1, 1]) + ).reshape([B, n, n - 1]) + + if self.loss_cal_mode == "memory-efficient": + loss = 0.0 + for k in range(n - 1): + pos_logits = logits[:, :, k] # [B, n] + pos_label_diffs = label_diffs[:, :, k] # [B, n] + neg_mask = (label_diffs >= pos_label_diffs.unsqueeze(-1)).astype( + "float32" + ) # [B, n, n - 1] + pos_log_probs = pos_logits - paddle.log( + (neg_mask * exp_logits).sum(axis=-1) + ) # [B, n] + loss += -pos_log_probs.sum() + loss /= B * n * (n - 1) + elif self.loss_cal_mode == "computation-efficient": + neg_mask = (label_diffs.unsqueeze(-2) >= label_diffs.unsqueeze(-1)).astype( + "float32" + ) # [B, n, n - 1, n - 1] + pos_log_probs = logits - paddle.log( + (neg_mask * exp_logits.unsqueeze(-2).tile([1, 1, n - 1, 1])).sum( + axis=-1 + ) + ) # [B, n, n - 1] + loss = -pos_log_probs.mean() + else: + raise NotImplementedError + + return loss + + def forward(self, features, labels): + # features: [B, T_o, H, W, C_o] + # labels: [B, T_o, H, W, C_l] + + B, T_o, H, W, C_o = features.shape + _, _, _, _, C_l = labels.shape + + loss = None + if self.rank_mode == "batch": + features = features.reshape([B, -1, C_o]).transpose([1, 0, 2]) + labels = labels.reshape([B, -1, C_l]).transpose([1, 0, 2]) + loss = self.cal_loss(features, labels) + elif self.rank_mode == "batch+T+H+W": + feat = features.transpose([0, 2, 3, 1, 4]).reshape([-1, T_o, C_o]) + label = labels.transpose([0, 2, 3, 1, 4]).reshape([-1, T_o, C_l]) + loss_T = self.cal_loss(feat, label) + + feat = features.transpose([0, 1, 3, 2, 4]).reshape([-1, H, C_o]) + label = labels.transpose([0, 1, 3, 2, 4]).reshape([-1, H, C_l]) + loss_H = self.cal_loss(feat, label) + + feat = features.reshape([-1, W, C_o]) + label = labels.reshape([-1, W, C_l]) + loss_W = self.cal_loss(feat, label) + + feat = features.transpose([1, 2, 3, 0, 4]).reshape([-1, B, C_o]) + label = labels.transpose([1, 2, 3, 0, 4]).reshape([-1, B, C_l]) + loss_batch = self.cal_loss(feat, label) + + loss = loss_T + loss_H + loss_W + loss_batch + else: + raise NotImplementedError + + loss = self.rnc_weight * loss + + return loss diff --git a/examples/smc_reac/ppsci/arch/fno_block.py b/examples/smc_reac/ppsci/arch/fno_block.py new file mode 100644 index 0000000000..df40c36a0e --- /dev/null +++ b/examples/smc_reac/ppsci/arch/fno_block.py @@ -0,0 +1,1269 @@ +import itertools +from typing import Dict +from typing import List +from typing import Optional +from typing import Tuple +from typing import Union + +import omegaconf +import paddle +import paddle.nn.functional as F +from paddle import nn + +from ppsci.utils import initializer +from ppsci.utils import logger + +einsum_symbols = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ" + + +class DomainPadding(nn.Layer): + """Applies domain padding scaled automatically to the input's resolution + + Args: + domain_padding (Union[float, List[float]]): Typically, between zero and one, percentage of padding to use. + padding_mode (str, optional): Whether to pad on both sides, by default + 'one-sided'.Options are 'symmetric' or 'one-sided'。 Defaults to "one-sided". + output_scaling_factor (Union[int, List[int]], optional): Scaling factor for the + output. Defaults to 1. + """ + + def __init__( + self, + domain_padding: Union[float, List[float]], + padding_mode: str = "one-sided", + output_scaling_factor: Union[int, List[int]] = 1, + ): + super().__init__() + self.domain_padding = domain_padding + self.padding_mode = padding_mode.lower() + if output_scaling_factor is None: + output_scaling_factor = 1 + self.output_scaling_factor: Union[int, List[int]] = output_scaling_factor + + # dict(f'{resolution}'=padding) such that padded = F.pad(x, indices) + self._padding = dict() + + # dict(f'{resolution}'=indices_to_unpad) such that unpadded = x[indices] + self._unpad_indices = dict() + + def forward(self, x): + self.pad(x) + + def pad(self, x): + """Take an input and pad it by the desired fraction + + The amount of padding will be automatically scaled with the resolution + """ + resolution = x.shape[2:] + + if isinstance(self.domain_padding, (float, int)): + self.domain_padding = [float(self.domain_padding)] * len(resolution) + + assert len(self.domain_padding) == len(resolution), ( + "domain_padding length must match the number of spatial/time dimensions " + "(excluding batch, ch)" + ) + + output_scaling_factor = self.output_scaling_factor + if not isinstance(self.output_scaling_factor, list): + # if unset by the user, scaling_factor will be 1 be default, + # so `output_scaling_factor` should never be None. + output_scaling_factor: List[float] = validate_scaling_factor( + self.output_scaling_factor, len(resolution), n_layers=None + ) + + try: + padding = self._padding[f"{resolution}"] + return F.pad(x, padding, mode="constant") + except KeyError: + padding = [round(p * r) for (p, r) in zip(self.domain_padding, resolution)] + + output_pad = padding + output_pad = [ + round(i * j) for (i, j) in zip(output_scaling_factor, output_pad) + ] + + # padding is being applied in reverse order + # (so we must reverse the padding list) + padding = padding[::-1] + + # the F.pad(x, padding) function pads the tensor 'x' in reverse order of the "padding" list i.e. the last axis of tensor 'x' will be padded by the amount mention at the first position of the 'padding' vector. The details about F.pad can be found here: + # https://www.paddlepaddle.org.cn/documentation/docs/zh/api/paddle/nn/functional/pad_cn.html + + if self.padding_mode == "symmetric": + # Pad both sides + unpad_list = list() + for p in output_pad: + if p == 0: + padding_end = None + padding_start = None + else: + padding_end = p + padding_start = -p + unpad_list.append(slice(padding_end, padding_start, None)) + + unpad_indices = (Ellipsis,) + tuple( + [slice(p, -p, None) for p in padding] + ) + padding = [i for p in padding for i in (p, p)] + + elif self.padding_mode == "one-sided": + # One-side padding + unpad_list = list() + for p in output_pad: + if p == 0: + padding_start = None + else: + padding_start = -p + unpad_list.append(slice(None, padding_start, None)) + unpad_indices = (Ellipsis,) + tuple(unpad_list) + padding = [i for p in padding for i in (0, p)] + else: + raise ValueError(f"Got self.padding_mode = {self.padding_mode}") + + self._padding[f"{resolution}"] = padding + + padded = F.pad(x, padding, mode="constant") + output_shape = padded.shape[2:] + output_shape = [ + round(i * j) for (i, j) in zip(output_scaling_factor, output_shape) + ] + + self._unpad_indices[f"{[i for i in output_shape]}"] = unpad_indices + + return padded + + def unpad(self, x): + """Remove the padding from padding inputs""" + unpad_indices = self._unpad_indices[f"{x.shape[2:]}"] + + return x[unpad_indices] + + +class SoftGating(nn.Layer): + """Applies soft-gating by weighting the channels of the given input + + Given an input x of size `(batch-size, channels, height, width)`, + this returns `x * w ` + where w is of shape `(1, channels, 1, 1)` + + Args: + in_features (int): The number of input features. + out_features (int, optional): Number of output features. Defaults to None. + n_dim (int, optional): Dimensionality of the input (excluding batch-size and channels). + ``n_dim=2`` corresponds to having Module2D. Defaults to 2. + bias (bool, optional): Whether to use bias. Defaults to False. + """ + + def __init__( + self, in_features, out_features: int = None, n_dim: int = 2, bias: bool = False + ): + super().__init__() + if out_features is not None and in_features != out_features: + raise ValueError( + f"Got in_features = {in_features} and out_features = {out_features}" + "but these two must be the same for soft-gating" + ) + self.in_features = in_features + self.out_features = out_features + + self.weight = self.create_parameter( + shape=(1, self.in_features, *(1,) * n_dim), + default_initializer=nn.initializer.Constant(1.0), + ) + if bias: + self.bias = self.create_parameter( + shape=(1, self.in_features, *(1,) * n_dim), + default_initializer=nn.initializer.Constant(1.0), + ) + else: + self.bias = None + + def forward(self, x): + """Applies soft-gating to a batch of activations""" + if self.bias is not None: + return self.weight * x + self.bias + else: + return self.weight * x + + +def skip_connection( + in_features, + out_features, + n_dim: int = 2, + bias: bool = False, + type: str = "soft-gating", +): + """A wrapper for several types of skip connections. + Returns an nn.Module skip connections, one of {'identity', 'linear', soft-gating'} + + Args: + in_features (int): Number of input features. + out_features (int): Number of output features. + n_dim (int, optional): Dimensionality of the input (excluding batch-size and channels). + ``n_dim=2`` corresponds to having Module2D. . Defaults to 2. + bias (bool, optional): Whether to use a bias. Defaults to False. + type (str, optional): Kind of skip connection to use,{'identity', 'linear', soft-gating'}. + Defaults to "soft-gating". + """ + + if type.lower() == "soft-gating": + return SoftGating( + in_features=in_features, out_features=out_features, bias=bias, n_dim=n_dim + ) + elif type.lower() == "linear": + return getattr(nn, f"Conv{n_dim}D")( + in_channels=in_features, + out_channels=out_features, + kernel_size=1, + bias_attr=bias, + ) + elif type.lower() == "identity": + return nn.Identity() + else: + raise ValueError( + f"Got skip-connection type = {type}, expected one of {'soft-gating', 'linear', 'identity'}." + ) + + +class AdaIN(nn.Layer): + def __init__(self, embed_dim, in_channels, mlp=None, eps=1e-5): + super().__init__() + self.in_channels = in_channels + self.embed_dim = embed_dim + self.eps = eps + + if mlp is None: + mlp = nn.Sequential( + nn.Linear(embed_dim, 512), nn.GELU(), nn.Linear(512, 2 * in_channels) + ) + self.mlp = mlp + + self.embedding = None + + def set_embedding(self, x): + self.embedding = x.reshape( + self.embed_dim, + ) + + def forward(self, x): + assert ( + self.embedding is not None + ), "AdaIN: update embedding before running forward" + + weight, bias = paddle.split( + self.mlp(self.embedding), + self.embedding.shape[0] // self.in_channels, + axis=0, + ) + + return nn.functional.group_norm(x, self.in_channels, self.eps, weight, bias) + + +class MLP(nn.Layer): + """A Multi-Layer Perceptron, with arbitrary number of layers + + Args: + in_channels (int): The number of input channels. + out_channels (int, optional): The number of output channels. Defaults to None. + hidden_channels (int, optional): The number of hidden channels. Defaults to None. + n_layers (int, optional): The number of layers. Defaults to 2. + n_dim (int, optional): The type of convolution,2D or 3D. Defaults to 2. + non_linearity (nn.functional, optional): The activation function. Defaults to F.gelu. + dropout (float, optional): The ratio of dropout. Defaults to 0.0. + """ + + def __init__( + self, + in_channels: int, + out_channels: int = None, + hidden_channels: int = None, + n_layers: int = 2, + n_dim: int = 2, + non_linearity: nn.functional = F.gelu, + dropout: float = 0.0, + **kwargs, + ): + super().__init__() + self.n_layers = n_layers + self.in_channels = in_channels + self.out_channels = in_channels if out_channels is None else out_channels + self.hidden_channels = ( + in_channels if hidden_channels is None else hidden_channels + ) + self.non_linearity = non_linearity + self.dropout = ( + nn.LayerList([nn.Dropout(dropout) for _ in range(n_layers)]) + if dropout > 0.0 + else None + ) + + Conv = getattr(nn, f"Conv{n_dim}D") + self.fcs = nn.LayerList() + for i in range(n_layers): + if i == 0 and i == (n_layers - 1): + self.fcs.append(Conv(self.in_channels, self.out_channels, 1)) + elif i == 0: + self.fcs.append(Conv(self.in_channels, self.hidden_channels, 1)) + elif i == (n_layers - 1): + self.fcs.append(Conv(self.hidden_channels, self.out_channels, 1)) + else: + self.fcs.append(Conv(self.hidden_channels, self.hidden_channels, 1)) + + def forward(self, x): + for i, fc in enumerate(self.fcs): + x = fc(x) + if i < self.n_layers - 1: + x = self.non_linearity(x) + if self.dropout is not None: + x = self.dropout[i](x) + return x + + +def _contract_dense(x, weight, separable=False): + order = len(x.shape) + x_syms = list(einsum_symbols[:order]) + + # in_channels, out_channels, x, y... + weight_syms = list(x_syms[1:]) # no batch-size + + # batch-size, out_channels, x, y... + if separable: + out_syms = [x_syms[0]] + list(weight_syms) + else: + weight_syms.insert(1, einsum_symbols[order]) # outputs + out_syms = list(weight_syms) + out_syms[0] = x_syms[0] + + eq = "".join(x_syms) + "," + "".join(weight_syms) + "->" + "".join(out_syms) + # For the darcy flow, the only einsum is abcd,becd->aecd, where x and weights are shaped [32,32,8,8] + if not isinstance(weight, paddle.Tensor): + weight = paddle.to_tensor(weight) + + return paddle.einsum(eq, x, weight) + + +def _contract_dense_trick(x, weight_real, weight_imag, separable=False): + # the same as above function, but do the complex multiplication manually to avoid the einsum bug in paddle + order = len(x.shape) + # batch-size, in_channels, x, y... + x_syms = list(einsum_symbols[:order]) + + # in_channels, out_channels, x, y... + weight_syms = list(x_syms[1:]) # no batch-size + + # batch-size, out_channels, x, y... + if separable: + out_syms = [x_syms[0]] + list(weight_syms) + else: + weight_syms.insert(1, einsum_symbols[order]) # outputs + out_syms = list(weight_syms) + out_syms[0] = x_syms[0] + + eq = "".join(x_syms) + "," + "".join(weight_syms) + "->" + "".join(out_syms) + + o1_real = paddle.einsum(eq, x.real(), weight_real) - paddle.einsum( + eq, x.imag(), weight_imag + ) + o1_imag = paddle.einsum(eq, x.imag(), weight_real) + paddle.einsum( + eq, x.real(), weight_imag + ) + x = paddle.complex(o1_real, o1_imag) + return x + + +def _contract_dense_separable(x, weight, separable=True): + if not separable: + raise ValueError("This function is only for separable=True") + return x * weight + + +def get_contract_fun( + weight, implementation: str = "reconstructed", separable: bool = False +): + """Generic ND implementation of Fourier Spectral Conv contraction. + + Args: + weight (paddle.tensor): FactorizedTensor. + implementation (str, optional): {'reconstructed', 'factorized'}. + whether to reconstruct the weight and do a forward pass (reconstructed) + or contract directly the factors of the factorized weight with the input (factorized). Defaults to "reconstructed". + separable (bool, optional): Whether to use the separable implementation of contraction. This + arg is only checked when `implementation=reconstructed`. Defaults to False. + + Returns: + function : (x, weight) -> x * weight in Fourier space. + """ + + if implementation == "reconstructed": + if separable: + return _contract_dense_separable + else: + return _contract_dense_trick + elif implementation == "factorized": + if isinstance(weight, paddle.Tensor): + return _contract_dense_trick + + else: + raise ValueError( + f'Got implementation={implementation}, expected "reconstructed" or "factorized"' + ) + + +Number = Union[float, int] + + +def validate_scaling_factor( + scaling_factor: Union[None, Number, List[Number], List[List[Number]]], + n_dim: int, + n_layers: Optional[int] = None, +) -> Union[None, List[float], List[List[float]]]: + """Validates the format and dimensionality of the scaling factor. + + Args: + scaling_factor (Union[None, Number, List[Number], List[List[Number]]]): The + scaling factor. + n_dim (int): The required number of dimensions for expanding `scaling_factor`. + n_layers (Optional[int], optional): The number of layers for the returned + nested list. If None, return a single list (rather than a list of lists) + with `factor` repeated `dim` times. Defaults to None. + """ + + if scaling_factor is None: + return None + if isinstance(scaling_factor, (float, int)): + if n_layers is None: + return [float(scaling_factor)] * n_dim + + return [[float(scaling_factor)] * n_dim] * n_layers + if ( + isinstance(scaling_factor, list) + and len(scaling_factor) > 0 + and all([isinstance(s, (float, int)) for s in scaling_factor]) + ): + + return [[float(s)] * n_dim for s in scaling_factor] + + if ( + isinstance(scaling_factor, list) + and len(scaling_factor) > 0 + and all( + [isinstance(s, (omegaconf.listconfig.ListConfig)) for s in scaling_factor] + ) + ): + s_sub_pass = True + for s in scaling_factor: + if all([isinstance(s_sub, (int, float)) for s_sub in s]): + pass + else: + s_sub_pass = False + if s_sub_pass: + return scaling_factor + + return None + + +def resample(x, res_scale, axis, output_shape=None): + """A module for generic n-dimentional interpolation (Fourier resampling). + + Args: + x (paddle.Tensor): Input activation of size (batch_size, channels, d1, ..., dN). + res_scale (optional[int,tuple]): Scaling factor along each of the dimensions in + 'axis' parameter. If res_scale is scaler, then isotropic scaling is performed. + axis (int): Axis or dimensions along which interpolation will be performed. + output_shape (optional[None ,tuple[int]]): The output shape. Defaults to None. + """ + + if isinstance(res_scale, (float, int)): + if axis is None: + axis = list(range(2, x.ndim)) + res_scale = [res_scale] * len(axis) + elif isinstance(axis, int): + axis = [axis] + res_scale = [res_scale] + else: + res_scale = [res_scale] * len(axis) + else: + assert len(res_scale) == len(axis), "length of res_scale and axis are not same" + + old_size = x.shape[-len(axis) :] + if output_shape is None: + new_size = tuple([int(round(s * r)) for (s, r) in zip(old_size, res_scale)]) + else: + new_size = output_shape + + if len(axis) == 1: + return F.interpolate(x, size=new_size[0], mode="linear", align_corners=True) + if len(axis) == 2: + return F.interpolate(x, size=new_size, mode="bicubic", align_corners=True) + + X = paddle.fft.rfftn(x.astype("float32"), norm="forward", axes=axis) + + new_fft_size = list(new_size) + new_fft_size[-1] = new_fft_size[-1] // 2 + 1 # Redundant last coefficient + new_fft_size_c = [min(i, j) for (i, j) in zip(new_fft_size, X.shape[-len(axis) :])] + out_fft = paddle.zeros( + [x.shape[0], x.shape[1], *new_fft_size], dtype=paddle.complex64 + ) + + mode_indexing = [((None, m // 2), (-m // 2, None)) for m in new_fft_size_c[:-1]] + [ + ((None, new_fft_size_c[-1]),) + ] + for i, boundaries in enumerate(itertools.product(*mode_indexing)): + + idx_tuple = [slice(None), slice(None)] + [slice(*b) for b in boundaries] + + out_fft[idx_tuple] = X[idx_tuple] + y = paddle.fft.irfftn(out_fft, s=new_size, norm="forward", axes=axis) + + return y + + +class FactorizedTensor(nn.Layer): + def __init__(self, shape, init_scale): + super().__init__() + self.shape = shape + self.init_scale = init_scale + self.real = self.create_parameter( + shape=shape, + ) + self.real = initializer.normal_(self.real, 0, init_scale) + self.imag = self.create_parameter(shape=shape) + self.imag = initializer.normal_(self.imag, 0, init_scale) + + def __repr__(self): + return f"FactorizedTensor(shape={self.shape})" + + @property + def data(self): + return paddle.complex(self.real, self.imag) + + +class FactorizedSpectralConv(nn.Layer): + """Generic N-Dimensional Fourier Neural Operator + + Args: + in_channels (int): Number of input channels. + out_channels (int): Number of output channels. + n_modes (Tuple[int, ...]): Number of modes to use for contraction in Fourier domain during training. + .. warning:: + We take care of the redundancy in the Fourier modes, therefore, for an input + of size I_1, ..., I_N, please provide modes M_K that are I_1 < M_K <= I_N + We will automatically keep the right amount of modes: specifically, for the + last mode only, if you specify M_N modes we will use M_N // 2 + 1 modes + as the real FFT is redundant along that last dimension. + + .. note:: + Provided modes should be even integers. odd numbers will be rounded to the closest even number. + This can be updated dynamically during training. + max_n_modes (int, optional): * If not None, **maximum** number of modes to keep + in Fourier Layer, along each dim + The number of modes (`n_modes`) cannot be increased beyond that. + * If None, all the n_modes are used. Defaults to None. + bias (bool, optional): Whether to use bias in the layers. Defaults to True. + n_layers (int, optional): Number of Fourier Layers. Defaults to 1. + separable (bool, optional): Whether to use separable Fourier Conv. Defaults to False. + output_scaling_factor (Optional[Union[Number, List[Number]]], optional): Scaling factor for the + output. Defaults to None. + rank (float, optional): Rank of the tensor factorization of the Fourier weights. Defaults to 0.5. + factorization (str, optional): Tensor factorization of the parameters weight to use. + * If None, a dense tensor parametrizes the Spectral convolutions + * Otherwise, the specified tensor factorization is used. Defaults to None. + implementation (str, optional): If factorization is not None, forward mode to use. + * `reconstructed` : the full weight tensor is reconstructed from the + factorization and used for the forward pass + * `factorized` : the input is directly contracted with the factors of + the decomposition. Defaults to "reconstructed". + joint_factorization (bool, optional): Whether all the Fourier Layers should be parametrized by a + single tensor. Defaults to False. + init_std (str, optional): The std to use for the init. Defaults to "auto". + fft_norm (str, optional):The normalization mode for the FFT. Defaults to "backward". + """ + + def __init__( + self, + in_channels: int, + out_channels: int, + n_modes: Tuple[int, ...], + max_n_modes: int = None, + bias: bool = True, + n_layers: int = 1, + separable: bool = False, + output_scaling_factor: Optional[Union[Number, List[Number]]] = None, + rank: float = 0.5, + factorization: str = None, + implementation: str = "reconstructed", + joint_factorization: bool = False, + init_std: str = "auto", + fft_norm: str = "backward", + ): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.joint_factorization = joint_factorization + self.n_modes = n_modes + + self.order = len(self.n_modes) + if max_n_modes is None: + max_n_modes = self.n_modes + elif isinstance(max_n_modes, int): + max_n_modes = [max_n_modes] + self.max_n_modes = max_n_modes + self.rank = rank + self.factorization = factorization + self.n_layers = n_layers + self.implementation = implementation + + self.output_scaling_factor: Union[ + None, List[List[float]] + ] = validate_scaling_factor(output_scaling_factor, self.order, n_layers) + + if init_std == "auto": + init_std = (2 / (in_channels + out_channels)) ** 0.5 + else: + init_std = init_std + self.fft_norm = fft_norm + if factorization is None: + factorization = "Dense" + if not factorization.lower().startswith("complex"): + factorization = f"Complex{factorization}" + if separable: + if in_channels != out_channels: + raise ValueError( + f"To use separable Fourier Conv, in_channels must be equal to out_channels, but got in_channels={in_channels} and out_channels={out_channels}" + ) + weight_shape = (in_channels, *max_n_modes) + else: + weight_shape = (in_channels, out_channels, *max_n_modes) + self.separable = separable + if joint_factorization: + self.weight = paddle.create_parameter( + shape=(n_layers, *weight_shape), + dtype="float32", + ) + self.weight = initializer.normal_(self.weight, 0, init_std) + else: + self.weight = nn.LayerList( + [ + FactorizedTensor(weight_shape, init_scale=init_std) + for _ in range(n_layers) + ] + ) + + self._contract = get_contract_fun( + self.weight[0].data, implementation=implementation, separable=separable + ) + if bias: + shape = (n_layers, self.out_channels) + (1,) * self.order + init_bias = init_std * paddle.randn(shape) + self.bias = paddle.create_parameter( + shape=shape, + dtype=(init_bias.dtype), + default_initializer=nn.initializer.Assign(init_bias), + ) + self.bias.stop_gradient = False + else: + self.bias = None + + @property + def n_modes(self): + return self._n_modes + + @n_modes.setter + def n_modes(self, n_modes): + if isinstance(n_modes, int): # Should happen for 1D FNO only + n_modes = [n_modes] + else: + n_modes = list(n_modes) + # The last mode has a redundacy as we use real FFT + # As a design choice we do the operation here to avoid users dealing with the +1 + n_modes[-1] = n_modes[-1] // 2 + 1 + self._n_modes = n_modes + + def transform(self, x, layer_index=0, output_shape=None): + in_shape = list(x.shape[2:]) + + if self.output_scaling_factor is not None and output_shape is None: + out_shape = tuple( + [ + round(s * r) + for (s, r) in zip(in_shape, self.output_scaling_factor[layer_index]) + ] + ) + elif output_shape is not None: + out_shape = output_shape + else: + out_shape = in_shape + if in_shape == out_shape: + return x + else: + return resample( + x, + 1.0, + list(range(2, x.ndim)), + output_shape=out_shape, + ) + + def forward( + self, + x: paddle.Tensor, + indices: int = 0, + output_shape: Optional[Tuple[int]] = None, + ): + batchsize, channels, *mode_sizes = x.shape + fft_size = list(mode_sizes) + fft_size[-1] = fft_size[-1] // 2 + 1 + fft_dims = list(range(-self.order, 0)) + + x = paddle.fft.rfftn(x=x, norm=self.fft_norm, axes=fft_dims) + + if self.order > 1: + x = paddle.fft.fftshift(x=x, axes=fft_dims[:-1]) + + out_fft = paddle.zeros( + shape=[batchsize, self.out_channels, *fft_size], dtype=paddle.complex64 + ) + + starts = [ + (max_modes - min(size, n_mode)) + for size, n_mode, max_modes in zip(fft_size, self.n_modes, self.max_n_modes) + ] + slices_w = [slice(None), slice(None)] + slices_w += [ + (slice(start // 2, -start // 2) if start else slice(start, None)) + for start in starts[:-1] + ] + slices_w += [slice(None, -starts[-1]) if starts[-1] else slice(None)] + + w_real = self.weight[indices].real[ + slices_w[0], slices_w[1], slices_w[2], slices_w[3] + ] + w_imag = self.weight[indices].imag[ + slices_w[0], slices_w[1], slices_w[2], slices_w[3] + ] + + starts = [ + (size - min(size, n_mode)) + for (size, n_mode) in zip(list(x.shape[2:]), list(w_real.shape[2:])) + ] + slices_x = [slice(None), slice(None)] # Batch_size, channels + slices_x += [ + slice(start // 2, -start // 2) if start else slice(start, None) + for start in starts[:-1] + ] + slices_x += [ + slice(None, -starts[-1]) if starts[-1] else slice(None) + ] # The last mode already has redundant half removed + idx_tuple = slices_x + if len(idx_tuple) == 4: + out_fft[ + idx_tuple[0], idx_tuple[1], idx_tuple[2], idx_tuple[3] + ] = self._contract( + x[idx_tuple[0], idx_tuple[1], idx_tuple[2], idx_tuple[3]], + w_real, + w_imag, + separable=self.separable, + ) + elif len(idx_tuple) == 3: + out_fft[idx_tuple[0], idx_tuple[1], idx_tuple[2]] = self._contract( + x[idx_tuple[0], idx_tuple[1], idx_tuple[2]], + w_real, + w_imag, + separable=self.separable, + ) + else: + raise ValueError("Not implemented") + + if self.output_scaling_factor is not None and output_shape is None: + mode_sizes = tuple( + [ + round(s * r) + for (s, r) in zip(mode_sizes, self.output_scaling_factor[indices]) + ] + ) + + if output_shape is not None: + mode_sizes = output_shape + + if self.order > 1: + out_fft = paddle.fft.fftshift(x=out_fft, axes=fft_dims[:-1]) + + x = paddle.fft.irfftn( + x=out_fft, s=mode_sizes, axes=fft_dims, norm=self.fft_norm + ) + if self.bias is not None: + x = x + self.bias[indices, ...] + return x + + +class FactorizedSpectralConv1d(FactorizedSpectralConv): + """1D Spectral Conv + + This is provided for reference only, + see :class:`FactorizedSpectralConv` for the preferred, general implementation + """ + + def forward(self, x, indices=0): + batchsize, channels, width = x.shape + + x = paddle.fft.rfft(x, norm=self.fft_norm) + + out_fft = paddle.zeros( + shape=[batchsize, self.out_channels, width // 2 + 1], dtype=paddle.complex64 + ) + + slices = ( + slice(None), # Equivalent to: [:, + slice(None), # ............... :, + slice(None, self.n_modes[0]), # :half_n_modes[0]] + ) + + w_real = self.weight[indices].real[slices[0], slices[1], slices[2]] + w_imag = self.weight[indices].imag[slices[0], slices[1], slices[2]] + + out_fft[slices[0], slices[1], slices[2]] = self._contract( + x[slices[0], slices[1], slices[2]], + w_real, + w_imag, + separable=self.separable, + ) + + if self.output_scaling_factor is not None: + width = round(width * self.output_scaling_factor[0]) + + x = paddle.fft.irfft(out_fft, n=width, norm=self.fft_norm) + + if self.bias is not None: + x = x + self.bias[indices, ...] + + return x + + +class FactorizedSpectralConv2d(FactorizedSpectralConv): + """2D Spectral Conv. + + This is provided for reference only, + see :class:`FactorizedSpectralConv` for the preferred, general implementation + """ + + def forward(self, x, indices=0): + batchsize, channels, height, width = x.shape + + x = paddle.fft.rfft2(x.float(), norm=self.fft_norm, axes=(-2, -1)) + + # The output will be of size (batch_size, self.out_channels, + # x.size(-2), x.size(-1)//2 + 1) + out_fft = paddle.zeros( + shape=[batchsize, self.out_channels, height, width // 2 + 1], + dtype=paddle.complex64, + ) + + slices0 = ( + slice(None), # Equivalent to: [:, + slice(None), # ............... :, + slice(self.n_modes[0] // 2), # :half_n_modes[0], + slice(self.n_modes[1]), # :half_n_modes[1]] + ) + slices1 = ( + slice(None), # Equivalent to: [:, + slice(None), # ...................... :, + slice(-self.n_modes[0] // 2, None), # -half_n_modes[0]:, + slice(self.n_modes[1]), # ...... :half_n_modes[1]] + ) + logger.message( + f"2D: {x[slices0].shape=}, {self._get_weight(indices)[slices0].shape=}, {self._get_weight(indices).shape=}" + ) + + w_real = self.weight[indices].real[ + slices1[0], slices1[1], slices1[2], slices1[3] + ] + w_imag = self.weight[indices].imag[ + slices1[0], slices1[1], slices1[2], slices1[3] + ] + + """Upper block (truncate high frequencies).""" + out_fft[slices0[0], slices0[1], slices0[2], slices0[3]] = self._contract( + x[slices0[0], slices0[1], slices0[2], slices0[3]], + w_real, + w_imag, + separable=self.separable, + ) + + w_real = self.weight[indices].real[ + slices0[0], slices0[1], slices0[2], slices0[3] + ] + w_imag = self.weight[indices].imag[ + slices0[0], slices0[1], slices0[2], slices0[3] + ] + + """Lower block""" + out_fft[slices1[0], slices1[1], slices1[2], slices1[3]] = self._contract( + x[slices1[0], slices1[1], slices1[2], slices1[3]], + w_real, + w_imag, + separable=self.separable, + ) + + if self.output_scaling_factor is not None: + width = round(width * self.output_scaling_factor[indices][0]) + height = round(height * self.output_scaling_factor[indices][1]) + + x = paddle.fft.irfft2( + out_fft, s=(height, width), axes=(-2, -1), norm=self.fft_norm + ) + + if self.bias is not None: + x = x + self.bias[indices, ...] + + return x + + +class FactorizedSpectralConv3d(FactorizedSpectralConv): + """3D Spectral Conv. + + This is provided for reference only, + see :class:`FactorizedSpectralConv` for the preferred, general implementation + """ + + def forward(self, x, indices=0): + batchsize, channels, height, width, depth = x.shape + + x = paddle.fft.rfftn(x.float(), norm=self.fft_norm, axes=[-3, -2, -1]) + + out_fft = paddle.zeros( + shape=[batchsize, self.out_channels, height, width, depth // 2 + 1], + dtype=paddle.complex64, + ) + + slices0 = ( + slice(None), # Equivalent to: [:, + slice(None), # ............... :, + slice(self.n_modes[0] // 2), # :half_n_modes[0], + slice(self.n_modes[1] // 2), # :half_n_modes[1], + slice(self.n_modes[2]), # :half_n_modes[2]] + ) + slices1 = ( + slice(None), # Equivalent to: [:, + slice(None), # ...................... :, + slice(self.n_modes[0] // 2), # ...... :half_n_modes[0], + slice(-self.n_modes[1] // 2, None), # -half_n_modes[1]:, + slice(self.n_modes[2]), # ...... :half_n_modes[0]] + ) + slices2 = ( + slice(None), # Equivalent to: [:, + slice(None), # ...................... :, + slice(-self.n_modes[0] // 2, None), # -half_n_modes[0]:, + slice(self.n_modes[1] // 2), # ...... :half_n_modes[1], + slice(self.n_modes[2]), # ...... :half_n_modes[2]] + ) + slices3 = ( + slice(None), # Equivalent to: [:, + slice(None), # ...................... :, + slice(-self.n_modes[0] // 2, None), # -half_n_modes[0], + slice(-self.n_modes[1] // 2, None), # -half_n_modes[1], + slice(self.n_modes[2]), # ...... :half_n_modes[2]] + ) + + w_real = self.weight[indices].real[ + slices3[0], slices3[1], slices3[2], slices3[3], slices3[4] + ] + w_imag = self.weight[indices].imag[ + slices3[0], slices3[1], slices3[2], slices3[3], slices3[4] + ] + + """Upper block -- truncate high frequencies.""" + out_fft[ + slices0[0], slices0[1], slices0[2], slices0[3], slices0[4] + ] = self._contract( + x[slices0[0], slices0[1], slices0[2], slices0[3], slices0[4]], + w_real, + w_imag, + separable=self.separable, + ) + + w_real = self.weight[indices].real[ + slices2[0], slices2[1], slices2[2], slices2[3], slices2[4] + ] + w_imag = self.weight[indices].imag[ + slices2[0], slices2[1], slices2[2], slices2[3], slices2[4] + ] + """Low-pass filter for indices 2 & 4, and high-pass filter for index 3.""" + out_fft[ + slices1[0], slices1[1], slices1[2], slices1[3], slices1[4] + ] = self._contract( + x[slices1[0], slices1[1], slices1[2], slices1[3], slices1[4]], + w_real, + w_imag, + separable=self.separable, + ) + + w_real = self.weight[indices].real[ + slices1[0], slices1[1], slices1[2], slices1[3], slices1[4] + ] + w_imag = self.weight[indices].imag[ + slices1[0], slices1[1], slices1[2], slices1[3], slices1[4] + ] + """Low-pass filter for indices 3 & 4, and high-pass filter for index 2.""" + out_fft[ + slices2[0], slices2[1], slices2[2], slices2[3], slices2[4] + ] = self._contract( + x[slices2[0], slices2[1], slices2[2], slices2[3], slices2[4]], + w_real, + w_imag, + separable=self.separable, + ) + + w_real = self.weight[indices].real[ + slices0[0], slices0[1], slices0[2], slices0[3], slices0[4] + ] + w_imag = self.weight[indices].imag[ + slices0[0], slices0[1], slices0[2], slices0[3], slices0[4] + ] + """Lower block -- low-cut filter in indices 2 & 3 + and high-cut filter in index 4.""" + out_fft[ + slices3[0], slices3[1], slices3[2], slices3[3], slices3[4] + ] = self._contract( + x[slices3[0], slices3[1], slices3[2], slices3[3], slices3[4]], + w_real, + w_imag, + separable=self.separable, + ) + + if self.output_scaling_factor is not None: + width = round(width * self.output_scaling_factor[0]) + height = round(height * self.output_scaling_factor[1]) + depth = round(depth * self.output_scaling_factor[2]) + + x = paddle.fft.irfftn( + out_fft, s=(height, width, depth), axes=[-3, -2, -1], norm=self.fft_norm + ) + + if self.bias is not None: + x = x + self.bias[indices, ...] + return x + + +class FNOBlocks(nn.Layer): + def __init__( + self, + in_channels: int, + out_channels: int, + n_modes: Tuple[int, ...], + output_scaling_factor: Optional[Union[Number, List[Number]]] = None, + n_layers: int = 1, + max_n_modes: int = None, + use_mlp: bool = False, + mlp: Optional[Dict[str, float]] = None, + non_linearity: nn.functional = F.gelu, + stabilizer: str = None, + norm: str = None, + ada_in_features: Optional[int] = None, + preactivation: bool = False, + fno_skip: str = "linear", + mlp_skip: str = "soft-gating", + separable: bool = False, + factorization: str = None, + rank: float = 1.0, + SpectralConv: FactorizedSpectralConv = FactorizedSpectralConv, + joint_factorization: bool = False, + implementation: str = "factorized", + fft_norm: str = "forward", + **kwargs, + ): + super().__init__() + if isinstance(n_modes, int): + n_modes = [n_modes] + self._n_modes = n_modes + self.n_dim = len(n_modes) + + self.max_n_modes = max_n_modes + self.in_channels = in_channels + self.out_channels = out_channels + self.n_layers = n_layers + self.joint_factorization = joint_factorization + self.non_linearity = non_linearity + self.rank = rank + self.factorization = factorization + self.fno_skip = fno_skip + self.mlp_skip = mlp_skip + self.use_mlp = use_mlp + self.fft_norm = fft_norm + self.implementation = implementation + self.separable = separable + self.preactivation = preactivation + self.ada_in_features = ada_in_features + self.stabilizer = stabilizer + self.norm = norm + + self.convs = SpectralConv( + self.in_channels, + self.out_channels, + self.n_modes, + output_scaling_factor=output_scaling_factor, + max_n_modes=max_n_modes, + rank=rank, + implementation=implementation, + separable=separable, + factorization=factorization, + joint_factorization=joint_factorization, + n_layers=n_layers, + ) + + self.fno_skips = nn.LayerList( + [ + skip_connection( + self.in_channels, + self.out_channels, + type=fno_skip, + n_dim=self.n_dim, + ) + for _ in range(n_layers) + ] + ) + + if use_mlp: + self.mlp = nn.LayerList( + [ + MLP( + in_channels=self.out_channels, + hidden_channels=int( + round(self.out_channels * mlp["expansion"]) + ), + dropout=mlp["dropout"], + n_dim=self.n_dim, + ) + for _ in range(n_layers) + ] + ) + self.mlp_skips = nn.LayerList( + [ + skip_connection( + self.in_channels, + self.out_channels, + type=mlp_skip, + n_dim=self.n_dim, + ) + for _ in range(n_layers) + ] + ) + else: + self.mlp = None + + # Each block will have 2 norms if we also use an MLP + self.n_norms = 1 if self.mlp is None else 2 + if norm is None: + self.norm = None + elif norm == "instance_norm": + self.norm = nn.LayerList( + [ + getattr(nn, f"InstanceNorm{self.n_dim}d")( + num_features=self.out_channels + ) + for _ in range(n_layers * self.n_norms) + ] + ) + elif norm == "group_norm": + self.norm = nn.LayerList( + [ + nn.GroupNorm(num_groups=1, num_channels=self.out_channels) + for _ in range(n_layers * self.n_norms) + ] + ) + elif norm == "ada_in": + self.norm = nn.LayerList( + [ + AdaIN(ada_in_features, out_channels) + for _ in range(n_layers * self.n_norms) + ] + ) + else: + raise ValueError( + f"Got {norm} but expected None or one of [instance_norm, group_norm, layer_norm]" + ) + + def forward(self, x, index=0, output_shape=None): + if self.preactivation: + return self.forward_with_preactivation(x, index, output_shape=output_shape) + else: + return self.forward_with_postactivation(x, index, output_shape=output_shape) + + def forward_with_postactivation(self, x, index=0, output_shape=None): + x_skip_fno = self.fno_skips[index](x) + x_skip_fno = self.convs.transform(x_skip_fno, index, output_shape=output_shape) + if self.mlp is not None: + x_skip_mlp = self.mlp_skips[index](x) + x_skip_mlp = self.convs.transform( + x_skip_mlp, index, output_shape=output_shape + ) + if self.stabilizer == "tanh": + x = paddle.tanh(x) + + x_fno = self.convs(x, index, output_shape=output_shape) + if self.norm is not None: + x_fno = self.norm[self.n_norms * index](x_fno) + + x = x_fno + x_skip_fno + + if (self.mlp is not None) or (index < (self.n_layers - 1)): + x = self.non_linearity(x) + + if self.mlp is not None: + x = self.mlp[index](x) + x_skip_mlp + + if self.norm is not None: + x = self.norm[self.n_norms * index + 1](x) + + if index < (self.n_layers - 1): + x = self.non_linearity(x) + + return x + + def forward_with_preactivation(self, x, index=0, output_shape=None): + # Apply non-linear activation (and norm) + # before this block's convolution/forward pass: + x = self.non_linearity(x) + + if self.norm is not None: + x = self.norm[self.n_norms * index](x) + + x_skip_fno = self.fno_skips[index](x) + x_skip_fno = self.convs.transform(x_skip_fno, index, output_shape=output_shape) + + if self.mlp is not None: + x_skip_mlp = self.mlp_skips[index](x) + x_skip_mlp = self.convs.transform( + x_skip_mlp, index, output_shape=output_shape + ) + + if self.stabilizer == "tanh": + x = paddle.tanh(x) + + x_fno = self.convs(x, index, output_shape=output_shape) + x = x_fno + x_skip_fno + + if self.mlp is not None: + if index < (self.n_layers - 1): + x = self.non_linearity(x) + + if self.norm is not None: + x = self.norm[self.n_norms * index + 1](x) + + x = self.mlp[index](x) + x_skip_mlp + + return x + + @property + def n_modes(self): + return self._n_modes + + @n_modes.setter + def n_modes(self, n_modes): + if isinstance(n_modes, int): # Should happen for 1D FNO only + n_modes = [n_modes] + else: + n_modes = list(n_modes) + # The last mode has a redundacy as we use real FFT + # As a design choice we do the operation here to avoid users dealing with the +1 + n_modes[-1] = n_modes[-1] // 2 + 1 + self._n_modes = n_modes diff --git a/examples/smc_reac/ppsci/arch/gan.py b/examples/smc_reac/ppsci/arch/gan.py new file mode 100644 index 0000000000..a673b8e440 --- /dev/null +++ b/examples/smc_reac/ppsci/arch/gan.py @@ -0,0 +1,400 @@ +# 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. + +from __future__ import annotations + +from typing import Dict +from typing import List +from typing import Tuple + +import paddle +import paddle.nn as nn + +from ppsci.arch import activation as act_mod +from ppsci.arch import base + + +class Conv2DBlock(nn.Layer): + def __init__( + self, + in_channel, + out_channel, + kernel_size, + stride, + use_bn, + act, + mean, + std, + value, + ): + super().__init__() + weight_attr = paddle.ParamAttr( + initializer=nn.initializer.Normal(mean=mean, std=std) + ) + bias_attr = paddle.ParamAttr(initializer=nn.initializer.Constant(value=value)) + self.conv_2d = nn.Conv2D( + in_channel, + out_channel, + kernel_size, + stride, + padding="SAME", + weight_attr=weight_attr, + bias_attr=bias_attr, + ) + self.bn = nn.BatchNorm2D(out_channel) if use_bn else None + self.act = act_mod.get_activation(act) if act else None + + def forward(self, x): + y = x + y = self.conv_2d(y) + if self.bn: + y = self.bn(y) + if self.act: + y = self.act(y) + return y + + +class VariantResBlock(nn.Layer): + def __init__( + self, + in_channel, + out_channels, + kernel_sizes, + strides, + use_bns, + acts, + mean, + std, + value, + ): + super().__init__() + self.conv_2d_0 = Conv2DBlock( + in_channel=in_channel, + out_channel=out_channels[0], + kernel_size=kernel_sizes[0], + stride=strides[0], + use_bn=use_bns[0], + act=acts[0], + mean=mean, + std=std, + value=value, + ) + self.conv_2d_1 = Conv2DBlock( + in_channel=out_channels[0], + out_channel=out_channels[1], + kernel_size=kernel_sizes[1], + stride=strides[1], + use_bn=use_bns[1], + act=acts[1], + mean=mean, + std=std, + value=value, + ) + + self.conv_2d_2 = Conv2DBlock( + in_channel=in_channel, + out_channel=out_channels[2], + kernel_size=kernel_sizes[2], + stride=strides[2], + use_bn=use_bns[2], + act=acts[2], + mean=mean, + std=std, + value=value, + ) + + self.act = act_mod.get_activation("relu") + + def forward(self, x): + y = x + y = self.conv_2d_0(y) + y = self.conv_2d_1(y) + short = self.conv_2d_2(x) + y = paddle.add(y, short) + y = self.act(y) + return y + + +class FCBlock(nn.Layer): + def __init__(self, in_channel, act, mean, std, value): + super().__init__() + self.flatten = nn.Flatten() + weight_attr = paddle.ParamAttr( + initializer=nn.initializer.Normal(mean=mean, std=std) + ) + bias_attr = paddle.ParamAttr(initializer=nn.initializer.Constant(value=value)) + self.linear = nn.Linear( + in_channel, + 1, + weight_attr=weight_attr, + bias_attr=bias_attr, + ) + self.act = act_mod.get_activation(act) if act else None + + def forward(self, x): + y = x + y = self.flatten(y) + y = self.linear(y) + if self.act: + y = self.act(y) + return y + + +class Generator(base.Arch): + """Generator Net of GAN. Attention, the net using a kind of variant of ResBlock which is + unique to "tempoGAN" example but not an open source network. + + Args: + input_keys (Tuple[str, ...]): Name of input keys, such as ("input1", "input2"). + output_keys (Tuple[str, ...]): Name of output keys, such as ("output1", "output2"). + in_channel (int): Number of input channels of the first conv layer. + out_channels_tuple (Tuple[Tuple[int, ...], ...]): Number of output channels of all conv layers, + such as [[out_res0_conv0, out_res0_conv1], [out_res1_conv0, out_res1_conv1]] + kernel_sizes_tuple (Tuple[Tuple[int, ...], ...]): Number of kernel_size of all conv layers, + such as [[kernel_size_res0_conv0, kernel_size_res0_conv1], [kernel_size_res1_conv0, kernel_size_res1_conv1]] + strides_tuple (Tuple[Tuple[int, ...], ...]): Number of stride of all conv layers, + such as [[stride_res0_conv0, stride_res0_conv1], [stride_res1_conv0, stride_res1_conv1]] + use_bns_tuple (Tuple[Tuple[bool, ...], ...]): Whether to use the batch_norm layer after each conv layer. + acts_tuple (Tuple[Tuple[str, ...], ...]): Whether to use the activation layer after each conv layer. If so, witch activation to use, + such as [[act_res0_conv0, act_res0_conv1], [act_res1_conv0, act_res1_conv1]] + + Examples: + >>> import ppsci + >>> in_channel = 1 + >>> rb_channel0 = (2, 8, 8) + >>> rb_channel1 = (128, 128, 128) + >>> rb_channel2 = (32, 8, 8) + >>> rb_channel3 = (2, 1, 1) + >>> out_channels_tuple = (rb_channel0, rb_channel1, rb_channel2, rb_channel3) + >>> kernel_sizes_tuple = (((5, 5), ) * 2 + ((1, 1), ), ) * 4 + >>> strides_tuple = ((1, 1, 1), ) * 4 + >>> use_bns_tuple = ((True, True, True), ) * 3 + ((False, False, False), ) + >>> acts_tuple = (("relu", None, None), ) * 4 + >>> model = ppsci.arch.Generator(("in",), ("out",), in_channel, out_channels_tuple, kernel_sizes_tuple, strides_tuple, use_bns_tuple, acts_tuple) + >>> batch_size = 4 + >>> height = 64 + >>> width = 64 + >>> input_data = paddle.randn([batch_size, in_channel, height, width]) + >>> input_dict = {'in': input_data} + >>> output_data = model(input_dict) + >>> print(output_data['out'].shape) + [4, 1, 64, 64] + """ + + def __init__( + self, + input_keys: Tuple[str, ...], + output_keys: Tuple[str, ...], + in_channel: int, + out_channels_tuple: Tuple[Tuple[int, ...], ...], + kernel_sizes_tuple: Tuple[Tuple[int, ...], ...], + strides_tuple: Tuple[Tuple[int, ...], ...], + use_bns_tuple: Tuple[Tuple[bool, ...], ...], + acts_tuple: Tuple[Tuple[str, ...], ...], + ): + super().__init__() + self.input_keys = input_keys + self.output_keys = output_keys + self.in_channel = in_channel + self.out_channels_tuple = out_channels_tuple + self.kernel_sizes_tuple = kernel_sizes_tuple + self.strides_tuple = strides_tuple + self.use_bns_tuple = use_bns_tuple + self.acts_tuple = acts_tuple + + self.init_blocks() + + def init_blocks(self): + blocks_list = [] + for i in range(len(self.out_channels_tuple)): + in_channel = ( + self.in_channel if i == 0 else self.out_channels_tuple[i - 1][-1] + ) + blocks_list.append( + VariantResBlock( + in_channel=in_channel, + out_channels=self.out_channels_tuple[i], + kernel_sizes=self.kernel_sizes_tuple[i], + strides=self.strides_tuple[i], + use_bns=self.use_bns_tuple[i], + acts=self.acts_tuple[i], + mean=0.0, + std=0.04, + value=0.1, + ) + ) + self.blocks = nn.LayerList(blocks_list) + + def forward_tensor(self, x): + y = x + for block in self.blocks: + y = block(y) + return y + + def forward(self, x): + if self._input_transform is not None: + x = self._input_transform(x) + + y = self.concat_to_tensor(x, self.input_keys, axis=-1) + y = self.forward_tensor(y) + y = self.split_to_dict(y, self.output_keys, axis=-1) + + if self._output_transform is not None: + y = self._output_transform(x, y) + return y + + +class Discriminator(base.Arch): + """Discriminator Net of GAN. + + Args: + input_keys (Tuple[str, ...]): Name of input keys, such as ("input1", "input2"). + output_keys (Tuple[str, ...]): Name of output keys, such as ("output1", "output2"). + in_channel (int): Number of input channels of the first conv layer. + out_channels (Tuple[int, ...]): Number of output channels of all conv layers, + such as (out_conv0, out_conv1, out_conv2). + fc_channel (int): Number of input features of linear layer. Number of output features of the layer + is set to 1 in this Net to construct a fully_connected layer. + kernel_sizes (Tuple[int, ...]): Number of kernel_size of all conv layers, + such as (kernel_size_conv0, kernel_size_conv1, kernel_size_conv2). + strides (Tuple[int, ...]): Number of stride of all conv layers, + such as (stride_conv0, stride_conv1, stride_conv2). + use_bns (Tuple[bool, ...]): Whether to use the batch_norm layer after each conv layer. + acts (Tuple[str, ...]): Whether to use the activation layer after each conv layer. If so, witch activation to use, + such as (act_conv0, act_conv1, act_conv2). + + Examples: + >>> import ppsci + >>> in_channel = 2 + >>> in_channel_tempo = 3 + >>> out_channels = (32, 64, 128, 256) + >>> fc_channel = 65536 + >>> kernel_sizes = ((4, 4), (4, 4), (4, 4), (4, 4)) + >>> strides = (2, 2, 2, 1) + >>> use_bns = (False, True, True, True) + >>> acts = ("leaky_relu", "leaky_relu", "leaky_relu", "leaky_relu", None) + >>> output_keys_disc = ("out_1", "out_2", "out_3", "out_4", "out_5", "out_6", "out_7", "out_8", "out_9", "out_10") + >>> model = ppsci.arch.Discriminator(("in_1","in_2"), output_keys_disc, in_channel, out_channels, fc_channel, kernel_sizes, strides, use_bns, acts) + >>> input_data = [paddle.to_tensor(paddle.randn([1, in_channel, 128, 128])),paddle.to_tensor(paddle.randn([1, in_channel, 128, 128]))] + >>> input_dict = {"in_1": input_data[0],"in_2": input_data[1]} + >>> out_dict = model(input_dict) + >>> for k, v in out_dict.items(): + ... print(k, v.shape) + out_1 [1, 32, 64, 64] + out_2 [1, 64, 32, 32] + out_3 [1, 128, 16, 16] + out_4 [1, 256, 16, 16] + out_5 [1, 1] + out_6 [1, 32, 64, 64] + out_7 [1, 64, 32, 32] + out_8 [1, 128, 16, 16] + out_9 [1, 256, 16, 16] + out_10 [1, 1] + """ + + def __init__( + self, + input_keys: Tuple[str, ...], + output_keys: Tuple[str, ...], + in_channel: int, + out_channels: Tuple[int, ...], + fc_channel: int, + kernel_sizes: Tuple[int, ...], + strides: Tuple[int, ...], + use_bns: Tuple[bool, ...], + acts: Tuple[str, ...], + ): + super().__init__() + self.input_keys = input_keys + self.output_keys = output_keys + self.in_channel = in_channel + self.out_channels = out_channels + self.fc_channel = fc_channel + self.kernel_sizes = kernel_sizes + self.strides = strides + self.use_bns = use_bns + self.acts = acts + + self.init_layers() + + def init_layers(self): + layers_list = [] + for i in range(len(self.out_channels)): + in_channel = self.in_channel if i == 0 else self.out_channels[i - 1] + layers_list.append( + Conv2DBlock( + in_channel=in_channel, + out_channel=self.out_channels[i], + kernel_size=self.kernel_sizes[i], + stride=self.strides[i], + use_bn=self.use_bns[i], + act=self.acts[i], + mean=0.0, + std=0.04, + value=0.1, + ) + ) + + layers_list.append( + FCBlock(self.fc_channel, self.acts[4], mean=0.0, std=0.04, value=0.1) + ) + self.layers = nn.LayerList(layers_list) + + def forward_tensor(self, x): + y = x + y_list = [] + for layer in self.layers: + y = layer(y) + y_list.append(y) + return y_list # y_conv1, y_conv2, y_conv3, y_conv4, y_fc(y_out) + + def forward(self, x): + if self._input_transform is not None: + x = self._input_transform(x) + + y_list = [] + # y1_conv1, y1_conv2, y1_conv3, y1_conv4, y1_fc, y2_conv1, y2_conv2, y2_conv3, y2_conv4, y2_fc + for k in x: + y_list.extend(self.forward_tensor(x[k])) + + y = self.split_to_dict(y_list, self.output_keys) + + if self._output_transform is not None: + y = self._output_transform(x, y) + + return y + + @staticmethod + def split_to_dict( + data_list: List[paddle.Tensor], keys: Tuple[str, ...] + ) -> Dict[str, paddle.Tensor]: + """Overwrite of split_to_dict() method belongs to Class base.Arch. + + Reason for overwriting is there is no concat_to_tensor() method called in "tempoGAN" example. + That is because input in "tempoGAN" example is not in a regular format, but a format like: + { + "input1": paddle.concat([in1, in2], axis=1), + "input2": paddle.concat([in1, in3], axis=1), + } + + Args: + data_list (List[paddle.Tensor]): The data to be split. It should be a list of tensor(s), but not a paddle.Tensor. + keys (Tuple[str, ...]): Keys of outputs. + + Returns: + Dict[str, paddle.Tensor]: Dict with split data. + """ + if len(keys) == 1: + return {keys[0]: data_list[0]} + return {key: data_list[i] for i, key in enumerate(keys)} diff --git a/examples/smc_reac/ppsci/arch/geofno.py b/examples/smc_reac/ppsci/arch/geofno.py new file mode 100644 index 0000000000..10fe8d08b5 --- /dev/null +++ b/examples/smc_reac/ppsci/arch/geofno.py @@ -0,0 +1,205 @@ +# Copyright (c) 2024 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. + +import paddle +import paddle.nn as nn +import paddle.nn.functional as F +import paddle.nn.initializer as Initializer + + +class SpectralConv1d(nn.Layer): + """ + 1D Fourier layer. It does FFT, linear transform, and Inverse FFT. + """ + + def __init__(self, in_channels, out_channels, modes1): + super(SpectralConv1d, self).__init__() + + self.in_channels = in_channels + self.out_channels = out_channels + # Number of Fourier modes to multiply, at most floor(N/2) + 1 + self.modes1 = modes1 + self.scale = 1 / (in_channels * out_channels) + + real = paddle.rand(shape=[in_channels, out_channels, modes1]) + real.stop_gradient = False + img = paddle.rand(shape=[in_channels, out_channels, modes1]) + img.stop_gradient = False + self.weights1_real = self.create_parameter( + [in_channels, out_channels, self.modes1], + attr=Initializer.Assign(self.scale * real), + ) + self.weights1_imag = self.create_parameter( + [in_channels, out_channels, self.modes1], + attr=Initializer.Assign(self.scale * img), + ) + + def compl_mul1d(self, op1, op2_real, op2_imag): + # (batch, in_channel, x ), (in_channel, out_channel, x) -> (batch, out_channel, x) + eq = "bix,iox->box" + op1_real = op1.real() + op1_imag = op1.imag() + result_real = paddle.unsqueeze( + paddle.einsum(eq, op1_real, op2_real) + - paddle.einsum(eq, op1_imag, op2_imag), + axis=-1, + ) + result_imag = paddle.unsqueeze( + paddle.einsum(eq, op1_real, op2_imag) + + paddle.einsum(eq, op1_imag, op2_real), + axis=-1, + ) + result_complex = paddle.as_complex( + paddle.concat([result_real, result_imag], axis=-1) + ) + return result_complex + + def forward(self, x, output_size=None): + batchsize = x.shape[0] + # Compute Fourier coefficients up to factor of e^(- something constant) + x_ft = paddle.fft.rfft(x) + + # Multiply relevant Fourier modes + out_ft_real = paddle.zeros( + [batchsize, self.out_channels, x.shape[-1] // 2 + 1], dtype="float32" + ) + out_ft_img = paddle.zeros( + [batchsize, self.out_channels, x.shape[-1] // 2 + 1], dtype="float32" + ) + out_ft = paddle.complex(out_ft_real, out_ft_img) + + out_ft[:, :, : self.modes1] = self.compl_mul1d( + x_ft[:, :, : self.modes1], self.weights1_real, self.weights1_imag + ) + + # Return to physical space + if output_size is None: + x = paddle.fft.irfft(out_ft, n=x.shape[-1]) + else: + x = paddle.fft.irfft(out_ft, n=output_size) + + return x + + +class FNO1d(nn.Layer): + """The overall network. It contains 4 layers of the Fourier layer. + 1. Lift the input to the desire channel dimension by self.fc0 . + 2. 4 layers of the integral operators u' = (W + K)(u). + W defined by self.w; K defined by self.conv . + 3. Project from the channel space to the output space by self.fc1 and self.fc2 . + + Args: + input_key (Tuple[str, ...], optional): Key to get the input tensor from the dict. Defaults to ("intput",). + output_key (Tuple[str, ...], optional): Key to save the output tensor into the dict. Defaults to ("output",). + modes (int, optional, optional): Number of Fourier modes to compute, it should be the same as + that in fft part of the code below. Defaults to 64. + width (int, optional, optional): Number of channels in each Fourier layer. Defaults to 64. + padding (int, optional, optional): How many zeros to pad to the input Tensor. Defaults to 100. + input_channel (int, optional, optional): Number of channels of the input tensor. Defaults to 2. + output_np (int, optional, optional): Number of points to sample the solution. Defaults to 2001. + + Examples: + >>> model = ppsci.arch.FNO1d() + >>> input_data = paddle.randn([100, 2001, 2]) + >>> input_dict = {"input": input_data} + >>> out_dict = model(input_dict) + >>> for k, v in out_dict.items(): + ... print(k, v.shape) + output [100, 1] + """ + + def __init__( + self, + input_key=("input",), + output_key=("output",), + modes=64, + width=64, + padding=100, + input_channel=2, + output_np=2001, + ): + super().__init__() + self.input_keys = input_key + self.output_keys = output_key + + self.output_np = output_np + self.modes1 = modes + self.width = width + self.padding = padding + self.fc0 = nn.Linear(input_channel, self.width) + + self.conv0 = SpectralConv1d(self.width, self.width, self.modes1) + self.conv1 = SpectralConv1d(self.width, self.width, self.modes1) + self.conv2 = SpectralConv1d(self.width, self.width, self.modes1) + self.conv3 = SpectralConv1d(self.width, self.width, self.modes1) + self.conv4 = SpectralConv1d(self.width, self.width, self.modes1) + + self.w0 = nn.Conv1D(self.width, self.width, 1) + self.w1 = nn.Conv1D(self.width, self.width, 1) + self.w2 = nn.Conv1D(self.width, self.width, 1) + self.w3 = nn.Conv1D(self.width, self.width, 1) + + self.fc1 = nn.Linear(self.width, 128) + self.fc2 = nn.Linear(128, 1) + + def _functional_pad(self, x, pad, mode="constant", value=0.0, data_format="NCL"): + if len(x.shape) * 2 == len(pad) and mode == "constant": + pad = ( + paddle.to_tensor(pad, dtype="float32") + .reshape((-1, 2)) + .flip([0]) + .flatten() + .tolist() + ) + return F.pad(x, pad, mode, value, data_format) + + def forward(self, x): + x = x[self.input_keys[0]] + # Dict + x = self.fc0(x) + x = paddle.transpose(x, perm=[0, 2, 1]) + # pad the domain if input is non-periodic + x = self._functional_pad(x, [0, self.padding]) + + x1 = self.conv0(x) + x2 = self.w0(x) + x = x1 + x2 + x = F.gelu(x=x, approximate=False) + + x1 = self.conv1(x) + x2 = self.w1(x) + x = x1 + x2 + x = F.gelu(x, approximate=False) + + x1 = self.conv2(x) + x2 = self.w2(x) + x = x1 + x2 + x = F.gelu(x, approximate=False) + + x1 = self.conv3(x) + x2 = self.w3(x) + x = x1 + x2 + x = F.gelu(x, approximate=False) + + x = x[..., : -self.padding] + x1 = self.conv4(x, self.output_np) + x2 = F.interpolate(x, size=[self.output_np], mode="linear", align_corners=True) + x = x1 + x2 + # Change the x-dimension to (batch, channel, 2001) + x = x.transpose(perm=[0, 2, 1]) + x = self.fc1(x) + x = F.gelu(x, approximate=False) + x = self.fc2(x) + + return {self.output_keys[0]: x} diff --git a/examples/smc_reac/ppsci/arch/graphcast.py b/examples/smc_reac/ppsci/arch/graphcast.py new file mode 100644 index 0000000000..79a1c0aeae --- /dev/null +++ b/examples/smc_reac/ppsci/arch/graphcast.py @@ -0,0 +1,492 @@ +# Copyright (c) 2024 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. + +from typing import TYPE_CHECKING +from typing import Dict +from typing import Tuple + +import paddle +import paddle.nn as nn + +from ppsci.arch import base + +if TYPE_CHECKING: + import ppsci.data.dataset.atmospheric_dataset as atmospheric_dataset + + +class ResidualConnection(nn.Layer): + def __init__(self, fn): + super().__init__() + self.fn = fn + + def forward(self, inputs): + return inputs + self.fn(inputs) + + +class GraphCastMLP(nn.Layer): + def __init__( + self, in_features, out_features, latent_features=None, layer_norm=True + ): + super().__init__() + + if latent_features is None: + latent_features = out_features + + self.mlp = nn.Sequential( + nn.Linear(in_features, latent_features, bias_attr=True), + nn.Silu(), + nn.Linear(latent_features, out_features, bias_attr=True), + ) + self.layer_norm = layer_norm + if layer_norm: + self.layer_norm = nn.LayerNorm(out_features) + + def forward(self, feat): + if self.layer_norm: + out = self.layer_norm(self.mlp(feat)) + else: + out = self.mlp(feat) + return out + + +class GraphCastGNN(nn.Layer): + def __init__( + self, + grid_node_num: int, + grid_node_emb_dim: int, + mesh_node_num: int, + mesh_node_emb_dim: int, + mesh_edge_emb_dim: int, + grid2mesh_edge_emb_dim: int, + mesh2grid_edge_emb_dim: int, + src_type: str = "mesh", + dst_type: str = "mesh", + ): + super().__init__() + + self.src = src_type + self.dst = dst_type + self.grid_node_num = grid_node_num + self.mesh_node_num = mesh_node_num + self.edge_in_dim = grid_node_emb_dim + mesh_node_emb_dim + + if src_type == "mesh" and dst_type == "mesh": + self.edge_in_dim += mesh_edge_emb_dim + self.edge_out_dim = mesh_edge_emb_dim + self.node_in_dim = mesh_node_emb_dim + mesh_edge_emb_dim + self.node_out_dim = mesh_node_emb_dim + elif src_type == "grid" and dst_type == "mesh": + self.edge_in_dim += grid2mesh_edge_emb_dim + self.edge_out_dim = grid2mesh_edge_emb_dim + self.node_in_dim = mesh_node_emb_dim + grid2mesh_edge_emb_dim + self.node_out_dim = mesh_node_emb_dim + elif src_type == "mesh" and dst_type == "grid": + self.edge_in_dim += mesh2grid_edge_emb_dim + self.edge_out_dim = mesh2grid_edge_emb_dim + self.node_in_dim = grid_node_emb_dim + mesh2grid_edge_emb_dim + self.node_out_dim = grid_node_emb_dim + else: + raise ValueError + + self.edge_layer = GraphCastMLP(self.edge_in_dim, self.edge_out_dim) + self.node_layer = GraphCastMLP(self.node_in_dim, self.node_out_dim) + + def forward(self, graph: "atmospheric_dataset.GraphGridMesh"): + if self.src == "mesh" and self.dst == "mesh": + edge_feats = graph.mesh_edge_feat + src_node_feats = graph.mesh_node_feat + dst_node_feats = graph.mesh_node_feat + src_idx = graph.mesh2mesh_src_index + dst_idx = graph.mesh2mesh_dst_index + dst_node_num = self.mesh_node_num + elif self.src == "grid" and self.dst == "mesh": + edge_feats = graph.grid2mesh_edge_feat + src_node_feats = graph.grid_node_feat + dst_node_feats = graph.mesh_node_feat + src_idx = graph.grid2mesh_src_index + dst_idx = graph.grid2mesh_dst_index + dst_node_num = self.mesh_node_num + elif self.src == "mesh" and self.dst == "grid": + edge_feats = graph.mesh2grid_edge_feat + src_node_feats = graph.mesh_node_feat + dst_node_feats = graph.grid_node_feat + src_idx = graph.mesh2grid_src_index + dst_idx = graph.mesh2grid_dst_index + dst_node_num = self.grid_node_num + + # update edge features + edge_feats_concat = paddle.concat( + [ + edge_feats, + paddle.gather(src_node_feats, src_idx), + paddle.gather(dst_node_feats, dst_idx), + ], + axis=-1, + ) + edge_feats_out = self.edge_layer(edge_feats_concat) + + _, batch_dim, _ = edge_feats_out.shape + + # update node features + edge_feats_scatter = paddle.zeros([dst_node_num, batch_dim, self.edge_out_dim]) + node_feats_concat = paddle.concat( + [ + dst_node_feats, + paddle.scatter( + edge_feats_scatter, dst_idx, edge_feats_out, overwrite=False + ), + ], + axis=-1, + ) + node_feats_out = self.node_layer(node_feats_concat) + + if self.src == "mesh" and self.dst == "mesh": + graph.mesh_edge_feat += edge_feats_out + graph.mesh_node_feat += node_feats_out + elif self.src == "grid" and self.dst == "mesh": + graph.grid2mesh_edge_feat += edge_feats_out + graph.mesh_node_feat += node_feats_out + elif self.src == "mesh" and self.dst == "grid": + graph.mesh2grid_edge_feat += edge_feats_out + graph.grid_node_feat += node_feats_out + + return graph + + +class GraphCastEmbedding(nn.Layer): + def __init__( + self, + grid_node_dim: int, + grid_node_emb_dim: int, + mesh_node_dim: int, + mesh_node_emb_dim: int, + mesh_edge_dim: int, + mesh_edge_emb_dim: int, + grid2mesh_edge_dim: int, + grid2mesh_edge_emb_dim: int, + mesh2grid_edge_dim: int, + mesh2grid_edge_emb_dim: int, + ): + super().__init__() + + self.grid_node_embedding = GraphCastMLP(grid_node_dim, grid_node_emb_dim) + self.mesh_node_embedding = GraphCastMLP(mesh_node_dim, mesh_node_emb_dim) + self.mesh_edge_embedding = GraphCastMLP(mesh_edge_dim, mesh_edge_emb_dim) + self.grid2mesh_edge_embedding = GraphCastMLP( + grid2mesh_edge_dim, grid2mesh_edge_emb_dim + ) + self.mesh2grid_edge_embedding = GraphCastMLP( + mesh2grid_edge_dim, mesh2grid_edge_emb_dim + ) + + def forward(self, graph: "atmospheric_dataset.GraphGridMesh"): + grid_node_emb = self.grid_node_embedding(graph.grid_node_feat) + mesh_node_emb = self.mesh_node_embedding(graph.mesh_node_feat) + mesh_edge_emb = self.mesh_edge_embedding(graph.mesh_edge_feat) + grid2mesh_edge_emb = self.grid2mesh_edge_embedding(graph.grid2mesh_edge_feat) + mesh2grid_edge_emb = self.mesh2grid_edge_embedding(graph.mesh2grid_edge_feat) + + graph.grid_node_feat = grid_node_emb + graph.mesh_node_feat = mesh_node_emb + graph.mesh_edge_feat = mesh_edge_emb + graph.grid2mesh_edge_feat = grid2mesh_edge_emb + graph.mesh2grid_edge_feat = mesh2grid_edge_emb + + return graph + + +class GraphCastGrid2Mesh(nn.Layer): + def __init__( + self, + grid_node_num: int, + grid_node_emb_dim: int, + mesh_node_num: int, + mesh_node_emb_dim: int, + mesh_edge_emb_dim: int, + grid2mesh_edge_emb_dim: int, + mesh2grid_edge_emb_dim: int, + ): + super().__init__() + self.grid2mesh_gnn = GraphCastGNN( + grid_node_num=grid_node_num, + grid_node_emb_dim=grid_node_emb_dim, + mesh_node_num=mesh_node_num, + mesh_node_emb_dim=mesh_node_emb_dim, + mesh_edge_emb_dim=mesh_edge_emb_dim, + grid2mesh_edge_emb_dim=grid2mesh_edge_emb_dim, + mesh2grid_edge_emb_dim=mesh2grid_edge_emb_dim, + src_type="grid", + dst_type="mesh", + ) + self.grid_node_layer = ResidualConnection( + GraphCastMLP(grid_node_emb_dim, grid_node_emb_dim) + ) + + def forward(self, graph: "atmospheric_dataset.GraphGridMesh"): + graph = self.grid2mesh_gnn(graph) + graph.grid_node_feat = self.grid_node_layer(graph.grid_node_feat) + return graph + + +class GraphCastMesh2Grid(nn.Layer): + def __init__( + self, + grid_node_num: int, + grid_node_emb_dim: int, + mesh_node_num: int, + mesh_node_emb_dim: int, + mesh_edge_emb_dim: int, + grid2mesh_edge_emb_dim: int, + mesh2grid_edge_emb_dim: int, + ): + super().__init__() + self.mesh2grid_gnn = GraphCastGNN( + grid_node_num=grid_node_num, + grid_node_emb_dim=grid_node_emb_dim, + mesh_node_num=mesh_node_num, + mesh_node_emb_dim=mesh_node_emb_dim, + mesh_edge_emb_dim=mesh_edge_emb_dim, + grid2mesh_edge_emb_dim=grid2mesh_edge_emb_dim, + mesh2grid_edge_emb_dim=mesh2grid_edge_emb_dim, + src_type="mesh", + dst_type="grid", + ) + self.mesh_node_layer = ResidualConnection( + GraphCastMLP(mesh_node_emb_dim, mesh_node_emb_dim) + ) + + def forward(self, graph: "atmospheric_dataset.GraphGridMesh"): + graph = self.mesh2grid_gnn(graph) + graph.mesh_node_feat = self.mesh_node_layer(graph.mesh_node_feat) + return graph + + +class GraphCastEncoder(nn.Layer): + def __init__( + self, + grid_node_num: int, + grid_node_dim: int, + grid_node_emb_dim: int, + mesh_node_num: int, + mesh_node_dim: int, + mesh_node_emb_dim: int, + mesh_edge_dim: int, + mesh_edge_emb_dim: int, + grid2mesh_edge_dim: int, + grid2mesh_edge_emb_dim: int, + mesh2grid_edge_dim: int, + mesh2grid_edge_emb_dim: int, + ): + super().__init__() + self.embedding = GraphCastEmbedding( + grid_node_dim=grid_node_dim, + grid_node_emb_dim=grid_node_emb_dim, + mesh_node_dim=mesh_node_dim, + mesh_node_emb_dim=mesh_node_emb_dim, + mesh_edge_dim=mesh_edge_dim, + mesh_edge_emb_dim=mesh_edge_emb_dim, + grid2mesh_edge_dim=grid2mesh_edge_dim, + grid2mesh_edge_emb_dim=grid2mesh_edge_emb_dim, + mesh2grid_edge_dim=mesh2grid_edge_dim, + mesh2grid_edge_emb_dim=mesh2grid_edge_emb_dim, + ) + self.grid2mesh_gnn = GraphCastGrid2Mesh( + grid_node_num=grid_node_num, + grid_node_emb_dim=grid_node_emb_dim, + mesh_node_num=mesh_node_num, + mesh_node_emb_dim=mesh_node_emb_dim, + mesh_edge_emb_dim=mesh_edge_emb_dim, + grid2mesh_edge_emb_dim=grid2mesh_edge_emb_dim, + mesh2grid_edge_emb_dim=mesh2grid_edge_emb_dim, + ) + + def forward(self, graph: "atmospheric_dataset.GraphGridMesh"): + graph = self.embedding(graph) + graph = self.grid2mesh_gnn(graph) + return graph + + +class GraphCastDecoder(nn.Layer): + def __init__( + self, + grid_node_num: int, + grid_node_emb_dim: int, + mesh_node_num: int, + mesh_node_emb_dim: int, + mesh_edge_emb_dim: int, + grid2mesh_edge_emb_dim: int, + mesh2grid_edge_emb_dim: int, + node_output_dim: int, + ): + super().__init__() + self.mesh2grid_gnn = GraphCastMesh2Grid( + grid_node_num=grid_node_num, + grid_node_emb_dim=grid_node_emb_dim, + mesh_node_num=mesh_node_num, + mesh_node_emb_dim=mesh_node_emb_dim, + mesh_edge_emb_dim=mesh_edge_emb_dim, + grid2mesh_edge_emb_dim=grid2mesh_edge_emb_dim, + mesh2grid_edge_emb_dim=mesh2grid_edge_emb_dim, + ) + self.grid_node_layer = GraphCastMLP( + grid_node_emb_dim, + node_output_dim, + latent_features=grid_node_emb_dim, + layer_norm=False, + ) + + def forward(self, graph: "atmospheric_dataset.GraphGridMesh"): + graph = self.mesh2grid_gnn(graph) + graph.grid_node_feat = self.grid_node_layer(graph.grid_node_feat) + return graph + + +class GraphCastProcessor(nn.Layer): + def __init__( + self, + grid_node_num: int, + grid_node_emb_dim: int, + mesh_node_num: int, + mesh_node_emb_dim: int, + mesh_edge_emb_dim: int, + grid2mesh_edge_emb_dim: int, + mesh2grid_edge_emb_dim: int, + gnn_msg_steps: int, + ): + super().__init__() + + self.processor = nn.Sequential() + for idx in range(gnn_msg_steps): + self.processor.add_sublayer( + f"{idx}", + GraphCastGNN( + grid_node_num=grid_node_num, + grid_node_emb_dim=grid_node_emb_dim, + mesh_node_num=mesh_node_num, + mesh_node_emb_dim=mesh_node_emb_dim, + mesh_edge_emb_dim=mesh_edge_emb_dim, + grid2mesh_edge_emb_dim=grid2mesh_edge_emb_dim, + mesh2grid_edge_emb_dim=mesh2grid_edge_emb_dim, + src_type="mesh", + dst_type="mesh", + ), + ) + + def forward(self, graph: "atmospheric_dataset.GraphGridMesh"): + graph = self.processor(graph) + return graph + + +class GraphCastNet(base.Arch): + """GraphCast Network + + Args: + input_keys (Tuple[str, ...]): Name of input keys. + output_keys (Tuple[str, ...]): Name of output keys. + grid_node_num (int): Number of grid nodes. + grid_node_dim (int): Dimension of grid nodes. + grid_node_emb_dim (int): Dimension of emdding grid nodes. + mesh_node_num (int): Number of mesh nodes. + mesh_node_dim (int): Dimension of mesh nodes. + mesh_node_emb_dim (int): Dimension of emdding mesh nodes. + mesh_edge_dim (int): Dimension of mesh edges. + mesh_edge_emb_dim (int): Dimension of emdding mesh edges. + grid2mesh_edge_dim (int): Dimension of mesh edges in Grid2Mesh GNN. + grid2mesh_edge_emb_dim (int): Dimension of emdding mesh edges in Grid2Mesh GNN. + mesh2grid_edge_dim (int): Dimension of mesh edges in Mesh2Grid GNN. + mesh2grid_edge_emb_dim (int): Dimension of emdding mesh edges in Mesh2Grid GNN. + gnn_msg_steps (int): Step of gnn messages. + node_output_dim (int): Dimension of output nodes. + """ + + def __init__( + self, + input_keys: Tuple[str, ...], + output_keys: Tuple[str, ...], + grid_node_num: int, + grid_node_dim: int, + grid_node_emb_dim: int, + mesh_node_num: int, + mesh_node_dim: int, + mesh_node_emb_dim: int, + mesh_edge_dim: int, + mesh_edge_emb_dim: int, + grid2mesh_edge_dim: int, + grid2mesh_edge_emb_dim: int, + mesh2grid_edge_dim: int, + mesh2grid_edge_emb_dim: int, + gnn_msg_steps: int, + node_output_dim: int, + ): + super().__init__() + self.input_keys = input_keys + self.output_keys = output_keys + self.graphcast = nn.Sequential( + ( + "encoder", + GraphCastEncoder( + grid_node_num=grid_node_num, + grid_node_dim=grid_node_dim, + grid_node_emb_dim=grid_node_emb_dim, + mesh_node_num=mesh_node_num, + mesh_node_dim=mesh_node_dim, + mesh_node_emb_dim=mesh_node_emb_dim, + mesh_edge_dim=mesh_edge_dim, + mesh_edge_emb_dim=mesh_edge_emb_dim, + grid2mesh_edge_dim=grid2mesh_edge_dim, + grid2mesh_edge_emb_dim=grid2mesh_edge_emb_dim, + mesh2grid_edge_dim=mesh2grid_edge_dim, + mesh2grid_edge_emb_dim=mesh2grid_edge_emb_dim, + ), + ), + ( + "processor", + GraphCastProcessor( + grid_node_num=grid_node_num, + grid_node_emb_dim=grid_node_emb_dim, + mesh_node_num=mesh_node_num, + mesh_node_emb_dim=mesh_node_emb_dim, + mesh_edge_emb_dim=mesh_edge_emb_dim, + grid2mesh_edge_emb_dim=grid2mesh_edge_emb_dim, + mesh2grid_edge_emb_dim=mesh2grid_edge_emb_dim, + gnn_msg_steps=gnn_msg_steps, + ), + ), + ( + "decoder", + GraphCastDecoder( + grid_node_num=grid_node_num, + grid_node_emb_dim=grid_node_emb_dim, + mesh_node_num=mesh_node_num, + mesh_node_emb_dim=mesh_node_emb_dim, + mesh_edge_emb_dim=mesh_edge_emb_dim, + grid2mesh_edge_emb_dim=grid2mesh_edge_emb_dim, + mesh2grid_edge_emb_dim=mesh2grid_edge_emb_dim, + node_output_dim=node_output_dim, + ), + ), + ) + + def forward( + self, x: Dict[str, "atmospheric_dataset.GraphGridMesh"] + ) -> Dict[str, paddle.Tensor]: + if self._input_transform is not None: + x = self._input_transform(x) + + graph = x[self.input_keys[0]] + y = self.graphcast(graph) + + if self._output_transform is not None: + y = self._output_transform(x, y) + return {self.output_keys[0]: y} diff --git a/examples/smc_reac/ppsci/arch/he_deeponets.py b/examples/smc_reac/ppsci/arch/he_deeponets.py new file mode 100644 index 0000000000..811da0d1b1 --- /dev/null +++ b/examples/smc_reac/ppsci/arch/he_deeponets.py @@ -0,0 +1,197 @@ +# 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. + +from __future__ import annotations + +from typing import Tuple +from typing import Union + +import paddle +import paddle.nn as nn + +from ppsci.arch import activation as act_mod +from ppsci.arch import base +from ppsci.arch import mlp + + +class HEDeepONets(base.Arch): + """Physical information deep operator networks. + + Args: + heat_input_keys (Tuple[str, ...]): Name of input data for heat boundary. + cold_input_keys (Tuple[str, ...]): Name of input data for cold boundary. + trunk_input_keys (Tuple[str, ...]): Name of input data for trunk net. + output_keys (Tuple[str, ...]): Output name of predicted temperature. + heat_num_loc (int): Number of sampled input data for heat boundary. + cold_num_loc (int): Number of sampled input data for cold boundary. + num_features (int): Number of features extracted from heat boundary, same for cold boundary and trunk net. + branch_num_layers (int): Number of hidden layers of branch net. + trunk_num_layers (int): Number of hidden layers of trunk net. + branch_hidden_size (Union[int, Tuple[int, ...]]): Number of hidden size of branch net. + An integer for all layers, or list of integer specify each layer's size. + trunk_hidden_size (Union[int, Tuple[int, ...]]): Number of hidden size of trunk net. + An integer for all layers, or list of integer specify each layer's size. + branch_skip_connection (bool, optional): Whether to use skip connection for branch net. Defaults to False. + trunk_skip_connection (bool, optional): Whether to use skip connection for trunk net. Defaults to False. + branch_activation (str, optional): Name of activation function for branch net. Defaults to "tanh". + trunk_activation (str, optional): Name of activation function for trunk net. Defaults to "tanh". + branch_weight_norm (bool, optional): Whether to apply weight norm on parameter(s) for branch net. Defaults to False. + trunk_weight_norm (bool, optional): Whether to apply weight norm on parameter(s) for trunk net. Defaults to False. + use_bias (bool, optional): Whether to add bias on predicted G(u)(y). Defaults to True. + + Examples: + >>> import ppsci + >>> model = ppsci.arch.HEDeepONets( + ... ('qm_h',), + ... ('qm_c',), + ... ("x",'t'), + ... ("T_h",'T_c','T_w'), + ... 1, + ... 1, + ... 100, + ... 9, + ... 6, + ... 256, + ... 128, + ... branch_activation="swish", + ... trunk_activation="swish", + ... use_bias=True, + ... ) + """ + + def __init__( + self, + heat_input_keys: Tuple[str, ...], + cold_input_keys: Tuple[str, ...], + trunk_input_keys: Tuple[str, ...], + output_keys: Tuple[str, ...], + heat_num_loc: int, + cold_num_loc: int, + num_features: int, + branch_num_layers: int, + trunk_num_layers: int, + branch_hidden_size: Union[int, Tuple[int, ...]], + trunk_hidden_size: Union[int, Tuple[int, ...]], + branch_skip_connection: bool = False, + trunk_skip_connection: bool = False, + branch_activation: str = "tanh", + trunk_activation: str = "tanh", + branch_weight_norm: bool = False, + trunk_weight_norm: bool = False, + use_bias: bool = True, + ): + super().__init__() + self.trunk_input_keys = trunk_input_keys + self.heat_input_keys = heat_input_keys + self.cold_input_keys = cold_input_keys + self.input_keys = ( + self.trunk_input_keys + self.heat_input_keys + self.cold_input_keys + ) + self.output_keys = output_keys + self.num_features = num_features + + self.heat_net = mlp.MLP( + self.heat_input_keys, + ("h",), + branch_num_layers, + branch_hidden_size, + branch_activation, + branch_skip_connection, + branch_weight_norm, + input_dim=heat_num_loc, + output_dim=num_features * len(self.output_keys), + ) + + self.cold_net = mlp.MLP( + self.cold_input_keys, + ("c",), + branch_num_layers, + branch_hidden_size, + branch_activation, + branch_skip_connection, + branch_weight_norm, + input_dim=cold_num_loc, + output_dim=num_features * len(self.output_keys), + ) + + self.trunk_net = mlp.MLP( + self.trunk_input_keys, + ("t",), + trunk_num_layers, + trunk_hidden_size, + trunk_activation, + trunk_skip_connection, + trunk_weight_norm, + input_dim=len(self.trunk_input_keys), + output_dim=num_features * len(self.output_keys), + ) + self.trunk_act = act_mod.get_activation(trunk_activation) + self.heat_act = act_mod.get_activation(branch_activation) + self.cold_act = act_mod.get_activation(branch_activation) + + self.use_bias = use_bias + if use_bias: + # register bias to parameter for updating in optimizer and storage + self.b = self.create_parameter( + shape=(len(self.output_keys),), + attr=nn.initializer.Constant(0.0), + ) + + def forward(self, x): + if self._input_transform is not None: + x = self._input_transform(x) + + # Branch net to encode the input function + heat_features = self.heat_net(x)[self.heat_net.output_keys[0]] + cold_features = self.cold_net(x)[self.cold_net.output_keys[0]] + # Trunk net to encode the domain of the output function + y_features = self.trunk_net(x)[self.trunk_net.output_keys[0]] + y_features = self.trunk_act(y_features) + # Dot product + G_u_h = paddle.sum( + heat_features[:, : self.num_features] + * y_features[:, : self.num_features] + * cold_features[:, : self.num_features], + axis=1, + keepdim=True, + ) + G_u_c = paddle.sum( + heat_features[:, self.num_features : 2 * self.num_features] + * y_features[:, self.num_features : 2 * self.num_features] + * cold_features[:, self.num_features : 2 * self.num_features], + axis=1, + keepdim=True, + ) + G_u_w = paddle.sum( + heat_features[:, 2 * self.num_features :] + * y_features[:, 2 * self.num_features :] + * cold_features[:, 2 * self.num_features :], + axis=1, + keepdim=True, + ) + # Add bias + if self.use_bias: + G_u_h += self.b[0] + G_u_c += self.b[1] + G_u_w += self.b[2] + + result_dict = { + self.output_keys[0]: G_u_h, + self.output_keys[1]: G_u_c, + self.output_keys[2]: G_u_w, + } + if self._output_transform is not None: + result_dict = self._output_transform(x, result_dict) + + return result_dict diff --git a/examples/smc_reac/ppsci/arch/ifm_mlp.py b/examples/smc_reac/ppsci/arch/ifm_mlp.py new file mode 100644 index 0000000000..4685b34be6 --- /dev/null +++ b/examples/smc_reac/ppsci/arch/ifm_mlp.py @@ -0,0 +1,540 @@ +# Copyright (c) 2024 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. + +from typing import List +from typing import Tuple + +import numpy as np +import paddle +import paddle.nn as nn +import paddle.nn.functional as F + +import ppsci +from ppsci.arch import base + + +def init_parameter_uniform( + parameter: paddle.base.framework.EagerParamBase, n: int +) -> None: + ppsci.utils.initializer.uniform_(parameter, -1 / np.sqrt(n), 1 / np.sqrt(n)) + + +# inputs, hidden_units, outputs, d_out, sigma, dp_ratio, first_omega_0, hidden_omega_0, reg +class IFMMLP(base.Arch): + """Understanding the limitations of deep models for molecular property prediction: Insights and solutions. + [Xia, Jun, et al. Advances in Neural Information Processing Systems 36 (2023): 64774-64792.]https://openreview.net/forum?id=NLFqlDeuzt) + + Code reference: https://github.com/junxia97/IFM + + Args: + input_keys (Tuple[str, ...]): Name of input keys, such as ("input", ). + output_keys (Tuple[str, ...]): Name of output keys, such as ("pred", ). + hidden_units (List[int]): Units num in hidden layers. + embed_name (str): Embed name used in arch, such as "IMF", "None". + inputs (int): Input dim. + outputs (int): Output dim. + d_out (int): Embedding output dim for some architecture. + sigma (float): Hyper parameter for some architecture. + dp_ratio (float): Dropout ratio. + reg (bool): Regularization flag. + first_omega_0 (float): Frequency factor used in first layer. + hidden_omega_0 (float): Frequency factor used in hidden layer. + """ + + def __init__( + self, + input_keys: Tuple[str, ...], + output_keys: Tuple[str, ...], + hidden_units: List[int], + embed_name: str, + inputs: int, + outputs: int, + d_out: int, + sigma: float, + dp_ratio: float, + reg: bool, + first_omega_0: float, + hidden_omega_0: float, + ): + super().__init__() + self.input_keys = input_keys + self.output_keys = output_keys + + # initialization + if embed_name == "None": + my_model = MyDNN( + inputs=inputs, + hidden_units=hidden_units, + dp_ratio=dp_ratio, + outputs=outputs, + reg=reg, + ) + elif embed_name == "LE": + my_model = LE_DNN( + inputs=inputs, + hidden_units=hidden_units, + d_out=d_out + 1, + dp_ratio=dp_ratio, + outputs=outputs, + reg=reg, + ) + elif embed_name == "LSIM": + my_model = LSIM_DNN( + inputs=inputs, + hidden_units=hidden_units, + d_out=d_out + 1, + sigma=sigma, + dp_ratio=dp_ratio, + outputs=outputs, + reg=reg, + ) + elif embed_name == "IFM": + my_model = IFM_DNN( + inputs=inputs, + hidden_units=hidden_units, + outputs=outputs, + dp_ratio=dp_ratio, + first_omega_0=first_omega_0, + hidden_omega_0=hidden_omega_0, + reg=reg, + ) + elif embed_name == "GM": + my_model = GM_DNN( + inputs=inputs, + hidden_units=hidden_units, + d_out=d_out + 1, + sigma=sigma + 1, + dp_ratio=dp_ratio, + outputs=outputs, + reg=reg, + ) + elif embed_name == "SIM": + my_model = SIM_DNN( + inputs=inputs, + hidden_units=hidden_units, + d_out=d_out + 1, + sigma=sigma + 1, + dp_ratio=dp_ratio, + outputs=outputs, + reg=reg, + ) + else: + raise ValueError("Invalid Embedding Name") + + self.model = my_model + + def forward(self, x): + Xs = x[self.input_keys[0]] + ret = self.model(Xs) + return {self.output_keys[0]: ret} + + +class MyDNN(nn.Layer): + """ + Args: + inputs (int): Input dim. + hidden_units (List[int]): Units num in hidden layers. + outputs (int): Output dim. + dp_ratio (float): Dropout ratio. + reg (bool): Regularization flag. + """ + + def __init__(self, inputs, hidden_units, outputs, dp_ratio, reg): + super(MyDNN, self).__init__() + # parameters + self.reg = reg + + # layers + self.hidden1 = nn.Linear(inputs, hidden_units[0]) + self.dropout1 = nn.Dropout(dp_ratio) + + self.hidden2 = nn.Linear(hidden_units[0], hidden_units[1]) + self.dropout2 = nn.Dropout(dp_ratio) + + self.hidden3 = nn.Linear(hidden_units[1], hidden_units[2]) + self.dropout3 = nn.Dropout(dp_ratio) + + if reg: + self.output = nn.Linear(hidden_units[2], 1) + else: + self.output = nn.Linear(hidden_units[2], outputs) + + def forward(self, x): + x = self.hidden1(x) + x = F.relu(self.dropout1(x)) + + x = self.hidden2(x) + x = F.relu(self.dropout2(x)) + + x = self.hidden3(x) + x = F.relu(self.dropout3(x)) + + return self.output(x) + + +class LE(nn.Layer): + def __init__(self, n_tokens: int, d_out: int): + super().__init__() + self.weight = self.create_parameter([n_tokens, 1, d_out]) + self.bias = self.create_parameter([n_tokens, d_out]) + self.reset_parameters() + + def reset_parameters(self) -> None: + d_out = self.weight.shape[-1] + init_parameter_uniform(self.weight, d_out) + init_parameter_uniform(self.bias, d_out) + + def forward(self, x: paddle.Tensor) -> paddle.Tensor: + """ + x: (n_batch, n_features, d_in) + returns: (n_batch, n_features, d_out) + """ + x = x.unsqueeze(-1) + x = (x.unsqueeze(-2) @ self.weight[None]).squeeze(-2) + x = x + self.bias[None] + return x + + +class PLE(nn.Layer): + def __init__(self, n_num_features: int, d_out: int, sigma: float): + super().__init__() + self.d_out = d_out + self.sigma = sigma + self.coefficients = self.create_parameter([n_num_features, d_out]) + self.reset_parameters() + + def reset_parameters(self) -> None: + ppsci.utils.initializer.normal_(self.coefficients, 0.0, self.sigma) + + def forward(self, x: paddle.Tensor) -> paddle.Tensor: + x = 2 * np.pi * self.coefficients[None] * x[..., None] + return paddle.concat([paddle.cos(x), paddle.sin(x)], -1) + + +class LE_DNN(nn.Layer): + def __init__(self, inputs, hidden_units, outputs, d_out, dp_ratio, reg): + super(LE_DNN, self).__init__() + # parameters + self.reg = reg + # layers + self.hidden1 = nn.Linear(inputs * d_out, hidden_units[0]) + self.dropout1 = nn.Dropout(dp_ratio) + + self.hidden2 = nn.Linear(hidden_units[0], hidden_units[1]) + self.dropout2 = nn.Dropout(dp_ratio) + + self.hidden3 = nn.Linear(hidden_units[1], hidden_units[2]) + self.dropout3 = nn.Dropout(dp_ratio) + + if reg: + self.output = nn.Linear(hidden_units[2], 1) + else: + self.output = nn.Linear(hidden_units[2], outputs) + self.embedding = LE(inputs, d_out) + + def forward(self, x): + x = self.embedding(x).view([x.size(0), -1]) + x = self.hidden1(x) + x = F.relu(self.dropout1(x)) + + x = self.hidden2(x) + x = F.relu(self.dropout2(x)) + + x = self.hidden3(x) + x = F.relu(self.dropout3(x)) + + return self.output(x) + + +class LSIM_DNN(nn.Layer): + def __init__(self, inputs, hidden_units, outputs, d_out, sigma, dp_ratio, reg): + super(LSIM_DNN, self).__init__() + # parameters + self.reg = reg + # layers + self.hidden1 = nn.Linear(inputs, hidden_units[0]) + self.dropout1 = nn.Dropout(dp_ratio) + + self.hidden2 = nn.Linear(hidden_units[0], hidden_units[1]) + self.dropout2 = nn.Dropout(dp_ratio) + + self.hidden3 = nn.Linear(hidden_units[1], hidden_units[2]) + self.dropout3 = nn.Dropout(dp_ratio) + + if reg: + self.output = nn.Linear(hidden_units[2], 1) + else: + self.output = nn.Linear(hidden_units[2], outputs) + self.embedding = PLE(inputs, d_out, sigma) + self.linear = nn.Linear(d_out * 2, inputs) + + def forward(self, x): + x = self.embedding(x).sum(1) + x = F.relu(self.linear(x)) + x = self.hidden1(x) + x = F.relu(self.dropout1(x)) + + x = self.hidden2(x) + x = F.relu(self.dropout2(x)) + + x = self.hidden3(x) + x = F.relu(self.dropout3(x)) + + return self.output(x) + + +class GaussianEncoding(nn.Layer): + def __init__(self, n_num_features: int, d_out: int, sigma: float) -> None: + super().__init__() + self.d_out = d_out + self.sigma = sigma + self.n_num_features = n_num_features + self.size = (d_out, n_num_features) + self.B = paddle.randn(self.size) * sigma + + def forward(self, x: paddle.Tensor) -> paddle.Tensor: + """ + x: (n_batch, n_features) + returns: (n_batch, n_features * 2 * d_out) + """ + self.B = self.B.to(x.place) + xp = 2 * np.pi * x @ self.B.T + return paddle.concat((paddle.cos(xp), paddle.sin(xp)), axis=-1) + + +class GM_DNN(nn.Layer): + """ + Args: + inputs (int): Input dim. + hidden_units (List[int]): Units num in hidden layers. + outputs (int): Output dim. + d_out (int): Embedding output dim for some architecture. + sigma (float): Hyper parameter for some architecture. + dp_ratio (float): Dropout ratio. + reg (bool): Regularization flag. + """ + + def __init__(self, inputs, hidden_units, outputs, d_out, sigma, dp_ratio, reg): + super(GM_DNN, self).__init__() + # parameters + self.reg = reg + self.d_out = d_out + self.sigma = sigma + # layers + self.hidden1 = nn.Linear(d_out * 2, hidden_units[0]) + self.dropout1 = nn.Dropout(dp_ratio) + + self.hidden2 = nn.Linear(hidden_units[0], hidden_units[1]) + self.dropout2 = nn.Dropout(dp_ratio) + + self.hidden3 = nn.Linear(hidden_units[1], hidden_units[2]) + self.dropout3 = nn.Dropout(dp_ratio) + + if reg: + self.output = nn.Linear(hidden_units[2], 1) + else: + self.output = nn.Linear(hidden_units[2], outputs) + + self.embedding = GaussianEncoding(inputs, d_out, sigma) + + def forward(self, x): + x = self.embedding(x) + x = self.hidden1(x) + x = F.relu(self.dropout1(x)) + + x = self.hidden2(x) + x = F.relu(self.dropout2(x)) + + x = self.hidden3(x) + x = F.relu(self.dropout3(x)) + + return self.output(x) + + +class SineLayer(nn.Layer): + # If is_first=True, omega_0 is a frequency factor which simply multiplies the activations before the + # nonlinearity. Different signals may require different omega_0 in the first layer - this is a + # hyperparameter. + + # If is_first=False, then the weights will be divided by omega_0 so as to keep the magnitude of + # activations constant, but boost gradients to the weight matrix + + def __init__( + self, in_features, out_features, bias=True, is_first=False, omega_0=30 + ): + super().__init__() + self.omega_0 = omega_0 + self.is_first = is_first + + self.in_features = in_features + self.linear = nn.Linear(in_features, out_features, bias_attr=bias) + + self.init_weights() + + def init_weights(self): + with paddle.no_grad(): + if self.is_first: + self.linear.weight.uniform_(-1 / self.in_features, 1 / self.in_features) + else: + self.linear.weight.uniform_( + -np.sqrt(6 / self.in_features) / self.omega_0, + np.sqrt(6 / self.in_features) / self.omega_0, + ) + + def forward(self, input): + return paddle.sin(self.omega_0 * self.linear(input)) + + def forward_with_intermediate(self, input): + # For visualization of activation distributions + intermediate = self.omega_0 * self.linear(input) + return paddle.sin(intermediate), intermediate + + +class IFM_DNN(nn.Layer): + """ + Args: + inputs (int): Input dim. + hidden_units (List[int]): Units num in hidden layers. + outputs (int): Output dim. + dp_ratio (float): Dropout ratio. + first_omega_0 (float): Frequency factor used in first layer. + hidden_omega_0 (float): Frequency factor used in hidden layer. + reg (bool): Regularization flag. + """ + + def __init__( + self, + inputs, + hidden_units, + outputs, + dp_ratio, + first_omega_0, + hidden_omega_0, + reg, + ): + super(IFM_DNN, self).__init__() + # parameters + self.reg = reg + # layers + self.hidden1 = SineLayer( + inputs, hidden_units[0], is_first=True, omega_0=first_omega_0 + ) + self.dropout1 = nn.Dropout(dp_ratio) + + self.hidden2 = SineLayer( + hidden_units[0], hidden_units[1], is_first=False, omega_0=hidden_omega_0 + ) + self.dropout2 = nn.Dropout(dp_ratio) + + self.hidden3 = SineLayer( + hidden_units[1], hidden_units[2], is_first=False, omega_0=hidden_omega_0 + ) + self.dropout3 = nn.Dropout(dp_ratio) + + if reg: + self.output = nn.Linear(hidden_units[2], 1) + with paddle.no_grad(): + self.output.weight.uniform_( + -np.sqrt(6 / hidden_units[2]) / hidden_omega_0, + np.sqrt(6 / hidden_units[2]) / hidden_omega_0, + ) + else: + self.output = nn.Linear(hidden_units[2], outputs) + with paddle.no_grad(): + self.output.weight.uniform_( + -np.sqrt(6 / hidden_units[2]) / hidden_omega_0, + np.sqrt(6 / hidden_units[2]) / hidden_omega_0, + ) + + def forward(self, x): + + x = self.hidden1(x) + x = F.relu(self.dropout1(x)) + # x = self.dropout1(x) + + x = self.hidden2(x) + x = F.relu(self.dropout2(x)) + # x = self.dropout2(x) + + x = self.hidden3(x) + x = F.relu(self.dropout3(x)) + # x = self.dropout3(x) + + return self.output(x) + + +class SIM_encoding(nn.Layer): + def __init__(self, n_num_features: int, d_out: int, sigma: float): + super().__init__() + self.d_out = d_out + self.sigma = sigma + self.n_num_features = n_num_features + self.coeffs = 2 * np.pi * sigma ** (paddle.arange(d_out) / d_out) + + def forward(self, x: paddle.Tensor) -> paddle.Tensor: + """ + x: (n_batch, n_features) + returns: (n_batch, n_features * 2 * d_out) + """ + xp = self.coeffs.to(x.device) * paddle.unsqueeze(x, -1) + xp_cat = paddle.concat((paddle.cos(xp), paddle.sin(xp)), axis=-1) + return xp_cat.flatten(-2, -1) + + +class SIM_DNN(nn.Layer): + """ + Args: + inputs (int): Input dim. + hidden_units (List[int]): Units num in hidden layers. + outputs (int): Output dim. + d_out (int): Embedding output dim for some architecture. + sigma (float): Hyper parameter for some architecture. + dp_ratio (float): Dropout ratio. + reg (bool): Regularization flag. + """ + + def __init__(self, inputs, hidden_units, outputs, d_out, sigma, dp_ratio, reg): + super(SIM_DNN, self).__init__() + # parameters + self.reg = reg + self.d_out = d_out + self.sigma = sigma + # layers + self.hidden1 = nn.Linear(d_out * 2 * inputs, hidden_units[0]) + self.dropout1 = nn.Dropout(dp_ratio) + + self.hidden2 = nn.Linear(hidden_units[0], hidden_units[1]) + self.dropout2 = nn.Dropout(dp_ratio) + + self.hidden3 = nn.Linear(hidden_units[1], hidden_units[2]) + self.dropout3 = nn.Dropout(dp_ratio) + + if reg: + self.output = nn.Linear(hidden_units[2], 1) + else: + self.output = nn.Linear(hidden_units[2], outputs) + self.embedding = SIM_encoding(inputs, d_out, sigma) + + def forward(self, x): + x = self.embedding(x) + x = self.hidden1(x) + x = F.relu(self.dropout1(x)) + + x = self.hidden2(x) + x = F.relu(self.dropout2(x)) + + x = self.hidden3(x) + x = F.relu(self.dropout3(x)) + + return self.output(x) diff --git a/examples/smc_reac/ppsci/arch/kan.py b/examples/smc_reac/ppsci/arch/kan.py new file mode 100644 index 0000000000..2c46d60c64 --- /dev/null +++ b/examples/smc_reac/ppsci/arch/kan.py @@ -0,0 +1,385 @@ +# Copyright (c) 2025 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. + +import math +from typing import Callable +from typing import Tuple + +import paddle + +from ppsci.arch import base +from ppsci.utils import initializer + +""" +This is the paddle implementation of Korogonov-Arnold-Network (KAN) +which is based on the torch implementation [efficient-kan] by Blealtan and akkashdash +please refer to their work (https://github.com/Blealtan/efficient-kan) +Authors: guhaohao0991(guhaohao@baidu.com) +Date: 2025/04/ +""" + + +class KANLinear(paddle.nn.Layer): + def __init__( + self, + in_features: int, + out_features: int, + grid_size: int = 5, + spline_order: int = 3, + scale_noise: float = 0.1, + scale_base: float = 1.0, + scale_spline: float = 1.0, + enable_standalone_scale_spline: bool = True, + base_activation: Callable[[paddle.Tensor], paddle.Tensor] = paddle.nn.Silu, + grid_eps: float = 0.02, + grid_range: Tuple[float, float] = (-1, 1), + ): + super().__init__() + self.in_features = in_features + self.out_features = out_features + self.grid_size = grid_size + self.spline_order = spline_order + + h = (grid_range[1] - grid_range[0]) / grid_size + grid = ( + ( + paddle.arange(start=-spline_order, end=grid_size + spline_order + 1) * h + + grid_range[0] + ) + .expand(shape=[in_features, -1]) + .contiguous() + ) + self.register_buffer(name="grid", tensor=grid) + + self.base_weight = self.create_parameter( + shape=[out_features, in_features], + default_initializer=paddle.nn.initializer.Assign( + paddle.empty(shape=[out_features, in_features]) + ), + ) + self.spline_weight = self.create_parameter( + shape=[out_features, in_features, grid_size + spline_order], + default_initializer=paddle.nn.initializer.Assign( + paddle.empty( + shape=[out_features, in_features, grid_size + spline_order] + ) + ), + ) + + if enable_standalone_scale_spline: + self.spline_scaler = self.create_parameter( + shape=[out_features, in_features], + default_initializer=paddle.nn.initializer.Assign( + paddle.empty(shape=[out_features, in_features]) + ), + ) + + self.scale_noise = scale_noise + self.scale_base = scale_base + self.scale_spline = scale_spline + self.enable_standalone_scale_spline = enable_standalone_scale_spline + self.base_activation = base_activation() + self.grid_eps = grid_eps + + self.reset_parameters() + + def reset_parameters(self): + self.base_weight = initializer.kaiming_uniform_( + tensor=self.base_weight, + a=math.sqrt(5) * self.scale_base, + nonlinearity="leaky_relu", + ) + with paddle.no_grad(): + noise = ( + ( + paddle.rand( + shape=[self.grid_size + 1, self.in_features, self.out_features] + ) + - 1 / 2 + ) + * self.scale_noise + / self.grid_size + ) + + paddle.assign( + (self.scale_spline if not self.enable_standalone_scale_spline else 1.0) + * self.curve2coeff( + self.grid.T[self.spline_order : -self.spline_order], noise + ), + output=self.spline_weight.data, + ) + + if self.enable_standalone_scale_spline: + self.spline_scaler = initializer.kaiming_uniform_( + tensor=self.spline_scaler, + a=math.sqrt(5) * self.scale_spline, + nonlinearity="leaky_relu", + ) + + def b_splines(self, x: paddle.Tensor): + """ + Compute the B-spline bases for the given input tensor. + + Args: + x (paddle.Tensor): Input tensor of shape (batch_size, in_features). + + Returns: + paddle.Tensor: B-spline bases tensor of shape (batch_size, in_features, grid_size + spline_order). + """ + assert x.dim() == 2 and x.shape[1] == self.in_features + grid: paddle.Tensor = self.grid + x = x.unsqueeze(axis=-1) + bases = ((x >= grid[:, :-1]) & (x < grid[:, 1:])).to(x.dtype) + + for k in range(1, self.spline_order + 1): + bases = (x - grid[:, : -(k + 1)]) / ( + grid[:, k:-1] - grid[:, : -(k + 1)] + ) * bases[:, :, :-1] + (grid[:, k + 1 :] - x) / ( + grid[:, k + 1 :] - grid[:, 1:-k] + ) * bases[ + :, :, 1: + ] + + assert tuple(bases.shape) == ( + x.shape[0], + self.in_features, + self.grid_size + self.spline_order, + ) + + return bases.contiguous() + + def curve2coeff(self, x: paddle.Tensor, y: paddle.Tensor): + """ + Compute the coefficients of the curve that interpolates the given points. + + Args: + x (paddle.Tensor): Input tensor of shape (batch_size, in_features). + y (paddle.Tensor): Output tensor of shape (batch_size, in_features, out_features). + + Returns: + paddle.Tensor: Coefficients tensor of shape (out_features, in_features, grid_size + spline_order). + """ + assert x.dim() == 2 and x.shape[1] == self.in_features + assert tuple(y.shape) == (x.shape[0], self.in_features, self.out_features) + + A = self.b_splines(x).transpose( + perm=dim2perm(self.b_splines(x).ndim, 0, 1) + ) # [in_features, batch_size, grid_size + spline_order] + B = y.transpose( + perm=dim2perm(y.ndim, 0, 1) + ) # [in_features, batch_size, out_features] + solution = paddle.linalg.lstsq(x=A, y=B)[ + 0 + ] # [in_features, grid_size + spline_order, out_features] + if A.shape[0] == 1: + solution = solution.unsqueeze(axis=0) + # print("A shape: ", A.shape, "B shape: ", B.shape, "Solution shape: ", solution.shape) + result = solution.transpose([2, 0, 1]) + assert tuple(result.shape) == ( + self.out_features, + self.in_features, + self.grid_size + self.spline_order, + ) + + return result.contiguous() + + @property + def scaled_spline_weight(self): + return self.spline_weight * ( + self.spline_scaler.unsqueeze(axis=-1) + if self.enable_standalone_scale_spline + else 1.0 + ) + + def forward(self, x: paddle.Tensor): + assert x.dim() == 2 and x.shape[1] == self.in_features + + base_output = paddle.nn.functional.linear( + x=self.base_activation(x), weight=self.base_weight.T + ) + + spline_output = paddle.nn.functional.linear( + x=self.b_splines(x).reshape([x.shape[0], -1]).contiguous(), + weight=self.scaled_spline_weight.reshape( + [self.out_features, -1] + ).T.contiguous(), + ) + # cant calculate 1st order derivation using view + # spline_output = paddle.nn.functional.linear( + # x=self.b_splines(x).view(x.shape[0], -1), + # weight=self.scaled_spline_weight.view(self.out_features, -1).T) + + return base_output + spline_output + + @paddle.no_grad() + def update_grid(self, x: paddle.Tensor, margin=0.01): + assert x.dim() == 2 and x.shape[1] == self.in_features + batch = x.shape[0] + + splines = self.b_splines(x) # [batch, in, coeff] + splines = splines.transpose(perm=[1, 0, 2]) # [in, batch, coeff] + orig_coeff = self.scaled_spline_weight # [out, in, coeff] + orig_coeff = orig_coeff.transpose(perm=[1, 2, 0]) # [in, coeff, out] + unreduced_spline_output = paddle.bmm( + x=splines, y=orig_coeff + ) # [in, batch, out] + unreduced_spline_output = unreduced_spline_output.transpose( + perm=[1, 0, 2] + ) # [batch, in, out] + + # sort each channel individually to collect data distribution + x_sorted = (paddle.sort(x=x, axis=0), paddle.argsort(x=x, axis=0))[0] + grid_adaptive = x_sorted[ + paddle.linspace( + start=0, stop=batch - 1, num=self.grid_size + 1, dtype="int64" + ) + ] + uniform_step = (x_sorted[-1] - x_sorted[0] + 2 * margin) / self.grid_size + grid_uniform = ( + paddle.arange(dtype="float32", end=self.grid_size + 1).unsqueeze(axis=1) + * uniform_step + + x_sorted[0] + - margin + ) + + grid = self.grid_eps * grid_uniform + (1 - self.grid_eps) * grid_adaptive + grid = paddle.concat( + x=[ + grid[:1] + - uniform_step + * paddle.arange( + start=self.spline_order, end=0, step=-1, dtype="float32" + ).unsqueeze(axis=1), + grid, + grid[-1:] + + uniform_step + * paddle.arange( + start=1, end=self.spline_order + 1, dtype="float32" + ).unsqueeze(axis=1), + ], + axis=0, + ) + + paddle.assign(grid.T, output=self.grid) + paddle.assign( + self.curve2coeff(x, unreduced_spline_output), output=self.spline_weight.data + ) + + def regularization_loss(self, regularize_activation=1.0, regularize_entropy=1.0): + """ + Compute the regularization loss. + + L1 and the entropy loss is for the feature selection, i.e., let the weight of the activation function be small. + """ + l1_fake = self.spline_weight.abs().mean(axis=-1) + regularization_loss_activation = l1_fake.sum() + p = l1_fake / regularization_loss_activation + regularization_loss_entropy = -paddle.sum(x=p * p.log()) + return ( + regularize_activation * regularization_loss_activation + + regularize_entropy * regularization_loss_entropy + ) + + +class KAN(base.Arch): + """Kolmogorov-Arnold Network (KAN). + + Args: + layers_hidden (Tuple[int, ...]): The number of hidden neurons in each layer. + input_keys (Tuple[str, ...]): The keys of the input dictionary. + output_keys (Tuple[str, ...]): The keys of the output dictionary. + grid_size (int): The size of the grid used by the spline basis functions. Default: 5. + spline_order (int): The order of the spline basis functions. Default: 3. + scale_noise (float): The scaling factor for the noise added to the weights of the KAN-linear layers. Default: 0.1. + scale_base (float): The scaling factor for the base activation output. Default: 1.0. + scale_spline (float): The scaling factor for the b-spline output. Default: 1.0. + base_activation (Callable[[paddle.Tensor], paddle.Tensor]): The base activation function. Default: paddle.nn.Silu. + grid_eps (float): The epsilon value used to initialize the grid. Default: 0.02. + grid_range (Tuple[float, float]): The domain range of the grid for b-spline interpolation. Default: (-1, 1). + + Examples: + >>> import paddle + >>> import ppsci + >>> model = ppsci.arch.KAN( + ... layers_hidden=(2, 5, 5, 1), + ... input_keys=("x", "y"), + ... output_keys=("z"), + ... grid_size=5, + ... spline_order=3 + >>> ) + >>> input_dict = {"x": paddle.rand([64, 1]), + ... "y": paddle.rand([64, 1])} + >>> output_dict = model(input_dict) + >>> print(output_dict["z"].shape) + [64, 1] + """ + + def __init__( + self, + layers_hidden: Tuple[int, ...], + input_keys: Tuple[str, ...], + output_keys: Tuple[str, ...], + grid_size: int = 5, + spline_order: int = 3, + scale_noise: float = 0.1, + scale_base: float = 1.0, + scale_spline: float = 1.0, + base_activation: Callable[[paddle.Tensor], paddle.Tensor] = paddle.nn.Silu, + grid_eps: float = 0.02, + grid_range: Tuple[float, float] = (-1, 1), + ): + super().__init__() + self.input_keys = input_keys + self.output_keys = output_keys + self.grid_size = grid_size + self.spline_order = spline_order + self.layers = paddle.nn.LayerList() + for in_features, out_features in zip(layers_hidden, layers_hidden[1:]): + self.layers.append( + KANLinear( + in_features, + out_features, + grid_size=grid_size, + spline_order=spline_order, + scale_noise=scale_noise, + scale_base=scale_base, + scale_spline=scale_spline, + base_activation=base_activation, + grid_eps=grid_eps, + grid_range=grid_range, + ) + ) + + def forward(self, x_dict, update_grid=False): + x = self.concat_to_tensor(x_dict, self.input_keys, axis=-1) + for index, layer in enumerate(self.layers): + if update_grid: + layer.update_grid(x) + x = layer(x) + if index < len(self.layers) - 1: + x = paddle.nn.functional.tanh(x=x) + out_dic = self.split_to_dict(x, self.output_keys, axis=-1) + return out_dic + + def regularization_loss(self, regularize_activation=1.0, regularize_entropy=1.0): + return sum( + layer.regularization_loss(regularize_activation, regularize_entropy) + for layer in self.layers + ) + + +def dim2perm(ndim, dim0, dim1): + perm = list(range(ndim)) + perm[dim0], perm[dim1] = perm[dim1], perm[dim0] + return perm diff --git a/examples/smc_reac/ppsci/arch/lno.py b/examples/smc_reac/ppsci/arch/lno.py new file mode 100644 index 0000000000..d600c5d028 --- /dev/null +++ b/examples/smc_reac/ppsci/arch/lno.py @@ -0,0 +1,312 @@ +# Copyright (c) 2024 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. + +from __future__ import annotations + +import operator +from functools import reduce +from typing import Optional +from typing import Tuple + +import numpy as np +import paddle +import paddle.nn as nn + +from ppsci.arch import activation as act_mod +from ppsci.arch import base +from ppsci.utils import initializer + + +class Laplace(nn.Layer): + """Generic N-Dimensional Laplace Operator with Pole-Residue Method. + + Args: + in_channels (int): Number of input channels of the first layer. + out_channels (int): Number of output channels of the last layer. + modes (Tuple[int, ...]): Number of modes to use for contraction in Laplace domain during training. + T (paddle.Tensor): Linspace of time dimension. + data (Tuple[paddle.Tensor, ...]): Linspaces of other dimensions. + """ + + def __init__( + self, + in_channels: int, + out_channels: int, + modes: Tuple[int, ...], + T: paddle.Tensor, + data: Tuple[paddle.Tensor, ...], + ): + super().__init__() + self.char1 = "pqr" + self.char2 = "mnk" + self.modes = modes + self.scale = 1 / (in_channels * out_channels) + self.dims = len(modes) + + self.weights_pole_real = nn.ParameterList() + self.weights_pole_imag = nn.ParameterList() + for i in range(self.dims): + weight_real = self._init_weights( + self.create_parameter((in_channels, out_channels, modes[i], 1)) + ) + weight_imag = self._init_weights( + self.create_parameter((in_channels, out_channels, modes[i], 1)) + ) + self.weights_pole_real.append(weight_real) + self.weights_pole_imag.append(weight_imag) + + residues_shape = (in_channels, out_channels) + modes + (1,) + self.weights_residue_real = self._init_weights( + self.create_parameter(residues_shape) + ) + self.weights_residue_imag = self._init_weights( + self.create_parameter(residues_shape) + ) + + self.initialize_lambdas(T, data) + self.get_einsum_eqs() + + def _init_weights(self, weight) -> paddle.Tensor: + return initializer.uniform_(weight, a=0, b=self.scale) + + def initialize_lambdas(self, T, data) -> None: + self.t_lst = (T,) + data + self.lambdas = [] + for i in range(self.dims): + t_i = self.t_lst[i] + self.register_buffer(f"t_{i}", t_i) + dt = (t_i[0, 1] - t_i[0, 0]).item() + omega = paddle.fft.fftfreq(n=tuple(t_i.shape)[1], d=dt) * 2 * np.pi * 1.0j + lambda_ = omega.reshape([*omega.shape, 1, 1, 1]) + self.register_buffer(f"lambda_{i}", lambda_) + self.lambdas.append(lambda_) + + def get_einsum_eqs(self) -> None: + terms_eq = [] + terms_x2_eq = [] + for i in range(self.dims): + term_eq = self.char1[i] + "io" + self.char2[i] + terms_eq.append(term_eq) + term_x2_eq = "io" + self.char2[i] + self.char1[i] + terms_x2_eq.append(term_x2_eq) + self.eq1 = ( + "bi" + + "".join(self.char1) + + "," + + "io" + + "".join(self.char2) + + "," + + ",".join(terms_eq) + + "->" + + "bo" + + "".join(self.char1) + ) + self.eq2 = ( + "bi" + + "".join(self.char1) + + "," + + "io" + + "".join(self.char2) + + "," + + ",".join(terms_eq) + + "->" + + "bo" + + "".join(self.char2) + ) + self.eq_x2 = ( + "bi" + + "".join(self.char2) + + "," + + ",".join(terms_x2_eq) + + "->bo" + + "".join(self.char1) + ) + + def output_PR(self, alpha) -> Tuple[paddle.Tensor, paddle.Tensor]: + weights_residue = paddle.as_complex( + paddle.concat( + [self.weights_residue_real, self.weights_residue_imag], axis=-1 + ) + ) + self.weights_pole = [] + terms = [] + for i in range(self.dims): + weights_pole = paddle.as_complex( + paddle.concat( + [self.weights_pole_real[i], self.weights_pole_imag[i]], axis=-1 + ) + ) + self.weights_pole.append(weights_pole) + sub = paddle.subtract(self.lambdas[i], weights_pole) + terms.append(paddle.divide(paddle.to_tensor(1, dtype=sub.dtype), sub)) + + output_residue1 = paddle.einsum(self.eq1, alpha, weights_residue, *terms) + output_residue2 = (-1) ** self.dims * paddle.einsum( + self.eq2, alpha, weights_residue, *terms + ) + return output_residue1, output_residue2 + + def forward(self, x): + alpha = paddle.fft.fftn(x=x, axes=[-3, -2, -1]) + output_residue1, output_residue2 = self.output_PR(alpha) + + x1 = paddle.fft.ifftn( + x=output_residue1, s=(x.shape[-3], x.shape[-2], x.shape[-1]) + ) + x1 = paddle.real(x=x1) + + exp_terms = [] + for i in range(self.dims): + term = paddle.einsum( + "io" + + self.char2[i] + + ",d" + + self.char1[i] + + "->io" + + self.char2[i] + + self.char1[i], + self.weights_pole[i], + self.t_lst[i].astype(paddle.complex64).reshape([1, -1]), + ) + exp_terms.append(paddle.exp(term)) + + x2 = paddle.einsum(self.eq_x2, output_residue2, *exp_terms) + x2 = paddle.real(x2) + x2 = x2 / reduce(operator.mul, x.shape[-3:], 1) + return x1 + x2 + + +class LNO(base.Arch): + """Laplace Neural Operator net. + + Args: + input_keys (Tuple[str, ...]): Name of input keys, such as ("input1", "input2"). + output_keys (Tuple[str, ...]): Name of output keys, such as ("output1", "output2"). + width (int): Tensor width of Laplace Layer. + modes (Tuple[int, ...]): Number of modes to use for contraction in Laplace domain during training. + T (paddle.Tensor): Linspace of time dimension. + data (Tuple[paddle.Tensor, ...]): Linspaces of other dimensions. + in_features (int, optional): Number of input channels of the first layer.. Defaults to 1. + hidden_features (int, optional): Number of channels of the fully-connected layer. Defaults to 64. + activation (str, optional): The activation function. Defaults to "sin". + use_norm (bool, optional): Whether to use normalization layers. Defaults to True. + use_grid (bool, optional): Whether to create grid. Defaults to False. + """ + + def __init__( + self, + input_keys: Tuple[str, ...], + output_keys: Tuple[str, ...], + width: int, + modes: Tuple[int, ...], + T: paddle.Tensor, + data: Optional[Tuple[paddle.Tensor, ...]] = None, + in_features: int = 1, + hidden_features: int = 64, + activation: str = "sin", + use_norm: bool = True, + use_grid: bool = False, + ): + super().__init__() + self.input_keys = input_keys + self.output_keys = output_keys + self.width = width + self.modes = modes + self.dims = len(modes) + assert self.dims <= 3, "Only 3 dims and lower of modes are supported now." + + if data is None: + data = () + assert ( + self.dims == len(data) + 1 + ), f"Dims of modes is {self.dims} but only {len(data)} dims(except T) of data received." + + self.fc0 = nn.Linear(in_features=in_features, out_features=self.width) + self.laplace = Laplace(self.width, self.width, self.modes, T, data) + self.conv = getattr(nn, f"Conv{self.dims}D")( + in_channels=self.width, + out_channels=self.width, + kernel_size=1, + data_format="NCDHW", + ) + if use_norm: + self.norm = getattr(nn, f"InstanceNorm{self.dims}D")( + num_features=self.width, + weight_attr=False, + bias_attr=False, + ) + self.fc1 = nn.Linear(in_features=self.width, out_features=hidden_features) + self.fc2 = nn.Linear(in_features=hidden_features, out_features=1) + self.act = act_mod.get_activation(activation) + + self.use_norm = use_norm + self.use_grid = use_grid + + def get_grid(self, shape): + batchsize, size_t, size_x, size_y = shape[0], shape[1], shape[2], shape[3] + gridt = paddle.linspace(0, 1, size_t) + gridt = gridt.reshape([1, size_t, 1, 1, 1]).tile( + [batchsize, 1, size_x, size_y, 1] + ) + gridx = paddle.linspace(0, 1, size_x) + gridx = gridx.reshape([1, 1, size_x, 1, 1]).tile( + [batchsize, size_t, 1, size_y, 1] + ) + gridy = paddle.linspace(0, 1, size_y) + gridy = gridy.reshape([1, 1, 1, size_y, 1]).tile( + [batchsize, size_t, size_x, 1, 1] + ) + return paddle.concat([gridt, gridx, gridy], axis=-1) + + def transpoe_to_NCDHW(self, x): + perm = [0, self.dims + 1] + list(range(1, self.dims + 1)) + return paddle.transpose(x, perm=perm) + + def transpoe_to_NDHWC(self, x): + perm = [0] + list(range(2, self.dims + 2)) + [1] + return paddle.transpose(x, perm=perm) + + def forward_tensor(self, x): + if self.use_grid: + grid = self.get_grid(x.shape) + x = paddle.concat([x, grid], axis=-1) + x = self.fc0(x) + x = self.transpoe_to_NCDHW(x) + + if self.use_norm: + x1 = self.norm(self.laplace(self.norm(x))) + else: + x1 = self.laplace(x) + + x2 = self.conv(x) + x = x1 + x2 + + x = self.transpoe_to_NDHWC(x) + + x = self.fc1(x) + x = self.act(x) + x = self.fc2(x) + return x + + def forward(self, x): + if self._input_transform is not None: + x = self._input_transform(x) + + y = self.concat_to_tensor(x, self.input_keys, axis=-1) + y = self.forward_tensor(y) + y = self.split_to_dict(y, self.output_keys, axis=-1) + + if self._output_transform is not None: + y = self._output_transform(x, y) + return y diff --git a/examples/smc_reac/ppsci/arch/mlp.py b/examples/smc_reac/ppsci/arch/mlp.py new file mode 100644 index 0000000000..ef2d6e9e44 --- /dev/null +++ b/examples/smc_reac/ppsci/arch/mlp.py @@ -0,0 +1,828 @@ +# 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. + +from __future__ import annotations + +from typing import Dict +from typing import Optional +from typing import Tuple +from typing import Union + +import numpy as np +import paddle +import paddle.nn as nn + +from ppsci.arch import activation as act_mod +from ppsci.arch import base +from ppsci.utils import initializer + + +class WeightNormLinear(nn.Layer): + def __init__(self, in_features: int, out_features: int, bias: bool = True) -> None: + super().__init__() + self.in_features = in_features + self.out_features = out_features + self.weight_v = self.create_parameter((in_features, out_features)) + self.weight_g = self.create_parameter((out_features,)) + if bias: + self.bias = self.create_parameter((out_features,)) + else: + self.bias = None + self._init_weights() + + def _init_weights(self) -> None: + initializer.xavier_uniform_(self.weight_v) + initializer.constant_(self.weight_g, 1.0) + if self.bias is not None: + initializer.constant_(self.bias, 0.0) + + def forward(self, input): + norm = self.weight_v.norm(p=2, axis=0, keepdim=True) + weight = self.weight_g * self.weight_v / norm + return nn.functional.linear(input, weight, self.bias) + + +class RandomWeightFactorization(nn.Layer): + def __init__( + self, + in_features: int, + out_features: int, + bias: bool = True, + mean: float = 0.5, + std: float = 0.1, + ): + super().__init__() + self.in_features = in_features + self.out_features = out_features + self.weight_v = self.create_parameter((in_features, out_features)) + self.weight_g = self.create_parameter((out_features,)) + if bias: + self.bias = self.create_parameter((out_features,)) + else: + self.bias = None + + self._init_weights(mean, std) + + def _init_weights(self, mean, std): + with paddle.no_grad(): + initializer.glorot_normal_(self.weight_v) + + nn.initializer.Normal(mean, std)(self.weight_g) + paddle.assign(paddle.exp(self.weight_g), self.weight_g) + paddle.assign(self.weight_v / self.weight_g, self.weight_v) + if self.bias is not None: + initializer.constant_(self.bias, 0.0) + + self.weight_g.stop_gradient = False + self.weight_v.stop_gradient = False + self.bias.stop_gradient = False + + def forward(self, input): + return nn.functional.linear(input, self.weight_g * self.weight_v, self.bias) + + +class PeriodEmbedding(nn.Layer): + def __init__(self, periods: Dict[str, Tuple[float, bool]]): + super().__init__() + self.freqs_dict = { + k: self.create_parameter( + [], + attr=paddle.ParamAttr(trainable=trainable), + default_initializer=nn.initializer.Constant(2 * np.pi / float(p)), + ) # mu = 2*pi / period for sin/cos function + for k, (p, trainable) in periods.items() + } + self.freqs = nn.ParameterList(list(self.freqs_dict.values())) + + def forward(self, x: Dict[str, paddle.Tensor]): + y = {k: v for k, v in x.items()} # shallow copy to avoid modifying input dict + + for k, w in self.freqs_dict.items(): + y[k] = paddle.concat([paddle.cos(w * x[k]), paddle.sin(w * x[k])], axis=-1) + + return y + + +class FourierEmbedding(nn.Layer): + def __init__(self, in_features, out_features, scale): + super().__init__() + if out_features % 2 != 0: + raise ValueError(f"out_features must be even, but got {out_features}.") + + self.kernel = self.create_parameter( + [in_features, out_features // 2], + default_initializer=nn.initializer.Normal(std=scale), + ) + + def forward(self, x: paddle.Tensor): + y = paddle.concat( + [ + paddle.cos(x @ self.kernel), + paddle.sin(x @ self.kernel), + ], + axis=-1, + ) + return y + + +class MLP(base.Arch): + """Multi layer perceptron network. + + Args: + input_keys (Tuple[str, ...]): Name of input keys, such as ("x", "y", "z"). + output_keys (Tuple[str, ...]): Name of output keys, such as ("u", "v", "w"). + num_layers (int): Number of hidden layers. + hidden_size (Union[int, Tuple[int, ...]]): Number of hidden size. + An integer for all layers, or list of integer specify each layer's size. + activation (str, optional): Name of activation function. Defaults to "tanh". + skip_connection (bool, optional): Whether to use skip connection. Defaults to False. + weight_norm (bool, optional): Whether to apply weight norm on parameter(s). Defaults to False. + input_dim (Optional[int]): Number of input's dimension. Defaults to None. + output_dim (Optional[int]): Number of output's dimension. Defaults to None. + periods (Optional[Dict[int, Tuple[float, bool]]]): Period of each input key, + input in given channel will be period embedded if specified, each tuple of + periods list is [period, trainable]. Defaults to None. + fourier (Optional[Dict[str, Union[float, int]]]): Random fourier feature embedding, + e.g. {'dim': 256, 'scale': 1.0}. Defaults to None. + random_weight (Optional[Dict[str, float]]): Mean and std of random weight + factorization layer, e.g. {"mean": 0.5, "std: 0.1"}. Defaults to None. + + Examples: + >>> import paddle + >>> import ppsci + >>> model = ppsci.arch.MLP( + ... input_keys=("x", "y"), + ... output_keys=("u", "v"), + ... num_layers=5, + ... hidden_size=128 + ... ) + >>> input_dict = {"x": paddle.rand([64, 1]), + ... "y": paddle.rand([64, 1])} + >>> output_dict = model(input_dict) + >>> print(output_dict["u"].shape) + [64, 1] + >>> print(output_dict["v"].shape) + [64, 1] + """ + + def __init__( + self, + input_keys: Tuple[str, ...], + output_keys: Tuple[str, ...], + num_layers: int, + hidden_size: Union[int, Tuple[int, ...]], + activation: str = "tanh", + skip_connection: bool = False, + weight_norm: bool = False, + input_dim: Optional[int] = None, + output_dim: Optional[int] = None, + periods: Optional[Dict[int, Tuple[float, bool]]] = None, + fourier: Optional[Dict[str, Union[float, int]]] = None, + random_weight: Optional[Dict[str, float]] = None, + ): + super().__init__() + self.input_keys = input_keys + self.output_keys = output_keys + self.linears = [] + self.acts = [] + self.periods = periods + self.fourier = fourier + if periods: + self.period_emb = PeriodEmbedding(periods) + + if isinstance(hidden_size, (tuple, list)): + if num_layers is not None: + raise ValueError( + "num_layers should be None when hidden_size is specified" + ) + elif isinstance(hidden_size, int): + if not isinstance(num_layers, int): + raise ValueError( + "num_layers should be an int when hidden_size is an int" + ) + hidden_size = [hidden_size] * num_layers + else: + raise ValueError( + f"hidden_size should be list of int or int, but got {type(hidden_size)}" + ) + + # initialize FC layer(s) + cur_size = len(self.input_keys) if input_dim is None else input_dim + if input_dim is None and periods: + # period embedded channel(s) will be doubled automatically + # if input_dim is not specified + cur_size += len(periods) + + if fourier: + self.fourier_emb = FourierEmbedding( + cur_size, fourier["dim"], fourier["scale"] + ) + cur_size = fourier["dim"] + + for i, _size in enumerate(hidden_size): + if weight_norm: + self.linears.append(WeightNormLinear(cur_size, _size)) + elif random_weight: + self.linears.append( + RandomWeightFactorization( + cur_size, + _size, + mean=random_weight["mean"], + std=random_weight["std"], + ) + ) + else: + self.linears.append(nn.Linear(cur_size, _size)) + + # initialize activation function + self.acts.append( + act_mod.get_activation(activation) + if activation != "stan" + else act_mod.get_activation(activation)(_size) + ) + # special initialization for certain activation + # TODO: Adapt code below to a more elegant style + if activation == "siren": + if i == 0: + act_mod.Siren.init_for_first_layer(self.linears[-1]) + else: + act_mod.Siren.init_for_hidden_layer(self.linears[-1]) + + cur_size = _size + + self.linears = nn.LayerList(self.linears) + self.acts = nn.LayerList(self.acts) + if random_weight: + self.last_fc = RandomWeightFactorization( + cur_size, + len(self.output_keys) if output_dim is None else output_dim, + mean=random_weight["mean"], + std=random_weight["std"], + ) + else: + self.last_fc = nn.Linear( + cur_size, + len(self.output_keys) if output_dim is None else output_dim, + ) + + self.skip_connection = skip_connection + + def forward_tensor(self, x): + y = x + skip = None + for i, linear in enumerate(self.linears): + y = linear(y) + if self.skip_connection and i % 2 == 0: + if skip is not None: + skip = y + y = y + skip + else: + skip = y + y = self.acts[i](y) + + y = self.last_fc(y) + + return y + + def forward(self, x): + if self._input_transform is not None: + x = self._input_transform(x) + + if self.periods: + x = self.period_emb(x) + + y = self.concat_to_tensor(x, self.input_keys, axis=-1) + + if self.fourier: + y = self.fourier_emb(y) + + y = self.forward_tensor(y) + y = self.split_to_dict(y, self.output_keys, axis=-1) + + if self._output_transform is not None: + y = self._output_transform(x, y) + return y + + +class ModifiedMLP(base.Arch): + """Modified Multi layer perceptron network. + + Understanding and mitigating gradient pathologies in physics-informed + neural networks. https://arxiv.org/pdf/2001.04536.pdf. + + Args: + input_keys (Tuple[str, ...]): Name of input keys, such as ("x", "y", "z"). + output_keys (Tuple[str, ...]): Name of output keys, such as ("u", "v", "w"). + num_layers (int): Number of hidden layers. + hidden_size (int): Number of hidden size, an integer for all layers. + activation (str, optional): Name of activation function. Defaults to "tanh". + skip_connection (bool, optional): Whether to use skip connection. Defaults to False. + weight_norm (bool, optional): Whether to apply weight norm on parameter(s). Defaults to False. + input_dim (Optional[int]): Number of input's dimension. Defaults to None. + output_dim (Optional[int]): Number of output's dimension. Defaults to None. + + Examples: + >>> import paddle + >>> import ppsci + >>> model = ppsci.arch.ModifiedMLP( + ... input_keys=("x", "y"), + ... output_keys=("u", "v"), + ... num_layers=5, + ... hidden_size=128 + ... ) + >>> input_dict = {"x": paddle.rand([64, 1]), + ... "y": paddle.rand([64, 1])} + >>> output_dict = model(input_dict) + >>> print(output_dict["u"].shape) + [64, 1] + >>> print(output_dict["v"].shape) + [64, 1] + """ + + def __init__( + self, + input_keys: Tuple[str, ...], + output_keys: Tuple[str, ...], + num_layers: int, + hidden_size: int, + activation: str = "tanh", + skip_connection: bool = False, + weight_norm: bool = False, + input_dim: Optional[int] = None, + output_dim: Optional[int] = None, + periods: Optional[Dict[int, Tuple[float, bool]]] = None, + fourier: Optional[Dict[str, Union[float, int]]] = None, + random_weight: Optional[Dict[str, float]] = None, + ): + super().__init__() + self.input_keys = input_keys + self.output_keys = output_keys + self.linears = [] + self.acts = [] + self.periods = periods + self.fourier = fourier + if periods: + self.period_emb = PeriodEmbedding(periods) + if isinstance(hidden_size, int): + if not isinstance(num_layers, int): + raise ValueError("num_layers should be an int") + hidden_size = [hidden_size] * num_layers + else: + raise ValueError(f"hidden_size should be int, but got {type(hidden_size)}") + + # initialize FC layer(s) + cur_size = len(self.input_keys) if input_dim is None else input_dim + if input_dim is None and periods: + # period embedded channel(s) will be doubled automatically + # if input_dim is not specified + cur_size += len(periods) + + if fourier: + self.fourier_emb = FourierEmbedding( + cur_size, fourier["dim"], fourier["scale"] + ) + cur_size = fourier["dim"] + + self.embed_u = nn.Sequential( + ( + WeightNormLinear(cur_size, hidden_size[0]) + if weight_norm + else ( + nn.Linear(cur_size, hidden_size[0]) + if random_weight is None + else RandomWeightFactorization( + cur_size, + hidden_size[0], + mean=random_weight["mean"], + std=random_weight["std"], + ) + ) + ), + ( + act_mod.get_activation(activation) + if activation != "stan" + else act_mod.get_activation(activation)(hidden_size[0]) + ), + ) + self.embed_v = nn.Sequential( + ( + WeightNormLinear(cur_size, hidden_size[0]) + if weight_norm + else ( + nn.Linear(cur_size, hidden_size[0]) + if random_weight is None + else RandomWeightFactorization( + cur_size, + hidden_size[0], + mean=random_weight["mean"], + std=random_weight["std"], + ) + ) + ), + ( + act_mod.get_activation(activation) + if activation != "stan" + else act_mod.get_activation(activation)(hidden_size[0]) + ), + ) + + for i, _size in enumerate(hidden_size): + if weight_norm: + self.linears.append(WeightNormLinear(cur_size, _size)) + elif random_weight: + self.linears.append( + RandomWeightFactorization( + cur_size, + _size, + mean=random_weight["mean"], + std=random_weight["std"], + ) + ) + else: + self.linears.append(nn.Linear(cur_size, _size)) + + # initialize activation function + self.acts.append( + act_mod.get_activation(activation) + if activation != "stan" + else act_mod.get_activation(activation)(_size) + ) + # special initialization for certain activation + # TODO: Adapt code below to a more elegant style + if activation == "siren": + if i == 0: + act_mod.Siren.init_for_first_layer(self.linears[-1]) + else: + act_mod.Siren.init_for_hidden_layer(self.linears[-1]) + + cur_size = _size + + self.linears = nn.LayerList(self.linears) + self.acts = nn.LayerList(self.acts) + if random_weight: + self.last_fc = RandomWeightFactorization( + cur_size, + len(self.output_keys) if output_dim is None else output_dim, + mean=random_weight["mean"], + std=random_weight["std"], + ) + else: + self.last_fc = nn.Linear( + cur_size, + len(self.output_keys) if output_dim is None else output_dim, + ) + + self.skip_connection = skip_connection + + def forward_tensor(self, x): + u = self.embed_u(x) + v = self.embed_v(x) + + y = x + skip = None + for i, linear in enumerate(self.linears): + y = linear(y) + y = self.acts[i](y) + y = y * u + (1 - y) * v + if self.skip_connection and i % 2 == 0: + if skip is not None: + skip = y + y = y + skip + else: + skip = y + + y = self.last_fc(y) + + return y + + def forward(self, x): + x_identity = x + if self._input_transform is not None: + x = self._input_transform(x) + + if self.periods: + x = self.period_emb(x) + + y = self.concat_to_tensor(x, self.input_keys, axis=-1) + + if self.fourier: + y = self.fourier_emb(y) + + y = self.forward_tensor(y) + y = self.split_to_dict(y, self.output_keys, axis=-1) + + if self._output_transform is not None: + y = self._output_transform(x_identity, y) + return y + + +class PirateNetBlock(nn.Layer): + r"""Basic block of PirateNet. + + $$ + \begin{align*} + \Phi(\mathbf{x})=\left[\begin{array}{l} + \cos (\mathbf{B} \mathbf{x}) \\ + \sin (\mathbf{B} \mathbf{x}) + \end{array}\right] \\ + \mathbf{f}^{(l)} & =\sigma\left(\mathbf{W}_1^{(l)} \mathbf{x}^{(l)}+\mathbf{b}_1^{(l)}\right) \\ + \mathbf{z}_1^{(l)} & =\mathbf{f}^{(l)} \odot \mathbf{U}+\left(1-\mathbf{f}^{(l)}\right) \odot \mathbf{V} \\ + \mathbf{g}^{(l)} & =\sigma\left(\mathbf{W}_2^{(l)} \mathbf{z}_1^{(l)}+\mathbf{b}_2^{(l)}\right) \\ + \mathbf{z}_2^{(l)} & =\mathbf{g}^{(l)} \odot \mathbf{U}+\left(1-\mathbf{g}^{(l)}\right) \odot \mathbf{V} \\ + \mathbf{h}^{(l)} & =\sigma\left(\mathbf{W}_3^{(l)} \mathbf{z}_2^{(l)}+\mathbf{b}_3^{(l)}\right) \\ + \mathbf{x}^{(l+1)} & =\alpha^{(l)} \cdot \mathbf{h}^{(l)}+\left(1-\alpha^{(l)}\right) \cdot \mathbf{x}^{(l)} + \end{align*} + $$ + + Args: + input_dim (int): Input dimension. + embed_dim (int): Embedding dimension. + activation (str, optional): Name of activation function. Defaults to "tanh". + random_weight (Optional[Dict[str, float]]): Mean and std of random weight + factorization layer, e.g. {"mean": 0.5, "std: 0.1"}. Defaults to None. + """ + + def __init__( + self, + input_dim: int, + embed_dim: int, + activation: str = "tanh", + random_weight: Optional[Dict[str, float]] = None, + ): + super().__init__() + self.linear1 = ( + nn.Linear(input_dim, embed_dim) + if random_weight is None + else RandomWeightFactorization( + input_dim, + embed_dim, + mean=random_weight["mean"], + std=random_weight["std"], + ) + ) + self.linear2 = ( + nn.Linear(embed_dim, embed_dim) + if random_weight is None + else RandomWeightFactorization( + embed_dim, + embed_dim, + mean=random_weight["mean"], + std=random_weight["std"], + ) + ) + self.linear3 = ( + nn.Linear(embed_dim, embed_dim) + if random_weight is None + else RandomWeightFactorization( + embed_dim, + embed_dim, + mean=random_weight["mean"], + std=random_weight["std"], + ) + ) + self.alpha = self.create_parameter( + [ + 1, + ], + default_initializer=nn.initializer.Constant(0), + ) + self.act1 = ( + act_mod.get_activation(activation) + if activation != "stan" + else act_mod.get_activation(activation)(embed_dim) + ) + self.act2 = ( + act_mod.get_activation(activation) + if activation != "stan" + else act_mod.get_activation(activation)(embed_dim) + ) + self.act3 = ( + act_mod.get_activation(activation) + if activation != "stan" + else act_mod.get_activation(activation)(embed_dim) + ) + + def forward(self, x, u, v): + f = self.act1(self.linear1(x)) + z1 = f * u + (1 - f) * v + g = self.act2(self.linear2(z1)) + z2 = g * u + (1 - g) * v + h = self.act3(self.linear3(z2)) + out = self.alpha * h + (1 - self.alpha) * x + return out + + +class PirateNet(base.Arch): + r"""PirateNet. + + [PIRATENETS: PHYSICS-INFORMED DEEP LEARNING WITHRESIDUAL ADAPTIVE NETWORKS](https://arxiv.org/pdf/2402.00326.pdf) + + $$ + \begin{align*} + \Phi(\mathbf{x}) &= \left[\begin{array}{l} + \cos (\mathbf{B} \mathbf{x}) \\ + \sin (\mathbf{B} \mathbf{x}) + \end{array}\right] \\ + \mathbf{f}^{(l)} &= \sigma\left(\mathbf{W}_1^{(l)} \mathbf{x}^{(l)}+\mathbf{b}_1^{(l)}\right) \\ + \mathbf{z}_1^{(l)} &= \mathbf{f}^{(l)} \odot \mathbf{U}+\left(1-\mathbf{f}^{(l)}\right) \odot \mathbf{V} \\ + \mathbf{g}^{(l)} &= \sigma\left(\mathbf{W}_2^{(l)} \mathbf{z}_1^{(l)}+\mathbf{b}_2^{(l)}\right) \\ + \mathbf{z}_2^{(l)} &= \mathbf{g}^{(l)} \odot \mathbf{U}+\left(1-\mathbf{g}^{(l)}\right) \odot \mathbf{V} \\ + \mathbf{h}^{(l)} &= \sigma\left(\mathbf{W}_3^{(l)} \mathbf{z}_2^{(l)}+\mathbf{b}_3^{(l)}\right) \\ + \mathbf{x}^{(l+1)} &= \text{PirateBlock}^{(l)}\left(\mathbf{x}^{(l)}\right), l=1...L-1\\ + \mathbf{u}_\theta &= \mathbf{W}^{(L+1)} \mathbf{x}^{(L)} + \end{align*} + $$ + + Args: + input_keys (Tuple[str, ...]): Name of input keys, such as ("x", "y", "z"). + output_keys (Tuple[str, ...]): Name of output keys, such as ("u", "v", "w"). + num_blocks (int): Number of PirateBlocks. + hidden_size (Union[int, Tuple[int, ...]]): Number of hidden size. + An integer for all layers, or list of integer specify each layer's size. + activation (str, optional): Name of activation function. Defaults to "tanh". + weight_norm (bool, optional): Whether to apply weight norm on parameter(s). Defaults to False. + input_dim (Optional[int]): Number of input's dimension. Defaults to None. + output_dim (Optional[int]): Number of output's dimension. Defaults to None. + periods (Optional[Dict[int, Tuple[float, bool]]]): Period of each input key, + input in given channel will be period embedded if specified, each tuple of + periods list is [period, trainable]. Defaults to None. + fourier (Optional[Dict[str, Union[float, int]]]): Random fourier feature embedding, + e.g. {'dim': 256, 'scale': 1.0}. Defaults to None. + random_weight (Optional[Dict[str, float]]): Mean and std of random weight + factorization layer, e.g. {"mean": 0.5, "std: 0.1"}. Defaults to None. + + Examples: + >>> import paddle + >>> import ppsci + >>> model = ppsci.arch.PirateNet( + ... input_keys=("x", "y"), + ... output_keys=("u", "v"), + ... num_blocks=3, + ... hidden_size=256, + ... fourier={'dim': 256, 'scale': 1.0}, + ... ) + >>> input_dict = {"x": paddle.rand([64, 1]), + ... "y": paddle.rand([64, 1])} + >>> output_dict = model(input_dict) + >>> print(output_dict["u"].shape) + [64, 1] + >>> print(output_dict["v"].shape) + [64, 1] + """ + + def __init__( + self, + input_keys: Tuple[str, ...], + output_keys: Tuple[str, ...], + num_blocks: int, + hidden_size: int, + activation: str = "tanh", + weight_norm: bool = False, + input_dim: Optional[int] = None, + output_dim: Optional[int] = None, + periods: Optional[Dict[int, Tuple[float, bool]]] = None, + fourier: Optional[Dict[str, Union[float, int]]] = None, + random_weight: Optional[Dict[str, float]] = None, + ): + super().__init__() + self.input_keys = input_keys + self.output_keys = output_keys + self.blocks = [] + self.periods = periods + self.fourier = fourier + if periods: + self.period_emb = PeriodEmbedding(periods) + + if isinstance(hidden_size, int): + if not isinstance(num_blocks, int): + raise ValueError("num_blocks should be an int") + hidden_size = [hidden_size] * num_blocks + else: + raise ValueError(f"hidden_size should be int, but got {type(hidden_size)}") + + # initialize FC layer(s) + cur_size = len(self.input_keys) if input_dim is None else input_dim + if input_dim is None and periods: + # period embedded channel(s) will be doubled automatically + # if input_dim is not specified + cur_size += len(periods) + + if fourier: + self.fourier_emb = FourierEmbedding( + cur_size, fourier["dim"], fourier["scale"] + ) + cur_size = fourier["dim"] + else: + self.linear_emb = nn.Linear(cur_size, hidden_size[0]) + cur_size = hidden_size[0] + + self.embed_u = nn.Sequential( + ( + WeightNormLinear(cur_size, hidden_size[0]) + if weight_norm + else ( + nn.Linear(cur_size, hidden_size[0]) + if random_weight is None + else RandomWeightFactorization( + cur_size, + hidden_size[0], + mean=random_weight["mean"], + std=random_weight["std"], + ) + ) + ), + ( + act_mod.get_activation(activation) + if activation != "stan" + else act_mod.get_activation(activation)(hidden_size[0]) + ), + ) + self.embed_v = nn.Sequential( + ( + WeightNormLinear(cur_size, hidden_size[0]) + if weight_norm + else ( + nn.Linear(cur_size, hidden_size[0]) + if random_weight is None + else RandomWeightFactorization( + cur_size, + hidden_size[0], + mean=random_weight["mean"], + std=random_weight["std"], + ) + ) + ), + ( + act_mod.get_activation(activation) + if activation != "stan" + else act_mod.get_activation(activation)(hidden_size[0]) + ), + ) + + for i, _size in enumerate(hidden_size): + self.blocks.append( + PirateNetBlock( + cur_size, + _size, + activation=activation, + random_weight=random_weight, + ) + ) + cur_size = _size + + self.blocks = nn.LayerList(self.blocks) + if random_weight: + self.last_fc = RandomWeightFactorization( + cur_size, + len(self.output_keys) if output_dim is None else output_dim, + mean=random_weight["mean"], + std=random_weight["std"], + ) + else: + self.last_fc = nn.Linear( + cur_size, + len(self.output_keys) if output_dim is None else output_dim, + ) + + def forward_tensor(self, x): + u = self.embed_u(x) + v = self.embed_v(x) + + y = x + for i, block in enumerate(self.blocks): + y = block(y, u, v) + + y = self.last_fc(y) + return y + + def forward(self, x): + if self._input_transform is not None: + x = self._input_transform(x) + + if self.periods: + x = self.period_emb(x) + + y = self.concat_to_tensor(x, self.input_keys, axis=-1) + + if self.fourier: + y = self.fourier_emb(y) + else: + y = self.linear_emb(y) + + y = self.forward_tensor(y) + y = self.split_to_dict(y, self.output_keys, axis=-1) + + if self._output_transform is not None: + y = self._output_transform(x, y) + return y diff --git a/examples/smc_reac/ppsci/arch/model_list.py b/examples/smc_reac/ppsci/arch/model_list.py new file mode 100644 index 0000000000..f5f7feeb8b --- /dev/null +++ b/examples/smc_reac/ppsci/arch/model_list.py @@ -0,0 +1,72 @@ +# 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. + +from __future__ import annotations + +from typing import Tuple + +from paddle import nn + +from ppsci.arch import base + + +class ModelList(base.Arch): + """ModelList layer which wrap more than one model that shares inputs. + + Args: + model_list (Tuple[base.Arch, ...]): Model(s) nested in tuple. + + Examples: + >>> import paddle + >>> import ppsci + >>> model1 = ppsci.arch.MLP(("x", "y"), ("u", "v"), 10, 128) + >>> model2 = ppsci.arch.MLP(("x", "y"), ("w", "p"), 5, 128) + >>> model = ppsci.arch.ModelList((model1, model2)) + >>> input_dict = {"x": paddle.rand([64, 64, 1]),"y": paddle.rand([64, 64, 1])} + >>> output_dict = model(input_dict) + >>> for k, v in output_dict.items(): + ... print(k, v.shape) + u [64, 64, 1] + v [64, 64, 1] + w [64, 64, 1] + p [64, 64, 1] + """ + + def __init__( + self, + model_list: Tuple[base.Arch, ...], + ): + super().__init__() + self.input_keys = sum([model.input_keys for model in model_list], ()) + self.input_keys = set(self.input_keys) + + output_keys_set = set() + for model in model_list: + if len(output_keys_set & set(model.output_keys)): + raise ValueError( + "output_keys of model from model_list should be unique," + f"but got duplicate keys: {output_keys_set & set(model.output_keys)}" + ) + output_keys_set = output_keys_set | set(model.output_keys) + self.output_keys = tuple(output_keys_set) + + self.model_list = nn.LayerList(model_list) + + def forward(self, x): + y_all = {} + for model in self.model_list: + y = model(x) + y_all.update(y) + + return y_all diff --git a/examples/smc_reac/ppsci/arch/moflow_basic.py b/examples/smc_reac/ppsci/arch/moflow_basic.py new file mode 100644 index 0000000000..68f10efafb --- /dev/null +++ b/examples/smc_reac/ppsci/arch/moflow_basic.py @@ -0,0 +1,297 @@ +# Copyright (c) 2024 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. + +# Copyright 2020 Chengxi Zang + +import numpy as np +import paddle +import paddle.nn as nn +from scipy import linalg as la + + +# logabs = lambda x: paddle.log(x=paddle.abs(x=x)) +def logabs(x): + return paddle.log(paddle.abs(x)) + + +class ActNorm(nn.Layer): + def __init__(self, in_channel, logdet=True): + super().__init__() + self.loc = self.create_parameter( + [1, in_channel, 1, 1], + default_initializer=nn.initializer.Constant(value=0.0), + ) + + self.scale = self.create_parameter( + [1, in_channel, 1, 1], + default_initializer=nn.initializer.Constant(value=1.0), + ) + + self.register_buffer( + name="initialized", tensor=paddle.to_tensor(data=0, dtype="uint8") + ) + self.logdet = logdet + + def initialize(self, input): + with paddle.no_grad(): + flatten = input.transpose(perm=[1, 0, 2, 3]).reshape( + [tuple(input.shape)[1], -1] + ) + mean = ( + flatten.mean(axis=1) + .unsqueeze(axis=1) + .unsqueeze(axis=2) + .unsqueeze(axis=3) + .transpose(perm=[1, 0, 2, 3]) + ) + std = ( + flatten.std(axis=1) + .unsqueeze(axis=1) + .unsqueeze(axis=2) + .unsqueeze(axis=3) + .transpose(perm=[1, 0, 2, 3]) + ) + paddle.assign(-mean, output=self.loc.data) + paddle.assign(1 / (std + 1e-06), output=self.scale.data) + + def forward(self, input): + _, _, height, width = tuple(input.shape) + if self.initialized.item() == 0: + self.initialize(input) + self.initialized.fill_(value=1) + log_abs = logabs(self.scale) + logdet = height * width * paddle.sum(x=log_abs) + if self.logdet: + return self.scale * (input + self.loc), logdet + else: + return self.scale * (input + self.loc) + + def reverse(self, output): + return output / self.scale - self.loc + + +class ActNorm2D(nn.Layer): + def __init__(self, in_dim, logdet=True): + super().__init__() + self.loc = self.create_parameter( + [1, in_dim, 1], + default_initializer=nn.initializer.Constant(value=0.0), + ) + + self.scale = self.create_parameter( + [1, in_dim, 1], + default_initializer=nn.initializer.Constant(value=1.0), + ) + + self.register_buffer( + name="initialized", tensor=paddle.to_tensor(data=0, dtype="uint8") + ) + self.logdet = logdet + + def initialize(self, input): + with paddle.no_grad(): + flatten = input.transpose(perm=[1, 0, 2]).reshape( + [tuple(input.shape)[1], -1] + ) + mean = ( + flatten.mean(axis=1) + .unsqueeze(axis=1) + .unsqueeze(axis=2) + .transpose(perm=[1, 0, 2]) + ) + std = ( + flatten.std(axis=1) + .unsqueeze(axis=1) + .unsqueeze(axis=2) + .transpose(perm=[1, 0, 2]) + ) + paddle.assign(-mean, output=self.loc.data) + paddle.assign(1 / (std + 1e-06), output=self.scale.data) + + def forward(self, input): + _, _, height = tuple(input.shape) + if self.initialized.item() == 0: + self.initialize(input) + self.initialized.fill_(value=1) + log_abs = logabs(self.scale) + logdet = height * paddle.sum(x=log_abs) + if self.logdet: + return self.scale * (input + self.loc), logdet + else: + return self.scale * (input + self.loc) + + def reverse(self, output): + return output / self.scale - self.loc + + +class InvConv2d(nn.Layer): + def __init__(self, in_channel): + super().__init__() + weight = paddle.randn([in_channel, in_channel]) + q, _ = paddle.linalg.qr(weight) + weight = q.unsqueeze(2).unsqueeze(3) + self.weight = paddle.create_parameter( + weight.shape, + weight.numpy().dtype, + default_initializer=nn.initializer.Assign(weight), + ) + + def forward(self, input): + _, _, height, width = tuple(input.shape) + out = nn.functional.conv2d(x=input, weight=self.weight) + res = paddle.linalg.slogdet(self.weight.squeeze().astype(dtype="float64")) + logdet = height * width * (res[0], res[1])[1].astype(dtype="float32") + return out, logdet + + def reverse(self, output): + return nn.functional.conv2d( + x=output, + weight=self.weight.squeeze().inverse().unsqueeze(axis=2).unsqueeze(axis=3), + ) + + +class InvConv2dLU(nn.Layer): + def __init__(self, in_channel): + super().__init__() + weight = np.random.randn(in_channel, in_channel) + q, _ = la.qr(weight) + w_p, w_l, w_u = la.lu(q.astype(np.float32)) + w_s = np.diag(w_u) + w_u = np.triu(w_u, 1) + u_mask = np.triu(np.ones_like(w_u), 1) + l_mask = u_mask.T + w_p = paddle.to_tensor(data=w_p) + w_l = paddle.to_tensor(data=w_l) + w_s = paddle.to_tensor(data=w_s) + w_u = paddle.to_tensor(data=w_u) + self.register_buffer(name="w_p", tensor=w_p) + self.register_buffer(name="u_mask", tensor=paddle.to_tensor(data=u_mask)) + self.register_buffer(name="l_mask", tensor=paddle.to_tensor(data=l_mask)) + self.register_buffer(name="s_sign", tensor=paddle.sign(x=w_s)) + self.register_buffer( + name="l_eye", tensor=paddle.eye(num_rows=tuple(l_mask.shape)[0]) + ) + self.w_l = paddle.create_parameter( + w_l.shape, + w_l.numpy().dtype, + default_initializer=nn.initializer.Assign(w_l), + ) + + self.w_s = paddle.create_parameter( + logabs(w_s).shape, + logabs(w_s).numpy().dtype, + default_initializer=nn.initializer.Assign(logabs(w_s)), + ) + + self.w_u = paddle.create_parameter( + w_u.shape, + w_u.numpy().dtype, + default_initializer=nn.initializer.Assign(w_u), + ) + + def forward(self, input): + _, _, height, width = tuple(input.shape) + weight = self.calc_weight() + out = nn.functional.conv2d(x=input, weight=weight) + logdet = height * width * paddle.sum(x=self.w_s) + return out, logdet + + def calc_weight(self): + weight = ( + self.w_p + @ (self.w_l * self.l_mask + self.l_eye) + @ ( + self.w_u * self.u_mask + + paddle.diag(x=self.s_sign * paddle.exp(x=self.w_s)) + ) + ) + return weight.unsqueeze(axis=2).unsqueeze(axis=3) + + def reverse(self, output): + weight = self.calc_weight() + return nn.functional.conv2d( + x=output, + weight=weight.squeeze().inverse().unsqueeze(axis=2).unsqueeze(axis=3), + ) + + +class GraphLinear(nn.Layer): + """Graph Linear layer. + This function assumes its input is 3-dimensional. Or 4-dim or whatever, only last dim are changed + Differently from :class:`nn.Linear`, it applies an affine + transformation to the third axis of input `x`. + Warning: original Chainer.link.Link use i.i.d. Gaussian initialization as default, + while default nn.Linear initialization using init.kaiming_uniform_ + """ + + def __init__(self, in_size, out_size, bias=True): + super(GraphLinear, self).__init__() + self.in_size = in_size + self.out_size = out_size + self.linear = nn.Linear( + in_features=in_size, out_features=out_size, bias_attr=bias + ) + + def forward(self, x): + """Forward propagation. + Args: + x (:class:`chainer.Variable`, or :class:`numpy.ndarray` ): + Input array that should be a float array whose ``ndim`` is 3. + + It represents a minibatch of atoms, each of which consists + of a sequence of molecules. Each molecule is represented + by integer IDs. The first axis is an index of atoms + (i.e. minibatch dimension) and the second one an index + of molecules. + + Returns: + class:`chainer.Variable`: + A 3-dimeisional array. + + """ + h = x + h = h.reshape([-1, tuple(x.shape)[-1]]) + h = self.linear(h) + h = h.reshape(tuple(tuple(x.shape)[:-1] + (self.out_size,))) + return h + + +class GraphConv(nn.Layer): + """ + graph convolution over batch and multi-graphs + Args: + in_channels: e.g. 8 + out_channels: e.g. 64 + num_edge_type (types of edges/bonds): e.g. 4 + return: + class:`chainer.Variable`: + """ + + def __init__(self, in_channels, out_channels, num_edge_type=4): + super(GraphConv, self).__init__() + self.graph_linear_self = GraphLinear(in_channels, out_channels) + self.graph_linear_edge = GraphLinear(in_channels, out_channels * num_edge_type) + self.num_edge_type = num_edge_type + self.in_ch = in_channels + self.out_ch = out_channels + + def forward(self, adj, h): + mb, node, ch = tuple(h.shape) + hs = self.graph_linear_self(h) + m = self.graph_linear_edge(h) + m = m.reshape([mb, node, self.out_ch, self.num_edge_type]) + m = m.transpose(perm=[0, 3, 1, 2]) + hr = paddle.matmul(x=adj, y=m) + hr = hr.sum(axis=1) + return hs + hr diff --git a/examples/smc_reac/ppsci/arch/moflow_glow.py b/examples/smc_reac/ppsci/arch/moflow_glow.py new file mode 100644 index 0000000000..5fbeb71520 --- /dev/null +++ b/examples/smc_reac/ppsci/arch/moflow_glow.py @@ -0,0 +1,477 @@ +# Copyright (c) 2024 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. + +# Copyright 2020 Chengxi Zang + +import warnings + +import paddle +import paddle.nn as nn + +from ppsci.arch.moflow_basic import ActNorm +from ppsci.arch.moflow_basic import ActNorm2D +from ppsci.arch.moflow_basic import GraphConv +from ppsci.arch.moflow_basic import GraphLinear +from ppsci.arch.moflow_basic import InvConv2d +from ppsci.arch.moflow_basic import InvConv2dLU + +warnings.filterwarnings( + "ignore", message="when training, we now always track global mean and variance." +) + + +class AffineCoupling(nn.Layer): + def __init__(self, in_channel, hidden_channels, affine=True, mask_swap=False): + super(AffineCoupling, self).__init__() + self.affine = affine + self.layers = nn.LayerList() + self.norms = nn.LayerList() + self.mask_swap = mask_swap + last_h = in_channel // 2 + if affine: + vh = tuple(hidden_channels) + (in_channel,) + else: + vh = tuple(hidden_channels) + (in_channel // 2,) + for h in vh: + self.layers.append( + nn.Conv2D(in_channels=last_h, out_channels=h, kernel_size=3, padding=1) + ) + self.norms.append(nn.BatchNorm2D(num_features=h)) + last_h = h + + def forward(self, input): + in_a, in_b = input.chunk(chunks=2, axis=1) + if self.mask_swap: + in_a, in_b = in_b, in_a + if self.affine: + s, t = self._s_t_function(in_a) + out_b = (in_b + t) * s + logdet = paddle.sum( + x=paddle.log(x=paddle.abs(x=s)).reshape([tuple(input.shape)[0], -1]), + axis=1, + ) + else: + _, t = self._s_t_function(in_a) + out_b = in_b + t + logdet = None + if self.mask_swap: + result = paddle.concat(x=[out_b, in_a], axis=1) + else: + result = paddle.concat(x=[in_a, out_b], axis=1) + return result, logdet + + def reverse(self, output): + out_a, out_b = output.chunk(chunks=2, axis=1) + if self.mask_swap: + out_a, out_b = out_b, out_a + if self.affine: + s, t = self._s_t_function(out_a) + in_b = out_b / s - t + else: + _, t = self._s_t_function(out_a) + in_b = out_b - t + if self.mask_swap: + result = paddle.concat(x=[in_b, out_a], axis=1) + else: + result = paddle.concat(x=[out_a, in_b], axis=1) + return result + + def _s_t_function(self, x): + h = x + for i in range(len(self.layers) - 1): + h = self.layers[i](h) + h = self.norms[i](h) + h = nn.functional.relu(x=h) + h = self.layers[-1](h) + s = None + if self.affine: + log_s, t = h.chunk(chunks=2, axis=1) + s = nn.functional.sigmoid(x=log_s) + else: + t = h + return s, t + + +class GraphAffineCoupling(nn.Layer): + def __init__(self, n_node, in_dim, hidden_dim_dict, masked_row, affine=True): + super(GraphAffineCoupling, self).__init__() + self.n_node = n_node + self.in_dim = in_dim + self.hidden_dim_dict = hidden_dim_dict + self.masked_row = masked_row + self.affine = affine + self.hidden_dim_gnn = hidden_dim_dict["gnn"] + self.hidden_dim_linear = hidden_dim_dict["linear"] + self.net = nn.LayerList() + self.norm = nn.LayerList() + last_dim = in_dim + for out_dim in self.hidden_dim_gnn: + self.net.append(GraphConv(last_dim, out_dim)) + self.norm.append(nn.BatchNorm1D(num_features=n_node)) + last_dim = out_dim + self.net_lin = nn.LayerList() + self.norm_lin = nn.LayerList() + for out_dim in self.hidden_dim_linear: + self.net_lin.append(GraphLinear(last_dim, out_dim)) + self.norm_lin.append(nn.BatchNorm1D(num_features=n_node)) + last_dim = out_dim + if affine: + self.net_lin.append(GraphLinear(last_dim, in_dim * 2)) + else: + self.net_lin.append(GraphLinear(last_dim, in_dim)) + self.scale = paddle.create_parameter( + paddle.zeros(shape=[1]).shape, + paddle.zeros(shape=[1]).numpy().dtype, + default_initializer=nn.initializer.Assign(paddle.zeros(shape=[1])), + ) + + mask = paddle.ones(shape=[n_node, in_dim]) + mask[masked_row, :] = 0 + self.register_buffer(name="mask", tensor=mask) + + def forward(self, adj, input): + masked_x = self.mask * input + s, t = self._s_t_function(adj, masked_x) + if self.affine: + out = masked_x + (1 - self.mask) * (input + t) * s + logdet = paddle.sum( + x=paddle.log(x=paddle.abs(x=s)).reshape([tuple(input.shape)[0], -1]), + axis=1, + ) + else: + out = masked_x + t * (1 - self.mask) + logdet = None + return out, logdet + + def reverse(self, adj, output): + masked_y = self.mask * output + s, t = self._s_t_function(adj, masked_y) + if self.affine: + input = masked_y + (1 - self.mask) * (output / s - t) + else: + input = masked_y + (1 - self.mask) * (output - t) + return input + + def _s_t_function(self, adj, x): + s = None + h = x + for i in range(len(self.net)): + h = self.net[i](adj, h) + h = self.norm[i](h) + h = nn.functional.relu(x=h) + for i in range(len(self.net_lin) - 1): + h = self.net_lin[i](h) + h = self.norm_lin[i](h) + h = nn.functional.relu(x=h) + h = self.net_lin[-1](h) + if self.affine: + log_s, t = h.chunk(chunks=2, axis=-1) + s = nn.functional.sigmoid(x=log_s) + else: + t = h + return s, t + + +class Flow(nn.Layer): + def __init__( + self, in_channel, hidden_channels, affine=True, conv_lu=2, mask_swap=False + ): + super(Flow, self).__init__() + self.actnorm = ActNorm(in_channel) + if conv_lu == 0: + self.invconv = InvConv2d(in_channel) + elif conv_lu == 1: + self.invconv = InvConv2dLU(in_channel) + elif conv_lu == 2: + self.invconv = None + else: + raise ValueError( + "conv_lu in {0,1,2}, 0:InvConv2d, 1:InvConv2dLU, 2:none-just swap to update in coupling" + ) + self.coupling = AffineCoupling( + in_channel, hidden_channels, affine=affine, mask_swap=mask_swap + ) + + def forward(self, input): + out, logdet = self.actnorm(input) + if self.invconv: + out, det1 = self.invconv(out) + else: + det1 = 0 + out, det2 = self.coupling(out) + logdet = logdet + det1 + if det2 is not None: + logdet = logdet + det2 + return out, logdet + + def reverse(self, output): + input = self.coupling.reverse(output) + if self.invconv: + input = self.invconv.reverse(input) + input = self.actnorm.reverse(input) + return input + + +class FlowOnGraph(nn.Layer): + def __init__(self, n_node, in_dim, hidden_dim_dict, masked_row, affine=True): + super(FlowOnGraph, self).__init__() + self.n_node = n_node + self.in_dim = in_dim + self.hidden_dim_dict = hidden_dim_dict + self.masked_row = masked_row + self.affine = affine + self.actnorm = ActNorm2D(in_dim=n_node) + self.coupling = GraphAffineCoupling( + n_node, in_dim, hidden_dim_dict, masked_row, affine=affine + ) + + def forward(self, adj, input): + out, logdet = self.actnorm(input) + det1 = 0 + out, det2 = self.coupling(adj, out) + logdet = logdet + det1 + if det2 is not None: + logdet = logdet + det2 + return out, logdet + + def reverse(self, adj, output): + input = self.coupling.reverse(adj, output) + input = self.actnorm.reverse(input) + return input + + +class Block(nn.Layer): + def __init__( + self, in_channel, n_flow, squeeze_fold, hidden_channels, affine=True, conv_lu=2 + ): + super(Block, self).__init__() + self.squeeze_fold = squeeze_fold + squeeze_dim = in_channel * self.squeeze_fold * self.squeeze_fold + self.flows = nn.LayerList() + for i in range(n_flow): + if conv_lu in (0, 1): + self.flows.append( + Flow( + squeeze_dim, + hidden_channels, + affine=affine, + conv_lu=conv_lu, + mask_swap=False, + ) + ) + else: + self.flows.append( + Flow( + squeeze_dim, + hidden_channels, + affine=affine, + conv_lu=2, + mask_swap=bool(i % 2), + ) + ) + + def forward(self, input): + out = self._squeeze(input) + logdet = 0 + for flow in self.flows: + out, det = flow(out) + logdet = logdet + det + out = self._unsqueeze(out) + return out, logdet + + def reverse(self, output): + input = self._squeeze(output) + for flow in self.flows[::-1]: + input = flow.reverse(input) + unsqueezed = self._unsqueeze(input) + return unsqueezed + + def _squeeze(self, x): + """Trade spatial extent for channels. In forward direction, convert each + 1x4x4 volume of input into a 4x1x1 volume of output. + + Args: + x (paddle.Tensor): Input to squeeze or unsqueeze. + reverse (bool): Reverse the operation, i.e., unsqueeze. + + Returns: + x (paddle.Tensor): Squeezed or unsqueezed tensor. + """ + assert len(tuple(x.shape)) == 4 + b_size, n_channel, height, width = tuple(x.shape) + fold = self.squeeze_fold + squeezed = x.reshape( + [b_size, n_channel, height // fold, fold, width // fold, fold] + ) + squeezed = squeezed.transpose(perm=[0, 1, 3, 5, 2, 4]).contiguous() + out = squeezed.reshape( + [b_size, n_channel * fold * fold, height // fold, width // fold] + ) + return out + + def _unsqueeze(self, x): + assert len(tuple(x.shape)) == 4 + b_size, n_channel, height, width = tuple(x.shape) + fold = self.squeeze_fold + unsqueezed = x.reshape( + [b_size, n_channel // (fold * fold), fold, fold, height, width] + ) + unsqueezed = unsqueezed.transpose(perm=[0, 1, 4, 2, 5, 3]).contiguous() + out = unsqueezed.reshape( + [b_size, n_channel // (fold * fold), height * fold, width * fold] + ) + return out + + +class BlockOnGraph(nn.Layer): + def __init__( + self, + n_node, + in_dim, + hidden_dim_dict, + n_flow, + mask_row_size=1, + mask_row_stride=1, + affine=True, + ): + """ + + :param n_node: + :param in_dim: + :param hidden_dim: + :param n_flow: + :param mask_row_size: number of rows to be masked for update + :param mask_row_stride: number of steps between two masks' firs row + :param affine: + """ + super(BlockOnGraph, self).__init__() + assert 0 < mask_row_size < n_node + self.flows = nn.LayerList() + for i in range(n_flow): + start = i * mask_row_stride + masked_row = [(r % n_node) for r in range(start, start + mask_row_size)] + self.flows.append( + FlowOnGraph( + n_node, + in_dim, + hidden_dim_dict, + masked_row=masked_row, + affine=affine, + ) + ) + + def forward(self, adj, input): + out = input + logdet = 0 + for flow in self.flows: + out, det = flow(adj, out) + logdet = logdet + det + return out, logdet + + def reverse(self, adj, output): + input = output + for flow in self.flows[::-1]: + input = flow.reverse(adj, input) + return input + + +class Glow(nn.Layer): + def __init__( + self, + in_channel, + n_flow, + n_block, + squeeze_fold, + hidden_channel, + affine=True, + conv_lu=2, + ): + super(Glow, self).__init__() + self.blocks = nn.LayerList() + n_channel = in_channel + for i in range(n_block): + self.blocks.append( + Block( + n_channel, + n_flow, + squeeze_fold, + hidden_channel, + affine=affine, + conv_lu=conv_lu, + ) + ) + + def forward(self, input): + logdet = 0 + out = input + for block in self.blocks: + out, det = block(out) + logdet = logdet + det + return out, logdet + + def reverse(self, z): + h = z + for i, block in enumerate(self.blocks[::-1]): + h = block.reverse(h) + return h + + +class GlowOnGraph(nn.Layer): + def __init__( + self, + n_node, + in_dim, + hidden_dim_dict, + n_flow, + n_block, + mask_row_size_list=[2], + mask_row_stride_list=[1], + affine=True, + ): + super(GlowOnGraph, self).__init__() + assert len(mask_row_size_list) == n_block or len(mask_row_size_list) == 1 + assert len(mask_row_stride_list) == n_block or len(mask_row_stride_list) == 1 + if len(mask_row_size_list) == 1: + mask_row_size_list = mask_row_size_list * n_block + if len(mask_row_stride_list) == 1: + mask_row_stride_list = mask_row_stride_list * n_block + self.blocks = nn.LayerList() + for i in range(n_block): + mask_row_size = mask_row_size_list[i] + mask_row_stride = mask_row_stride_list[i] + self.blocks.append( + BlockOnGraph( + n_node, + in_dim, + hidden_dim_dict, + n_flow, + mask_row_size, + mask_row_stride, + affine=affine, + ) + ) + + def forward(self, adj, x): + logdet = 0 + out = x + for block in self.blocks: + out, det = block(adj, out) + logdet = logdet + det + return out, logdet + + def reverse(self, adj, z): + input = z + for i, block in enumerate(self.blocks[::-1]): + input = block.reverse(adj, input) + return input diff --git a/examples/smc_reac/ppsci/arch/moflow_net.py b/examples/smc_reac/ppsci/arch/moflow_net.py new file mode 100644 index 0000000000..c2f88607bc --- /dev/null +++ b/examples/smc_reac/ppsci/arch/moflow_net.py @@ -0,0 +1,335 @@ +# Copyright (c) 2024 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. + +# Copyright 2020 Chengxi Zang +from __future__ import annotations + +import math +from typing import Dict +from typing import Tuple + +import paddle + +from ppsci.arch import base +from ppsci.arch.moflow_glow import Glow +from ppsci.arch.moflow_glow import GlowOnGraph + + +def gaussian_nll(x, mean, ln_var, reduce="sum"): + """Computes the negative log-likelihood of a Gaussian distribution. + + Given two variable ``mean`` representing :math:`\\mu` and ``ln_var`` + representing :math:`\\log(\\sigma^2)`, this function computes in + elementwise manner the negative log-likelihood of :math:`x` on a + Gaussian distribution :math:`N(\\mu, S)`, + + .. math:: + + -\\log N(x; \\mu, \\sigma^2) = + \\log\\left(\\sqrt{(2\\pi)^D |S|}\\right) + + \\frac{1}{2}(x - \\mu)^\\top S^{-1}(x - \\mu), + + where :math:`D` is a dimension of :math:`x` and :math:`S` is a diagonal + matrix where :math:`S_{ii} = \\sigma_i^2`. + + The output is a variable whose value depends on the value of + the option ``reduce``. If it is ``'no'``, it holds the elementwise + loss values. If it is ``'sum'``, loss values are summed up. + + Args: + x (:class:`~chainer.Variable` or :ref:`ndarray`): Input variable. + mean (:class:`~chainer.Variable` or :ref:`ndarray`): A variable + representing mean of a Gaussian distribution, :math:`\\mu`. + ln_var (:class:`~chainer.Variable` or :ref:`ndarray`): A variable + representing logarithm of variance of a Gaussian distribution, + :math:`\\log(\\sigma^2)`. + reduce (str): Reduction option. Its value must be either + ``'sum'`` or ``'no'``. Otherwise, :class:`ValueError` is raised. + + Returns: + ~chainer.Variable: + A variable representing the negative log-likelihood. + If ``reduce`` is ``'no'``, the output variable holds array + whose shape is same as one of (hence both of) input variables. + If it is ``'sum'``, the output variable holds a scalar value. + + """ + if reduce not in ("sum", "no"): + raise ValueError( + "only 'sum' and 'no' are valid for 'reduce', but '%s' is given" % reduce + ) + x_prec = paddle.exp(x=-ln_var) + x_diff = x - mean + x_power = x_diff * x_diff * x_prec * -0.5 + loss = (ln_var + math.log(2 * math.pi)) / 2 - x_power + if reduce == "sum": + return loss.sum() + else: + return loss + + +def rescale_adj(adj, type="all"): + if type == "view": + out_degree = adj.sum(axis=-1) + out_degree_sqrt_inv = out_degree.pow(y=-1) + out_degree_sqrt_inv[out_degree_sqrt_inv == float("inf")] = 0 + adj_prime = out_degree_sqrt_inv.unsqueeze(axis=-1) * adj + else: + num_neighbors = adj.sum(axis=(1, 2)).astype(dtype="float32") + num_neighbors_inv = num_neighbors.pow(y=-1) + num_neighbors_inv[num_neighbors_inv == float("inf")] = 0 + adj_prime = num_neighbors_inv[:, None, None, :] * adj + return adj_prime + + +class MoFlowNet(base.Arch): + """ + Args: + input_keys (Tuple[str, ...]): Name of input keys, such as ("nodes","edges",). + output_keys (Tuple[str, ...]): Name of output keys, such as ("output","sum_log_det"). + hyper_params (object): More parameters derived from hyper_params for easy use. + """ + + def __init__( + self, + input_keys: Tuple[str, ...], + output_keys: Tuple[str, ...], + hyper_params: None, + ): + super(MoFlowNet, self).__init__() + self.input_keys = input_keys + self.output_keys = output_keys + self.hyper_params = hyper_params + self.b_n_type = hyper_params.b_n_type + self.a_n_node = hyper_params.a_n_node + self.a_n_type = hyper_params.a_n_type + self.b_size = self.a_n_node * self.a_n_node * self.b_n_type + self.a_size = self.a_n_node * self.a_n_type + self.noise_scale = hyper_params.noise_scale + if hyper_params.learn_dist: + self.ln_var = paddle.create_parameter( + paddle.zeros(shape=[1]).shape, + paddle.zeros(shape=[1]).numpy().dtype, + default_initializer=paddle.nn.initializer.Assign( + paddle.zeros(shape=[1]) + ), + ) + + else: + self.register_buffer(name="ln_var", tensor=paddle.zeros(shape=[1])) + self.bond_model = Glow( + in_channel=hyper_params.b_n_type, + n_flow=hyper_params.b_n_flow, + n_block=hyper_params.b_n_block, + squeeze_fold=hyper_params.b_n_squeeze, + hidden_channel=hyper_params.b_hidden_ch, + affine=hyper_params.b_affine, + conv_lu=hyper_params.b_conv_lu, + ) + self.atom_model = GlowOnGraph( + n_node=hyper_params.a_n_node, + in_dim=hyper_params.a_n_type, + hidden_dim_dict={ + "gnn": hyper_params.a_hidden_gnn, + "linear": hyper_params.a_hidden_lin, + }, + n_flow=hyper_params.a_n_flow, + n_block=hyper_params.a_n_block, + mask_row_size_list=hyper_params.mask_row_size_list, + mask_row_stride_list=hyper_params.mask_row_stride_list, + affine=hyper_params.a_affine, + ) + + def forward(self, x): + h = x[self.input_keys[0]] + adj = x[self.input_keys[1]] + adj_normalized = rescale_adj(adj).to(adj) + + if self.training: + if self.noise_scale == 0: + h = h / 2.0 - 0.5 + paddle.rand(shape=h.shape, dtype=h.dtype) * 0.4 + else: + h = h + paddle.rand(shape=h.shape, dtype=h.dtype) * self.noise_scale + h, sum_log_det_jacs_x = self.atom_model(adj_normalized, h) + if self.training: + if self.noise_scale == 0: + adj = ( + adj / 2.0 + - 0.5 + + paddle.rand(shape=adj.shape, dtype=adj.dtype) * 0.4 + ) + else: + adj = ( + adj + + paddle.rand(shape=adj.shape, dtype=adj.dtype) * self.noise_scale + ) + adj_h, sum_log_det_jacs_adj = self.bond_model(adj) + out = [h, adj_h] + result_dict = { + self.output_keys[0]: out, + self.output_keys[1]: [sum_log_det_jacs_x, sum_log_det_jacs_adj], + } + + return result_dict + + def reverse(self, z, true_adj=None): + """ + Returns a molecule, given its latent vector. + + Args: + z: latent vector. Shape: [B, N*N*M + N*T] (100,369) 369=9*9 * 4 + 9*5 + B = Batch size, N = number of atoms, M = number of bond types, + T = number of atom types (Carbon, Oxygen etc.) + true_adj: used for testing. An adjacency matrix of a real molecule + + return: + adjacency matrix and feature matrix of a molecule + """ + batch_size = tuple(z.shape)[0] + with paddle.no_grad(): + z_x = z[:, : self.a_size] + z_adj = z[:, self.a_size :] + if true_adj is None: + h_adj = z_adj.reshape( + [batch_size, self.b_n_type, self.a_n_node, self.a_n_node] + ) + h_adj = self.bond_model.reverse(h_adj) + if self.noise_scale == 0: + h_adj = (h_adj + 0.5) * 2 + adj = h_adj + adj = adj + adj.transpose(perm=[0, 1, 3, 2]) + adj = adj / 2 + adj = paddle.nn.functional.softmax(adj, axis=1) + max_bond = adj.max(axis=1).reshape( + [batch_size, -1, self.a_n_node, self.a_n_node] + ) + adj = paddle.floor(x=adj / max_bond) + else: + adj = true_adj + h_x = z_x.reshape([batch_size, self.a_n_node, self.a_n_type]) + adj_normalized = rescale_adj(adj).to(h_x) + h_x = self.atom_model.reverse(adj_normalized, h_x) + if self.noise_scale == 0: + h_x = (h_x + 0.5) * 2 + return adj, h_x + + def log_prob_loss(self, output_dict: Dict, *args): + losses = 0 + z = output_dict[self.output_keys[0]] + logdet = output_dict[self.output_keys[1]] + z[0] = z[0].reshape([tuple(z[0].shape)[0], -1]) + z[1] = z[1].reshape([tuple(z[1].shape)[0], -1]) + logdet[0] = logdet[0] - self.a_size * math.log(2.0) + logdet[1] = logdet[1] - self.b_size * math.log(2.0) + if len(self.ln_var) == 1: + ln_var_adj = self.ln_var * paddle.ones(shape=[self.b_size]).to(z[0]) + ln_var_x = self.ln_var * paddle.ones(shape=[self.a_size]).to(z[0]) + else: + ln_var_adj = self.ln_var[0] * paddle.ones(shape=[self.b_size]).to(z[0]) + ln_var_x = self.ln_var[1] * paddle.ones(shape=[self.a_size]).to(z[0]) + nll_adj = paddle.mean( + paddle.sum( + gaussian_nll( + z[1], + paddle.zeros(shape=self.b_size).to(z[0]), + ln_var_adj, + reduce="no", + ), + axis=1, + ) + - logdet[1] + ) + nll_adj = nll_adj / (self.b_size * math.log(2.0)) + nll_x = paddle.mean( + paddle.sum( + gaussian_nll( + z[0], + paddle.zeros(shape=self.a_size).to(z[0]), + ln_var_x, + reduce="no", + ), + axis=1, + ) + - logdet[0] + ) + nll_x = nll_x / (self.a_size * math.log(2.0)) + if nll_x.item() < 0: + print(f"nll_x: {nll_x.item()}") + losses = nll_x + nll_adj + return {"total_loss": losses} + + def save_hyperparams(self, path): + self.hyper_params.save(path) + + +class MoFlowProp(base.Arch): + """ + Args: + input_keys (Tuple[str, ...]): Name of input keys, such as ("nodes","edges",). + output_keys (Tuple[str, ...]): Name of output keys, such as ("output","sum_log_det"). + model (MoFlowNet): pre-trained model. + hidden_size (int): Hidden dimension list for output regression. + """ + + def __init__( + self, + input_keys: Tuple[str, ...], + output_keys: Tuple[str, ...], + model: MoFlowNet, + hidden_size, + ): + super(MoFlowProp, self).__init__() + self.input_keys = input_keys + self.output_keys = output_keys + self.model = model + self.latent_size = model.b_size + model.a_size + self.hidden_size = hidden_size + vh = (self.latent_size,) + tuple(hidden_size) + (1,) + modules = [] + for i in range(len(vh) - 1): + modules.append(paddle.nn.Linear(in_features=vh[i], out_features=vh[i + 1])) + if i < len(vh) - 2: + modules.append(paddle.nn.Tanh()) + self.propNN = paddle.nn.Sequential(*modules) + + def encode(self, x): + with paddle.no_grad(): + self.model.eval() + output_dict = self.model(x) + z = output_dict["output"] + sum_log_det_jacs = output_dict["sum_log_det"] + h = paddle.concat( + [ + z[0].reshape([tuple(z[0].shape)[0], -1]), + z[1].reshape([tuple(z[1].shape)[0], -1]), + ], + axis=1, + ) + return h, sum_log_det_jacs + + def reverse(self, z): + with paddle.no_grad(): + self.model.eval() + adj, x = self.model.reverse(z, true_adj=None) + return adj, x + + def forward(self, x): + h, sum_log_det_jacs = self.encode(x) + output = self.propNN(h) + result_dict = { + self.output_keys[0]: [h, output], + self.output_keys[1]: sum_log_det_jacs, + } + + return result_dict diff --git a/examples/smc_reac/ppsci/arch/nowcastnet.py b/examples/smc_reac/ppsci/arch/nowcastnet.py new file mode 100644 index 0000000000..bc7538ad91 --- /dev/null +++ b/examples/smc_reac/ppsci/arch/nowcastnet.py @@ -0,0 +1,639 @@ +# 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. + +import collections +from typing import Tuple + +import paddle +from paddle import nn + +from ppsci.arch import base + + +class NowcastNet(base.Arch): + """The NowcastNet model. + + Args: + input_keys (Tuple[str, ...]): Name of input keys, such as ("input",). + output_keys (Tuple[str, ...]): Name of output keys, such as ("output",). + input_length (int, optional): Input length. Defaults to 9. + total_length (int, optional): Total length. Defaults to 29. + image_height (int, optional): Image height. Defaults to 512. + image_width (int, optional): Image width. Defaults to 512. + image_ch (int, optional): Image channel. Defaults to 2. + ngf (int, optional): Noise Projector input length. Defaults to 32. + + Examples: + >>> import ppsci + >>> model = ppsci.arch.NowcastNet(("input", ), ("output", )) + >>> input_data = paddle.rand([1, 9, 512, 512, 2]) + >>> input_dict = {"input": input_data} + >>> output_dict = model(input_dict) + >>> print(output_dict["output"].shape) + [1, 20, 512, 512, 1] + """ + + def __init__( + self, + input_keys: Tuple[str, ...], + output_keys: Tuple[str, ...], + input_length: int = 9, + total_length: int = 29, + image_height: int = 512, + image_width: int = 512, + image_ch: int = 2, + ngf: int = 32, + ): + super().__init__() + self.input_keys = input_keys + self.output_keys = output_keys + + self.input_length = input_length + self.total_length = total_length + self.image_height = image_height + self.image_width = image_width + self.image_ch = image_ch + self.ngf = ngf + + configs = collections.namedtuple( + "Object", ["ngf", "evo_ic", "gen_oc", "ic_feature"] + ) + configs.ngf = self.ngf + configs.evo_ic = self.total_length - self.input_length + configs.gen_oc = self.total_length - self.input_length + configs.ic_feature = self.ngf * 10 + + self.pred_length = self.total_length - self.input_length + self.evo_net = Evolution_Network(self.input_length, self.pred_length, base_c=32) + self.gen_enc = Generative_Encoder(self.total_length, base_c=self.ngf) + self.gen_dec = Generative_Decoder(configs) + self.proj = Noise_Projector(self.ngf) + sample_tensor = paddle.zeros(shape=[1, 1, self.image_height, self.image_width]) + self.grid = make_grid(sample_tensor) + + @staticmethod + def split_to_dict(data_tensors: Tuple[paddle.Tensor, ...], keys: Tuple[str, ...]): + return {key: data_tensors[i] for i, key in enumerate(keys)} + + def forward(self, x): + if self._input_transform is not None: + x = self._input_transform(x) + + x_tensor = self.concat_to_tensor(x, self.input_keys) + + y = [] + out = self.forward_tensor(x_tensor) + y.append(out) + y = self.split_to_dict(y, self.output_keys) + + if self._output_transform is not None: + y = self._output_transform(x, y) + return y + + def forward_tensor(self, x): + all_frames = x[:, :, :, :, :1] + frames = all_frames.transpose(perm=[0, 1, 4, 2, 3]) + batch = frames.shape[0] + height = frames.shape[3] + width = frames.shape[4] + # Input Frames + input_frames = frames[:, : self.input_length] + input_frames = input_frames.reshape((batch, self.input_length, height, width)) + # Evolution Network + intensity, motion = self.evo_net(input_frames) + motion_ = motion.reshape((batch, self.pred_length, 2, height, width)) + intensity_ = intensity.reshape((batch, self.pred_length, 1, height, width)) + series = [] + last_frames = all_frames[:, self.input_length - 1 : self.input_length, :, :, 0] + grid = self.grid.tile((batch, 1, 1, 1)) + for i in range(self.pred_length): + last_frames = warp( + last_frames, motion_[:, i], grid, mode="nearest", padding_mode="border" + ) + last_frames = last_frames + intensity_[:, i] + series.append(last_frames) + evo_result = paddle.concat(x=series, axis=1) + evo_result = evo_result / 128 + # Generative Network + evo_feature = self.gen_enc(paddle.concat(x=[input_frames, evo_result], axis=1)) + noise = paddle.randn(shape=[batch, self.ngf, height // 32, width // 32]) + noise = self.proj(noise) + ngf = noise.shape[1] + noise_feature = ( + noise.reshape((batch, -1, 4, 4, 8, 8)) + .transpose(perm=[0, 1, 4, 5, 2, 3]) + .reshape((batch, ngf // 16, height // 8, width // 8)) + ) + feature = paddle.concat(x=[evo_feature, noise_feature], axis=1) + gen_result = self.gen_dec(feature, evo_result) + return gen_result.unsqueeze(axis=-1) + + +class Evolution_Network(nn.Layer): + def __init__(self, n_channels, n_classes, base_c=64, bilinear=True): + super().__init__() + self.n_channels = n_channels + self.n_classes = n_classes + self.bilinear = bilinear + base_c = base_c + self.inc = DoubleConv(n_channels, base_c) + self.down1 = Down(base_c * 1, base_c * 2) + self.down2 = Down(base_c * 2, base_c * 4) + self.down3 = Down(base_c * 4, base_c * 8) + factor = 2 if bilinear else 1 + self.down4 = Down(base_c * 8, base_c * 16 // factor) + self.up1 = Up(base_c * 16, base_c * 8 // factor, bilinear) + self.up2 = Up(base_c * 8, base_c * 4 // factor, bilinear) + self.up3 = Up(base_c * 4, base_c * 2 // factor, bilinear) + self.up4 = Up(base_c * 2, base_c * 1, bilinear) + self.outc = OutConv(base_c * 1, n_classes) + param1 = paddle.zeros(shape=[1, n_classes, 1, 1]) + gamma = self.create_parameter( + shape=param1.shape, + dtype=param1.dtype, + default_initializer=nn.initializer.Assign(param1), + ) + gamma.stop_gradient = False + self.gamma = gamma + self.up1_v = Up(base_c * 16, base_c * 8 // factor, bilinear) + self.up2_v = Up(base_c * 8, base_c * 4 // factor, bilinear) + self.up3_v = Up(base_c * 4, base_c * 2 // factor, bilinear) + self.up4_v = Up(base_c * 2, base_c * 1, bilinear) + self.outc_v = OutConv(base_c * 1, n_classes * 2) + + def forward(self, x): + x1 = self.inc(x) + x2 = self.down1(x1) + x3 = self.down2(x2) + x4 = self.down3(x3) + x5 = self.down4(x4) + x = self.up1(x5, x4) + x = self.up2(x, x3) + x = self.up3(x, x2) + x = self.up4(x, x1) + x = self.outc(x) * self.gamma + v = self.up1_v(x5, x4) + v = self.up2_v(v, x3) + v = self.up3_v(v, x2) + v = self.up4_v(v, x1) + v = self.outc_v(v) + return x, v + + +class DoubleConv(nn.Layer): + def __init__(self, in_channels, out_channels, kernel=3, mid_channels=None): + super().__init__() + if not mid_channels: + mid_channels = out_channels + self.double_conv = nn.Sequential( + nn.BatchNorm2D(num_features=in_channels), + nn.ReLU(), + nn.utils.spectral_norm( + layer=nn.Conv2D( + in_channels=in_channels, + out_channels=mid_channels, + kernel_size=kernel, + padding=kernel // 2, + ) + ), + nn.BatchNorm2D(num_features=mid_channels), + nn.ReLU(), + nn.utils.spectral_norm( + layer=nn.Conv2D( + in_channels=mid_channels, + out_channels=out_channels, + kernel_size=kernel, + padding=kernel // 2, + ) + ), + ) + self.single_conv = nn.Sequential( + nn.BatchNorm2D(num_features=in_channels), + nn.utils.spectral_norm( + layer=nn.Conv2D( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=kernel, + padding=kernel // 2, + ) + ), + ) + + def forward(self, x): + shortcut = self.single_conv(x) + x = self.double_conv(x) + x = x + shortcut + return x + + +class Down(nn.Layer): + def __init__(self, in_channels, out_channels, kernel=3): + super().__init__() + self.maxpool_conv = nn.Sequential( + nn.MaxPool2D(kernel_size=2), + DoubleConv(in_channels, out_channels, kernel), + ) + + def forward(self, x): + x = self.maxpool_conv(x) + return x + + +class Up(nn.Layer): + def __init__(self, in_channels, out_channels, bilinear=True, kernel=3): + super().__init__() + if bilinear: + self.up = nn.Upsample(scale_factor=2, mode="bilinear", align_corners=True) + self.conv = DoubleConv( + in_channels, out_channels, kernel=kernel, mid_channels=in_channels // 2 + ) + else: + self.up = nn.Conv2DTranspose( + in_channels=in_channels, + out_channels=in_channels // 2, + kernel_size=2, + stride=2, + ) + self.conv = DoubleConv(in_channels, out_channels, kernel) + + def forward(self, x1, x2): + x1 = self.up(x1) + # input is CHW + diffY = x2.shape[2] - x1.shape[2] + diffX = x2.shape[3] - x1.shape[3] + x1 = nn.functional.pad( + x1, [diffX // 2, diffX - diffX // 2, diffY // 2, diffY - diffY // 2] + ) + x = paddle.concat(x=[x2, x1], axis=1) + return self.conv(x) + + +class Up_S(nn.Layer): + def __init__(self, in_channels, out_channels, bilinear=True, kernel=3): + super().__init__() + if bilinear: + self.up = nn.Upsample(scale_factor=2, mode="bilinear", align_corners=True) + self.conv = DoubleConv( + in_channels, out_channels, kernel=kernel, mid_channels=in_channels + ) + else: + self.up = nn.Conv2DTranspose( + in_channels=in_channels, + out_channels=in_channels, + kernel_size=2, + stride=2, + ) + self.conv = DoubleConv(in_channels, out_channels, kernel) + + def forward(self, x): + x = self.up(x) + return self.conv(x) + + +class OutConv(nn.Layer): + def __init__(self, in_channels, out_channels): + super().__init__() + self.conv = nn.Conv2D( + in_channels=in_channels, out_channels=out_channels, kernel_size=1 + ) + + def forward(self, x): + return self.conv(x) + + +class Generative_Encoder(nn.Layer): + def __init__(self, n_channels, base_c=64): + super().__init__() + base_c = base_c + self.inc = DoubleConv(n_channels, base_c, kernel=3) + self.down1 = Down(base_c * 1, base_c * 2, 3) + self.down2 = Down(base_c * 2, base_c * 4, 3) + self.down3 = Down(base_c * 4, base_c * 8, 3) + + def forward(self, x): + x = self.inc(x) + x = self.down1(x) + x = self.down2(x) + x = self.down3(x) + return x + + +class Generative_Decoder(nn.Layer): + def __init__(self, opt): + super().__init__() + self.opt = opt + nf = opt.ngf + ic = opt.ic_feature + self.fc = nn.Conv2D( + in_channels=ic, out_channels=8 * nf, kernel_size=3, padding=1 + ) + self.head_0 = GenBlock(8 * nf, 8 * nf, opt) + self.G_middle_0 = GenBlock(8 * nf, 4 * nf, opt, double_conv=True) + self.G_middle_1 = GenBlock(4 * nf, 4 * nf, opt, double_conv=True) + self.up_0 = GenBlock(4 * nf, 2 * nf, opt) + self.up_1 = GenBlock(2 * nf, 1 * nf, opt, double_conv=True) + self.up_2 = GenBlock(1 * nf, 1 * nf, opt, double_conv=True) + final_nc = nf * 1 + self.conv_img = nn.Conv2D( + in_channels=final_nc, out_channels=self.opt.gen_oc, kernel_size=3, padding=1 + ) + self.up = nn.Upsample(scale_factor=2) + + def forward(self, x, evo): + x = self.fc(x) + x = self.head_0(x, evo) + x = self.up(x) + x = self.G_middle_0(x, evo) + x = self.G_middle_1(x, evo) + x = self.up(x) + x = self.up_0(x, evo) + x = self.up(x) + x = self.up_1(x, evo) + x = self.up_2(x, evo) + x = self.conv_img(nn.functional.leaky_relu(x=x, negative_slope=0.2)) + return x + + +class GenBlock(nn.Layer): + def __init__(self, fin, fout, opt, use_se=False, dilation=1, double_conv=False): + super().__init__() + self.learned_shortcut = fin != fout + fmiddle = min(fin, fout) + self.opt = opt + self.double_conv = double_conv + self.pad = nn.Pad2D(padding=dilation, mode="reflect") + self.conv_0 = nn.Conv2D( + in_channels=fin, + out_channels=fmiddle, + kernel_size=3, + padding=0, + dilation=dilation, + ) + self.conv_1 = nn.Conv2D( + in_channels=fmiddle, + out_channels=fout, + kernel_size=3, + padding=0, + dilation=dilation, + ) + if self.learned_shortcut: + self.conv_s = nn.Conv2D( + in_channels=fin, out_channels=fout, kernel_size=1, bias_attr=False + ) + self.conv_0 = nn.utils.spectral_norm(layer=self.conv_0) + self.conv_1 = nn.utils.spectral_norm(layer=self.conv_1) + if self.learned_shortcut: + self.conv_s = nn.utils.spectral_norm(layer=self.conv_s) + ic = opt.evo_ic + self.norm_0 = SPADE(fin, ic) + self.norm_1 = SPADE(fmiddle, ic) + if self.learned_shortcut: + self.norm_s = SPADE(fin, ic) + + def forward(self, x, evo): + x_s = self.shortcut(x, evo) + dx = self.conv_0(self.pad(self.actvn(self.norm_0(x, evo)))) + if self.double_conv: + dx = self.conv_1(self.pad(self.actvn(self.norm_1(dx, evo)))) + out = x_s + dx + return out + + def shortcut(self, x, evo): + if self.learned_shortcut: + x_s = self.conv_s(self.norm_s(x, evo)) + else: + x_s = x + return x_s + + def actvn(self, x): + return nn.functional.leaky_relu(x=x, negative_slope=0.2) + + +class SPADE(nn.Layer): + def __init__(self, norm_nc, label_nc): + super().__init__() + ks = 3 + self.param_free_norm = nn.InstanceNorm2D( + num_features=norm_nc, weight_attr=False, bias_attr=False, momentum=1 - 0.1 + ) + nhidden = 64 + ks = 3 + pw = ks // 2 + self.mlp_shared = nn.Sequential( + nn.Pad2D(padding=pw, mode="reflect"), + nn.Conv2D( + in_channels=label_nc, out_channels=nhidden, kernel_size=ks, padding=0 + ), + nn.ReLU(), + ) + self.pad = nn.Pad2D(padding=pw, mode="reflect") + self.mlp_gamma = nn.Conv2D( + in_channels=nhidden, out_channels=norm_nc, kernel_size=ks, padding=0 + ) + self.mlp_beta = nn.Conv2D( + in_channels=nhidden, out_channels=norm_nc, kernel_size=ks, padding=0 + ) + + def forward(self, x, evo): + normalized = self.param_free_norm(x) + evo = nn.functional.adaptive_avg_pool2d(x=evo, output_size=x.shape[2:]) + actv = self.mlp_shared(evo) + gamma = self.mlp_gamma(self.pad(actv)) + beta = self.mlp_beta(self.pad(actv)) + out = normalized * (1 + gamma) + beta + return out + + +class Noise_Projector(nn.Layer): + def __init__(self, input_length): + super().__init__() + self.input_length = input_length + self.conv_first = nn.utils.spectral_norm( + nn.Conv2D( + in_channels=self.input_length, + out_channels=self.input_length * 2, + kernel_size=3, + padding=1, + ) + ) + self.L1 = ProjBlock(self.input_length * 2, self.input_length * 4) + self.L2 = ProjBlock(self.input_length * 4, self.input_length * 8) + self.L3 = ProjBlock(self.input_length * 8, self.input_length * 16) + self.L4 = ProjBlock(self.input_length * 16, self.input_length * 32) + + def forward(self, x): + x = self.conv_first(x) + x = self.L1(x) + x = self.L2(x) + x = self.L3(x) + x = self.L4(x) + return x + + +class ProjBlock(nn.Layer): + def __init__(self, in_channel, out_channel): + super().__init__() + self.one_conv = nn.utils.spectral_norm( + nn.Conv2D( + in_channels=in_channel, + out_channels=out_channel - in_channel, + kernel_size=1, + padding=0, + ) + ) + self.double_conv = nn.Sequential( + nn.utils.spectral_norm( + nn.Conv2D( + in_channels=in_channel, + out_channels=out_channel, + kernel_size=3, + padding=1, + ) + ), + nn.ReLU(), + nn.utils.spectral_norm( + nn.Conv2D( + in_channels=out_channel, + out_channels=out_channel, + kernel_size=3, + padding=1, + ) + ), + ) + + def forward(self, x): + x1 = paddle.concat(x=[x, self.one_conv(x)], axis=1) + x2 = self.double_conv(x) + output = x1 + x2 + return output + + +def make_grid(input): + B, C, H, W = input.shape + xx = paddle.arange(start=0, end=W).reshape((1, -1)).tile((H, 1)) + yy = paddle.arange(start=0, end=H).reshape((-1, 1)).tile((1, W)) + xx = xx.reshape((1, 1, H, W)).tile((B, 1, 1, 1)) + yy = yy.reshape((1, 1, H, W)).tile((B, 1, 1, 1)) + grid = paddle.concat(x=(xx, yy), axis=1).astype(dtype=paddle.get_default_dtype()) + return grid + + +def warp(input, flow, grid, mode="bilinear", padding_mode="zeros"): + B, C, H, W = input.shape + vgrid = grid + flow + vgrid[:, 0, :, :] = 2.0 * vgrid[:, 0, :, :].clone() / max(W - 1, 1) - 1.0 + vgrid[:, 1, :, :] = 2.0 * vgrid[:, 1, :, :].clone() / max(H - 1, 1) - 1.0 + vgrid = vgrid.transpose(perm=[0, 2, 3, 1]) + output = nn.functional.grid_sample( + x=input.cpu(), + grid=vgrid.cpu(), + padding_mode=padding_mode, + mode=mode, + align_corners=True, + ) + return output.cuda() + + +def l2normalize(v, eps=1e-12): + return v / (v.norm() + eps) + + +class spectral_norm(nn.Layer): + def __init__(self, module, name="weight", power_iterations=1): + super().__init__() + self.module = module + self.name = name + self.power_iterations = power_iterations + if not self._made_params(): + self._make_params() + + def _update_u_v(self): + u = getattr(self.module, self.name + "_u") + v = getattr(self.module, self.name + "_v") + w = getattr(self.module, self.name + "_bar") + height = w.detach().shape[0] + for _ in range(self.power_iterations): + v = l2normalize( + paddle.mv( + x=paddle.t(input=w.reshape((height, -1)).detach()), vec=u.detach() + ) + ) + u = l2normalize( + paddle.mv(x=w.reshape((height, -1)).detach(), vec=v.detach()) + ) + sigma = u.dot(y=w.reshape((height, -1)).mv(vec=v)) + setattr(self.module, self.name, w / sigma.expand_as(y=w)) + + def _made_params(self): + try: + _ = getattr(self.module, self.name + "_u") + _ = getattr(self.module, self.name + "_v") + _ = getattr(self.module, self.name + "_bar") + return True + except AttributeError: + return False + + def _make_params(self): + w = getattr(self.module, self.name) + height = w.detach().shape[0] + width = w.reshape((height, -1)).detach().shape[1] + + tmp_w = paddle.normal(shape=[height]) + out_0 = paddle.create_parameter( + shape=tmp_w.shape, + dtype=tmp_w.numpy().dtype, + default_initializer=nn.initializer.Assign(tmp_w), + ) + out_0.stop_gradient = True + u = out_0 + + tmp_w = paddle.normal(shape=[width]) + out_1 = paddle.create_parameter( + shape=tmp_w.shape, + dtype=tmp_w.numpy().dtype, + default_initializer=nn.initializer.Assign(tmp_w), + ) + out_1.stop_gradient = True + v = out_1 + u = l2normalize(u) + v = l2normalize(v) + tmp_w = w.detach() + out_2 = paddle.create_parameter( + shape=tmp_w.shape, + dtype=tmp_w.numpy().dtype, + default_initializer=nn.initializer.Assign(tmp_w), + ) + out_2.stop_gradient = False + w_bar = out_2 + del self.module._parameters[self.name] + + u = create_param(u) + v = create_param(v) + self.module.add_parameter(name=self.name + "_u", parameter=u) + self.module.add_parameter(name=self.name + "_v", parameter=v) + self.module.add_parameter(name=self.name + "_bar", parameter=w_bar) + + def forward(self, *args): + self._update_u_v() + return self.module.forward(*args) + + +def create_param(x): + param = paddle.create_parameter( + shape=x.shape, + dtype=x.dtype, + default_initializer=nn.initializer.Assign(x), + ) + param.stop_gradient = x.stop_gradient + return param diff --git a/examples/smc_reac/ppsci/arch/paddle_harmonics/legendre.py b/examples/smc_reac/ppsci/arch/paddle_harmonics/legendre.py new file mode 100644 index 0000000000..376599719e --- /dev/null +++ b/examples/smc_reac/ppsci/arch/paddle_harmonics/legendre.py @@ -0,0 +1,176 @@ +# coding=utf-8 + +# SPDX-FileCopyrightText: Copyright (c) 2022 The torch-harmonics Authors. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +""" +Code below is heavily based on [torch-harmonics](https://github.com/NVIDIA/torch-harmonics/blob/main/torch_harmonics/legendre.py) +""" + +import numpy as np + + +def clm(l, m): + """ + Defines the normalization factor to orthonormalize the Spherical Harmonics + """ + return np.sqrt((2 * l + 1) / 4 / np.pi) * np.sqrt( + np.math.factorial(l - m) / np.math.factorial(l + m) + ) + + +def legpoly(mmax, lmax, x, norm="ortho", inverse=False, csphase=True): + r""" + Computes the values of (-1)^m c^l_m P^l_m(x) at the positions specified by x. + The resulting tensor has shape (mmax, lmax, len(x)). The Condon-Shortley Phase (-1)^m + can be turned off optionally. + + method of computation follows + [1] Schaeffer, N.; Efficient spherical harmonic transforms aimed at pseudospectral numerical simulations, G3: Geochemistry, Geophysics, Geosystems. + [2] Rapp, R.H.; A Fortran Program for the Computation of Gravimetric Quantities from High Degree Spherical Harmonic Expansions, Ohio State University Columbus; report; 1982; + https://apps.dtic.mil/sti/citations/ADA123406 + [3] Schrama, E.; Orbit integration based upon interpolated gravitational gradients + """ + # compute the tensor P^m_n: + nmax = max(mmax, lmax) + vdm = np.zeros((nmax, nmax, len(x)), dtype=np.float64) + + norm_factor = 1.0 if norm == "ortho" else np.sqrt(4 * np.pi) + norm_factor = 1.0 / norm_factor if inverse else norm_factor + + # initial values to start the recursion + vdm[0, 0, :] = norm_factor / np.sqrt(4 * np.pi) + + # fill the diagonal and the lower diagonal + for l in range(1, nmax): + vdm[l - 1, l, :] = np.sqrt(2 * l + 1) * x * vdm[l - 1, l - 1, :] + vdm[l, l, :] = ( + np.sqrt((2 * l + 1) * (1 + x) * (1 - x) / 2 / l) * vdm[l - 1, l - 1, :] + ) + + # fill the remaining values on the upper triangle and multiply b + for l in range(2, nmax): + for m in range(0, l - 1): + vdm[m, l, :] = ( + x + * np.sqrt((2 * l - 1) / (l - m) * (2 * l + 1) / (l + m)) + * vdm[m, l - 1, :] + - np.sqrt( + (l + m - 1) + / (l - m) + * (2 * l + 1) + / (2 * l - 3) + * (l - m - 1) + / (l + m) + ) + * vdm[m, l - 2, :] + ) + + if norm == "schmidt": + for l in range(0, nmax): + if inverse: + vdm[:, l, :] = vdm[:, l, :] * np.sqrt(2 * l + 1) + else: + vdm[:, l, :] = vdm[:, l, :] / np.sqrt(2 * l + 1) + + vdm = vdm[:mmax, :lmax] + + if csphase: + for m in range(1, mmax, 2): + vdm[m] *= -1 + + return vdm + + +def _precompute_legpoly(mmax, lmax, t, norm="ortho", inverse=False, csphase=True): + r""" + Computes the values of (-1)^m c^l_m P^l_m(\cos \theta) at the positions specified by t (theta). + The resulting tensor has shape (mmax, lmax, len(x)). The Condon-Shortley Phase (-1)^m + can be turned off optionally. + + method of computation follows + [1] Schaeffer, N.; Efficient spherical harmonic transforms aimed at pseudospectral numerical simulations, G3: Geochemistry, Geophysics, Geosystems. + [2] Rapp, R.H.; A Fortran Program for the Computation of Gravimetric Quantities from High Degree Spherical Harmonic Expansions, Ohio State University Columbus; report; 1982; + https://apps.dtic.mil/sti/citations/ADA123406 + [3] Schrama, E.; Orbit integration based upon interpolated gravitational gradients + """ + return legpoly(mmax, lmax, np.cos(t), norm=norm, inverse=inverse, csphase=csphase) + + +def _precompute_dlegpoly(mmax, lmax, t, norm="ortho", inverse=False, csphase=True): + r""" + Computes the values of the derivatives $\frac{d}{d \theta} P^m_l(\cos \theta)$ + at the positions specified by t (theta), as well as $\frac{1}{\sin \theta} P^m_l(\cos \theta)$, + needed for the computation of the vector spherical harmonics. The resulting tensor has shape + (2, mmax, lmax, len(t)). + + computation follows + [2] Wang, B., Wang, L., Xie, Z.; Accurate calculation of spherical and vector spherical harmonic expansions via spectral element grids; Adv Comput Math. + """ + pct = _precompute_legpoly( + mmax + 1, lmax + 1, t, norm=norm, inverse=inverse, csphase=False + ) + + dpct = np.zeros((2, mmax, lmax, len(t)), dtype=np.float64) + + # fill the derivative terms wrt theta + for l in range(0, lmax): + + # m = 0 + dpct[0, 0, l] = -np.sqrt(l * (l + 1)) * pct[1, l] + + # 0 < m < l + for m in range(1, min(l, mmax)): + dpct[0, m, l] = 0.5 * ( + np.sqrt((l + m) * (l - m + 1)) * pct[m - 1, l] + - np.sqrt((l - m) * (l + m + 1)) * pct[m + 1, l] + ) + + # m == l + if mmax > l: + dpct[0, l, l] = np.sqrt(l / 2) * pct[l - 1, l] + + # fill the - 1j m P^m_l / sin(phi). as this component is purely imaginary, + # we won't store it explicitly in a complex array + for m in range(1, min(l + 1, mmax)): + # this component is implicitly complex + # we do not divide by m here as this cancels with the derivative of the exponential + dpct[1, m, l] = ( + 0.5 + * np.sqrt((2 * l + 1) / (2 * l + 3)) + * ( + np.sqrt((l - m + 1) * (l - m + 2)) * pct[m - 1, l + 1] + + np.sqrt((l + m + 1) * (l + m + 2)) * pct[m + 1, l + 1] + ) + ) + + if csphase: + for m in range(1, mmax, 2): + dpct[:, m] *= -1 + + return dpct diff --git a/examples/smc_reac/ppsci/arch/paddle_harmonics/quadrature.py b/examples/smc_reac/ppsci/arch/paddle_harmonics/quadrature.py new file mode 100644 index 0000000000..25de0e1436 --- /dev/null +++ b/examples/smc_reac/ppsci/arch/paddle_harmonics/quadrature.py @@ -0,0 +1,156 @@ +# coding=utf-8 + +# SPDX-FileCopyrightText: Copyright (c) 2022 The torch-harmonics Authors. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +""" +Code below is heavily based on [torch-harmonics](https://github.com/NVIDIA/torch-harmonics/blob/main/torch_harmonics/quadrature.py) +""" + +import numpy as np + + +def legendre_gauss_weights(n, a=-1.0, b=1.0): + r""" + Helper routine which returns the Legendre-Gauss nodes and weights + on the interval [a, b] + """ + xlg, wlg = np.polynomial.legendre.leggauss(n) + xlg = (b - a) * 0.5 * xlg + (b + a) * 0.5 + wlg = wlg * (b - a) * 0.5 + + return xlg, wlg + + +def lobatto_weights(n, a=-1.0, b=1.0, tol=1e-16, maxiter=100): + r""" + Helper routine which returns the Legendre-Gauss-Lobatto nodes and weights + on the interval [a, b] + """ + wlg = np.zeros((n,)) + tlg = np.zeros((n,)) + tmp = np.zeros((n,)) + + # Vandermonde Matrix + vdm = np.zeros((n, n)) + + # initialize Chebyshev nodes as first guess + for i in range(n): + tlg[i] = -np.cos(np.pi * i / (n - 1)) + + tmp = 2.0 + + for i in range(maxiter): + tmp = tlg + + vdm[:, 0] = 1.0 + vdm[:, 1] = tlg + + for k in range(2, n): + vdm[:, k] = ( + (2 * k - 1) * tlg * vdm[:, k - 1] - (k - 1) * vdm[:, k - 2] + ) / k + + tlg = tmp - (tlg * vdm[:, n - 1] - vdm[:, n - 2]) / (n * vdm[:, n - 1]) + + if max(abs(tlg - tmp).flatten()) < tol: + break + + wlg = 2.0 / ((n * (n - 1)) * (vdm[:, n - 1] ** 2)) + + # rescale + tlg = (b - a) * 0.5 * tlg + (b + a) * 0.5 + wlg = wlg * (b - a) * 0.5 + + return tlg, wlg + + +def clenshaw_curtiss_weights(n, a=-1.0, b=1.0): + r""" + Computation of the Clenshaw-Curtis quadrature nodes and weights. + This implementation follows + + [1] Joerg Waldvogel, Fast Construction of the Fejer and Clenshaw-Curtis Quadrature Rules; BIT Numerical Mathematics, Vol. 43, No. 1, pp. 001–018. + """ + assert n > 1 + + tcc = np.cos(np.linspace(np.pi, 0, n)) + + if n == 2: + wcc = np.array([1.0, 1.0]) + else: + + n1 = n - 1 + N = np.arange(1, n1, 2) + l = len(N) + m = n1 - l + + v = np.concatenate([2 / N / (N - 2), 1 / N[-1:], np.zeros(m)]) + v = 0 - v[:-1] - v[-1:0:-1] + + g0 = -np.ones(n1) + g0[l] = g0[l] + n1 + g0[m] = g0[m] + n1 + g = g0 / (n1**2 - 1 + (n1 % 2)) + wcc = np.fft.ifft(v + g).real + wcc = np.concatenate((wcc, wcc[:1])) + + # rescale + tcc = (b - a) * 0.5 * tcc + (b + a) * 0.5 + wcc = wcc * (b - a) * 0.5 + + return tcc, wcc + + +def fejer2_weights(n, a=-1.0, b=1.0): + r""" + Computation of the Fejer quadrature nodes and weights. + This implementation follows + + [1] Joerg Waldvogel, Fast Construction of the Fejer and Clenshaw-Curtis Quadrature Rules; BIT Numerical Mathematics, Vol. 43, No. 1, pp. 001–018. + """ + assert n > 2 + + tcc = np.cos(np.linspace(np.pi, 0, n)) + + n1 = n - 1 + N = np.arange(1, n1, 2) + l = len(N) + m = n1 - l + + v = np.concatenate([2 / N / (N - 2), 1 / N[-1:], np.zeros(m)]) + v = 0 - v[:-1] - v[-1:0:-1] + + wcc = np.fft.ifft(v).real + wcc = np.concatenate((wcc, wcc[:1])) + + # rescale + tcc = (b - a) * 0.5 * tcc + (b + a) * 0.5 + wcc = wcc * (b - a) * 0.5 + + return tcc, wcc diff --git a/examples/smc_reac/ppsci/arch/paddle_harmonics/random_fields.py b/examples/smc_reac/ppsci/arch/paddle_harmonics/random_fields.py new file mode 100644 index 0000000000..8fad28cf26 --- /dev/null +++ b/examples/smc_reac/ppsci/arch/paddle_harmonics/random_fields.py @@ -0,0 +1,148 @@ +# coding=utf-8 + +# SPDX-FileCopyrightText: Copyright (c) 2022 The paddle-harmonics Authors. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +""" +Code below is heavily based on [torch-harmonics](https://github.com/NVIDIA/torch-harmonics/blob/main/torch_harmonics/random_fields.py) +""" + +import paddle +from paddle import nn + +from ppsci.arch.paddle_harmonics.sht import InverseRealSHT + + +class GaussianRandomFieldS2(nn.Layer): + r""" + A mean-zero Gaussian Random Field on the sphere with Matern covariance: + C = sigma^2 (-Lap + tau^2 I)^(-alpha). + + Lap is the Laplacian on the sphere, I the identity operator, + and sigma, tau, alpha are scalar parameters. + + Note: C is trace-class on L^2 if and only if alpha > 1. + + Args: + nlat (int): Number of latitudinal modes.longitudinal modes are 2*nlat. + alpha (float, optional): Regularity parameter. Larger means smoother. Defaults to 2.0. + tau (float, optional): Lenght-scale parameter. Larger means more scales. Defaults to 3.0. + sigma (float, optional): Scale parameter. Larger means bigger. + If None, sigma = tau**(0.5*(2*alpha - 2.0)). Defaults to None. + radius (float, optional): Radius of the sphere. Defaults to 1.0. + grid (str, optional): Grid type. Currently supports "equiangular" and + "legendre-gauss". Defaults to "equiangular". + dtype (paddle.dtype, optional): Numerical type for the calculations. Defaults to paddle.float32. + """ + + def __init__( + self, + nlat, + alpha: float = 2.0, + tau: float = 3.0, + sigma: float = None, + radius: float = 1.0, + grid: str = "equiangular", + dtype: paddle.dtype = paddle.float32, + ): + + super().__init__() + + # Number of latitudinal modes. + self.nlat = nlat + + # Default value of sigma if None is given. + if sigma is None: + assert alpha > 1.0, f"Alpha must be greater than one, got {alpha}." + sigma = tau ** (0.5 * (2 * alpha - 2.0)) + + # Inverse SHT + self.isht = InverseRealSHT( + self.nlat, 2 * self.nlat, grid=grid, norm="backward" + ).to(dtype=dtype) + + # Square root of the eigenvalues of C. + sqrt_eig = ( + paddle.to_tensor([j * (j + 1) for j in range(self.nlat)]) + .reshape([self.nlat, 1]) + .tile([1, self.nlat + 1]) + ) + sqrt_eig = paddle.tril( + sigma * (((sqrt_eig / radius**2) + tau**2) ** (-alpha / 2.0)) + ) + sqrt_eig[0, 0] = 0.0 + sqrt_eig = sqrt_eig.unsqueeze(0) + self.register_buffer("sqrt_eig", sqrt_eig) + + # Save mean and var of the standard Gaussian. + # Need these to re-initialize distribution on a new device. + mean = paddle.to_tensor([0.0]).astype(dtype) + var = paddle.to_tensor([1.0]).astype(dtype) + self.register_buffer("mean", mean) + self.register_buffer("var", var) + + # Standard normal noise sampler. + self.gaussian_noise = paddle.distribution.Normal(self.mean, self.var) + + def forward(self, N, xi=None): + """Sample random functions from a spherical GRF. + + Args: + N (int): Number of functions to sample. + xi (paddle.Tensor, optional): Noise is a complex tensor of size (N, nlat, nlat+1). + If None, new Gaussian noise is sampled. + If xi is provided, N is ignored.. Defaults to None. + + Returns: + u (paddle.Tensor): N random samples from the GRF returned as a + tensor of size (N, nlat, 2*nlat) on a equiangular grid. + """ + + # Sample Gaussian noise. + if xi is None: + xi = self.gaussian_noise.sample((N, self.nlat, self.nlat + 1, 2)).squeeze() + xi = paddle.as_complex(xi) + + # Karhunen-Loeve expansion. + u = self.isht(xi * self.sqrt_eig) + + return u + + # Override cuda and to methods so sampler gets initialized with mean + # and variance on the correct device. + def cuda(self, *args, **kwargs): + super().cuda(*args, **kwargs) + self.gaussian_noise = paddle.distribution.Normal(self.mean, self.var) + + return self + + def to(self, *args, **kwargs): + super().to(*args, **kwargs) + self.gaussian_noise = paddle.distribution.Normal(self.mean, self.var) + + return self diff --git a/examples/smc_reac/ppsci/arch/paddle_harmonics/sht.py b/examples/smc_reac/ppsci/arch/paddle_harmonics/sht.py new file mode 100644 index 0000000000..bf5e685a04 --- /dev/null +++ b/examples/smc_reac/ppsci/arch/paddle_harmonics/sht.py @@ -0,0 +1,461 @@ +# coding=utf-8 + +# SPDX-FileCopyrightText: Copyright (c) 2022 The torch-harmonics Authors. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +""" +Code below is heavily based on [torch-harmonics](https://github.com/NVIDIA/torch-harmonics/blob/main/torch_harmonics/sht.py) +""" + +import math + +import numpy as np +import paddle +import paddle.fft +import paddle.nn as nn + +from ppsci.arch.paddle_harmonics.legendre import _precompute_dlegpoly +from ppsci.arch.paddle_harmonics.legendre import _precompute_legpoly +from ppsci.arch.paddle_harmonics.quadrature import clenshaw_curtiss_weights +from ppsci.arch.paddle_harmonics.quadrature import legendre_gauss_weights +from ppsci.arch.paddle_harmonics.quadrature import lobatto_weights + + +class RealSHT(nn.Layer): + """ + Defines a module for computing the forward (real-valued) SHT. + Precomputes Legendre Gauss nodes, weights and associated Legendre polynomials on these nodes. + The SHT is applied to the last two dimensions of the input + + [1] Schaeffer, N. Efficient spherical harmonic transforms aimed at pseudospectral numerical simulations, G3: Geochemistry, Geophysics, Geosystems. + [2] Wang, B., Wang, L., Xie, Z.; Accurate calculation of spherical and vector spherical harmonic expansions via spectral element grids; Adv Comput Math. + + Initializes the SHT Layer, precomputing the necessary quadrature weights. + + Args: + nlat (int): Input grid resolution in the latitudinal direction. + nlon (int): Input grid resolution in the longitudinal direction. + lmax (int, optional): The max input grid resolution in the latitudinal direction. Defaults to None. + mmax (int, optional): The max input grid resolution in the longitudinal direction. Defaults to None. + grid (str, optional): Grid in the latitude direction (for now only tensor product grids are supported). + Defaults to "lobatto". + norm (str, optional): The type of normalization to use. Defaults to "ortho". + csphase (bool, optional): Whether to apply the complex-conjugate symmetry phase factor. Defaults to True. + """ + + def __init__( + self, + nlat, + nlon, + lmax=None, + mmax=None, + grid="lobatto", + norm="ortho", + csphase=True, + ): + super().__init__() + + self.nlat = nlat + self.nlon = nlon + self.grid = grid + self.norm = norm + self.csphase = csphase + + # compute quadrature points + if self.grid == "legendre-gauss": + cost, w = legendre_gauss_weights(nlat, -1, 1) + self.lmax = lmax or self.nlat + elif self.grid == "lobatto": + cost, w = lobatto_weights(nlat, -1, 1) + self.lmax = lmax or self.nlat - 1 + elif self.grid == "equiangular": + cost, w = clenshaw_curtiss_weights(nlat, -1, 1) + # cost, w = fejer2_weights(nlat, -1, 1) + self.lmax = lmax or self.nlat + else: + raise (ValueError("Unknown quadrature mode")) + + # apply cosine transform and flip them + tq = np.flip(np.arccos(cost)) + + # determine the dimensions + self.mmax = mmax or self.nlon // 2 + 1 + + # combine quadrature weights with the legendre weights + weights = paddle.to_tensor(w) + pct = _precompute_legpoly( + self.mmax, self.lmax, tq, norm=self.norm, csphase=self.csphase + ) + pct = paddle.to_tensor(pct) + self.weights = paddle.einsum("mlk,k->mlk", pct, weights) + + def extra_repr(self): + return f"nlat={self.nlat}, nlon={self.nlon},\n lmax={self.lmax}, mmax={self.mmax},\n grid={self.grid}, csphase={self.csphase}" + + def forward(self, x: paddle.Tensor): + + assert x.shape[-2] == self.nlat + assert x.shape[-1] == self.nlon + + # apply real fft in the longitudinal direction + x = ( + 2.0 + * paddle.to_tensor(math.pi) + * paddle.fft.rfft(x, axis=-1, norm="forward") + ) + + # do the Legendre-Gauss quadrature + x = paddle.as_real(x) + # distributed contraction: fork + out_shape = list(x.shape) + out_shape[-3] = self.lmax + out_shape[-2] = self.mmax + xout = paddle.zeros(out_shape, dtype=x.dtype) + + # contraction + xout[..., 0] = paddle.einsum( + "...km,mlk->...lm", x[..., : self.mmax, 0], self.weights.astype(x.dtype) + ) + xout[..., 1] = paddle.einsum( + "...km,mlk->...lm", x[..., : self.mmax, 1], self.weights.astype(x.dtype) + ) + x = paddle.as_complex(xout) + + return x + + +class InverseRealSHT(nn.Layer): + """ + Defines a module for computing the inverse (real-valued) SHT. + Precomputes Legendre Gauss nodes, weights and associated Legendre polynomials on these nodes. + nlat, nlon: Output dimensions + lmax, mmax: Input dimensions (spherical coefficients). For convenience, these are inferred from the output dimensions + + [1] Schaeffer, N. Efficient spherical harmonic transforms aimed at pseudospectral numerical simulations, G3: Geochemistry, Geophysics, Geosystems. + [2] Wang, B., Wang, L., Xie, Z.; Accurate calculation of spherical and vector spherical harmonic expansions via spectral element grids; Adv Comput Math. + """ + + def __init__( + self, + nlat, + nlon, + lmax=None, + mmax=None, + grid="lobatto", + norm="ortho", + csphase=True, + ): + + super().__init__() + + self.nlat = nlat + self.nlon = nlon + self.grid = grid + self.norm = norm + self.csphase = csphase + + # compute quadrature points + if self.grid == "legendre-gauss": + cost, _ = legendre_gauss_weights(nlat, -1, 1) + self.lmax = lmax or self.nlat + elif self.grid == "lobatto": + cost, _ = lobatto_weights(nlat, -1, 1) + self.lmax = lmax or self.nlat - 1 + elif self.grid == "equiangular": + cost, _ = clenshaw_curtiss_weights(nlat, -1, 1) + self.lmax = lmax or self.nlat + else: + raise (ValueError("Unknown quadrature mode")) + + # apply cosine transform and flip them + t = np.flip(np.arccos(cost)) + + # determine the dimensions + self.mmax = mmax or self.nlon // 2 + 1 + + pct = _precompute_legpoly( + self.mmax, self.lmax, t, norm=self.norm, inverse=True, csphase=self.csphase + ) + self.pct = paddle.to_tensor(pct) + + def extra_repr(self): + return f"nlat={self.nlat}, nlon={self.nlon},\n lmax={self.lmax}, mmax={self.mmax},\n grid={self.grid}, csphase={self.csphase}" + + def forward(self, x: paddle.Tensor): + + assert x.shape[-2] == self.lmax + assert x.shape[-1] == self.mmax + + # Evaluate associated Legendre functions on the output nodes + x = paddle.as_real(x) + + rl = paddle.einsum("...lm, mlk->...km", x[..., 0], self.pct.astype(x.dtype)) + im = paddle.einsum("...lm, mlk->...km", x[..., 1], self.pct.astype(x.dtype)) + xs = paddle.stack((rl, im), -1) + + # apply the inverse (real) FFT + x = paddle.as_complex(xs) + x = paddle.fft.irfft(x, n=self.nlon, axis=-1, norm="forward") + + return x + + +class RealVectorSHT(nn.Layer): + """ + Defines a module for computing the forward (real) vector SHT. + Precomputes Legendre Gauss nodes, weights and associated Legendre polynomials on these nodes. + The SHT is applied to the last three dimensions of the input. + + [1] Schaeffer, N. Efficient spherical harmonic transforms aimed at pseudospectral numerical simulations, G3: Geochemistry, Geophysics, Geosystems. + [2] Wang, B., Wang, L., Xie, Z.; Accurate calculation of spherical and vector spherical harmonic expansions via spectral element grids; Adv Comput Math. + + Initializes the vector SHT Layer, precomputing the necessary quadrature weights. + """ + + def __init__( + self, + nlat, + nlon, + lmax=None, + mmax=None, + grid="lobatto", + norm="ortho", + csphase=True, + ): + super().__init__() + + self.nlat = nlat + self.nlon = nlon + self.grid = grid + self.norm = norm + self.csphase = csphase + + # compute quadrature points + if self.grid == "legendre-gauss": + cost, w = legendre_gauss_weights(nlat, -1, 1) + self.lmax = lmax or self.nlat + elif self.grid == "lobatto": + cost, w = lobatto_weights(nlat, -1, 1) + self.lmax = lmax or self.nlat - 1 + elif self.grid == "equiangular": + cost, w = clenshaw_curtiss_weights(nlat, -1, 1) + # cost, w = fejer2_weights(nlat, -1, 1) + self.lmax = lmax or self.nlat + else: + raise (ValueError("Unknown quadrature mode")) + + # apply cosine transform and flip them + tq = np.flip(np.arccos(cost)) + + # determine the dimensions + self.mmax = mmax or self.nlon // 2 + 1 + + weights = paddle.to_tensor(w) + dpct = _precompute_dlegpoly( + self.mmax, self.lmax, tq, norm=self.norm, csphase=self.csphase + ) + dpct = paddle.to_tensor(dpct) + + # combine integration weights, normalization factor in to one: + l = paddle.arange(0, self.lmax) + norm_factor = 1.0 / l / (l + 1) + norm_factor[0] = 1.0 + weights = paddle.einsum("dmlk,k,l->dmlk", dpct, weights, norm_factor) + # since the second component is imaginary, we need to take complex conjugation into account + weights[1] = -1 * weights[1] + + self.weights = weights + + def extra_repr(self): + return f"nlat={self.nlat}, nlon={self.nlon},\n lmax={self.lmax}, mmax={self.mmax},\n grid={self.grid}, csphase={self.csphase}" + + def forward(self, x: paddle.Tensor): + + assert len(x.shape) >= 3 + + # apply real fft in the longitudinal direction + x = 2.0 * paddle.to_tensor(np.pi) * paddle.fft.rfft(x, axis=-1, norm="forward") + + # do the Legendre-Gauss quadrature + x = paddle.as_real(x) + + # distributed contraction: fork + out_shape = list(x.shape) + out_shape[-3] = self.lmax + out_shape[-2] = self.mmax + xout = paddle.zeros(out_shape, dtype=x.dtype) + + # contraction - spheroidal component + # real component + xout[..., 0, :, :, 0] = paddle.einsum( + "...km,mlk->...lm", + x[..., 0, :, : self.mmax, 0], + self.weights[0].astype(x.dtype), + ) - paddle.einsum( + "...km,mlk->...lm", + x[..., 1, :, : self.mmax, 1], + self.weights[1].astype(x.dtype), + ) + + # iamg component + xout[..., 0, :, :, 1] = paddle.einsum( + "...km,mlk->...lm", + x[..., 0, :, : self.mmax, 1], + self.weights[0].astype(x.dtype), + ) + paddle.einsum( + "...km,mlk->...lm", + x[..., 1, :, : self.mmax, 0], + self.weights[1].astype(x.dtype), + ) + + # contraction - toroidal component + # real component + xout[..., 1, :, :, 0] = -paddle.einsum( + "...km,mlk->...lm", + x[..., 0, :, : self.mmax, 1], + self.weights[1].astype(x.dtype), + ) - paddle.einsum( + "...km,mlk->...lm", + x[..., 1, :, : self.mmax, 0], + self.weights[0].astype(x.dtype), + ) + # imag component + xout[..., 1, :, :, 1] = paddle.einsum( + "...km,mlk->...lm", + x[..., 0, :, : self.mmax, 0], + self.weights[1].astype(x.dtype), + ) - paddle.einsum( + "...km,mlk->...lm", + x[..., 1, :, : self.mmax, 1], + self.weights[0].astype(x.dtype), + ) + + return paddle.as_complex(xout) + + +class InverseRealVectorSHT(nn.Layer): + """ + Defines a module for computing the inverse (real-valued) vector SHT. + Precomputes Legendre Gauss nodes, weights and associated Legendre polynomials on these nodes. + + [1] Schaeffer, N. Efficient spherical harmonic transforms aimed at pseudospectral numerical simulations, G3: Geochemistry, Geophysics, Geosystems. + [2] Wang, B., Wang, L., Xie, Z.; Accurate calculation of spherical and vector spherical harmonic expansions via spectral element grids; Adv Comput Math. + """ + + def __init__( + self, + nlat, + nlon, + lmax=None, + mmax=None, + grid="lobatto", + norm="ortho", + csphase=True, + ): + + super().__init__() + + self.nlat = nlat + self.nlon = nlon + self.grid = grid + self.norm = norm + self.csphase = csphase + + # compute quadrature points + if self.grid == "legendre-gauss": + cost, _ = legendre_gauss_weights(nlat, -1, 1) + self.lmax = lmax or self.nlat + elif self.grid == "lobatto": + cost, _ = lobatto_weights(nlat, -1, 1) + self.lmax = lmax or self.nlat - 1 + elif self.grid == "equiangular": + cost, _ = clenshaw_curtiss_weights(nlat, -1, 1) + self.lmax = lmax or self.nlat + else: + raise (ValueError("Unknown quadrature mode")) + + # apply cosine transform and flip them + t = np.flip(np.arccos(cost)) + + # determine the dimensions + self.mmax = mmax or self.nlon // 2 + 1 + + dpct = _precompute_dlegpoly( + self.mmax, self.lmax, t, norm=self.norm, inverse=True, csphase=self.csphase + ) + self.dpct = paddle.to_tensor(dpct) + + def extra_repr(self): + return f"nlat={self.nlat}, nlon={self.nlon},\n lmax={self.lmax}, mmax={self.mmax},\n grid={self.grid}, csphase={self.csphase}" + + def forward(self, x: paddle.Tensor): + + assert x.shape[-2] == self.lmax + assert x.shape[-1] == self.mmax + + # Evaluate associated Legendre functions on the output nodes + x = paddle.as_real(x) + + # contraction - spheroidal component + # real component + srl = paddle.einsum( + "...lm,mlk->...km", x[..., 0, :, :, 0], self.dpct[0].astype(x.dtype) + ) - paddle.einsum( + "...lm,mlk->...km", x[..., 1, :, :, 1], self.dpct[1].astype(x.dtype) + ) + # iamg component + sim = paddle.einsum( + "...lm,mlk->...km", x[..., 0, :, :, 1], self.dpct[0].astype(x.dtype) + ) + paddle.einsum( + "...lm,mlk->...km", x[..., 1, :, :, 0], self.dpct[1].astype(x.dtype) + ) + + # contraction - toroidal component + # real component + trl = -paddle.einsum( + "...lm,mlk->...km", x[..., 0, :, :, 1], self.dpct[1].astype(x.dtype) + ) - paddle.einsum( + "...lm,mlk->...km", x[..., 1, :, :, 0], self.dpct[0].astype(x.dtype) + ) + # imag component + tim = paddle.einsum( + "...lm,mlk->...km", x[..., 0, :, :, 0], self.dpct[1].astype(x.dtype) + ) - paddle.einsum( + "...lm,mlk->...km", x[..., 1, :, :, 1], self.dpct[0].astype(x.dtype) + ) + + # reassemble + s = paddle.stack((srl, sim), -1) + t = paddle.stack((trl, tim), -1) + xs = paddle.stack((s, t), -4) + + # apply the inverse (real) FFT + x = paddle.as_complex(xs) + x = paddle.fft.irfft(x, n=self.nlon, axis=-1, norm="forward") + + return x diff --git a/examples/smc_reac/ppsci/arch/phycrnet.py b/examples/smc_reac/ppsci/arch/phycrnet.py new file mode 100644 index 0000000000..c72583ebf9 --- /dev/null +++ b/examples/smc_reac/ppsci/arch/phycrnet.py @@ -0,0 +1,540 @@ +# 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. + +from typing import Tuple + +import numpy as np +import paddle +import paddle.nn as nn +from paddle.nn import utils + +from ppsci.arch import base + +# define the high-order finite difference kernels +LALP_OP = [ + [ + [ + [0, 0, -1 / 12, 0, 0], + [0, 0, 4 / 3, 0, 0], + [-1 / 12, 4 / 3, -5, 4 / 3, -1 / 12], + [0, 0, 4 / 3, 0, 0], + [0, 0, -1 / 12, 0, 0], + ] + ] +] + +PARTIAL_Y = [ + [ + [ + [0, 0, 0, 0, 0], + [0, 0, 0, 0, 0], + [1 / 12, -8 / 12, 0, 8 / 12, -1 / 12], + [0, 0, 0, 0, 0], + [0, 0, 0, 0, 0], + ] + ] +] + +PARTIAL_X = [ + [ + [ + [0, 0, 1 / 12, 0, 0], + [0, 0, -8 / 12, 0, 0], + [0, 0, 0, 0, 0], + [0, 0, 8 / 12, 0, 0], + [0, 0, -1 / 12, 0, 0], + ] + ] +] + + +# specific parameters for burgers equation +def _initialize_weights(module): + if isinstance(module, nn.Conv2D): + c = 1.0 # 0.5 + initializer = nn.initializer.Uniform( + -c * np.sqrt(1 / (3 * 3 * 320)), c * np.sqrt(1 / (3 * 3 * 320)) + ) + initializer(module.weight) + elif isinstance(module, nn.Linear): + initializer = nn.initializer.Constant(0.0) + initializer(module.bias) + + +class PhyCRNet(base.Arch): + """Physics-informed convolutional-recurrent neural networks. + + Args: + input_channels (int): The input channels. + hidden_channels (Tuple[int, ...]): The hidden channels. + input_kernel_size (Tuple[int, ...]): The input kernel size(s). + input_stride (Tuple[int, ...]): The input stride(s). + input_padding (Tuple[int, ...]): The input padding(s). + dt (float): The dt parameter. + num_layers (Tuple[int, ...]): The number of layers. + upscale_factor (int): The upscale factor. + step (int, optional): The step(s). Defaults to 1. + effective_step (Tuple[int, ...], optional): The effective step. Defaults to (1, ). + + Examples: + >>> import ppsci + >>> model = ppsci.arch.PhyCRNet( + ... input_channels=2, + ... hidden_channels=[8, 32, 128, 128], + ... input_kernel_size=[4, 4, 4, 3], + ... input_stride=[2, 2, 2, 1], + ... input_padding=[1, 1, 1, 1], + ... dt=0.002, + ... num_layers=[3, 1], + ... upscale_factor=8 + ... ) + """ + + def __init__( + self, + input_channels: int, + hidden_channels: Tuple[int, ...], + input_kernel_size: Tuple[int, ...], + input_stride: Tuple[int, ...], + input_padding: Tuple[int, ...], + dt: float, + num_layers: Tuple[int, ...], + upscale_factor: int, + step: int = 1, + effective_step: Tuple[int, ...] = (1,), + ): + super(PhyCRNet, self).__init__() + + # input channels of layer includes input_channels and hidden_channels of cells + self.input_channels = [input_channels] + hidden_channels + self.hidden_channels = hidden_channels + self.input_kernel_size = input_kernel_size + self.input_stride = input_stride + self.input_padding = input_padding + self.step = step + self.effective_step = effective_step + self._all_layers = [] + self.dt = dt + self.upscale_factor = upscale_factor + + # number of layers + self.num_encoder = num_layers[0] + self.num_convlstm = num_layers[1] + + # encoder - downsampling + self.encoder = nn.LayerList( + [ + encoder_block( + input_channels=self.input_channels[i], + hidden_channels=self.hidden_channels[i], + input_kernel_size=self.input_kernel_size[i], + input_stride=self.input_stride[i], + input_padding=self.input_padding[i], + ) + for i in range(self.num_encoder) + ] + ) + + # ConvLSTM + self.convlstm = nn.LayerList( + [ + ConvLSTMCell( + input_channels=self.input_channels[i], + hidden_channels=self.hidden_channels[i], + input_kernel_size=self.input_kernel_size[i], + input_stride=self.input_stride[i], + input_padding=self.input_padding[i], + ) + for i in range(self.num_encoder, self.num_encoder + self.num_convlstm) + ] + ) + + # output layer + self.output_layer = nn.Conv2D( + 2, 2, kernel_size=5, stride=1, padding=2, padding_mode="circular" + ) + + # pixelshuffle - upscale + self.pixelshuffle = nn.PixelShuffle(self.upscale_factor) + + # initialize weights + self.apply(_initialize_weights) + initializer_0 = nn.initializer.Constant(0.0) + initializer_0(self.output_layer.bias) + self.enable_transform = True + + def forward(self, x): + if self.enable_transform: + if self._input_transform is not None: + x = self._input_transform(x) + output_x = x + + self.initial_state = x["initial_state"] + x = x["input"] + internal_state = [] + outputs = [] + second_last_state = [] + + for step in range(self.step): + xt = x + + # encoder + for encoder in self.encoder: + x = encoder(x) + + # convlstm + for i, lstm in enumerate(self.convlstm, self.num_encoder): + if step == 0: + (h, c) = lstm.init_hidden_tensor( + prev_state=self.initial_state[i - self.num_encoder] + ) + internal_state.append((h, c)) + + # one-step forward + (h, c) = internal_state[i - self.num_encoder] + x, new_c = lstm(x, h, c) + internal_state[i - self.num_encoder] = (x, new_c) + + # output + x = self.pixelshuffle(x) + x = self.output_layer(x) + + # residual connection + x = xt + self.dt * x + + if step == (self.step - 2): + second_last_state = internal_state.copy() + + if step in self.effective_step: + outputs.append(x) + + result_dict = {"outputs": outputs, "second_last_state": second_last_state} + if self.enable_transform: + if self._output_transform is not None: + result_dict = self._output_transform(output_x, result_dict) + return result_dict + + +class ConvLSTMCell(nn.Layer): + """Convolutional LSTM""" + + def __init__( + self, + input_channels, + hidden_channels, + input_kernel_size, + input_stride, + input_padding, + hidden_kernel_size=3, + num_features=4, + ): + super(ConvLSTMCell, self).__init__() + + self.input_channels = input_channels + self.hidden_channels = hidden_channels + self.hidden_kernel_size = hidden_kernel_size # Page 9, The convolutional operations in ConvLSTM have 3x3 kernels. + self.input_kernel_size = input_kernel_size + self.input_stride = input_stride + self.input_padding = input_padding + self.num_features = ( + num_features # Page 10, block of different dense layers {4, 3, 4} + ) + + # padding for hidden state + self.padding = int((self.hidden_kernel_size - 1) / 2) + + self.Wxi = nn.Conv2D( + self.input_channels, + self.hidden_channels, + self.input_kernel_size, + self.input_stride, + self.input_padding, + bias_attr=None, + padding_mode="circular", + ) + + self.Whi = nn.Conv2D( + self.hidden_channels, + self.hidden_channels, + self.hidden_kernel_size, + 1, + padding=1, + bias_attr=False, + padding_mode="circular", + ) + + self.Wxf = nn.Conv2D( + self.input_channels, + self.hidden_channels, + self.input_kernel_size, + self.input_stride, + self.input_padding, + bias_attr=None, + padding_mode="circular", + ) + + self.Whf = nn.Conv2D( + self.hidden_channels, + self.hidden_channels, + self.hidden_kernel_size, + 1, + padding=1, + bias_attr=False, + padding_mode="circular", + ) + + self.Wxc = nn.Conv2D( + self.input_channels, + self.hidden_channels, + self.input_kernel_size, + self.input_stride, + self.input_padding, + bias_attr=None, + padding_mode="circular", + ) + + self.Whc = nn.Conv2D( + self.hidden_channels, + self.hidden_channels, + self.hidden_kernel_size, + 1, + padding=1, + bias_attr=False, + padding_mode="circular", + ) + + self.Wxo = nn.Conv2D( + self.input_channels, + self.hidden_channels, + self.input_kernel_size, + self.input_stride, + self.input_padding, + bias_attr=None, + padding_mode="circular", + ) + + self.Who = nn.Conv2D( + self.hidden_channels, + self.hidden_channels, + self.hidden_kernel_size, + 1, + padding=1, + bias_attr=False, + padding_mode="circular", + ) + + initializer_0 = nn.initializer.Constant(0.0) + initializer_1 = nn.initializer.Constant(1.0) + + initializer_0(self.Wxi.bias) + initializer_0(self.Wxf.bias) + initializer_0(self.Wxc.bias) + initializer_1(self.Wxo.bias) + + def forward(self, x, h, c): + ci = nn.functional.sigmoid(self.Wxi(x) + self.Whi(h)) + cf = nn.functional.sigmoid(self.Wxf(x) + self.Whf(h)) + cc = cf * c + ci * paddle.tanh(self.Wxc(x) + self.Whc(h)) + co = nn.functional.sigmoid(self.Wxo(x) + self.Who(h)) + ch = co * paddle.tanh(cc) + return ch, cc + + def init_hidden_tensor(self, prev_state): + return ((prev_state[0]).cuda(), (prev_state[1]).cuda()) + + +class encoder_block(nn.Layer): + """Encoder with CNN""" + + def __init__( + self, + input_channels, + hidden_channels, + input_kernel_size, + input_stride, + input_padding, + ): + super(encoder_block, self).__init__() + + self.input_channels = input_channels + self.hidden_channels = hidden_channels + self.input_kernel_size = input_kernel_size + self.input_stride = input_stride + self.input_padding = input_padding + + self.conv = utils.weight_norm( + nn.Conv2D( + self.input_channels, + self.hidden_channels, + self.input_kernel_size, + self.input_stride, + self.input_padding, + bias_attr=None, + padding_mode="circular", + ) + ) + + self.act = nn.ReLU() + + initializer_0 = nn.initializer.Constant(0.0) + initializer_0(self.conv.bias) + + def forward(self, x): + return self.act(self.conv(x)) + + +class Conv2DDerivative(nn.Layer): + def __init__(self, der_filter, resol, kernel_size=3, name=""): + super(Conv2DDerivative, self).__init__() + + self.resol = resol # constant in the finite difference + self.name = name + self.input_channels = 1 + self.output_channels = 1 + self.kernel_size = kernel_size + + self.padding = int((kernel_size - 1) / 2) + self.filter = nn.Conv2D( + self.input_channels, + self.output_channels, + self.kernel_size, + 1, + padding=0, + bias_attr=False, + ) + + # Fixed gradient operator + self.filter.weight = self.create_parameter( + shape=self.filter.weight.shape, + dtype=self.filter.weight.dtype, + default_initializer=nn.initializer.Assign( + paddle.to_tensor( + der_filter, dtype=paddle.get_default_dtype(), stop_gradient=True + ) + ), + ) + self.filter.weight.stop_gradient = True + + def forward(self, input): + derivative = self.filter(input) + return derivative / self.resol + + +class Conv1DDerivative(nn.Layer): + def __init__(self, der_filter, resol, kernel_size=3, name=""): + super(Conv1DDerivative, self).__init__() + + self.resol = resol # $\delta$*constant in the finite difference + self.name = name + self.input_channels = 1 + self.output_channels = 1 + self.kernel_size = kernel_size + + self.padding = int((kernel_size - 1) / 2) + self.filter = nn.Conv1D( + self.input_channels, + self.output_channels, + self.kernel_size, + 1, + padding=0, + bias_attr=False, + ) + + # Fixed gradient operator + self.filter.weight = self.create_parameter( + shape=self.filter.weight.shape, + dtype=self.filter.weight.dtype, + default_initializer=nn.initializer.Assign( + paddle.to_tensor( + der_filter, dtype=paddle.get_default_dtype(), stop_gradient=True + ) + ), + ) + self.filter.weight.stop_gradient = True + + def forward(self, input): + derivative = self.filter(input) + return derivative / self.resol + + +class loss_generator(nn.Layer): + """Loss generator for physics loss""" + + def __init__(self, dt, dx): + """Construct the derivatives, X = Width, Y = Height""" + super(loss_generator, self).__init__() + + # spatial derivative operator + self.laplace = Conv2DDerivative( + der_filter=LALP_OP, resol=(dx**2), kernel_size=5, name="laplace_operator" + ) + + self.dx = Conv2DDerivative( + der_filter=PARTIAL_X, resol=(dx * 1), kernel_size=5, name="dx_operator" + ) + + self.dy = Conv2DDerivative( + der_filter=PARTIAL_Y, resol=(dx * 1), kernel_size=5, name="dy_operator" + ) + + # temporal derivative operator + self.dt = Conv1DDerivative( + der_filter=[[[-1, 0, 1]]], resol=(dt * 2), kernel_size=3, name="partial_t" + ) + + def get_phy_Loss(self, output): + # spatial derivatives + laplace_u = self.laplace(output[1:-1, 0:1, :, :]) # [t,c,h,w] + laplace_v = self.laplace(output[1:-1, 1:2, :, :]) + + u_x = self.dx(output[1:-1, 0:1, :, :]) + u_y = self.dy(output[1:-1, 0:1, :, :]) + v_x = self.dx(output[1:-1, 1:2, :, :]) + v_y = self.dy(output[1:-1, 1:2, :, :]) + + # temporal derivative - u + u = output[:, 0:1, 2:-2, 2:-2] + lent = u.shape[0] + lenx = u.shape[3] + leny = u.shape[2] + u_conv1d = u.transpose((2, 3, 1, 0)) # [height(Y), width(X), c, step] + u_conv1d = u_conv1d.reshape((lenx * leny, 1, lent)) + u_t = self.dt(u_conv1d) # lent-2 due to no-padding + u_t = u_t.reshape((leny, lenx, 1, lent - 2)) + u_t = u_t.transpose((3, 2, 0, 1)) # [step-2, c, height(Y), width(X)] + + # temporal derivative - v + v = output[:, 1:2, 2:-2, 2:-2] + v_conv1d = v.transpose((2, 3, 1, 0)) # [height(Y), width(X), c, step] + v_conv1d = v_conv1d.reshape((lenx * leny, 1, lent)) + v_t = self.dt(v_conv1d) # lent-2 due to no-padding + v_t = v_t.reshape((leny, lenx, 1, lent - 2)) + v_t = v_t.transpose((3, 2, 0, 1)) # [step-2, c, height(Y), width(X)] + + u = output[1:-1, 0:1, 2:-2, 2:-2] # [t, c, height(Y), width(X)] + v = output[1:-1, 1:2, 2:-2, 2:-2] # [t, c, height(Y), width(X)] + + assert laplace_u.shape == u_t.shape + assert u_t.shape == v_t.shape + assert laplace_u.shape == u.shape + assert laplace_v.shape == v.shape + + # Reynolds number + R = 200.0 + + # 2D burgers eqn + f_u = u_t + u * u_x + v * u_y - (1 / R) * laplace_u + f_v = v_t + u * v_x + v * v_y - (1 / R) * laplace_v + + return f_u, f_v diff --git a/examples/smc_reac/ppsci/arch/phylstm.py b/examples/smc_reac/ppsci/arch/phylstm.py new file mode 100644 index 0000000000..a840d3f7ad --- /dev/null +++ b/examples/smc_reac/ppsci/arch/phylstm.py @@ -0,0 +1,239 @@ +# 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. + +import paddle +import paddle.nn as nn + +from ppsci.arch import base + + +class DeepPhyLSTM(base.Arch): + """DeepPhyLSTM init function. + + Args: + input_size (int): The input size. + output_size (int): The output size. + hidden_size (int, optional): The hidden size. Defaults to 100. + model_type (int, optional): The model type, value is 2 or 3, 2 indicates having two sub-models, 3 indicates having three submodels. Defaults to 2. + + Examples: + >>> import paddle + >>> import ppsci + >>> # model_type is `2` + >>> model = ppsci.arch.DeepPhyLSTM( + ... input_size=16, + ... output_size=1, + ... hidden_size=100, + ... model_type=2) + >>> out = model( + ... {"ag":paddle.rand([64, 16, 16]), + ... "ag_c":paddle.rand([64, 16, 16]), + ... "phi":paddle.rand([1, 16, 16])}) + >>> for k, v in out.items(): + ... print(f"{k} {v.dtype} {v.shape}") + eta_pred paddle.float32 [64, 16, 1] + eta_dot_pred paddle.float32 [64, 16, 1] + g_pred paddle.float32 [64, 16, 1] + eta_t_pred_c paddle.float32 [64, 16, 1] + eta_dot_pred_c paddle.float32 [64, 16, 1] + lift_pred_c paddle.float32 [64, 16, 1] + >>> # model_type is `3` + >>> model = ppsci.arch.DeepPhyLSTM( + ... input_size=16, + ... output_size=1, + ... hidden_size=100, + ... model_type=3) + >>> out = model( + ... {"ag":paddle.rand([64, 16, 1]), + ... "ag_c":paddle.rand([64, 16, 1]), + ... "phi":paddle.rand([1, 16, 16])}) + >>> for k, v in out.items(): + ... print(f"{k} {v.dtype} {v.shape}") + eta_pred paddle.float32 [64, 16, 1] + eta_dot_pred paddle.float32 [64, 16, 1] + g_pred paddle.float32 [64, 16, 1] + eta_t_pred_c paddle.float32 [64, 16, 1] + eta_dot_pred_c paddle.float32 [64, 16, 1] + lift_pred_c paddle.float32 [64, 16, 1] + g_t_pred_c paddle.float32 [64, 16, 1] + g_dot_pred_c paddle.float32 [64, 16, 1] + """ + + def __init__(self, input_size, output_size, hidden_size=100, model_type=2): + super().__init__() + self.input_size = input_size + self.output_size = output_size + self.hidden_size = hidden_size + self.model_type = model_type + + if self.model_type == 2: + self.lstm_model = nn.Sequential( + nn.LSTM(input_size, hidden_size), + nn.ReLU(), + nn.LSTM(hidden_size, hidden_size), + nn.ReLU(), + nn.LSTM(hidden_size, hidden_size), + nn.ReLU(), + nn.Linear(hidden_size, hidden_size), + nn.Linear(hidden_size, 3 * output_size), + ) + + self.lstm_model_f = nn.Sequential( + nn.LSTM(3 * output_size, hidden_size), + nn.ReLU(), + nn.LSTM(hidden_size, hidden_size), + nn.ReLU(), + nn.LSTM(hidden_size, hidden_size), + nn.ReLU(), + nn.Linear(hidden_size, hidden_size), + nn.Linear(hidden_size, output_size), + ) + elif self.model_type == 3: + self.lstm_model = nn.Sequential( + nn.LSTM(1, hidden_size), + nn.ReLU(), + nn.LSTM(hidden_size, hidden_size), + nn.ReLU(), + nn.LSTM(hidden_size, hidden_size), + nn.ReLU(), + nn.Linear(hidden_size, 3 * output_size), + ) + + self.lstm_model_f = nn.Sequential( + nn.LSTM(3 * output_size, hidden_size), + nn.ReLU(), + nn.LSTM(hidden_size, hidden_size), + nn.ReLU(), + nn.LSTM(hidden_size, hidden_size), + nn.ReLU(), + nn.Linear(hidden_size, output_size), + ) + + self.lstm_model_g = nn.Sequential( + nn.LSTM(2 * output_size, hidden_size), + nn.ReLU(), + nn.LSTM(hidden_size, hidden_size), + nn.ReLU(), + nn.LSTM(hidden_size, hidden_size), + nn.ReLU(), + nn.Linear(hidden_size, output_size), + ) + else: + raise ValueError(f"model_type should be 2 or 3, but got {model_type}") + + def forward(self, x): + if self._input_transform is not None: + x = self._input_transform(x) + + if self.model_type == 2: + result_dict = self._forward_type_2(x) + elif self.model_type == 3: + result_dict = self._forward_type_3(x) + if self._output_transform is not None: + result_dict = self._output_transform(x, result_dict) + return result_dict + + def _forward_type_2(self, x): + output = x["ag"] + for layer in self.lstm_model: + output = layer(output) + if isinstance(output, tuple): + output = output[0] + + eta_pred = output[:, :, 0 : self.output_size] + eta_dot_pred = output[:, :, self.output_size : 2 * self.output_size] + g_pred = output[:, :, 2 * self.output_size :] + + # for ag_c + output_c = x["ag_c"] + for layer in self.lstm_model: + output_c = layer(output_c) + if isinstance(output_c, tuple): + output_c = output_c[0] + + eta_pred_c = output_c[:, :, 0 : self.output_size] + eta_dot_pred_c = output_c[:, :, self.output_size : 2 * self.output_size] + g_pred_c = output_c[:, :, 2 * self.output_size :] + eta_t_pred_c = paddle.matmul(x["phi"], eta_pred_c) + eta_tt_pred_c = paddle.matmul(x["phi"], eta_dot_pred_c) + eta_dot1_pred_c = eta_dot_pred_c[:, :, 0:1] + tmp = paddle.concat([eta_pred_c, eta_dot1_pred_c, g_pred_c], 2) + f = tmp + for layer in self.lstm_model_f: + f = layer(f) + if isinstance(f, tuple): + f = f[0] + + lift_pred_c = eta_tt_pred_c + f + + return { + "eta_pred": eta_pred, + "eta_dot_pred": eta_dot_pred, + "g_pred": g_pred, + "eta_t_pred_c": eta_t_pred_c, + "eta_dot_pred_c": eta_dot_pred_c, + "lift_pred_c": lift_pred_c, + } + + def _forward_type_3(self, x): + # physics informed neural networks + output = x["ag"] + for layer in self.lstm_model: + output = layer(output) + if isinstance(output, tuple): + output = output[0] + + eta_pred = output[:, :, 0 : self.output_size] + eta_dot_pred = output[:, :, self.output_size : 2 * self.output_size] + g_pred = output[:, :, 2 * self.output_size :] + + output_c = x["ag_c"] + for layer in self.lstm_model: + output_c = layer(output_c) + if isinstance(output_c, tuple): + output_c = output_c[0] + + eta_pred_c = output_c[:, :, 0 : self.output_size] + eta_dot_pred_c = output_c[:, :, self.output_size : 2 * self.output_size] + g_pred_c = output_c[:, :, 2 * self.output_size :] + + eta_t_pred_c = paddle.matmul(x["phi"], eta_pred_c) + eta_tt_pred_c = paddle.matmul(x["phi"], eta_dot_pred_c) + g_t_pred_c = paddle.matmul(x["phi"], g_pred_c) + + f = paddle.concat([eta_pred_c, eta_dot_pred_c, g_pred_c], 2) + for layer in self.lstm_model_f: + f = layer(f) + if isinstance(f, tuple): + f = f[0] + + lift_pred_c = eta_tt_pred_c + f + + eta_dot1_pred_c = eta_dot_pred_c[:, :, 0:1] + g_dot_pred_c = paddle.concat([eta_dot1_pred_c, g_pred_c], 2) + for layer in self.lstm_model_g: + g_dot_pred_c = layer(g_dot_pred_c) + if isinstance(g_dot_pred_c, tuple): + g_dot_pred_c = g_dot_pred_c[0] + + return { + "eta_pred": eta_pred, + "eta_dot_pred": eta_dot_pred, + "g_pred": g_pred, + "eta_t_pred_c": eta_t_pred_c, + "eta_dot_pred_c": eta_dot_pred_c, + "lift_pred_c": lift_pred_c, + "g_t_pred_c": g_t_pred_c, + "g_dot_pred_c": g_dot_pred_c, + } diff --git a/examples/smc_reac/ppsci/arch/physx_transformer.py b/examples/smc_reac/ppsci/arch/physx_transformer.py new file mode 100644 index 0000000000..267fb458c6 --- /dev/null +++ b/examples/smc_reac/ppsci/arch/physx_transformer.py @@ -0,0 +1,407 @@ +# 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. + +""" +Code below is heavily based on [transformer-physx](https://github.com/zabaras/transformer-physx) +""" + +from __future__ import annotations + +from typing import Optional +from typing import Tuple + +import paddle +import paddle.nn.functional as F +from paddle import nn +from paddle.nn.initializer import Constant +from paddle.nn.initializer import Normal + +from ppsci.arch import base +from ppsci.arch.embedding_koopman import CylinderEmbedding + +zeros_ = Constant(value=0.0) +ones_ = Constant(value=1.0) + + +class MaskedAttention(nn.Layer): + """Masked self-attention module. + + Args: + embed_dim (int): The expected feature size in the input and output. + num_ctx (int): Context length of block. + num_heads (int): The number of heads in multi-head attention. + attn_drop (float, optional): The dropout probability used on attention + weights to drop some attention targets. Defaults to 0. + proj_drop (float, optional): The dropout probability used on output. Defaults to 0. + scale (bool, optional): Whether to scale attention weights. Defaults to False. + """ + + def __init__( + self, + embed_dim: int, + num_ctx: int, + num_heads: int, + attn_drop: float = 0.0, + proj_drop: float = 0.0, + scale: bool = False, + ): + super().__init__() + self.register_buffer( + "bias", + paddle.tril(paddle.ones((num_ctx, num_ctx), dtype="int32")).reshape( + [1, 1, num_ctx, num_ctx] + ), + ) + + self.register_buffer("masked_bias", paddle.to_tensor(-1e4)) + self.num_heads = num_heads + self.split_size = embed_dim + self.scale = scale + + self.qkv_proj = nn.Linear(embed_dim, embed_dim * 3) + self.out_proj = nn.Linear(embed_dim, embed_dim) + self.attn_drop = nn.Dropout(attn_drop) + self.proj_drop = nn.Dropout(proj_drop) + + def _attn( + self, + query, + key, + value, + attention_mask=None, + head_mask=None, + output_attentions=False, + ): + attn = paddle.matmul(query, key) + if self.scale: + attn = attn / (float(value.shape[-1]) ** 0.5) + + nd, ns = attn.shape[-2], attn.shape[-1] + mask = self.bias[:, :, ns - nd : ns, :ns] + attn = paddle.where(mask > 0, attn, self.masked_bias.cast(attn.dtype)) + + if attention_mask is not None: + attn = attn + attention_mask + + attn = F.softmax(attn, axis=-1) + attn = self.attn_drop(attn) + + if head_mask is not None: + attn = attn * head_mask + + outputs = [paddle.matmul(attn, value)] + if output_attentions: + outputs.append(attn) + return outputs + + def merge_heads(self, x): + x = x.transpose([0, 2, 1, 3]) + new_x_shape = x.shape[:-2] + [ + x.shape[-2] * x.shape[-1], + ] + return x.reshape(new_x_shape) + + def split_heads(self, x, k=False): + new_x_shape = x.shape[:-1] + [self.num_heads, x.shape[-1] // self.num_heads] + x = x.reshape(new_x_shape) + if k: + return x.transpose([0, 2, 3, 1]) + return x.transpose([0, 2, 1, 3]) + + def forward( + self, + x, + layer_past=None, + attention_mask=None, + head_mask=None, + output_attentions=False, + ): + x = self.qkv_proj(x) + query, key, value = x.split(x.shape[2] // self.split_size, axis=2) + query = self.split_heads(query) + key = self.split_heads(key, k=True) + value = self.split_heads(value) + # Concat previous key and value tensors + if layer_past is not None: + past_key, past_value = layer_past[0].transpose([0, 1, 3, 2]), layer_past[1] + key = paddle.concat((past_key, key), axis=-1) + value = paddle.concat((past_value, value), axis=-2) + + attn_outputs = self._attn( + query, key, value, attention_mask, head_mask, output_attentions + ) + output = attn_outputs[0] + output = self.merge_heads(output) + output = self.out_proj(output) + output = self.proj_drop(output) + + outputs = [output] + attn_outputs[1:] + return outputs + + +class MLP(nn.Layer): + """Multi layer perceptron module used in Transformer. + + Args: + in_features (int): Number of the input features. + hidden_features (Optional[int]): Number of the hidden size. Defaults to None. + out_features (Optional[int]): Number of the output features. Defaults to None. + drop (float, optional): Probability of dropout the units. Defaults to 0. + """ + + def __init__( + self, + in_features: int, + hidden_features: Optional[int] = None, + out_features: Optional[int] = None, + drop: float = 0.0, + ): + super().__init__() + out_features = out_features or in_features + hidden_features = hidden_features or in_features + + self.fc1 = nn.Linear(in_features, hidden_features) + self.act = nn.GELU(approximate=True) + self.fc2 = nn.Linear(hidden_features, out_features) + self.drop = nn.Dropout(drop) + + def forward(self, x): + x = self.fc1(x) + x = self.act(x) + x = self.fc2(x) + x = self.drop(x) + return x + + +class Block(nn.Layer): + """Transformer decoder block consisting of layer norm, + masked self-attention, layer norm and fully connected layer. + + Args: + num_ctx (int): Context length of block + embed_size (int): The number of embedding size. + num_heads (int): The number of heads in multi-head attention. + attn_pdrop (float): The dropout probability used on attention + weights to drop some attention targets. + resid_pdrop (float): The dropout probability used on output. + scale (bool, optional): Scaled self-attention calculation. Defaults to False. + """ + + def __init__( + self, + num_ctx: int, + embed_size: int, + num_heads: int, + attn_pdrop: float, + resid_pdrop: float, + scale: bool = False, + ): + super().__init__() + self.ln_1 = nn.LayerNorm(embed_size) + self.attn = MaskedAttention( + embed_size, num_ctx, num_heads, attn_pdrop, resid_pdrop, scale + ) + self.ln_2 = nn.LayerNorm(embed_size) + self.mlp = MLP(embed_size, 4 * embed_size, resid_pdrop) + + def forward( + self, + x, + layer_past=None, + attention_mask=None, + head_mask=None, + output_attentions=False, + ): + # Evaluate attention heads + output_attn = self.attn.forward( + self.ln_1(x), + layer_past=layer_past, + attention_mask=attention_mask, + head_mask=head_mask, + output_attentions=output_attentions, + ) + x = x + output_attn[0] + m = self.mlp(self.ln_2(x)) + x = x + m + outputs = [x] + output_attn[1:] + return outputs + + +class PhysformerGPT2(base.Arch): + """Transformer decoder model for modeling physics. + + Args: + input_keys (Tuple[str, ...]): Input keys, such as ("embeds",). + output_keys (Tuple[str, ...]): Output keys, such as ("pred_embeds",). + num_layers (int): Number of transformer layers. + num_ctx (int): Context length of block. + embed_size (int): The number of embedding size. + num_heads (int): The number of heads in multi-head attention. + embd_pdrop (float, optional): The dropout probability used on embedding features. Defaults to 0.0. + attn_pdrop (float, optional): The dropout probability used on attention weights. Defaults to 0.0. + resid_pdrop (float, optional): The dropout probability used on block outputs. Defaults to 0.0. + initializer_range (float, optional): Initializer range of linear layer. Defaults to 0.05. + embedding_model (Optional[base.Arch]): Embedding model, If this parameter is set, + the embedding model will map the input data to the embedding space and the + output data to the physical space. Defaults to None. + + Examples: + >>> import paddle + >>> import ppsci + >>> model = ppsci.arch.PhysformerGPT2(("embeds", ), ("pred_embeds", ), 6, 16, 128, 4) + >>> data = paddle.to_tensor(paddle.randn([10, 16, 128])) + >>> inputs = {"embeds": data} + >>> outputs = model(inputs) + >>> print(outputs["pred_embeds"].shape) + [10, 16, 128] + """ + + def __init__( + self, + input_keys: Tuple[str, ...], + output_keys: Tuple[str, ...], + num_layers: int, + num_ctx: int, + embed_size: int, + num_heads: int, + embd_pdrop: float = 0.0, + attn_pdrop: float = 0.0, + resid_pdrop: float = 0.0, + initializer_range: float = 0.05, + embedding_model: Optional[base.Arch] = None, + ): + super().__init__() + self.input_keys = input_keys + self.output_keys = output_keys + + self.num_layers = num_layers + self.num_ctx = num_ctx + self.embed_size = embed_size + self.num_heads = num_heads + self.embd_pdrop = embd_pdrop + self.attn_pdrop = attn_pdrop + self.resid_pdrop = resid_pdrop + self.initializer_range = initializer_range + + self.drop = nn.Dropout(embd_pdrop) + self.blocks = nn.LayerList( + [ + Block( + num_ctx, embed_size, num_heads, attn_pdrop, resid_pdrop, scale=True + ) + for _ in range(num_layers) + ] + ) + self.ln = nn.LayerNorm(embed_size) + self.linear = nn.Linear(embed_size, embed_size) + + self.apply(self._init_weights) + self.embedding_model = embedding_model + + def _init_weights(self, module): + if isinstance(module, nn.Linear): + normal_ = Normal(mean=0.0, std=self.initializer_range) + normal_(module.weight) + if module.bias is not None: + zeros_(module.bias) + elif isinstance(module, nn.LayerNorm): + zeros_(module.bias) + ones_(module.weight) + + def get_position_embed(self, x): + B, N, _ = x.shape + position_ids = paddle.arange(0, N, dtype=paddle.get_default_dtype()).reshape( + [1, N, 1] + ) + position_ids = position_ids.repeat_interleave(B, axis=0) + + position_embeds = paddle.zeros_like(x) + i = paddle.arange(0, self.embed_size // 2).unsqueeze(0).unsqueeze(0) + position_embeds[:, :, ::2] = paddle.sin( + position_ids / 10000 ** (2 * i / self.embed_size) + ) + position_embeds[:, :, 1::2] = paddle.cos( + position_ids / 10000 ** (2 * i / self.embed_size) + ) + return position_embeds + + def _generate_time_series(self, x, max_length): + cur_len = x.shape[1] + if cur_len >= max_length: + raise ValueError( + f"max_length({max_length}) should be larger than " + f"the length of input context({cur_len})" + ) + + while cur_len < max_length: + model_inputs = x[:, -1:] + outputs = self.forward_tensor(model_inputs) + next_output = outputs[0][:, -1:] + x = paddle.concat([x, next_output], axis=1) + cur_len = cur_len + 1 + return x + + @paddle.no_grad() + def generate(self, x, max_length=256): + if max_length <= 0: + raise ValueError( + f"max_length({max_length}) should be a strictly positive integer." + ) + outputs = self._generate_time_series(x, max_length) + return outputs + + def forward_tensor(self, x): + position_embeds = self.get_position_embed(x) + # Combine input embedding, position embedding + hidden_states = x + position_embeds + hidden_states = self.drop(hidden_states) + + # Loop through transformer self-attention layers + for block in self.blocks: + block_outputs = block(hidden_states) + hidden_states = block_outputs[0] + outputs = self.linear(self.ln(hidden_states)) + return (outputs,) + + def forward_eval(self, x): + input_embeds = x[:, :1] + outputs = self.generate(input_embeds) + return (outputs[:, 1:],) + + @staticmethod + def split_to_dict(data_tensors, keys): + return {key: data_tensors[i] for i, key in enumerate(keys)} + + def forward(self, x): + if self._input_transform is not None: + x = self._input_transform(x) + x_tensor = self.concat_to_tensor(x, self.input_keys, axis=-1) + if self.embedding_model is not None: + if isinstance(self.embedding_model, CylinderEmbedding): + x_tensor = self.embedding_model.encoder(x_tensor, x["visc"]) + else: + x_tensor = self.embedding_model.encoder(x_tensor) + + if self.training: + y = self.forward_tensor(x_tensor) + else: + y = self.forward_eval(x_tensor) + + if self.embedding_model is not None: + y = (self.embedding_model.decoder(y[0]),) + + y = self.split_to_dict(y, self.output_keys) + if self._output_transform is not None: + y = self._output_transform(x, y) + return y diff --git a/examples/smc_reac/ppsci/arch/regdgcnn.py b/examples/smc_reac/ppsci/arch/regdgcnn.py new file mode 100644 index 0000000000..c39a5b0540 --- /dev/null +++ b/examples/smc_reac/ppsci/arch/regdgcnn.py @@ -0,0 +1,250 @@ +# Copyright (c) 2024 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. + +""" +Created on Mon May 29 22:18:28 2023 + +@author: Mohamed Elrefaie, mohamed.elrefaie@mit.edu, mohamed.elrefaie@tum.de + +This module is part of the research presented in the paper +"DrivAerNet: A Parametric Car Dataset for Data-driven Aerodynamic Design and Graph-Based Drag Prediction". +It extends the work by introducing a Deep Graph Convolutional Neural Network (RegDGCNN) model for Regression Tasks, +specifically designed for processing 3D point cloud data of car models from the DrivAerNet dataset. + +The RegDGCNN model utilizes a series of graph-based convolutional layers to effectively capture the complex geometric +and topological structure of 3D car models, facilitating advanced aerodynamic analyses and predictions. +The model architecture incorporates several techniques, including dynamic graph construction, +EdgeConv operations, and global feature aggregation, to robustly learn from graphs and point cloud data. + +Parts of this code are modified from the original version authored by Yue Wang +""" + +from __future__ import annotations + +from typing import Dict +from typing import Tuple + +import paddle + + +def transpose_aux_func(dims, dim0, dim1): + perm = list(range(dims)) + perm[dim0], perm[dim1] = perm[dim1], perm[dim0] + return perm + + +def knn(x, k): + """ + Computes the k-nearest neighbors for each point in x. + + Args: + x (paddle.Tensor): The input tensor of shape (batch_size, num_dims, num_points). + k (int): The number of nearest neighbors to find. + + Returns: + paddle.Tensor: Indices of the k-nearest neighbors for each point, shape (batch_size, num_points, k). + """ + inner = -2 * paddle.matmul( + x=x.transpose(perm=transpose_aux_func(x.ndim, 2, 1)), y=x + ) + xx = paddle.sum(x=x**2, axis=1, keepdim=True) + pairwise_distance = ( + -xx - inner - xx.transpose(perm=transpose_aux_func(xx.ndim, 2, 1)) + ) + idx = pairwise_distance.topk(k=k, axis=-1)[1] + return idx + + +def get_graph_feature(x, k=20, idx=None): + """ + Constructs local graph features for each point by finding its k-nearest neighbors and + concatenating the relative position vectors. + + Args: + x (paddle.Tensor): The input tensor of shape (batch_size, num_dims, num_points). + k (int): The number of neighbors to consider for graph construction. + idx (paddle.Tensor, optional): Precomputed k-nearest neighbor indices. + + Returns: + paddle.Tensor: The constructed graph features of shape (batch_size, 2*num_dims, num_points, k). + """ + batch_size = x.shape[0] + num_points = x.shape[2] + x = x.reshape([batch_size, -1, num_points]) + if idx is None: + idx = knn(x, k=k) + idx_base = paddle.arange(start=0, end=batch_size).reshape([-1, 1, 1]) * num_points + idx = idx + idx_base + idx = idx.reshape([-1]) + _, num_dims, _ = tuple(x.shape) + x = x.transpose(perm=transpose_aux_func(x.ndim, 2, 1)).contiguous() + feature = x.reshape([batch_size * num_points, -1])[idx, :] + feature = feature.reshape([batch_size, num_points, k, num_dims]) + x = x.reshape([batch_size, num_points, 1, num_dims]).tile(repeat_times=[1, 1, k, 1]) + feature = ( + paddle.concat(x=(feature - x, x), axis=3) + .transpose(perm=[0, 3, 1, 2]) + .contiguous() + ) + del x, idx, idx_base + paddle.device.cuda.empty_cache() + return feature + + +class RegDGCNN(paddle.nn.Layer): + """Deep Graph Convolutional Neural Network for Regression Tasks (RegDGCNN) designed to process 3D point cloud data. + + This network architecture extracts hierarchical features from point clouds using graph-based convolutions, + enabling effective learning of spatial structures. + + Args: + input_keys (Tuple[str, ...]): Keys for input data fields. + label_keys (Tuple[str, ...]): Keys for label data fields. + weight_keys (Tuple[str, ...]): Keys for weight data fields. + args (dict): Configuration parameters including: + - 'k' (int): Number of neighbors for graph convolution. + - 'emb_dims' (int): Embedding dimensions for feature aggregation. + - 'dropout' (float): Dropout rate for regularization. + output_channels (int, optional): Number of output channels. Defaults to 1. + """ + + def __init__( + self, + input_keys: Tuple[str, ...], + label_keys: Tuple[str, ...], + weight_keys: Tuple[str, ...], + args: dict, + output_channels=1, + ): + + super(RegDGCNN, self).__init__() + self.input_keys = input_keys + self.label_keys = label_keys + self.weight_keys = weight_keys + self.args = args + self.k = args["k"] + self.bn1 = paddle.nn.BatchNorm2D(num_features=256) + self.bn2 = paddle.nn.BatchNorm2D(num_features=512) + self.bn3 = paddle.nn.BatchNorm2D(num_features=512) + self.bn4 = paddle.nn.BatchNorm2D(num_features=1024) + self.bn5 = paddle.nn.BatchNorm1D(num_features=args["emb_dims"]) + self.conv1 = paddle.nn.Sequential( + paddle.nn.Conv2D( + in_channels=6, out_channels=256, kernel_size=1, bias_attr=False + ), + self.bn1, + paddle.nn.LeakyReLU(negative_slope=0.2), + ) + self.conv2 = paddle.nn.Sequential( + paddle.nn.Conv2D( + in_channels=256 * 2, out_channels=512, kernel_size=1, bias_attr=False + ), + self.bn2, + paddle.nn.LeakyReLU(negative_slope=0.2), + ) + self.conv3 = paddle.nn.Sequential( + paddle.nn.Conv2D( + in_channels=512 * 2, out_channels=512, kernel_size=1, bias_attr=False + ), + self.bn3, + paddle.nn.LeakyReLU(negative_slope=0.2), + ) + self.conv4 = paddle.nn.Sequential( + paddle.nn.Conv2D( + in_channels=512 * 2, out_channels=1024, kernel_size=1, bias_attr=False + ), + self.bn4, + paddle.nn.LeakyReLU(negative_slope=0.2), + ) + self.conv5 = paddle.nn.Sequential( + paddle.nn.Conv1D( + in_channels=2304, + out_channels=args["emb_dims"], + kernel_size=1, + bias_attr=False, + ), + self.bn5, + paddle.nn.LeakyReLU(negative_slope=0.2), + ) + self.linear1 = paddle.nn.Linear( + in_features=args["emb_dims"] * 2, out_features=128, bias_attr=False + ) + self.bn6 = paddle.nn.BatchNorm1D(num_features=128) + self.dp1 = paddle.nn.Dropout(p=args["dropout"]) + self.linear2 = paddle.nn.Linear(in_features=128, out_features=64) + self.bn7 = paddle.nn.BatchNorm1D(num_features=64) + self.dp2 = paddle.nn.Dropout(p=args["dropout"]) + self.linear3 = paddle.nn.Linear(in_features=64, out_features=32) + self.bn8 = paddle.nn.BatchNorm1D(num_features=32) + self.dp3 = paddle.nn.Dropout(p=args["dropout"]) + self.linear4 = paddle.nn.Linear(in_features=32, out_features=16) + self.bn9 = paddle.nn.BatchNorm1D(num_features=16) + self.dp4 = paddle.nn.Dropout(p=args["dropout"]) + self.linear5 = paddle.nn.Linear(in_features=16, out_features=output_channels) + + def forward(self, x: paddle.Tensor) -> Dict[str, paddle.Tensor]: + """ + Forward pass of the model to process input data and predict outputs. + + Args: + x (paddle.Tensor): Input tensor representing a batch of point clouds. + + Returns: + Dict[str, paddle.Tensor]: Model predictions for the input batch. + + """ + + x = x[self.input_keys[0]] + batch_size = x.shape[0] + x = x.transpose(perm=[0, 2, 1]) + + x = get_graph_feature(x, k=self.k) + x = self.conv1(x) + x1 = x.max(axis=-1, keepdim=False) + x = get_graph_feature(x1, k=self.k) + x = self.conv2(x) + x2 = x.max(axis=-1, keepdim=False) + x = get_graph_feature(x2, k=self.k) + x = self.conv3(x) + x3 = x.max(axis=-1, keepdim=False) + x = get_graph_feature(x3, k=self.k) + x = self.conv4(x) + x4 = x.max(axis=-1, keepdim=False) + x = paddle.concat(x=(x1, x2, x3, x4), axis=1) + x = self.conv5(x) + x1 = paddle.nn.functional.adaptive_max_pool1d(x=x, output_size=1).reshape( + [batch_size, -1] + ) + x2 = paddle.nn.functional.adaptive_avg_pool1d(x=x, output_size=1).reshape( + [batch_size, -1] + ) + x = paddle.concat(x=(x1, x2), axis=1) + x = paddle.nn.functional.leaky_relu( + x=self.bn6(self.linear1(x)), negative_slope=0.2 + ) + x = self.dp1(x) + x = paddle.nn.functional.leaky_relu( + x=self.bn7(self.linear2(x)), negative_slope=0.2 + ) + x = self.dp2(x) + x = paddle.nn.functional.leaky_relu( + x=self.bn8(self.linear3(x)), negative_slope=0.2 + ) + x = self.dp3(x) + x = paddle.nn.functional.leaky_relu( + x=self.bn9(self.linear4(x)), negative_slope=0.2 + ) + x = self.dp4(x) + x = self.linear5(x) + return {self.label_keys[0]: x} diff --git a/examples/smc_reac/ppsci/arch/regpointnet.py b/examples/smc_reac/ppsci/arch/regpointnet.py new file mode 100644 index 0000000000..2dee522326 --- /dev/null +++ b/examples/smc_reac/ppsci/arch/regpointnet.py @@ -0,0 +1,146 @@ +# Copyright (c) 2024 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. + +# Copyright 2024 Mohamed Elrefaie +""" +@author: Mohamed Elrefaie, mohamed.elrefaie@mit.edu mohamed.elrefaie@tum.de + +This module is part of the research presented in the paper: +"DrivAerNet++: A Large-Scale Multimodal Car Dataset with Computational Fluid Dynamics Simulations and Deep Learning Benchmarks". + +This module is used to define point-cloud models, includingPointNet +for the task of surrogate modeling of the aerodynamic drag. +""" + +from __future__ import annotations + +from typing import Dict +from typing import Tuple + +import paddle + + +class RegPointNet(paddle.nn.Layer): + """ + PointNet-based regression model for 3D point cloud data. + + This network architecture is designed to process 3D point cloud data using a series of convolutional layers, + followed by fully connected layers, enabling effective learning of spatial structures and features. + + Args: + input_keys (Tuple[str, ...]): Keys for input data fields. + output_keys (Tuple[str, ...]): Keys for output data fields. + weight_keys (Tuple[str, ...]): Keys for weight data fields. + args (dict): Configuration parameters including: + - 'emb_dims' (int): Dimensionality of the embedding space. + - 'dropout' (float): Dropout probability. + """ + + def __init__( + self, + input_keys: Tuple[str, ...], + output_keys: Tuple[str, ...], + weight_keys: Tuple[str, ...], + args, + ): + super().__init__() + self.input_keys = input_keys + self.output_keys = output_keys + self.weight_keys = weight_keys + self.args = args + self.conv1 = paddle.nn.Conv1D( + in_channels=3, out_channels=512, kernel_size=1, bias_attr=False + ) + self.conv2 = paddle.nn.Conv1D( + in_channels=512, out_channels=1024, kernel_size=1, bias_attr=False + ) + self.conv3 = paddle.nn.Conv1D( + in_channels=1024, out_channels=1024, kernel_size=1, bias_attr=False + ) + self.conv4 = paddle.nn.Conv1D( + in_channels=1024, out_channels=1024, kernel_size=1, bias_attr=False + ) + self.conv5 = paddle.nn.Conv1D( + in_channels=1024, out_channels=1024, kernel_size=1, bias_attr=False + ) + self.conv6 = paddle.nn.Conv1D( + in_channels=1024, + out_channels=args["emb_dims"], + kernel_size=1, + bias_attr=False, + ) + self.bn1 = paddle.nn.BatchNorm1D(num_features=512) + self.bn2 = paddle.nn.BatchNorm1D(num_features=1024) + self.bn3 = paddle.nn.BatchNorm1D(num_features=1024) + self.bn4 = paddle.nn.BatchNorm1D(num_features=1024) + self.bn5 = paddle.nn.BatchNorm1D(num_features=1024) + self.bn6 = paddle.nn.BatchNorm1D(num_features=args["emb_dims"]) + self.dropout_conv = paddle.nn.Dropout(p=args["dropout"]) + self.dropout_linear = paddle.nn.Dropout(p=args["dropout"]) + self.conv_shortcut = paddle.nn.Conv1D( + in_channels=3, out_channels=args["emb_dims"], kernel_size=1, bias_attr=False + ) + self.bn_shortcut = paddle.nn.BatchNorm1D(num_features=args["emb_dims"]) + self.linear1 = paddle.nn.Linear( + in_features=args["emb_dims"], out_features=512, bias_attr=False + ) + self.bn7 = paddle.nn.BatchNorm1D(num_features=512) + self.linear2 = paddle.nn.Linear( + in_features=512, out_features=256, bias_attr=False + ) + self.bn8 = paddle.nn.BatchNorm1D(num_features=256) + self.linear3 = paddle.nn.Linear(in_features=256, out_features=128) + self.bn9 = paddle.nn.BatchNorm1D(num_features=128) + self.linear4 = paddle.nn.Linear(in_features=128, out_features=64) + self.bn10 = paddle.nn.BatchNorm1D(num_features=64) + self.final_linear = paddle.nn.Linear(in_features=64, out_features=1) + + def forward(self, x: Dict[str, paddle.Tensor]) -> Dict[str, paddle.Tensor]: + """ + Forward pass of the network. + + Args: + x (Dict[str, paddle.Tensor]): Input tensor of shape (batch_size, 3, num_points). + + Returns: + Dict[str, paddle.Tensor]: A dictionary where the key is the first element of `self.output_keys` + and the value is the output tensor of the predicted scalar value. + """ + + x: paddle.Tensor = x[self.input_keys[0]] + + x_processed = x.transpose(perm=[0, 2, 1]) + + shortcut = self.bn_shortcut(self.conv_shortcut(x_processed)) + x = paddle.nn.functional.relu(x=self.bn1(self.conv1(x_processed))) + x = self.dropout_conv(x) + x = paddle.nn.functional.relu(x=self.bn2(self.conv2(x))) + x = self.dropout_conv(x) + x = paddle.nn.functional.relu(x=self.bn3(self.conv3(x))) + x = self.dropout_conv(x) + x = paddle.nn.functional.relu(x=self.bn4(self.conv4(x))) + x = self.dropout_conv(x) + x = paddle.nn.functional.relu(x=self.bn5(self.conv5(x))) + x = self.dropout_conv(x) + x = paddle.nn.functional.relu(x=self.bn6(self.conv6(x))) + x = x + shortcut + x = paddle.nn.functional.adaptive_max_pool1d(x=x, output_size=1).squeeze( + axis=-1 + ) + x = paddle.nn.functional.relu(x=self.bn7(self.linear1(x))) + x = paddle.nn.functional.relu(x=self.bn8(self.linear2(x))) + x = paddle.nn.functional.relu(x=self.bn9(self.linear3(x))) + x = paddle.nn.functional.relu(x=self.bn10(self.linear4(x))) + x = self.final_linear(x) + return {self.output_keys[0]: x} diff --git a/examples/smc_reac/ppsci/arch/sfnonet.py b/examples/smc_reac/ppsci/arch/sfnonet.py new file mode 100644 index 0000000000..c7695fb02d --- /dev/null +++ b/examples/smc_reac/ppsci/arch/sfnonet.py @@ -0,0 +1,568 @@ +from typing import Dict +from typing import List +from typing import Optional +from typing import Tuple +from typing import Union + +import paddle +import paddle.nn.functional as F +from paddle import nn + +from ppsci.arch import base +from ppsci.arch import fno_block +from ppsci.arch.paddle_harmonics import sht as paddle_sht +from ppsci.utils import initializer + +einsum_symbols = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ" + + +def _contract_dense(x, weight, separable=False, dhconv=True): + order = len(x.shape) + x_syms = list(einsum_symbols[:order]) + + # in_channels, out_channels, x, y... + weight_syms = list(x_syms[1:]) # no batch-size + + # batch-size, out_channels, x, y... + if separable: + out_syms = [x_syms[0]] + list(weight_syms) + else: + weight_syms.insert(1, einsum_symbols[order]) # outputs + out_syms = list(weight_syms) + out_syms[0] = x_syms[0] + + if dhconv: + weight_syms.pop() + + eq = "".join(x_syms) + "," + "".join(weight_syms) + "->" + "".join(out_syms) + # For the darcy flow, the only einsum is abcd,becd->aecd, where x and weights are shaped [32,32,8,8] + if not isinstance(weight, paddle.Tensor): + weight = paddle.to_tensor(weight) + + return paddle.einsum(eq, x, weight) + + +def _contract_dense_trick(x, weight_real, weight_imag, separable=False, dhconv=True): + # the same as above function, but do the complex multiplication manually to avoid the einsum bug in paddle + order = len(x.shape) + # batch-size, in_channels, x, y... + x_syms = list(einsum_symbols[:order]) + + # in_channels, out_channels, x, y... + weight_syms = list(x_syms[1:]) # no batch-size + + # batch-size, out_channels, x, y... + if separable: + out_syms = [x_syms[0]] + list(weight_syms) + else: + weight_syms.insert(1, einsum_symbols[order]) # outputs + out_syms = list(weight_syms) + out_syms[0] = x_syms[0] + + if dhconv: + weight_syms.pop() + + eq = "".join(x_syms) + "," + "".join(weight_syms) + "->" + "".join(out_syms) + + o1_real = paddle.einsum(eq, x.real(), weight_real) - paddle.einsum( + eq, x.imag(), weight_imag + ) + o1_imag = paddle.einsum(eq, x.imag(), weight_real) + paddle.einsum( + eq, x.real(), weight_imag + ) + x = paddle.complex(o1_real, o1_imag) + return x + + +def _contract_dense_separable(x, weight, separable=True): + if not separable: + raise ValueError("This function is only for separable=True") + return x * weight + + +def get_contract_fun(weight, implementation="reconstructed", separable=False): + """Generic ND implementation of Fourier Spectral Conv contraction. + + Args: + weight (FactorizedTensor): The factoriz Tensor. + implementation (str, optional): Whether to reconstruct the weight and do a forward pass (reconstructed) + or contract directly the factors of the factorized weight with the input (factorized). + {'reconstructed', 'factorized'} Defaults to "reconstructed". + separable (bool, optional): Whether to use the separable implementation of contraction. This arg is + only checked when `implementation=reconstructed`. Defaults to False. + """ + + if implementation == "reconstructed": + if separable: + return _contract_dense_separable + else: + return _contract_dense_trick + elif implementation == "factorized": + if isinstance(weight, paddle.Tensor): + return _contract_dense_trick + + else: + raise ValueError( + f'Got implementation={implementation}, expected "reconstructed" or "factorized"' + ) + + +class SHT(nn.Layer): + """A wrapper for the Spherical Harmonics transform + + Allows to call it with an interface similar to that of FFT + """ + + def __init__(self, dtype=paddle.float32): + super().__init__() + self.dtype = dtype + self._SHT_cache = nn.LayerDict() + self._iSHT_cache = nn.LayerDict() + + def sht(self, x, s=None, norm="ortho", grid="equiangular"): + *_, height, width = x.shape # height = latitude, width = longitude + if s is None: + if grid == "equiangular": + modes_width = height // 2 + else: + modes_width = height + modes_height = height + else: + modes_height, modes_width = s + + cache_key = f"{height}_{width}_{modes_height}_{modes_width}_{norm}_{grid}" + + try: + sht = self._SHT_cache[cache_key] + except KeyError: + sht = paddle_sht.RealSHT( + nlat=height, + nlon=width, + lmax=modes_height, + mmax=modes_width, + grid=grid, + norm=norm, + ).astype(dtype=self.dtype) + + self._SHT_cache[cache_key] = sht + + return sht(x) + + def isht(self, x, s=None, norm="ortho", grid="equiangular"): + *_, modes_height, modes_width = x.shape # height = latitude, width = longitude + if s is None: + if grid == "equiangular": + width = modes_width * 2 + else: + width = modes_width + height = modes_height + else: + height, width = s + + cache_key = f"{height}_{width}_{modes_height}_{modes_width}_{norm}_{grid}" + + try: + isht = self._iSHT_cache[cache_key] + except KeyError: + isht = paddle_sht.InverseRealSHT( + nlat=height, + nlon=width, + lmax=modes_height, + mmax=modes_width, + grid=grid, + norm=norm, + ).astype(dtype=self.dtype) + self._iSHT_cache[cache_key] = isht + + return isht(x) + + +Number = Union[int, float] + + +class SphericalConv(nn.Layer): + """Spherical Convolution, base class for the SFNO [1]. + .. [1] Spherical Fourier Neural Operators: Learning Stable Dynamics on the Sphere, + Boris Bonev, Thorsten Kurth, Christian Hundt, Jaideep Pathak, Maximilian Baust, Karthik Kashinath, Anima Anandkumar, + ICML 2023. + + Args: + in_channels (int): Number of input channels. + out_channels (int): Number of output channels. + n_modes (Tuple[int, ...]): Number of modes to use for contraction in Fourier domain during + training. + max_n_modes (int, optional): The maximum number of modes to use for contraction in Fourier domain during + training. Defaults to None. + bias (bool, optional): Whether to use bias in the layers. Defaults to True. + n_layers (int, optional): Number of Fourier Layers. Defaults to 1. + separable (bool, optional): Whether to use separable Fourier Conv. Defaults to False. + output_scaling_factor (Optional[Union[Number, List[Number]]], optional): Scaling factor for the + output. Defaults to None. + rank (float, optional): Rank of the tensor factorization of the Fourier weights. Defaults to 0.5. + factorization (str, optional): Tensor factorization of the parameters weight to use. Defaults to "dense". + implementation (str, optional): If factorization is not None, forward mode to use. Defaults to "reconstructed". + joint_factorization (bool, optional): Whether all the Fourier Layers should be parametrized by a + single tensor. Defaults to False. + init_std (str, optional): The std to use for the init. Defaults to "auto". + sht_norm (str, optional): The normalization mode of the SHT. Defaults to "ortho". + sht_grids (str, optional): The grid of the SHT. Defaults to "equiangular". + dtype (paddle.float32, optional): The data type. Defaults to paddle.float32. + """ + + def __init__( + self, + in_channels: int, + out_channels: int, + n_modes: Tuple[int, ...], + max_n_modes: int = None, + bias: bool = True, + n_layers: int = 1, + separable: bool = False, + output_scaling_factor: Optional[Union[Number, List[Number]]] = None, + rank: float = 0.5, + factorization: str = "dense", + implementation: str = "reconstructed", + joint_factorization: bool = False, + init_std: str = "auto", + sht_norm: str = "ortho", + sht_grids: str = "equiangular", + dtype: paddle.dtype = paddle.float32, + ): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + + self.dtype = dtype + + self.joint_factorization = joint_factorization + + if isinstance(n_modes, int): + n_modes = [n_modes] + self._n_modes = n_modes + self.order = len(n_modes) + + if max_n_modes is None: + max_n_modes = self.n_modes + elif isinstance(max_n_modes, int): + max_n_modes = [max_n_modes] + self.max_n_modes = max_n_modes + + self.rank = rank + self.factorization = factorization + self.n_layers = n_layers + self.implementation = implementation + + self.output_scaling_factor: Union[ + None, List[List[float]] + ] = fno_block.validate_scaling_factor( + output_scaling_factor, self.order, n_layers + ) + + if init_std == "auto": + init_std = (2 / (in_channels + out_channels)) ** 0.5 + else: + init_std = init_std + + if separable: + if in_channels != out_channels: + raise ValueError( + f"To use separable Fourier Conv, in_channels must be equal to out_channels, but got in_channels={in_channels} and out_channels={out_channels}" + ) + weight_shape = (in_channels, *self.n_modes[:-1]) + else: + weight_shape = (in_channels, out_channels, *self.n_modes[:-1]) + self.separable = separable + + if joint_factorization: + self.weight = paddle.create_parameter( + shape=(n_layers, *weight_shape), + dtype="float32", + ) + self.weight = initializer.normal_(self.weight, 0, init_std) + else: + self.weight = nn.LayerList( + [ + fno_block.FactorizedTensor(weight_shape, init_scale=init_std) + for _ in range(n_layers) + ] + ) + self._contract = get_contract_fun( + self.weight[0].data, implementation=implementation, separable=separable + ) + if bias: + shape = (n_layers, self.out_channels) + (1,) * self.order + init_bias = init_std * paddle.randn(shape) + self.bias = paddle.create_parameter( + shape=shape, + dtype=(init_bias.dtype), + default_initializer=nn.initializer.Assign(init_bias), + ) + self.bias.stop_gradient = False + else: + self.bias = None + + self.sht_norm = sht_norm + if isinstance(sht_grids, str): + sht_grids = [sht_grids] * (self.n_layers + 1) + self.sht_grids = sht_grids + self.sht_handle = SHT(dtype=self.dtype) + + @property + def n_modes(self): + return self._n_modes + + @n_modes.setter + def n_modes(self, n_modes): + if isinstance(n_modes, int): # Should happen for 1D FNO only + n_modes = [n_modes] + else: + n_modes = list(n_modes) + self._n_modes = n_modes + + def forward(self, x, indices=0, output_shape=None): + batchsize, channels, height, width = x.shape + + if self.output_scaling_factor is not None and output_shape is None: + scaling_factors = self.output_scaling_factor[indices] + height = round(height * scaling_factors[0]) + width = round(width * scaling_factors[1]) + elif output_shape is not None: + height, width = output_shape[0], output_shape[1] + + out_fft = self.sht_handle.sht( + x, + s=(self.n_modes[0], self.n_modes[1] // 2), + norm=self.sht_norm, + grid=self.sht_grids[indices], + ) + + w_real = self.weight[indices].real[:, :, : self.n_modes[0]] + w_imag = self.weight[indices].imag[:, :, : self.n_modes[0]] + + out_fft = self._contract( + out_fft[:, :, : self.n_modes[0], : self.n_modes[1] // 2], + w_real, + w_imag, + separable=self.separable, + dhconv=True, + ) + + x = self.sht_handle.isht( + out_fft, + s=(height, width), + norm=self.sht_norm, + grid=self.sht_grids[indices + 1], + ) + + if self.bias is not None: + x = x + self.bias[indices, ...] + + return x + + def transform(self, x, layer_index=0, output_shape=None): + *_, in_height, in_width = x.shape + + if self.output_scaling_factor is not None and output_shape is None: + height = round(in_height * self.output_scaling_factor[layer_index][0]) + width = round(in_width * self.output_scaling_factor[layer_index][1]) + elif output_shape is not None: + height, width = output_shape[0], output_shape[1] + else: + height, width = in_height, in_width + + # Return the identity if the resolution and grid of the input and output are the same + if ((in_height, in_width) == (height, width)) and ( + self.sht_grids[layer_index] == self.sht_grids[layer_index + 1] + ): + return x + else: + coefs = self.sht_handle.sht( + x, s=self.n_modes, norm=self.sht_norm, grid=self.sht_grids[layer_index] + ) + return self.sht_handle.isht( + coefs, + s=(height, width), + norm=self.sht_norm, + grid=self.sht_grids[layer_index + 1], + ) + + +class SFNONet(base.Arch): + """N-Dimensional Tensorized Fourier Neural Operator. + + Args: + input_keys (Tuple[str, ...]): Name of input keys, such as ("input",). + output_keys (Tuple[str, ...]): Name of output keys, such as ("output",). + n_modes (Tuple[int, ...]): Number of modes to keep in Fourier Layer, along each dimension + The dimensionality of the SFNO is inferred from ``len(n_modes)` + hidden_channels (int): Width of the FNO (i.e. number of channels) + in_channels (int, optional): Number of input channels. Defaults to 3. + out_channels (int, optional): Number of output channels. Defaults to 1. + lifting_channels (int, optional): Number of hidden channels of the lifting block of the FNO. + Defaults to 256. + projection_channels (int, optional): Number of hidden channels of the projection block of the FNO. + Defaults to 256. + n_layers (int, optional): Number of Fourier Layers. Defaults to 4. + use_mlp (bool, optional): Whether to use an MLP layer after each FNO block. Defaults to False. + mlp (Dict[str, float], optional): Parameters of the MLP. {'expansion': float, 'dropout': float}. + Defaults to None. + non_linearity (nn.functional, optional): Non-Linearity module to use. Defaults to F.gelu. + norm (str, optional): Normalization layer to use. Defaults to None. + ada_in_features (int,optional): The input channels of the adaptive normalization.Defaults to None. + preactivation (bool, optional): Whether to use resnet-style preactivation. Defaults to False. + fno_skip (str, optional): Type of skip connection to use,{'linear', 'identity', 'soft-gating'}. + Defaults to "soft-gating". + separable (bool, optional): Whether to use a depthwise separable spectral convolution. + Defaults to False. + factorization (str, optional): Tensor factorization of the parameters weight to use. + * If None, a dense tensor parametrizes the Spectral convolutions. + * Otherwise, the specified tensor factorization is used. Defaults to "Tucker". + rank (float, optional): Rank of the tensor factorization of the Fourier weights. Defaults to 1.0. + joint_factorization (bool, optional): Whether all the Fourier Layers should be parametrized by a + single tensor (vs one per layer). Defaults to False. + implementation (str, optional): {'factorized', 'reconstructed'}, optional. Defaults to "factorized". + If factorization is not None, forward mode to use:: + * `reconstructed` : the full weight tensor is reconstructed from the factorization and used for the forward pass. + * `factorized` : the input is directly contracted with the factors of the decomposition. + domain_padding (Optional[list], optional): Whether to use percentage of padding. Defaults to None. + domain_padding_mode (str, optional): {'symmetric', 'one-sided'}, optional + How to perform domain padding, by default 'one-sided'. Defaults to "one-sided". + fft_norm (str, optional): The normalization mode for the FFT. Defaults to "forward". + patching_levels (int, optional): Number of patching levels to use. Defaults to 0. + """ + + def __init__( + self, + input_keys: Tuple[str, ...], + output_keys: Tuple[str, ...], + n_modes: Tuple[int, ...], + hidden_channels: int, + in_channels: int = 3, + out_channels: int = 1, + lifting_channels: int = 256, + projection_channels: int = 256, + n_layers: int = 4, + use_mlp: bool = False, + mlp: Optional[Dict[str, float]] = None, + max_n_modes: int = None, + non_linearity: nn.functional = F.gelu, + stabilizer: str = None, + norm: str = None, + ada_in_features: Optional[int] = None, + preactivation: bool = False, + fno_skip: str = "linear", + mlp_skip: str = "soft-gating", + separable: bool = False, + factorization: str = None, + rank: float = 1.0, + joint_factorization: bool = False, + implementation: str = "factorized", + domain_padding: Optional[list] = None, + domain_padding_mode: str = "one-sided", + fft_norm: str = "forward", + patching_levels: int = 0, + **kwargs, + ): + super().__init__() + self.input_keys = input_keys + self.output_keys = output_keys + + self.n_dim = len(n_modes) + self.n_modes = n_modes + self.hidden_channels = hidden_channels + self.lifting_channels = lifting_channels + self.projection_channels = projection_channels + self.in_channels = in_channels + if patching_levels: + self.in_channels = self.in_channels * patching_levels + 1 + self.out_channels = out_channels + self.n_layers = n_layers + self.joint_factorization = joint_factorization + self.non_linearity = non_linearity + self.rank = rank + self.factorization = factorization + self.fno_skip = (fno_skip,) + self.mlp_skip = (mlp_skip,) + self.fft_norm = fft_norm + self.implementation = implementation + self.separable = separable + self.preactivation = preactivation + self.stabilizer = stabilizer + if domain_padding is not None and ( + (isinstance(domain_padding, list) and sum(domain_padding) > 0) + or (isinstance(domain_padding, (float, int)) and domain_padding > 0) + ): + self.domain_padding = fno_block.DomainPadding( + domain_padding=domain_padding, padding_mode=domain_padding_mode + ) + else: + self.domain_padding = None + self.domain_padding_mode = domain_padding_mode + + self.fno_blocks = fno_block.FNOBlocks( + in_channels=hidden_channels, + out_channels=hidden_channels, + n_modes=self.n_modes, + n_layers=n_layers, + max_n_modes=max_n_modes, + use_mlp=use_mlp, + mlp=mlp, + non_linearity=non_linearity, + stabilizer=stabilizer, + norm=norm, + ada_in_features=ada_in_features, + preactivation=preactivation, + fno_skip=fno_skip, + mlp_skip=mlp_skip, + separable=separable, + factorization=factorization, + rank=rank, + SpectralConv=SphericalConv, + joint_factorization=joint_factorization, + implementation=implementation, + fft_norm=fft_norm, + ) + # if lifting_channels is passed, make lifting an MLP + # with a hidden layer of size lifting_channels + if self.lifting_channels: + self.lifting = fno_block.MLP( + in_channels=in_channels, + out_channels=self.hidden_channels, + hidden_channels=self.lifting_channels, + n_layers=2, + n_dim=self.n_dim, + ) + # otherwise, make it a linear layer + else: + self.lifting = fno_block.MLP( + in_channels=in_channels, + out_channels=self.hidden_channels, + hidden_channels=self.hidden_channels, + n_layers=1, + n_dim=self.n_dim, + ) + self.projection = fno_block.MLP( + in_channels=self.hidden_channels, + out_channels=out_channels, + hidden_channels=self.projection_channels, + n_layers=2, + n_dim=self.n_dim, + non_linearity=non_linearity, + ) + + def forward(self, x): + """SFNO's forward pass""" + x = self.concat_to_tensor(x, self.input_keys) + + x = self.lifting(x) + if self.domain_padding is not None: + x = self.domain_padding.pad(x) + # x is 0.4 * [1, 32, 16, 16], passed + for index in range(self.n_layers): + x = self.fno_blocks(x, index) + + if self.domain_padding is not None: + x = self.domain_padding.unpad(x) + out = self.projection(x) + + return {self.output_keys[0]: out} diff --git a/examples/smc_reac/ppsci/arch/smc_reac.py b/examples/smc_reac/ppsci/arch/smc_reac.py new file mode 100644 index 0000000000..9e4d2595db --- /dev/null +++ b/examples/smc_reac/ppsci/arch/smc_reac.py @@ -0,0 +1,107 @@ +import paddle +from paddle import nn + +from ppsci.arch import base + + +class SuzukiMiyauraModel(base.Arch): + def __init__( + self, input_dim, hidden_dim, hidden_dim2, hidden_dim3, hidden_dim4, output_dim + ): + super().__init__() + + self.r1_fc = nn.Sequential( + nn.Linear(input_dim, hidden_dim), + nn.ReLU(), + nn.Linear(hidden_dim, hidden_dim2), + nn.ReLU(), + nn.Linear(hidden_dim2, hidden_dim3), + ) + + self.r2_fc = nn.Sequential( + nn.Linear(input_dim, hidden_dim), + nn.ReLU(), + nn.Linear(hidden_dim, hidden_dim2), + nn.ReLU(), + nn.Linear(hidden_dim2, hidden_dim3), + ) + + self.ligand_fc = nn.Sequential( + nn.Linear(input_dim, hidden_dim), + nn.ReLU(), + nn.Linear(hidden_dim, hidden_dim2), + nn.ReLU(), + nn.Linear(hidden_dim2, hidden_dim3), + ) + + self.base_fc = nn.Sequential( + nn.Linear(input_dim, hidden_dim), + nn.ReLU(), + nn.Linear(hidden_dim, hidden_dim2), + nn.ReLU(), + nn.Linear(hidden_dim2, hidden_dim3), + ) + + self.solvent_fc = nn.Sequential( + nn.Linear(input_dim, hidden_dim), + nn.ReLU(), + nn.Linear(hidden_dim, hidden_dim2), + nn.ReLU(), + nn.Linear(hidden_dim2, hidden_dim3), + nn.ReLU(), + ) + + self.weights = paddle.create_parameter( + shape=[5], + dtype="float32", + default_initializer=paddle.nn.initializer.Assign( + paddle.to_tensor([0.2, 0.2, 0.2, 0.2, 0.2]) + ), + ) + + self.fc_combined = nn.Sequential( + nn.Linear(hidden_dim3, hidden_dim4), + nn.ReLU(), + nn.Linear(hidden_dim4, output_dim), + ) + + def weighted_average(self, features, weights): + + weights = weights.clone().detach() + + weighted_sum = sum(f * w for f, w in zip(features, weights)) + + total_weight = weights.sum() + + return weighted_sum / total_weight + + def forward(self, x): + x = self.concat_to_tensor(x, ("v"), axis=-1) + + input_splits = paddle.split(x, num_or_sections=5, axis=1) + + r1_input, r2_input, ligand_input, base_input, solvent_input = input_splits + + r1_features = self.r1_fc(r1_input) + + r2_features = self.r2_fc(r2_input) + + ligand_features = self.ligand_fc(ligand_input) + + base_features = self.base_fc(base_input) + + solvent_features = self.solvent_fc(solvent_input) + + features = [ + r1_features, + r2_features, + ligand_features, + base_features, + solvent_features, + ] + + combined_features = self.weighted_average(features, self.weights) + + output = self.fc_combined(combined_features) + output = self.split_to_dict(output, ("u"), axis=-1) + return output diff --git a/examples/smc_reac/ppsci/arch/spinn.py b/examples/smc_reac/ppsci/arch/spinn.py new file mode 100644 index 0000000000..014446941f --- /dev/null +++ b/examples/smc_reac/ppsci/arch/spinn.py @@ -0,0 +1,180 @@ +# Copyright (c) 2024 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. + +from __future__ import annotations + +from typing import Dict +from typing import List +from typing import Optional +from typing import Tuple +from typing import Union + +import paddle +import paddle.nn as nn + +from ppsci.arch import base +from ppsci.arch.mlp import ModifiedMLP +from ppsci.utils import initializer + + +class SPINN(base.Arch): + """Separable Physics-Informed Neural Networks. + + Args: + input_keys (Tuple[str, ...]): Keys of input variables. + output_keys (Tuple[str, ...]): Keys of output variables. + r (int): Number of features for each output dimension. + num_layers (int): Number of layers. + hidden_size (Union[int, Tuple[int, ...]]): Size of hidden layer. + activation (str, optional): Name of activation function. + skip_connection (bool, optional): Whether to use skip connection. + weight_norm (bool, optional): Whether to use weight normalization. + periods (Optional[Dict[int, Tuple[float, bool]]], optional): Periodicity of input variables. + fourier (Optional[Dict[str, Union[float, int]]], optional): Frequency of input variables. + random_weight (Optional[Dict[str, float]], optional): Random weight of linear layer. + + Examples: + >>> from ppsci.arch import SPINN + >>> model = SPINN( + ... input_keys=('x', 'y', 'z'), + ... output_keys=('u', 'v'), + ... r=32, + ... num_layers=4, + ... hidden_size=32, + ... ) + >>> input_dict = {"x": paddle.rand([3, 1]), + ... "y": paddle.rand([4, 1]), + ... "z": paddle.rand([5, 1])} + >>> output_dict = model(input_dict) + >>> print(output_dict["u"].shape) + [3, 4, 5, 1] + >>> print(output_dict["v"].shape) + [3, 4, 5, 1] + """ + + def __init__( + self, + input_keys: Tuple[str, ...], + output_keys: Tuple[str, ...], + r: int, + num_layers: int, + hidden_size: Union[int, Tuple[int, ...]], + activation: str = "tanh", + skip_connection: bool = False, + weight_norm: bool = False, + periods: Optional[Dict[int, Tuple[float, bool]]] = None, + fourier: Optional[Dict[str, Union[float, int]]] = None, + random_weight: Optional[Dict[str, float]] = None, + ): + + super().__init__() + self.input_keys = input_keys + self.output_keys = output_keys + self.r = r + input_dim = len(self.input_keys) + + self.branch_nets = nn.LayerList() + for i in range(input_dim): + self.branch_nets.append( + ModifiedMLP( + input_keys=(input_keys[i],), + output_keys=("f",), + num_layers=num_layers, + hidden_size=hidden_size, + activation=activation, + skip_connection=skip_connection, + weight_norm=weight_norm, + output_dim=r * len(output_keys), + periods=periods, + fourier=fourier, + random_weight=random_weight, + ) + ) + + self._init_weights() + + def _init_weights(self): + for m in self.sublayers(True): + if isinstance(m, nn.Linear): + initializer.glorot_normal_(m.weight) + initializer.zeros_(m.bias) + + def _tensor_contraction(self, x: paddle.Tensor, y: paddle.Tensor) -> paddle.Tensor: + """Tensor contraction between two tensors along the last channel. + + Args: + x (Tensor): Input tensor with shape [*N, C]. + y (Tensor): Input tensor with shape [*M, C] + + Returns: + Tensor: Output tensor with shape [*N, *M, C]. + """ + x_ndim = x.ndim + y_ndim = y.ndim + out_dim = x_ndim + y_ndim - 1 + + # Align the dimensions of x and y to out_dim + if x_ndim < out_dim: + # Add singleton dimensions to x at the end of dimensions + x = x.unsqueeze([-2] * (out_dim - x_ndim)) + if y_ndim < out_dim: + # Add singleton dimensions to y at the begin of dimensions + y = y.unsqueeze([0] * (out_dim - y_ndim)) + + # Multiply x and y with implicit broadcasting + out = x * y + + return out + + def forward_tensor(self, *xs) -> List[paddle.Tensor]: + # forward each dim branch + feature_f = [] + for i, input_var in enumerate(xs): + input_i = {self.input_keys[i]: input_var} + output_f_i = self.branch_nets[i](input_i) + feature_f.append(output_f_i["f"]) # [B, r*output_dim] + + output = [] + for i, key in enumerate(self.output_keys): + st, ed = i * self.r, (i + 1) * self.r + # do tensor contraction and sum over all branch outputs + if ed - st == self.r: + output_i = feature_f[0] + else: + output_i = feature_f[0][:, st:ed] + + for j in range(1, len(self.input_keys)): + if ed - st == self.r: + output_ii = feature_f[j] + else: + output_ii = feature_f[j][:, st:ed] + output_i = self._tensor_contraction(output_i, output_ii) + + output_i = output_i.sum(-1, keepdim=True) + output.append(output_i) + + return output + + def forward(self, x): + if self._input_transform is not None: + x = self._input_transform(x) + + output = self.forward_tensor(*[x[key] for key in self.input_keys]) + + output = {key: output[i] for i, key in enumerate(self.output_keys)} + + if self._output_transform is not None: + output = self._output_transform(x, output) + + return output diff --git a/examples/smc_reac/ppsci/arch/tfnonet.py b/examples/smc_reac/ppsci/arch/tfnonet.py new file mode 100644 index 0000000000..91bcfd6f5c --- /dev/null +++ b/examples/smc_reac/ppsci/arch/tfnonet.py @@ -0,0 +1,514 @@ +from typing import Dict +from typing import Optional +from typing import Tuple +from typing import Union + +import paddle.nn.functional as F +from paddle import nn + +from ppsci.arch import base +from ppsci.arch import fno_block + + +class FNONet(base.Arch): + """N-Dimensional Tensorized Fourier Neural Operator. + + Args: + input_keys (Tuple[str, ...]): Name of input keys, such as ("input",). + output_keys (Tuple[str, ...]): Name of output keys, such as ("output",). + n_modes (Tuple[int, ...]): Number of modes to keep in Fourier Layer, along each dimension + The dimensionality of the TFNO is inferred from ``len(n_modes)` + hidden_channels (int): Width of the FNO (i.e. number of channels) + in_channels (int, optional): Number of input channels. Defaults to 3. + out_channels (int, optional): Number of output channels. Defaults to 1. + lifting_channels (int, optional): Number of hidden channels of the lifting block of the FNO. + Defaults to 256. + projection_channels (int, optional): Number of hidden channels of the projection block of the FNO. + Defaults to 256. + n_layers (int, optional): Number of Fourier Layers. Defaults to 4. + use_mlp (bool, optional): Whether to use an MLP layer after each FNO block. Defaults to False. + mlp (Dict[str, float], optional): Parameters of the MLP. {'expansion': float, 'dropout': float}. + Defaults to None. + non_linearity (nn.functional, optional): Non-Linearity module to use. Defaults to F.gelu. + norm (str, optional): Normalization layer to use. Defaults to None. + ada_in_features (int,optional): The input channels of the adaptive normalization.Defaults to None.s + preactivation (bool, optional): Whether to use resnet-style preactivation. Defaults to False. + skip (str, optional): Type of skip connection to use,{'linear', 'identity', 'soft-gating'}. + Defaults to "soft-gating". + separable (bool, optional): Whether to use a depthwise separable spectral convolution. + Defaults to False. + factorization (str, optional): Tensor factorization of the parameters weight to use. + * If None, a dense tensor parametrizes the Spectral convolutions. + * Otherwise, the specified tensor factorization is used. Defaults to "Tucker". + rank (float, optional): Rank of the tensor factorization of the Fourier weights. Defaults to 1.0. + joint_factorization (bool, optional): Whether all the Fourier Layers should be parametrized by a + single tensor (vs one per layer). Defaults to False. + implementation (str, optional): {'factorized', 'reconstructed'}, optional. Defaults to "factorized". + If factorization is not None, forward mode to use:: + * `reconstructed` : the full weight tensor is reconstructed from the factorization and used for the forward pass. + * `factorized` : the input is directly contracted with the factors of the decomposition. + domain_padding (Optional[Union[list,float,int]]): Whether to use percentage of padding. Defaults to + None. + domain_padding_mode (str, optional): {'symmetric', 'one-sided'}, optional + How to perform domain padding, by default 'one-sided'. Defaults to "one-sided". + fft_norm (str, optional): The normalization mode for the FFT. Defaults to "forward". + patching_levels (int, optional): Number of patching levels to use. Defaults to 0. + SpectralConv (nn.layer, optional): Spectral convolution layer to use. + Defaults to fno_block.FactorizedSpectralConv. + """ + + def __init__( + self, + input_keys: Tuple[str, ...], + output_keys: Tuple[str, ...], + n_modes: Tuple[int, ...], + hidden_channels: int, + in_channels: int = 3, + out_channels: int = 1, + lifting_channels: int = 256, + projection_channels: int = 256, + n_layers: int = 4, + use_mlp: bool = False, + mlp: Optional[Dict[str, float]] = None, + max_n_modes: int = None, + non_linearity: nn.functional = F.gelu, + stabilizer: str = None, + norm: str = None, + ada_in_features: Optional[int] = None, + preactivation: bool = False, + fno_skip: str = "linear", + mlp_skip: str = "soft-gating", + separable: bool = False, + factorization: str = None, + rank: float = 1.0, + joint_factorization: bool = False, + implementation: str = "factorized", + domain_padding: Optional[Union[list, float, int]] = None, + domain_padding_mode: str = "one-sided", + fft_norm: str = "forward", + patching_levels: int = 0, + SpectralConv: nn.Layer = fno_block.FactorizedSpectralConv, + **kwargs, + ): + super().__init__() + self.input_keys = input_keys + self.output_keys = output_keys + + self.n_dim = len(n_modes) + self.n_modes = n_modes + self.hidden_channels = hidden_channels + self.lifting_channels = lifting_channels + self.projection_channels = projection_channels + self.in_channels = in_channels + if patching_levels: + self.in_channels = self.in_channels * patching_levels + 1 + self.out_channels = out_channels + self.n_layers = n_layers + self.joint_factorization = joint_factorization + self.non_linearity = non_linearity + self.rank = rank + self.factorization = factorization + self.fno_skip = (fno_skip,) + self.mlp_skip = (mlp_skip,) + self.fft_norm = fft_norm + self.implementation = implementation + self.separable = separable + self.preactivation = preactivation + self.stabilizer = stabilizer + if domain_padding is not None and ( + (isinstance(domain_padding, list) and sum(domain_padding) > 0) + or (isinstance(domain_padding, (float, int)) and domain_padding > 0) + ): + self.domain_padding = fno_block.DomainPadding( + domain_padding=domain_padding, padding_mode=domain_padding_mode + ) + else: + self.domain_padding = None + self.domain_padding_mode = domain_padding_mode + self.fno_blocks = fno_block.FNOBlocks( + in_channels=hidden_channels, + out_channels=hidden_channels, + n_modes=self.n_modes, + n_layers=n_layers, + max_n_modes=max_n_modes, + use_mlp=use_mlp, + mlp=mlp, + non_linearity=non_linearity, + stabilizer=stabilizer, + norm=norm, + ada_in_features=ada_in_features, + preactivation=preactivation, + fno_skip=fno_skip, + mlp_skip=mlp_skip, + separable=separable, + factorization=factorization, + rank=rank, + SpectralConv=SpectralConv, + joint_factorization=joint_factorization, + implementation=implementation, + fft_norm=fft_norm, + ) + # if lifting_channels is passed, make lifting an MLP + # with a hidden layer of size lifting_channels + if self.lifting_channels: + self.lifting = fno_block.MLP( + in_channels=in_channels, + out_channels=self.hidden_channels, + hidden_channels=self.lifting_channels, + n_layers=2, + n_dim=self.n_dim, + ) + # otherwise, make it a linear layer + else: + self.lifting = fno_block.MLP( + in_channels=in_channels, + out_channels=self.hidden_channels, + hidden_channels=self.hidden_channels, + n_layers=1, + n_dim=self.n_dim, + ) + self.projection = fno_block.MLP( + in_channels=self.hidden_channels, + out_channels=out_channels, + hidden_channels=self.projection_channels, + n_layers=2, + n_dim=self.n_dim, + non_linearity=non_linearity, + ) + + def forward(self, x): + """TFNO's forward pass""" + x = self.concat_to_tensor(x, self.input_keys) + + x = self.lifting(x) + if self.domain_padding is not None: + x = self.domain_padding.pad(x) + # x is 0.4 * [1, 32, 16, 16], passed + for index in range(self.n_layers): + x = self.fno_blocks(x, index) + + if self.domain_padding is not None: + x = self.domain_padding.unpad(x) + out = self.projection(x) + return {self.output_keys[0]: out} + + +class TFNO1dNet(FNONet): + """1D Fourier Neural Operator. + + Args: + input_keys (Tuple[str, ...]): Name of input keys, such as ("input",). + output_keys (Tuple[str, ...]): Name of output keys, such as ("output",). + n_modes_height (Tuple[int, ...]): Number of Fourier modes to keep along the height, along each + dimension. + hidden_channels (int): Width of the FNO (i.e. number of channels). + in_channels (int, optional): Number of input channels. Defaults to 3. + out_channels (int, optional): Number of output channels. Defaults to 1. + lifting_channels (int, optional): Number of hidden channels of the lifting block of the FNO. + Defaults to 256. + projection_channels (int, optional): Number of hidden channels of the projection block of the FNO. + Defaults to 256. + n_layers (int, optional): Number of Fourier Layers. Defaults to 4. + use_mlp (bool, optional): Whether to use an MLP layer after each FNO block. Defaults to False. + mlp (dict[str, float], optional): Parameters of the MLP. {'expansion': float, 'dropout': float}. + Defaults to None. + non_linearity (nn.functional, optional): Non-Linearity module to use. Defaults to F.gelu. + norm (F.module, optional): Normalization layer to use. Defaults to None. + preactivation (bool, optional): Whether to use resnet-style preactivation. Defaults to False. + skip (str, optional): Type of skip connection to use,{'linear', 'identity', 'soft-gating'}. + Defaults to "soft-gating". + separable (bool, optional): Whether to use a depthwise separable spectral convolution. + Defaults to False. + factorization (str, optional): Tensor factorization of the parameters weight to use. + * If None, a dense tensor parametrizes the Spectral convolutions. + * Otherwise, the specified tensor factorization is used. Defaults to "Tucker". + rank (float, optional): Rank of the tensor factorization of the Fourier weights. Defaults to 1.0. + joint_factorization (bool, optional): Whether all the Fourier Layers should be parametrized by a + single tensor (vs one per layer). Defaults to False. + implementation (str, optional): {'factorized', 'reconstructed'}, optional. Defaults to "factorized". + If factorization is not None, forward mode to use:: + * `reconstructed` : the full weight tensor is reconstructed from the factorization and used for the forward pass. + * `factorized` : the input is directly contracted with the factors of the decomposition. + domain_padding (Optional[Union[list, float, int]], optional): Whether to use percentage of padding. + Defaults to None. + domain_padding_mode (str, optional): {'symmetric', 'one-sided'}, optional + How to perform domain padding, by default 'one-sided'. Defaults to "one-sided". + fft_norm (str, optional): The normalization mode for the FFT. Defaults to "forward". + patching_levels (int, optional): Number of patching levels to use. Defaults to 0. + SpectralConv (nn.layer, optional): Spectral convolution layer to use. + Defaults to fno_block.FactorizedSpectralConv. + """ + + def __init__( + self, + input_keys: Tuple[str, ...], + output_keys: Tuple[str, ...], + n_modes_height: Tuple[int, ...], + hidden_channels: int, + in_channels: int = 3, + out_channels: int = 1, + lifting_channels: int = 256, + projection_channels: int = 256, + n_layers: int = 4, + non_linearity: nn.functional = F.gelu, + use_mlp: bool = False, + mlp: Optional[Dict[str, float]] = None, + norm: str = None, + skip: str = "soft-gating", + separable: bool = False, + preactivation: bool = False, + factorization: str = "Tucker", + rank: float = 1.0, + joint_factorization: bool = False, + implementation: str = "factorized", + domain_padding: Optional[Union[list, float, int]] = None, + domain_padding_mode: str = "one-sided", + fft_norm: str = "forward", + patching_levels: int = 0, + SpectralConv: nn.Layer = fno_block.FactorizedSpectralConv, + **kwargs, + ): + super().__init__( + input_keys=input_keys, + output_keys=output_keys, + n_modes=(n_modes_height,), + hidden_channels=hidden_channels, + in_channels=in_channels, + out_channels=out_channels, + lifting_channels=lifting_channels, + projection_channels=projection_channels, + n_layers=n_layers, + non_linearity=non_linearity, + use_mlp=use_mlp, + mlp=mlp, + norm=norm, + skip=skip, + separable=separable, + preactivation=preactivation, + factorization=factorization, + rank=rank, + joint_factorization=joint_factorization, + implementation=implementation, + domain_padding=domain_padding, + domain_padding_mode=domain_padding_mode, + fft_norm=fft_norm, + patching_levels=patching_levels, + SpectralConv=SpectralConv, + ) + self.n_modes_height = n_modes_height + + +class TFNO2dNet(FNONet): + """2D Fourier Neural Operator. + + Args: + input_keys (Tuple[str, ...]): Name of input keys, such as ("input",). + output_keys (Tuple[str, ...]): Name of output keys, such as ("output",). + n_modes_height (int): Number of Fourier modes to keep along the height. + n_modes_width (int): Number of modes to keep in Fourier Layer, along the width. + hidden_channels (int): Width of the FNO (i.e. number of channels). + in_channels (int, optional): Number of input channels. Defaults to 3. + out_channels (int, optional): Number of output channels. Defaults to 1. + lifting_channels (int, optional): Number of hidden channels of the lifting block of the FNO. + Defaults to 256. + projection_channels (int, optional): Number of hidden channels of the projection block of the FNO. + Defaults to 256. + n_layers (int, optional): Number of Fourier Layers. Defaults to 4. + use_mlp (bool, optional): Whether to use an MLP layer after each FNO block. Defaults to False. + mlp (Dict[str, float], optional): Parameters of the MLP. {'expansion': float, 'dropout': float}. + Defaults to None. + non_linearity (nn.Layer, optional): Non-Linearity module to use. Defaults to F.gelu. + norm (F.module, optional): Normalization layer to use. Defaults to None. + preactivation (bool, optional): Whether to use resnet-style preactivation. Defaults to False. + skip (str, optional): Type of skip connection to use,{'linear', 'identity', 'soft-gating'}. + Defaults to "soft-gating". + separable (bool, optional): Whether to use a depthwise separable spectral convolution. + Defaults to False. + factorization (str, optional): Tensor factorization of the parameters weight to use. + * If None, a dense tensor parametrizes the Spectral convolutions. + * Otherwise, the specified tensor factorization is used. Defaults to "Tucker". + rank (float, optional): Rank of the tensor factorization of the Fourier weights. Defaults to 1.0. + joint_factorization (bool, optional): Whether all the Fourier Layers should be parametrized by a + single tensor (vs one per layer). Defaults to False. + implementation (str, optional): {'factorized', 'reconstructed'}, optional. Defaults to "factorized". + If factorization is not None, forward mode to use:: + * `reconstructed` : the full weight tensor is reconstructed from the factorization and used for the forward pass. + * `factorized` : the input is directly contracted with the factors of the decomposition. + domain_padding (Union[list,float,int], optional): Whether to use percentage of padding. Defaults to + None. + domain_padding_mode (str, optional): {'symmetric', 'one-sided'}, optional + How to perform domain padding, by default 'one-sided'. Defaults to "one-sided". + fft_norm (str, optional): The normalization mode for the FFT. Defaults to "forward". + patching_levels (int, optional): Number of patching levels to use. Defaults to 0. + SpectralConv (nn.layer, optional): Spectral convolution layer to use. + Defaults to fno_block.FactorizedSpectralConv. + """ + + def __init__( + self, + input_keys: Tuple[str, ...], + output_keys: Tuple[str, ...], + n_modes_height: int, + n_modes_width: int, + hidden_channels: int, + in_channels: int = 3, + out_channels: int = 1, + lifting_channels: int = 256, + projection_channels: int = 256, + n_layers: int = 4, + non_linearity: nn.functional = F.gelu, + use_mlp: bool = False, + mlp: Optional[Dict[str, float]] = None, + norm: str = None, + skip: str = "soft-gating", + separable: bool = False, + preactivation: bool = False, + factorization: str = "Tucker", + rank: float = 1.0, + joint_factorization: bool = False, + implementation: str = "factorized", + domain_padding: Optional[Union[list, float, int]] = None, + domain_padding_mode: str = "one-sided", + fft_norm: str = "forward", + patching_levels: int = 0, + SpectralConv: nn.layer = fno_block.FactorizedSpectralConv, + **kwargs, + ): + super().__init__( + input_keys=input_keys, + output_keys=output_keys, + n_modes=(n_modes_height, n_modes_width), + hidden_channels=hidden_channels, + in_channels=in_channels, + out_channels=out_channels, + lifting_channels=lifting_channels, + projection_channels=projection_channels, + n_layers=n_layers, + non_linearity=non_linearity, + use_mlp=use_mlp, + mlp=mlp, + norm=norm, + skip=skip, + separable=separable, + preactivation=preactivation, + factorization=factorization, + rank=rank, + joint_factorization=joint_factorization, + implementation=implementation, + domain_padding=domain_padding, + domain_padding_mode=domain_padding_mode, + fft_norm=fft_norm, + patching_levels=patching_levels, + SpectralConv=SpectralConv, + ) + self.n_modes_height = n_modes_height + self.n_modes_width = n_modes_width + + +class TFNO3dNet(FNONet): + """3D Fourier Neural Operator. + + Args: + input_keys (Tuple[str, ...]): Name of input keys, such as ("input",). + output_keys (Tuple[str, ...]): Name of output keys, such as ("output",). + n_modes_height (int): Number of Fourier modes to keep along the height. + n_modes_width (int): Number of modes to keep in Fourier Layer, along the width. + n_modes_depth (int): Number of Fourier modes to keep along the depth. + hidden_channels (int): Width of the FNO (i.e. number of channels). + in_channels (int, optional): Number of input channels. Defaults to 3. + out_channels (int, optional): Number of output channels. Defaults to 1. + lifting_channels (int, optional): Number of hidden channels of the lifting block of the FNO. + Defaults to 256. + projection_channels (int, optional): Number of hidden channels of the projection block of the FNO. + Defaults to 256. + n_layers (int, optional): Number of Fourier Layers. Defaults to 4. + use_mlp (bool, optional): Whether to use an MLP layer after each FNO block. Defaults to False. + mlp (Dict[str, float], optional): Parameters of the MLP. {'expansion': float, 'dropout': float}. + Defaults to None. + non_linearity (nn.Layer, optional): Non-Linearity module to use. Defaults to F.gelu. + norm (F.module, optional): Normalization layer to use. Defaults to None. + preactivation (bool, optional): Whether to use resnet-style preactivation. Defaults to False. + skip (str, optional): Type of skip connection to use,{'linear', 'identity', 'soft-gating'}. + Defaults to "soft-gating". + separable (bool, optional): Whether to use a depthwise separable spectral convolution. + Defaults to False. + factorization (str, optional): Tensor factorization of the parameters weight to use. + * If None, a dense tensor parametrizes the Spectral convolutions. + * Otherwise, the specified tensor factorization is used. Defaults to "Tucker". + rank (float, optional): Rank of the tensor factorization of the Fourier weights. Defaults to 1.0. + joint_factorization (bool, optional): Whether all the Fourier Layers should be parametrized by a + single tensor (vs one per layer). Defaults to False. + implementation (str, optional): {'factorized', 'reconstructed'}, optional. Defaults to "factorized". + If factorization is not None, forward mode to use:: + * `reconstructed` : the full weight tensor is reconstructed from the factorization and used for the forward pass. + * `factorized` : the input is directly contracted with the factors of the decomposition. + domain_padding (str, optional): Whether to use percentage of padding. Defaults to None. + domain_padding_mode (str, optional): {'symmetric', 'one-sided'}, optional + How to perform domain padding, by default 'one-sided'. Defaults to "one-sided". + fft_norm (str, optional): The normalization mode for the FFT. Defaults to "forward". + patching_levels (int, optional): Number of patching levels to use. Defaults to 0. + SpectralConv (nn.layer, optional): Spectral convolution layer to use. Defaults to fno_block. + FactorizedSpectralConv. + """ + + def __init__( + self, + input_keys: Tuple[str, ...], + output_keys: Tuple[str, ...], + n_modes_height: int, + n_modes_width: int, + n_modes_depth: int, + hidden_channels: int, + in_channels: int = 3, + out_channels: int = 1, + lifting_channels: int = 256, + projection_channels: int = 256, + n_layers: int = 4, + non_linearity: nn.functional = F.gelu, + use_mlp: bool = False, + mlp: Optional[Dict[str, float]] = None, + norm: str = None, + skip: str = "soft-gating", + separable: bool = False, + preactivation: bool = False, + factorization: str = "Tucker", + rank: float = 1.0, + joint_factorization: bool = False, + implementation: str = "factorized", + domain_padding: Optional[Union[list, float, int]] = None, + domain_padding_mode: str = "one-sided", + fft_norm: str = "forward", + patching_levels: int = 0, + SpectralConv: nn.layer = fno_block.FactorizedSpectralConv, + **kwargs, + ): + super().__init__( + input_keys=input_keys, + output_keys=output_keys, + n_modes=(n_modes_height, n_modes_width, n_modes_depth), + hidden_channels=hidden_channels, + in_channels=in_channels, + out_channels=out_channels, + lifting_channels=lifting_channels, + projection_channels=projection_channels, + n_layers=n_layers, + non_linearity=non_linearity, + use_mlp=use_mlp, + mlp=mlp, + norm=norm, + skip=skip, + separable=separable, + preactivation=preactivation, + factorization=factorization, + rank=rank, + joint_factorization=joint_factorization, + implementation=implementation, + domain_padding=domain_padding, + domain_padding_mode=domain_padding_mode, + fft_norm=fft_norm, + patching_levels=patching_levels, + SpectralConv=SpectralConv, + ) + self.n_modes_height = n_modes_height + self.n_modes_width = n_modes_width + self.n_modes_depth = n_modes_depth diff --git a/examples/smc_reac/ppsci/arch/tgcn.py b/examples/smc_reac/ppsci/arch/tgcn.py new file mode 100644 index 0000000000..5cf1ebc3eb --- /dev/null +++ b/examples/smc_reac/ppsci/arch/tgcn.py @@ -0,0 +1,200 @@ +from typing import Tuple + +import paddle as pp +import paddle.nn.functional as F +from numpy import ndarray +from paddle import nn +from paddle.nn.initializer import KaimingNormal + +from ppsci.arch.base import Arch + + +class graph_conv(nn.Layer): + def __init__(self, in_dim, out_dim, dropout, num_layer=2): + super(graph_conv, self).__init__() + self.mlp = nn.Conv2D( + (num_layer + 1) * in_dim, + out_dim, + kernel_size=(1, 1), + weight_attr=KaimingNormal(), + ) + self.dropout = dropout + self.num_layer = num_layer + + def forward(self, x, adj): + # B C N T + out = [x] + for _ in range(self.num_layer): + new_x = pp.matmul(adj, x) + out.append(new_x) + x = new_x + + h = pp.concat(out, axis=1) + h = self.mlp(h) + h = F.dropout(h, self.dropout, training=self.training) + return h + + +class tempol_conv(nn.Layer): + def __init__(self, in_dim, out_dim, hidden, num_layer=3, k_s=3, alpha=0.1): + super(tempol_conv, self).__init__() + self.leakyrelu = nn.LeakyReLU(alpha) + self.tc_convs = nn.LayerList() + self.num_layer = num_layer + for i in range(num_layer): + in_channels = in_dim if i == 0 else hidden + self.tc_convs.append( + nn.Conv2D( + in_channels=in_channels, + out_channels=hidden, + kernel_size=(1, k_s), + padding=(0, i + 1), + dilation=i + 1, + weight_attr=KaimingNormal(), + ) + ) + + self.mlp = nn.Conv2D( + in_channels=in_dim + hidden * num_layer, + out_channels=out_dim, + kernel_size=(1, 1), + weight_attr=KaimingNormal(), + ) + + def forward(self, x): + # B C N T + x_cat = [x] + for i in range(self.num_layer): + x = self.leakyrelu(self.tc_convs[i](x)) + x_cat.append(x) + tc_out = self.mlp(pp.concat(x_cat, axis=1)) + return tc_out + + +class TGCN(Arch): + """ + TGCN is a class that represents an Temporal Graph Convolutional Network model. + + Args: + input_keys (Tuple[str, ...]): A tuple of input keys. + output_keys (Tuple[str, ...]): A tuple of output keys. + adj (ndarray): The adjacency matrix of the graph. + in_dim (int): The dimension of the input data. + emb_dim (int): The dimension of the embedded space. + hidden (int): The dimension of the latent space. + gc_layer (int): The number of the graph convolutional layer. + tc_layer (int): The number of the temporal convolutional layer. + k_s (int): The kernel size of the temporal convolutional layer. + dropout (float): The dropout rate. + alpha (float): The negative slope of LeakyReLU. + input_len (int): The input timesteps. + label_len (int): The output timesteps. + + Examples: + >>> import paddle + >>> import ppsci + >>> model = ppsci.arch.TGCN( + ... input_keys=("input",), + ... output_keys=("label",), + ... adj=numpy.ones((307, 307), dtype=numpy.float32), + ... in_dim=1, + ... emb_dim=32 + ... hidden=64, + ... gc_layer=2, + ... tc_layer=2 + ... k_s=3, + ... dropout=0.25, + ... alpha=0.1, + ... input_len=12, + ... label_len=12, + ... ) + >>> input_dict = {"input": paddle.rand([64, 12, 307, 1]),} + >>> label_dict = model(input_dict) + >>> print(label_dict["label"].shape) + [64, 12, 307, 1] + """ + + def __init__( + self, + input_keys: Tuple[str, ...], + output_keys: Tuple[str, ...], + adj: ndarray, + in_dim: int, + emb_dim: int, + hidden: int, + gc_layer: int, + tc_layer: int, + k_s: int, + dropout: float, + alpha: float, + input_len: int, + label_len: int, + ): + super(TGCN, self).__init__() + + self.input_keys = input_keys + self.output_keys = output_keys + + self.register_buffer("adj", pp.to_tensor(data=adj)) + + self.emb_conv = nn.Conv2D( + in_channels=in_dim, + out_channels=emb_dim, + kernel_size=(1, 1), + weight_attr=KaimingNormal(), + ) + + self.tc1_conv = tempol_conv( + emb_dim, hidden, hidden, num_layer=tc_layer, k_s=k_s, alpha=alpha + ) + self.sc1_conv = graph_conv(hidden, hidden, dropout, num_layer=gc_layer) + self.bn1 = nn.BatchNorm2D(hidden) + + self.tc2_conv = tempol_conv( + hidden, hidden, hidden, num_layer=tc_layer, k_s=k_s, alpha=alpha + ) + self.sc2_conv = graph_conv(hidden, hidden, dropout, num_layer=gc_layer) + self.bn2 = nn.BatchNorm2D(hidden) + + self.end_conv_1 = nn.Conv2D( + in_channels=emb_dim + hidden + hidden, + out_channels=2 * hidden, + kernel_size=(1, 1), + weight_attr=KaimingNormal(), + ) + self.end_conv_2 = nn.Conv2D( + in_channels=2 * hidden, + out_channels=label_len, + kernel_size=(1, input_len), + weight_attr=KaimingNormal(), + ) + + def forward(self, raw): + # emb block + x = raw[self.input_keys[0]] + x = x.transpose(perm=[0, 3, 2, 1]) # B in_dim N T + emb_x = self.emb_conv(x) # B emd_dim N T + + # TC1 + tc1_out = self.tc1_conv(emb_x) # B hidden N T + + # SC1 + sc1_out = self.sc1_conv(tc1_out, self.adj) # B hidden N T + sc1_out = sc1_out + tc1_out + sc1_out = self.bn1(sc1_out) + + # TC2 + tc2_out = self.tc2_conv(sc1_out) # B hidden N T + + # SC2 + sc2_out = self.sc2_conv(tc2_out, self.adj) # B hidden N T + sc2_out = sc2_out + tc2_out + sc2_out = self.bn2(sc2_out) + + # readout block + x_out = F.relu(pp.concat((emb_x, sc1_out, sc2_out), axis=1)) + x_out = F.relu(self.end_conv_1(x_out)) + # transform + x_out = self.end_conv_2(x_out) # B T N 1 + + return {self.output_keys[0]: x_out} diff --git a/examples/smc_reac/ppsci/arch/transformer.py b/examples/smc_reac/ppsci/arch/transformer.py new file mode 100644 index 0000000000..a80eb46dcf --- /dev/null +++ b/examples/smc_reac/ppsci/arch/transformer.py @@ -0,0 +1,417 @@ +# Copyright (c) 2024 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. + +""" +Reference: https://github.com/omron-sinicx/transformer4sr +""" + +from __future__ import annotations + +import math +from typing import Callable +from typing import Tuple + +import paddle +import paddle.nn as nn + +from ppsci.arch import activation as act_mod +from ppsci.arch import base + + +def transpose_aux_func(dims, dim0, dim1): + perm = list(range(dims)) + perm[dim0], perm[dim1] = perm[dim1], perm[dim0] + return perm + + +class MultiHeadAttention(nn.Layer): + def __init__(self, heads, d_model): + super().__init__() + self.heads = heads + self.d_model = d_model + assert d_model % heads == 0 + self.d_k = d_model // heads + self.W_Q = nn.Linear(in_features=d_model, out_features=d_model) + self.W_K = nn.Linear(in_features=d_model, out_features=d_model) + self.W_V = nn.Linear(in_features=d_model, out_features=d_model) + self.W_O = nn.Linear(in_features=d_model, out_features=d_model) + + def scaled_dot_product_attention(self, Q, K, V, mask=None): + scores = paddle.matmul( + x=Q, y=K.transpose(perm=transpose_aux_func(K.ndim, -1, -2)) + ) / math.sqrt(self.d_k) + if mask is not None: + scores = paddle.where( + condition=mask, + x=paddle.to_tensor(data=[-1e9], dtype="float32"), + y=scores, + ) + weights = nn.functional.softmax(x=scores, axis=-1) + return paddle.matmul(x=weights, y=V) + + def forward(self, Q, K, V, mask=None): + Q_temp = paddle.reshape( + x=self.W_Q(Q), + shape=[i for i in tuple(Q.shape)[:-1]] + [self.heads] + [self.d_k], + ).transpose( + perm=transpose_aux_func( + paddle.reshape( + x=self.W_Q(Q), + shape=[i for i in tuple(Q.shape)[:-1]] + [self.heads] + [self.d_k], + ).ndim, + 1, + 2, + ) + ) + K_temp = paddle.reshape( + x=self.W_K(K), + shape=[i for i in tuple(K.shape)[:-1]] + [self.heads] + [self.d_k], + ).transpose( + perm=transpose_aux_func( + paddle.reshape( + x=self.W_K(K), + shape=[i for i in tuple(K.shape)[:-1]] + [self.heads] + [self.d_k], + ).ndim, + 1, + 2, + ) + ) + V_temp = paddle.reshape( + x=self.W_V(V), + shape=[i for i in tuple(V.shape)[:-1]] + [self.heads] + [self.d_k], + ).transpose( + perm=transpose_aux_func( + paddle.reshape( + x=self.W_V(V), + shape=[i for i in tuple(V.shape)[:-1]] + [self.heads] + [self.d_k], + ).ndim, + 1, + 2, + ) + ) + sdpa = self.scaled_dot_product_attention( + Q_temp, K_temp, V_temp, mask + ).transpose( + perm=transpose_aux_func( + self.scaled_dot_product_attention(Q_temp, K_temp, V_temp, mask).ndim, + 1, + 2, + ) + ) + sdpa = paddle.reshape( + x=sdpa, shape=[i for i in tuple(sdpa.shape)[:-2]] + [self.d_model] + ) + y_mha = self.W_O(sdpa) + return y_mha + + +class MLP(nn.Layer): + def __init__(self, list_dims, act="relu", dropout=0.0): + super().__init__() + self.layers = nn.LayerList() + for i in range(len(list_dims) - 1): + self.layers.append( + nn.Linear(in_features=list_dims[i], out_features=list_dims[i + 1]) + ) + self.layers.append(act_mod.get_activation(act) if act else None) + self.layers.append(nn.Dropout(p=dropout)) + + def forward(self, x): + y = x + for layer in self.layers: + y = layer(y) + return y + + +class EncoderLayerMix(nn.Layer): + def __init__(self, in_features, d_model, heads, act="relu", dropout=0.0): + super().__init__() + self.mlp = MLP([in_features, d_model, d_model], act="relu", dropout=dropout) + self.multihead_attention = MultiHeadAttention(heads, d_model) + self.dropout = nn.Dropout(p=dropout) + self.norm = nn.LayerNorm(normalized_shape=d_model) + + def forward(self, x): + y = x + y = paddle.flatten(y, start_axis=2) + y = self.mlp(y) + y = self.multihead_attention(y, y, y, mask=None) + y = self.dropout(y) + y = paddle.unsqueeze(y, axis=2) + y = x + y + y = self.norm(y) + return y + + +class Encoder(nn.Layer): + def __init__( + self, num_layers, num_var_max, d_model, heads, act="relu", dropout=0.0 + ): + super().__init__() + self.first_mlp = MLP([1, d_model, d_model], act="relu", dropout=dropout) + self.layers = nn.LayerList( + sublayers=[ + EncoderLayerMix( + d_model * num_var_max, d_model, heads, act="relu", dropout=dropout + ) + for _ in range(num_layers) + ] + ) + self.last_mlp = MLP([d_model, d_model], act="relu", dropout=dropout) + + def forward(self, x): + y = x + y = self.first_mlp(y) + for layer in self.layers: + y = layer(y) + y = self.last_mlp(y) + y = paddle.max(y, axis=1) + return y + + +class TokenEmbeddings(nn.Layer): + def __init__(self, vocab_size, seq_length, d_model, dropout=0.0): + super().__init__() + self.embed = nn.Embedding(num_embeddings=vocab_size, embedding_dim=d_model) + self.seq_length = seq_length + self.d_model = d_model + self.dropout = nn.Dropout(dropout) + self.get_pe_num() + + def get_pe_num(self): + self.pe = paddle.zeros(shape=[self.seq_length, self.d_model]) + numerator = paddle.arange( + self.seq_length, dtype=paddle.get_default_dtype() + ).unsqueeze(axis=1) + denominator = paddle.pow( + x=paddle.to_tensor(10e4, dtype=paddle.get_default_dtype()), + y=paddle.arange(self.d_model, step=2) / self.d_model, + ).unsqueeze(axis=0) + self.pe[:, 0::2] = paddle.sin(x=numerator / denominator) + self.pe[:, 1::2] = paddle.cos(x=numerator / denominator) + self.pe.stop_gradient = True + + def forward(self, x): + # embedding + y = x + y = self.embed(y) * math.sqrt(self.d_model) + # position encoding + y = self.dropout(y + self.pe) + return y + + +class DecoderLayer(nn.Layer): + def __init__(self, heads, d_model, act="relu", dropout=0.0): + super().__init__() + self.multihead_attention_1 = MultiHeadAttention(heads, d_model) + self.dropout_1 = nn.Dropout(p=dropout) + self.norm_1 = nn.LayerNorm(d_model) + + self.multihead_attention_2 = MultiHeadAttention(heads, d_model) + self.dropout_2 = nn.Dropout(p=dropout) + self.norm_2 = nn.LayerNorm(d_model) + + self.mlp = MLP([d_model, 2 * d_model, d_model], act="relu", dropout=dropout) + self.norm_3 = nn.LayerNorm(d_model) + + def forward(self, x_emb, x_enc, mask): + y_mha_1 = self.multihead_attention_1(x_emb, x_emb, x_emb, mask=mask) + y_mha_1 = self.dropout_1(y_mha_1) + y = y_mha_1 + x_emb + y = self.norm_1(y) + y_mha_2 = self.multihead_attention_2(y, x_enc, x_enc, mask=None) + y_mha_2 = self.dropout_2(y_mha_2) + y = y + y_mha_2 + y = self.norm_2(y) + y_mlp = self.mlp(y) + y = y + y_mlp + y = self.norm_3(y) + return y + + +class Decoder(nn.Layer): + def __init__( + self, + num_layers, + vocab_size, + seq_length, + d_model, + heads, + act="relu", + dropout=0.0, + ): + super().__init__() + self.token_embeddings = TokenEmbeddings( + vocab_size, seq_length, d_model, dropout + ) + self.dropout = nn.Dropout(p=dropout) + self.layers = nn.LayerList( + sublayers=[ + DecoderLayer(heads, d_model, act="relu", dropout=dropout) + for _ in range(num_layers) + ] + ) + + def forward(self, x_target, x_enc, mask): + y = x_target + y = self.token_embeddings(y) + y = self.dropout(y) + for layer in self.layers: + y = layer(y, x_enc, mask) + return y + + +class Transformer(base.Arch): + """A Kind of Transformer Model. + + Args: + input_keys (Tuple[str, ...]): Name of input keys, such as ("x", "y", "z"). + output_keys (Tuple[str, ...]): Name of output keys, such as ("u", "v", "w"). + num_var_max (int): Maximum number of variables. + vocab_size (int): Size of vocab. Size of unary operators = 1, binary operators = 2. + seq_length (int): Length of sequance. + d_model (int, optional): The innermost dimension of model. Defaults to 256. + heads (int, optional): The number of independent heads for the multi-head attention layers. Defaults to 4. + num_layers_enc (int, optional): The number of encoders. Defaults to 4. + num_layers_dec (int, optional): The number of decoders. Defaults to 8. + dropout (float, optional): Dropout regularization. Defaults to 0.0. + + Examples: + >>> import paddle + >>> import ppsci + >>> model = ppsci.arch.Transformer( + ... input_keys=("input", "target_seq"), + ... output_keys=("output",), + ... num_var_max=7, + ... vocab_size=20, + ... seq_length=30, + ... ) + >>> input_dict = {"input": paddle.rand([512, 50, 7, 1]), + ... "target_seq": paddle.rand([512, 30])} + >>> output_dict = model(input_dict) + >>> print(output_dict["output"].shape) + [512, 30, 20] + """ + + def __init__( + self, + input_keys: Tuple[str, ...], + output_keys: Tuple[str, ...], + num_var_max: int, + vocab_size: int, + seq_length: int, + d_model: int = 256, + heads: int = 4, + num_layers_enc: int = 4, + num_layers_dec: int = 8, + act: str = "relu", + dropout: float = 0.0, + ): + super().__init__() + self.input_keys = input_keys + self.output_keys = output_keys + self.num_var_max = num_var_max + self.vocab_size = vocab_size + self.seq_length = seq_length + self.d_model = d_model + self.heads = heads + self.num_layers_enc = num_layers_enc + self.num_layers_dec = num_layers_dec + self.act = act + self.dropout = dropout + + self.encoder = Encoder( + num_layers_enc, num_var_max, d_model, heads, act="relu", dropout=dropout + ) + self.decoder = Decoder( + num_layers_dec, + vocab_size, + seq_length, + d_model, + heads, + act="relu", + dropout=dropout, + ) + self.last_layer = paddle.nn.Linear(in_features=d_model, out_features=vocab_size) + + def get_mask(self, target_seq): + padding_mask = paddle.equal(target_seq, 0).unsqueeze(axis=1).unsqueeze(axis=1) + future_mask = paddle.triu( + paddle.ones(shape=[target_seq.shape[1], target_seq.shape[1]]), + diagonal=1, + ).astype(dtype="bool") + mask = paddle.logical_or(x=padding_mask, y=future_mask) + return mask + + def forward_tensor(self, x_lst): + y, target_seq = x_lst[0], x_lst[1] + mask = self.get_mask(target_seq) + y_enc = self.encoder(y) + y = self.decoder(target_seq, y_enc, mask) + y = self.last_layer(y) + return y + + def forward(self, x): + if self._input_transform is not None: + x = self._input_transform(x) + + x_lst = [x[key] for key in self.input_keys] # input, target_seq + y = self.forward_tensor(x_lst) + y = self.split_to_dict(y, self.output_keys, axis=-1) + + if self._output_transform is not None: + y = self._output_transform(x, y) + return y + + @paddle.no_grad() + def decode_process( + self, dataset: paddle.Tensor, complete_func: Callable + ) -> paddle.Tensor: + """Greedy decode with the Transformer model, decode until the equation tree is completed. + + Args: + dataset (paddle.Tensor): Tabular dataset. + complete_func (Callable): Function used to calculate whether inference is complete. + """ + encoder_output = self.encoder(dataset) + decoder_output = paddle.zeros( + shape=(dataset.shape[0], self.seq_length + 1), dtype=paddle.int64 + ) + decoder_output[:, 0] = 1 + is_complete = paddle.zeros(shape=dataset.shape[0], dtype=paddle.bool) + for n1 in range(self.seq_length): + padding_mask = ( + paddle.equal(x=decoder_output[:, :-1], y=0) + .unsqueeze(axis=1) + .unsqueeze(axis=1) + ) + future_mask = paddle.triu( + x=paddle.ones(shape=[self.seq_length, self.seq_length]), diagonal=1 + ).astype(dtype=paddle.bool) + mask_dec = paddle.logical_or(x=padding_mask, y=future_mask) + y_dec = self.decoder( + x_target=decoder_output[:, :-1], + x_enc=encoder_output, + mask=mask_dec, + ) + y_mlp = self.last_layer(y_dec) + # set value depending on complete condition + decoder_output[:, n1 + 1] = paddle.where( + is_complete, 0, paddle.argmax(y_mlp[:, n1], axis=-1) + ) + # set complete condition + for n2 in range(dataset.shape[0]): + if complete_func(decoder_output[n2, 1:]): + is_complete[n2] = True + return decoder_output diff --git a/examples/smc_reac/ppsci/arch/unetex.py b/examples/smc_reac/ppsci/arch/unetex.py new file mode 100644 index 0000000000..d0ba170464 --- /dev/null +++ b/examples/smc_reac/ppsci/arch/unetex.py @@ -0,0 +1,290 @@ +# 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. + +from typing import Optional +from typing import Tuple +from typing import Type + +import paddle +from paddle import nn + +from ppsci.arch import base + + +def create_layer( + in_channel, + out_channel, + kernel_size, + weight_norm=True, + batch_norm=True, + activation=nn.ReLU, + convolution=nn.Conv2D, +): + if kernel_size % 2 == 0: + raise ValueError("kernel_size should even number") + conv = convolution(in_channel, out_channel, kernel_size, padding=kernel_size // 2) + if weight_norm: + conv = nn.util.weight_norm(conv) + layer = [] + layer.append(conv) + if activation is not None: + layer.append(activation()) + if batch_norm: + layer.append(nn.BatchNorm2D(out_channel)) + return nn.Sequential(*layer) + + +def create_encoder_block( + in_channel, + out_channel, + kernel_size, + weight_norm=True, + batch_norm=True, + activation=nn.ReLU, + layers=2, +): + encoder = [] + encoder.append( + create_layer( + in_channel, + out_channel, + kernel_size, + weight_norm, + batch_norm, + activation, + nn.Conv2D, + ) + ) + for i in range(layers - 1): + encoder.append( + create_layer( + out_channel, + out_channel, + kernel_size, + weight_norm, + batch_norm, + activation, + nn.Conv2D, + ) + ) + return nn.Sequential(*encoder) + + +def create_decoder_block( + in_channel, + out_channel, + kernel_size, + weight_norm=True, + batch_norm=True, + activation=nn.ReLU, + layers=2, + final_layer=False, +): + decoder = [] + for i in range(layers): + _in = in_channel + _out = in_channel + _batch_norm = batch_norm + _activation = activation + if i == 0: + _in = in_channel * 2 + if i == layers - 1: + _out = out_channel + if final_layer: + _batch_norm = False + _activation = None + decoder.append( + create_layer( + _in, + _out, + kernel_size, + weight_norm, + _batch_norm, + _activation, + nn.Conv2DTranspose, + ) + ) + return nn.Sequential(*decoder) + + +def create_encoder( + in_channel, filters, kernel_size, wn=True, bn=True, activation=nn.ReLU, layers=2 +): + encoder = [] + for i in range(len(filters)): + encoder_layer = create_encoder_block( + in_channel if i == 0 else filters[i - 1], + filters[i], + kernel_size, + wn, + bn, + activation, + layers, + ) + encoder = encoder + [encoder_layer] + return nn.Sequential(*encoder) + + +def create_decoder( + out_channel, + filters, + kernel_size, + weight_norm=True, + batch_norm=True, + activation=nn.ReLU, + layers=2, +): + decoder = [] + for i in range(len(filters)): + if i == 0: + decoder_layer = create_decoder_block( + filters[i], + out_channel, + kernel_size, + weight_norm, + batch_norm, + activation, + layers, + final_layer=True, + ) + else: + decoder_layer = create_decoder_block( + filters[i], + filters[i - 1], + kernel_size, + weight_norm, + batch_norm, + activation, + layers, + final_layer=False, + ) + decoder = [decoder_layer] + decoder + return nn.Sequential(*decoder) + + +class UNetEx(base.Arch): + """U-Net Extension for CFD. + + Reference: [Ribeiro M D, Rehman A, Ahmed S, et al. DeepCFD: Efficient steady-state laminar flow approximation with deep convolutional neural networks[J]. arXiv preprint arXiv:2004.08826, 2020.](https://arxiv.org/abs/2004.08826) + + Args: + input_key (str): Name of function data for input. + output_key (str): Name of function data for output. + in_channel (int): Number of channels of input. + out_channel (int): Number of channels of output. + kernel_size (int, optional): Size of kernel of convolution layer. Defaults to 3. + filters (Tuple[int, ...], optional): Number of filters. Defaults to (16, 32, 64). + layers (int, optional): Number of encoders or decoders. Defaults to 3. + weight_norm (bool, optional): Whether use weight normalization layer. Defaults to True. + batch_norm (bool, optional): Whether add batch normalization layer. Defaults to True. + activation (Type[nn.Layer], optional): Name of activation function. Defaults to nn.ReLU. + final_activation (Optional[Type[nn.Layer]]): Name of final activation function. Defaults to None. + + Examples: + >>> import ppsci + >>> model = ppsci.arch.UNetEx( + ... input_key="input", + ... output_key="output", + ... in_channel=3, + ... out_channel=3, + ... kernel_size=5, + ... filters=(4, 4, 4, 4), + ... layers=3, + ... weight_norm=False, + ... batch_norm=False, + ... activation=None, + ... final_activation=None, + ... ) + >>> input_dict = {'input': paddle.rand([4, 3, 4, 4])} + >>> output_dict = model(input_dict) + >>> print(output_dict['output']) # doctest: +SKIP + >>> print(output_dict['output'].shape) + [4, 3, 4, 4] + """ + + def __init__( + self, + input_key: str, + output_key: str, + in_channel: int, + out_channel: int, + kernel_size: int = 3, + filters: Tuple[int, ...] = (16, 32, 64), + layers: int = 3, + weight_norm: bool = True, + batch_norm: bool = True, + activation: Type[nn.Layer] = nn.ReLU, + final_activation: Optional[Type[nn.Layer]] = None, + ): + if len(filters) == 0: + raise ValueError("The filters shouldn't be empty ") + + super().__init__() + self.input_keys = (input_key,) + self.output_keys = (output_key,) + self.final_activation = final_activation + self.encoder = create_encoder( + in_channel, + filters, + kernel_size, + weight_norm, + batch_norm, + activation, + layers, + ) + decoders = [ + create_decoder( + 1, filters, kernel_size, weight_norm, batch_norm, activation, layers + ) + for i in range(out_channel) + ] + self.decoders = nn.Sequential(*decoders) + + def encode(self, x): + tensors = [] + indices = [] + sizes = [] + for encoder in self.encoder: + x = encoder(x) + sizes.append(x.shape) + tensors.append(x) + x, ind = nn.functional.max_pool2d(x, 2, 2, return_mask=True) + indices.append(ind) + return x, tensors, indices, sizes + + def decode(self, x, tensors, indices, sizes): + y = [] + for _decoder in self.decoders: + _x = x + _tensors = tensors[:] + _indices = indices[:] + _sizes = sizes[:] + for decoder in _decoder: + tensor = _tensors.pop() + size = _sizes.pop() + indice = _indices.pop() + # upsample operations + _x = nn.functional.max_unpool2d(_x, indice, 2, 2, output_size=size) + _x = paddle.concat([tensor, _x], axis=1) + _x = decoder(_x) + y.append(_x) + return paddle.concat(y, axis=1) + + def forward(self, x): + x = x[self.input_keys[0]] + x, tensors, indices, sizes = self.encode(x) + x = self.decode(x, tensors, indices, sizes) + if self.final_activation is not None: + x = self.final_activation(x) + return {self.output_keys[0]: x} diff --git a/examples/smc_reac/ppsci/arch/unonet.py b/examples/smc_reac/ppsci/arch/unonet.py new file mode 100644 index 0000000000..c238a55be1 --- /dev/null +++ b/examples/smc_reac/ppsci/arch/unonet.py @@ -0,0 +1,289 @@ +from typing import Dict +from typing import Optional +from typing import Tuple +from typing import Union + +import paddle +import paddle.nn as nn +import paddle.nn.functional as F + +from ppsci.arch import base +from ppsci.arch import fno_block + + +class UNONet(base.Arch): + """N-Dimensional U-Shaped Neural Operator. + + Args: + input_keys (Tuple[str, ...]): Name of input keys, such as ("input",). + output_keys (Tuple[str, ...]): Name of output keys, such as ("output",). + in_channels (int, optional): Number of input channels. + out_channels (int, optional): Number of output channels. + hidden_channels (int): Width of the FNO (i.e. number of channels). + lifting_channels (int, optional): Number of hidden channels of the lifting block of the FNO. + Defaults to 256. + projection_channels (int, optional): Number of hidden channels of the projection block of the FNO. + Defaults to 256. + n_layers (int, optional): Number of Fourier Layers. Defaults to 4. + uno_out_channels (Tuple[int, ...], optional): Number of output channel of each Fourier Layers. + Eaxmple: For a Five layer UNO uno_out_channels can be [32,64,64,64,32].c + uno_n_modes (Tuple[Tuple[int, ...], ...]): Number of Fourier Modes to use in integral operation of each + Fourier Layers (along each dimension). + Example: For a five layer UNO with 2D input the uno_n_modes can be: [[5,5],[5,5],[5,5],[5,5],[5,5]]. Defaults to None. + uno_scalings (Tuple[Tuple[int, ...], ...]): Scaling Factors for each Fourier Layers. + Example: For a five layer UNO with 2D input, the uno_scalings can be : [[1.0,1.0],[0.5,0.5],[1,1],[1,1],[2,2]].Defaults to None. + horizontal_skips_map (Dict, optional): A map {...., b: a, ....} denoting horizontal skip connection + from a-th layer to b-th layer. If None default skip connection is applied. + Example: For a 5 layer UNO architecture, the skip connections can be horizontal_skips_map ={4:0,3:1}.Defaults to None. + incremental_n_modes (tuple[int],optional): Incremental number of modes to use in Fourier domain. + * If not None, this allows to incrementally increase the number of modes in Fourier domain + during training. Has to verify n <= N for (n, m) in zip(incremental_n_modes, n_modes). + * If None, all the n_modes are used. + This can be updated dynamically during training.Defaults to None. + use_mlp (bool, optional): Whether to use an MLP layer after each FNO block. Defaults to False. + mlp (Dict[str, float], optional): Parameters of the MLP. {'expansion': float, 'dropout': float}. + Defaults to None. + non_linearity (nn.functional, optional): Non-Linearity module to use. Defaults to F.gelu. + norm (str, optional): Normalization layer to use. Defaults to None. + ada_in_features (Optional[int],optional): The input channels of the adaptive normalization.Defaults to + None. + preactivation (bool, optional): Whether to use resnet-style preactivation. Defaults to False. + fno_skip (str, optional): Type of skip connection to use for fno_block. Defaults to "linear". + horizontal_skip (str, optional): Type of skip connection to use for horizontal skip. Defaults to + "linear". + mlp_skip (str, optional): Type of skip connection to use for mlp. Defaults to "soft-gating". + separable (bool, optional): Whether to use a depthwise separable spectral convolution. + Defaults to False. + factorization (str, optional): Tensor factorization of the parameters weight to use. + * If None, a dense tensor parametrizes the Spectral convolutions. + * Otherwise, the specified tensor factorization is used. Defaults to "Tucker". + rank (float, optional): Rank of the tensor factorization of the Fourier weights. Defaults to 1.0. + joint_factorization (bool, optional): Whether all the Fourier Layers should be parametrized by a + single tensor (vs one per layer). Defaults to False. + implementation (str, optional): {'factorized', 'reconstructed'}, optional. Defaults to "factorized". + If factorization is not None, forward mode to use:: + * `reconstructed` : the full weight tensor is reconstructed from the factorization and used for the forward pass. + * `factorized` : the input is directly contracted with the factors of the decomposition. + domain_padding (Optional[Union[list, float, int]], optional): Whether to use percentage of padding. + Defaults to None. + domain_padding_mode (str, optional): {'symmetric', 'one-sided'}, optional + How to perform domain padding, by default 'one-sided'. Defaults to "one-sided". + fft_norm (str, optional): The normalization mode for the FFT. Defaults to "forward". + patching_levels (int, optional): Number of patching levels to use. Defaults to 0. + """ + + def __init__( + self, + input_keys: Tuple[str, ...], + output_keys: Tuple[str, ...], + in_channels: int, + out_channels: int, + hidden_channels: int, + lifting_channels: int = 256, + projection_channels: int = 256, + n_layers: int = 4, + uno_out_channels: Tuple[int, ...] = None, + uno_n_modes: Tuple[Tuple[int, ...], ...] = None, + uno_scalings: Tuple[Tuple[int, ...], ...] = None, + horizontal_skips_map: Dict = None, + incremental_n_modes: Tuple[int, ...] = None, + use_mlp: bool = False, + mlp: Optional[Dict[str, float]] = None, + non_linearity: nn.functional = F.gelu, + norm: str = None, + ada_in_features: Optional[int] = None, + preactivation: bool = False, + fno_skip: str = "linear", + horizontal_skip: str = "linear", + mlp_skip: str = "soft-gating", + separable: bool = False, + factorization: str = None, + rank: float = 1.0, + joint_factorization: bool = False, + implementation: str = "factorized", + domain_padding: Optional[Union[list, float, int]] = None, + domain_padding_mode: str = "one-sided", + fft_norm: str = "forward", + patching_levels: int = 0, + **kwargs, + ): + super().__init__() + self.input_keys = input_keys + self.output_keys = output_keys + if uno_out_channels is None: + raise ValueError("uno_out_channels can not be None") + if uno_n_modes is None: + raise ValueError("uno_n_modes can not be None") + if uno_scalings is None: + raise ValueError("uno_scalings can not be None") + + if len(uno_out_channels) != n_layers: + raise ValueError("Output channels for all layers are not given") + + if len(uno_n_modes) != n_layers: + raise ValueError("Number of modes for all layers are not given") + + if len(uno_scalings) != n_layers: + raise ValueError("Scaling factor for all layers are not given") + + self.n_dim = len(uno_n_modes[0]) + self.uno_out_channels = uno_out_channels + self.uno_n_modes = uno_n_modes + self.uno_scalings = uno_scalings + + self.hidden_channels = hidden_channels + self.lifting_channels = lifting_channels + self.projection_channels = projection_channels + self.in_channels = in_channels + if patching_levels: + self.in_channels = self.in_channels * patching_levels + 1 + self.out_channels = out_channels + self.n_layers = n_layers + self.horizontal_skips_map = horizontal_skips_map + self.joint_factorization = joint_factorization + self.non_linearity = non_linearity + self.rank = rank + self.factorization = factorization + self.fno_skip = (fno_skip,) + self.mlp_skip = (mlp_skip,) + self.fft_norm = fft_norm + self.implementation = implementation + self.separable = separable + self.preactivation = preactivation + self._incremental_n_modes = incremental_n_modes + self.mlp = mlp + # constructing default skip maps + if self.horizontal_skips_map is None: + self.horizontal_skips_map = {} + for i in range( + 0, + n_layers // 2, + ): + # example, if n_layers = 5, then 4:0, 3:1 + self.horizontal_skips_map[n_layers - i - 1] = i + # self.uno_scalings may be a 1d list specifying uniform scaling factor at each layer + # or a 2d list, where each row specifies scaling factors along each dimension. + # To get the final (end to end) scaling factors we need to multiply + # the scaling factors (a list) of all layer. + + self.end_to_end_scaling_factor = [1] * len(self.uno_scalings[0]) + # multiplying scaling factors + for k in self.uno_scalings: + self.end_to_end_scaling_factor = [ + i * j for (i, j) in zip(self.end_to_end_scaling_factor, k) + ] + + # list with a single element is replaced by the scaler. + if len(self.end_to_end_scaling_factor) == 1: + self.end_to_end_scaling_factor = self.end_to_end_scaling_factor[0] + + if isinstance(self.end_to_end_scaling_factor, (float, int)): + self.end_to_end_scaling_factor = [ + self.end_to_end_scaling_factor + ] * self.n_dim + + if domain_padding is not None and ( + (isinstance(domain_padding, list) and sum(domain_padding) > 0) + or (isinstance(domain_padding, (float, int)) and domain_padding > 0) + ): + self.domain_padding = fno_block.DomainPadding( + domain_padding=domain_padding, padding_mode=domain_padding_mode + ) + else: + self.domain_padding = None + self.domain_padding_mode = domain_padding_mode + + self.lifting = fno_block.MLP( + in_channels=in_channels, + out_channels=self.hidden_channels, + hidden_channels=self.lifting_channels, + n_layers=2, + n_dim=self.n_dim, + ) + + self.fno_blocks = nn.LayerList([]) + self.horizontal_skips = nn.LayerDict({}) + prev_out = self.hidden_channels + for i in range(self.n_layers): + if i in self.horizontal_skips_map.keys(): + prev_out = ( + prev_out + self.uno_out_channels[self.horizontal_skips_map[i]] + ) + self.fno_blocks.append( + fno_block.FNOBlocks( + in_channels=prev_out, + out_channels=self.uno_out_channels[i], + n_modes=self.uno_n_modes[i], + use_mlp=use_mlp, + mlp=mlp, + output_scaling_factor=[self.uno_scalings[i]], + non_linearity=non_linearity, + norm=norm, + ada_in_features=ada_in_features, + preactivation=preactivation, + fno_skip=fno_skip, + mlp_skip=mlp_skip, + separable=separable, + incremental_n_modes=incremental_n_modes, + factorization=factorization, + rank=rank, + SpectralConv=fno_block.FactorizedSpectralConv, + joint_factorization=joint_factorization, + implementation=implementation, + fft_norm=fft_norm, + ) + ) + + if i in self.horizontal_skips_map.values(): + self.horizontal_skips[str(i)] = fno_block.skip_connection( + self.uno_out_channels[i], + self.uno_out_channels[i], + type=horizontal_skip, + n_dim=self.n_dim, + ) + prev_out = self.uno_out_channels[i] + + self.projection = fno_block.MLP( + in_channels=prev_out, + out_channels=out_channels, + hidden_channels=self.projection_channels, + n_layers=2, + n_dim=self.n_dim, + non_linearity=non_linearity, + ) + + def forward(self, x, **kwargs): + x = self.concat_to_tensor(x, self.input_keys) + x = self.lifting(x) + if self.domain_padding is not None: + x = self.domain_padding.pad(x) + output_shape = [ + int(round(i * j)) + for (i, j) in zip(x.shape[-self.n_dim :], self.end_to_end_scaling_factor) + ] + + skip_outputs = {} + cur_output = None + for layer_idx in range(self.n_layers): + if layer_idx in self.horizontal_skips_map.keys(): + skip_val = skip_outputs[self.horizontal_skips_map[layer_idx]] + output_scaling_factors = [ + m / n for (m, n) in zip(x.shape, skip_val.shape) + ] + output_scaling_factors = output_scaling_factors[-1 * self.n_dim :] + t = fno_block.resample( + skip_val, output_scaling_factors, list(range(-self.n_dim, 0)) + ) + x = paddle.concat([x, t], axis=1) + + if layer_idx == self.n_layers - 1: + cur_output = output_shape + x = self.fno_blocks[layer_idx](x, output_shape=cur_output) + if layer_idx in self.horizontal_skips_map.values(): + skip_outputs[layer_idx] = self.horizontal_skips[str(layer_idx)](x) + + if self.domain_padding is not None: + x = self.domain_padding.unpad(x) + + out = self.projection(x) + return {self.output_keys[0]: out} diff --git a/examples/smc_reac/ppsci/arch/uscnn.py b/examples/smc_reac/ppsci/arch/uscnn.py new file mode 100644 index 0000000000..5da8440ca4 --- /dev/null +++ b/examples/smc_reac/ppsci/arch/uscnn.py @@ -0,0 +1,124 @@ +# 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. + +from typing import Tuple +from typing import Union + +import numpy as np +from paddle import nn + +import ppsci +from ppsci.arch import base + + +class USCNN(base.Arch): + """Physics-informed convolutional neural networks. + + Args: + input_keys (Tuple[str, ...]): Name of input keys, such as ("coords"). + output_keys (Tuple[str, ...]):Name of output keys, such as ("outputV"). + hidden_size (Union[int, Tuple[int, ...]]): The hidden channel for convolutional layers + h (float): The spatial step + nx (int): the number of grids along x-axis + ny (int): The number of grids along y-axis + nvar_in (int, optional): input channel. Defaults to 1. + nvar_out (int, optional): Output channel. Defaults to 1. + pad_singleside (int, optional): Pad for hard boundary constraint. Defaults to 1. + k (int, optional): Kernel_size. Defaults to 5. + s (int, optional): Stride. Defaults to 1. + p (int, optional): Padding. Defaults to 2. + + Examples: + >>> import ppsci + >>> model = ppsci.arch.USCNN( + ... ["coords"], + ... ["outputV"], + ... [16, 32, 16], + ... h=0.01, + ... ny=19, + ... nx=84, + ... nvar_in=2, + ... nvar_out=1, + ... pad_singleside=1, + ... ) + """ + + def __init__( + self, + input_keys: Tuple[str, ...], + output_keys: Tuple[str, ...], + hidden_size: Union[int, Tuple[int, ...]], + h: float, + nx: int, + ny: int, + nvar_in: int = 1, + nvar_out: int = 1, + pad_singleside: int = 1, + k: int = 5, + s: int = 1, + p: int = 2, + ): + super().__init__() + self.input_keys = input_keys + self.output_keys = output_keys + self.nvar_in = nvar_in + self.nvar_out = nvar_out + self.k = k + self.s = s + self.p = p + self.deltaX = h + self.nx = nx + self.ny = ny + self.pad_singleside = pad_singleside + self.relu = nn.ReLU() + self.US = nn.Upsample(size=[self.ny - 2, self.nx - 2], mode="bicubic") + self.conv1 = nn.Conv2D( + self.nvar_in, hidden_size[0], kernel_size=k, stride=s, padding=p + ) + self.conv2 = nn.Conv2D( + hidden_size[0], hidden_size[1], kernel_size=k, stride=s, padding=p + ) + self.conv3 = nn.Conv2D( + hidden_size[1], hidden_size[2], kernel_size=k, stride=s, padding=p + ) + self.conv4 = nn.Conv2D( + hidden_size[2], self.nvar_out, kernel_size=k, stride=s, padding=p + ) + self.pixel_shuffle = nn.PixelShuffle(1) + self.apply(self.init_weights) + self.udfpad = nn.Pad2D( + [pad_singleside, pad_singleside, pad_singleside, pad_singleside], value=0 + ) + + def init_weights(self, m): + if isinstance(m, nn.Conv2D): + bound = 1 / np.sqrt(np.prod(m.weight.shape[1:])) + ppsci.utils.initializer.uniform_(m.weight, -bound, bound) + if m.bias is not None: + ppsci.utils.initializer.uniform_(m.bias, -bound, bound) + + def forward(self, x): + y = self.concat_to_tensor(x, self.input_keys, axis=-1) + y = self.US(y) + y = self.relu(self.conv1(y)) + y = self.relu(self.conv2(y)) + y = self.relu(self.conv3(y)) + y = self.pixel_shuffle(self.conv4(y)) + + y = self.udfpad(y) + y = y[:, 0, :, :].reshape([y.shape[0], 1, y.shape[2], y.shape[3]]) + y = self.split_to_dict(y, self.output_keys) + if self._output_transform is not None: + y = self._output_transform(x, y) + return y diff --git a/examples/smc_reac/ppsci/arch/vae.py b/examples/smc_reac/ppsci/arch/vae.py new file mode 100644 index 0000000000..2a05f0d648 --- /dev/null +++ b/examples/smc_reac/ppsci/arch/vae.py @@ -0,0 +1,103 @@ +# 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. + +from __future__ import annotations + +from typing import Tuple + +import paddle +import paddle.nn as nn + +from ppsci.arch import base + + +class AutoEncoder(base.Arch): + """ + AutoEncoder is a class that represents an autoencoder neural network model. + + Args: + input_keys (Tuple[str, ...]): A tuple of input keys. + output_keys (Tuple[str, ...]): A tuple of output keys. + input_dim (int): The dimension of the input data. + latent_dim (int): The dimension of the latent space. + hidden_dim (int): The dimension of the hidden layer. + + Examples: + >>> import paddle + >>> import ppsci + >>> model = ppsci.arch.AutoEncoder( + ... input_keys=("input1",), + ... output_keys=("mu", "log_sigma", "decoder_z",), + ... input_dim=100, + ... latent_dim=50, + ... hidden_dim=200 + ... ) + >>> input_dict = {"input1": paddle.rand([200, 100]),} + >>> output_dict = model(input_dict) + >>> print(output_dict["mu"].shape) + [200, 50] + >>> print(output_dict["log_sigma"].shape) + [200, 50] + >>> print(output_dict["decoder_z"].shape) + [200, 100] + """ + + def __init__( + self, + input_keys: Tuple[str, ...], + output_keys: Tuple[str, ...], + input_dim: int, + latent_dim: int, + hidden_dim: int, + ): + super(AutoEncoder, self).__init__() + self.input_keys = input_keys + self.output_keys = output_keys + # encoder + self._encoder_linear = nn.Sequential( + nn.Linear(input_dim, hidden_dim), + nn.Tanh(), + ) + self._encoder_mu = nn.Linear(hidden_dim, latent_dim) + self._encoder_log_sigma = nn.Linear(hidden_dim, latent_dim) + + self._decoder = nn.Sequential( + nn.Linear(latent_dim, hidden_dim), + nn.Tanh(), + nn.Linear(hidden_dim, input_dim), + ) + + def encoder(self, x): + h = self._encoder_linear(x) + mu = self._encoder_mu(h) + log_sigma = self._encoder_log_sigma(h) + return mu, log_sigma + + def decoder(self, x): + return self._decoder(x) + + def forward_tensor(self, x): + mu, log_sigma = self.encoder(x) + z = mu + paddle.randn(mu.shape) * paddle.exp(log_sigma) + return mu, log_sigma, self.decoder(z) + + def forward(self, x): + x = self.concat_to_tensor(x, self.input_keys, axis=-1) + mu, log_sigma, decoder_z = self.forward_tensor(x) + result_dict = { + self.output_keys[0]: mu, + self.output_keys[1]: log_sigma, + self.output_keys[2]: decoder_z, + } + return result_dict diff --git a/examples/smc_reac/ppsci/arch/velocitygan.py b/examples/smc_reac/ppsci/arch/velocitygan.py new file mode 100644 index 0000000000..28f44d14f1 --- /dev/null +++ b/examples/smc_reac/ppsci/arch/velocitygan.py @@ -0,0 +1,354 @@ +# 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. + +from math import ceil +from math import sqrt +from typing import Tuple + +import paddle + +import ppsci.utils.initializer as init +from ppsci.arch import base + + +class VelocityGenerator(base.Arch): + """The Generator Of VelocityGAN. + VelocityGAN is applied to full waveform inversion tasks, and the structure of the model + comes from https://arxiv.org/abs/2111.02926# + + Args: + input_keys (Tuple[str, ...]): Name of input keys, such as ("input",). + output_keys (Tuple[str, ...]): Name of output keys, such as ("output",). + dim1 (int, optional): Number of channels in the outermost layers of both encoder and decoder segments. Default is 32. + dim2 (int, optional): Number of channels in the second set of layers from the outermost in both encoder and decoder segments. Default is 64. + dim3 (int, optional): Number of channels in the intermediate layers. Default is 128. + dim4 (int, optional): Number of channels near the bottleneck, just before and after the deepest layer. Default is 256. + dim5 (int, optional): Number of channels at the bottleneck, the deepest layer in the network. Default is 512. + sample_spatial (float, optional): Spatial sampling rate of the input, used to dynamically calculate the kernel size in the last encoder layer. Default is 1.0. + + Examples: + >>> import ppsci + >>> import paddle + >>> model = ppsci.arch.VelocityGenerator(("input", ), ("output", )) + >>> input_dict = {"input": paddle.randn((1, 5, 1000, 70))} + >>> output_dict = model(input_dict) # doctest: +SKIP + >>> print(output_dict["output"].shape) # doctest: +SKIP + [1, 1, 70, 70] + """ + + def __init__( + self, + input_keys: Tuple[str, ...], + output_keys: Tuple[str, ...], + dim1: int = 32, + dim2: int = 64, + dim3: int = 128, + dim4: int = 256, + dim5: int = 512, + sample_spatial: float = 1.0, + ): + super().__init__() + self.input_keys = input_keys + self.output_keys = output_keys + self.generator = Generator( + dim1=dim1, + dim2=dim2, + dim3=dim3, + dim4=dim4, + dim5=dim5, + sample_spatial=sample_spatial, + ) + + def forward(self, x): + if self._input_transform is not None: + x = self._input_transform(x) + + y = self.concat_to_tensor(x, self.input_keys, axis=-1) + y = self.generator(y) + y = self.split_to_dict(y, self.output_keys, axis=-1) + + if self._output_transform is not None: + y = self._output_transform(x, y) + + return y + + +class VelocityDiscriminator(base.Arch): + """The Discriminator Of VelocityGAN. + VelocityGAN is applied to full waveform inversion tasks, and the structure of the model + comes from https://arxiv.org/abs/2111.02926# + + Args: + input_keys (Tuple[str, ...]): Name of input keys, such as ("input",). + output_keys (Tuple[str, ...]): Name of output keys, such as ("output",). + dim1 (int, optional): The number of output channels for convblock1_1 and convblock1_2. Default is 32. + dim2 (int, optional): The number of output channels for convblock2_1 and convblock2_2. Default is 64. + dim3 (int, optional): The number of output channels for convblock3_1 and convblock3_2. Default is 128. + dim4 (int, optional): The number of output channels for convblock4_1 and convblock4_2. Default is 256. + + Examples: + >>> import ppsci + >>> import paddle + >>> model = ppsci.arch.VelocityDiscriminator(("input", ), ("output", )) + >>> input_dict = {"input": paddle.randn((1, 1, 70, 70))} + >>> output_dict = model(input_dict) # doctest: +SKIP + >>> print(output_dict["output"].shape) # doctest: +SKIP + [1, 1] + """ + + def __init__( + self, + input_keys: Tuple[str, ...], + output_keys: Tuple[str, ...], + dim1: int = 32, + dim2: int = 64, + dim3: int = 128, + dim4: int = 256, + ): + super().__init__() + self.input_keys = input_keys + self.output_keys = output_keys + self.discriminator = Discriminator(dim1=dim1, dim2=dim2, dim3=dim3, dim4=dim4) + + def forward(self, x): + if self._input_transform is not None: + x = self._input_transform(x) + + y = self.concat_to_tensor(x, self.input_keys, axis=-1) + y = self.discriminator(y) + y = self.split_to_dict(y, self.output_keys, axis=-1) + + if self._output_transform is not None: + y = self._output_transform(x, y) + + return y + + +class Generator(paddle.nn.Layer): + """The specific implementation of the generator, which is encapsulated in the VelocityGenerator class. + + Args: + dim1 (int, optional): Number of channels in the outermost layers of both encoder and decoder segments. Default is 32. + dim2 (int, optional): Number of channels in the second set of layers from the outermost in both encoder and decoder segments. Default is 64. + dim3 (int, optional): Number of channels in the intermediate layers. Default is 128. + dim4 (int, optional): Number of channels near the bottleneck, just before and after the deepest layer. Default is 256. + dim5 (int, optional): Number of channels at the bottleneck, the deepest layer in the network. Default is 512. + sample_spatial (float, optional): Spatial sampling rate of the input, used to dynamically calculate the kernel size in the last encoder layer. Default is 1.0. + """ + + def __init__( + self, + dim1: int = 32, + dim2: int = 64, + dim3: int = 128, + dim4: int = 256, + dim5: int = 512, + sample_spatial: float = 1.0, + ): + super(Generator, self).__init__() + self.convblock1 = ConvBlock( + 5, dim1, kernel_size=(7, 1), stride=(2, 1), padding=(3, 0) + ) + self.convblock2_1 = ConvBlock( + dim1, dim2, kernel_size=(3, 1), stride=(2, 1), padding=(1, 0) + ) + self.convblock2_2 = ConvBlock(dim2, dim2, kernel_size=(3, 1), padding=(1, 0)) + self.convblock3_1 = ConvBlock( + dim2, dim2, kernel_size=(3, 1), stride=(2, 1), padding=(1, 0) + ) + self.convblock3_2 = ConvBlock(dim2, dim2, kernel_size=(3, 1), padding=(1, 0)) + self.convblock4_1 = ConvBlock( + dim2, dim3, kernel_size=(3, 1), stride=(2, 1), padding=(1, 0) + ) + self.convblock4_2 = ConvBlock(dim3, dim3, kernel_size=(3, 1), padding=(1, 0)) + self.convblock5_1 = ConvBlock(dim3, dim3, stride=2) + self.convblock5_2 = ConvBlock(dim3, dim3) + self.convblock6_1 = ConvBlock(dim3, dim4, stride=2) + self.convblock6_2 = ConvBlock(dim4, dim4) + self.convblock7_1 = ConvBlock(dim4, dim4, stride=2) + self.convblock7_2 = ConvBlock(dim4, dim4) + self.convblock8 = ConvBlock( + dim4, dim5, kernel_size=(8, ceil(70 * sample_spatial / 8)), padding=0 + ) + self.deconv1_1 = DeconvBlock(dim5, dim5, kernel_size=5) + self.deconv1_2 = ConvBlock(dim5, dim5) + self.deconv2_1 = DeconvBlock(dim5, dim4, kernel_size=4, stride=2, padding=1) + self.deconv2_2 = ConvBlock(dim4, dim4) + self.deconv3_1 = DeconvBlock(dim4, dim3, kernel_size=4, stride=2, padding=1) + self.deconv3_2 = ConvBlock(dim3, dim3) + self.deconv4_1 = DeconvBlock(dim3, dim2, kernel_size=4, stride=2, padding=1) + self.deconv4_2 = ConvBlock(dim2, dim2) + self.deconv5_1 = DeconvBlock(dim2, dim1, kernel_size=4, stride=2, padding=1) + self.deconv5_2 = ConvBlock(dim1, dim1) + self.deconv6 = ConvBlock_Tanh(dim1, 1) + self.initial_weight() + + def initial_weight(self): + for _, m in self.named_sublayers(): + if isinstance(m, paddle.nn.Conv2D) or isinstance( + m, paddle.nn.Conv2DTranspose + ): + init.kaiming_uniform_(m.weight, a=sqrt(5)) + if m.bias is not None: + fan_in, _ = init._calculate_fan_in_and_fan_out(m.weight) + bound = 1 / sqrt(fan_in) + init.uniform_(m.bias, -bound, bound) + + def forward(self, x): + x = self.convblock1(x) + x = self.convblock2_1(x) + x = self.convblock2_2(x) + x = self.convblock3_1(x) + x = self.convblock3_2(x) + x = self.convblock4_1(x) + x = self.convblock4_2(x) + x = self.convblock5_1(x) + x = self.convblock5_2(x) + x = self.convblock6_1(x) + x = self.convblock6_2(x) + x = self.convblock7_1(x) + x = self.convblock7_2(x) + x = self.convblock8(x) + x = self.deconv1_1(x) + x = self.deconv1_2(x) + x = self.deconv2_1(x) + x = self.deconv2_2(x) + x = self.deconv3_1(x) + x = self.deconv3_2(x) + x = self.deconv4_1(x) + x = self.deconv4_2(x) + x = self.deconv5_1(x) + x = self.deconv5_2(x) + x = paddle.nn.functional.pad(x, pad=[-5, -5, -5, -5], mode="constant", value=0) + x = self.deconv6(x) + return x + + +class Discriminator(paddle.nn.Layer): + """The specific implementation of the discriminator, which is encapsulated in the VelocityDiscriminator class. + + Args: + dim1 (int, optional): The number of output channels for convblock1_1 and convblock1_2. Default is 32. + dim2 (int, optional): The number of output channels for convblock2_1 and convblock2_2. Default is 64. + dim3 (int, optional): The number of output channels for convblock3_1 and convblock3_2. Default is 128. + dim4 (int, optional): The number of output channels for convblock4_1 and convblock4_2. Default is 256. + """ + + def __init__( + self, dim1: int = 32, dim2: int = 64, dim3: int = 128, dim4: int = 256 + ): + super(Discriminator, self).__init__() + self.convblock1_1 = ConvBlock(1, dim1, stride=2) + self.convblock1_2 = ConvBlock(dim1, dim1) + self.convblock2_1 = ConvBlock(dim1, dim2, stride=2) + self.convblock2_2 = ConvBlock(dim2, dim2) + self.convblock3_1 = ConvBlock(dim2, dim3, stride=2) + self.convblock3_2 = ConvBlock(dim3, dim3) + self.convblock4_1 = ConvBlock(dim3, dim4, stride=2) + self.convblock4_2 = ConvBlock(dim4, dim4) + self.convblock5 = ConvBlock(dim4, 1, kernel_size=5, padding=0) + self.initial_weight() + + def initial_weight(self): + for _, m in self.named_sublayers(): + if isinstance(m, paddle.nn.Conv2D): + init.kaiming_uniform_(m.weight, a=sqrt(5)) + if m.bias is not None: + fan_in, _ = init._calculate_fan_in_and_fan_out(m.weight) + bound = 1 / sqrt(fan_in) + init.uniform_(m.bias, -bound, bound) + + def forward(self, x): + x = self.convblock1_1(x) + x = self.convblock1_2(x) + x = self.convblock2_1(x) + x = self.convblock2_2(x) + x = self.convblock3_1(x) + x = self.convblock3_2(x) + x = self.convblock4_1(x) + x = self.convblock4_2(x) + x = self.convblock5(x) + x = x.reshape([x.shape[0], -1]) + return x + + +class ConvBlock(paddle.nn.Layer): + """A convolution block, including Conv2D, BatchNorm2D, and LeakyReLU""" + + def __init__( + self, in_fea, out_fea, kernel_size=3, stride=1, padding=1, relu_slop=0.2 + ): + super(ConvBlock, self).__init__() + layers = [ + paddle.nn.Conv2D( + in_channels=in_fea, + out_channels=out_fea, + kernel_size=kernel_size, + stride=stride, + padding=padding, + ), + paddle.nn.BatchNorm2D(out_fea), + paddle.nn.LeakyReLU(negative_slope=relu_slop), + ] + self.layers = paddle.nn.Sequential(*layers) + + def forward(self, x): + return self.layers(x) + + +class DeconvBlock(paddle.nn.Layer): + """A deconvolution block, including Conv2DTranspose, BatchNorm2D, and LeakyReLU""" + + def __init__( + self, in_fea, out_fea, kernel_size=2, stride=2, padding=0, output_padding=0 + ): + super(DeconvBlock, self).__init__() + layers = [ + paddle.nn.Conv2DTranspose( + in_channels=in_fea, + out_channels=out_fea, + kernel_size=kernel_size, + stride=stride, + padding=padding, + output_padding=output_padding, + ), + paddle.nn.BatchNorm2D(out_fea), + paddle.nn.LeakyReLU(negative_slope=0.2), + ] + self.layers = paddle.nn.Sequential(*layers) + + def forward(self, x): + return self.layers(x) + + +class ConvBlock_Tanh(paddle.nn.Layer): + """A convolution block, including Conv2D, BatchNorm2D, and Tanh""" + + def __init__(self, in_fea, out_fea, kernel_size=3, stride=1, padding=1): + super(ConvBlock_Tanh, self).__init__() + layers = [ + paddle.nn.Conv2D( + in_channels=in_fea, + out_channels=out_fea, + kernel_size=kernel_size, + stride=stride, + padding=padding, + ), + paddle.nn.BatchNorm2D(out_fea), + paddle.nn.Tanh(), + ] + self.layers = paddle.nn.Sequential(*layers) + + def forward(self, x): + return self.layers(x) diff --git a/examples/smc_reac/ppsci/autodiff/__init__.py b/examples/smc_reac/ppsci/autodiff/__init__.py new file mode 100644 index 0000000000..68a0570c24 --- /dev/null +++ b/examples/smc_reac/ppsci/autodiff/__init__.py @@ -0,0 +1,17 @@ +# 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. + +from ppsci.autodiff.ad import clear +from ppsci.autodiff.ad import hessian +from ppsci.autodiff.ad import jacobian diff --git a/examples/smc_reac/ppsci/autodiff/ad.py b/examples/smc_reac/ppsci/autodiff/ad.py new file mode 100644 index 0000000000..ba3afd14a4 --- /dev/null +++ b/examples/smc_reac/ppsci/autodiff/ad.py @@ -0,0 +1,341 @@ +# 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. + +""" +This module is adapted from [https://github.com/lululxvi/deepxde](https://github.com/lululxvi/deepxde) +""" + +from __future__ import annotations + +from typing import Callable +from typing import Dict +from typing import List +from typing import Optional +from typing import Union + +import paddle + + +class _Jacobian: + """Compute Jacobian matrix J: J[i][j] = dy_i/dx_j, where i = 0, ..., dim_y-1 and + j = 0, ..., dim_x - 1. + + It is lazy evaluation, i.e., it only computes J[i][j] when needed, and will cache + by output tensor(row index in jacobian matrix). + + Args: + ys (paddle.Tensor): Output Tensor of shape [batch_size, dim_y]. + xs (paddle.Tensor): Input Tensor of shape [batch_size, dim_x]. + """ + + def __init__( + self, + ys: "paddle.Tensor", + xs: "paddle.Tensor", + J: Optional[Dict[int, paddle.Tensor]] = None, + ): + self.ys = ys + self.xs = xs + + self.dim_y = ys.shape[1] + self.dim_x = xs.shape[1] + + self.J: Dict[int, paddle.Tensor] = {} if J is None else J + + def __call__( + self, + i: int = 0, + j: Optional[int] = None, + retain_graph: Optional[bool] = None, + create_graph: bool = True, + ) -> "paddle.Tensor": + """ + Returns J[`i`][`j`]. If `j` is ``None``, returns the gradient of y_i, i.e. J[i]. + """ + if not 0 <= i < self.dim_y: + raise ValueError(f"i({i}) should in range [0, {self.dim_y}).") + if j is not None and not 0 <= j < self.dim_x: + raise ValueError(f"j({j}) should in range [0, {self.dim_x}).") + # Compute J[i] + if i not in self.J: + y = self.ys[:, i : i + 1] if self.dim_y > 1 else self.ys + self.J[i] = paddle.grad( + y, self.xs, retain_graph=retain_graph, create_graph=create_graph + )[0] + + return self.J[i] if (j is None or self.dim_x == 1) else self.J[i][:, j : j + 1] + + +class Jacobians: + r"""Compute multiple Jacobians. + + $$ + \rm Jacobian(ys, xs, i, j) = \dfrac{\partial ys_i}{\partial xs_j} + $$ + + A new instance will be created for a new pair of (output, input). For the (output, + input) pair that has been computed before, it will reuse the previous instance, + rather than creating a new one. + """ + + def __init__(self): + self.Js = {} + + def __call__( + self, + ys: "paddle.Tensor", + xs: Union["paddle.Tensor", List["paddle.Tensor"]], + i: int = 0, + j: Optional[int] = None, + retain_graph: Optional[bool] = None, + create_graph: bool = True, + ) -> Union["paddle.Tensor", List["paddle.Tensor"]]: + """Compute jacobians for given ys and xs. + + Args: + ys (paddle.Tensor): Output tensor. + xs (Union[paddle.Tensor, List[paddle.Tensor]]): Input tensor(s). + i (int, optional): i-th output variable. Defaults to 0. + j (Optional[int]): j-th input variable. Defaults to None. + retain_graph (Optional[bool]): Whether to retain the forward graph which + is used to calculate the gradient. When it is True, the graph would + be retained, in which way users can calculate backward twice for the + same graph. When it is False, the graph would be freed. Default None, + which means it is equal to `create_graph`. + create_graph (bool, optional): Whether to create the gradient graphs of + the computing process. When it is True, higher order derivatives are + supported to compute; when it is False, the gradient graphs of the + computing process would be discarded. Default False. + + Returns: + paddle.Tensor: Jacobian matrix of ys[i] to xs[j]. + + Examples: + >>> import paddle + >>> import ppsci + >>> x = paddle.randn([4, 1]) + >>> x.stop_gradient = False + >>> y = x * x + >>> dy_dx = ppsci.autodiff.jacobian(y, x) + >>> print(dy_dx.shape) + [4, 1] + """ + if not isinstance(xs, (list, tuple)): + key = (ys, xs) + if key not in self.Js: + self.Js[key] = _Jacobian(ys, xs) + return self.Js[key](i, j, retain_graph, create_graph) + else: + xs_require = [xs[i] for i in range(len(xs)) if (ys, xs[i]) not in self.Js] + grads_require = paddle.grad( + ys, + xs_require, + create_graph=create_graph, + retain_graph=retain_graph, + ) + + idx = 0 + Js_list = [] + for k, xs_ in enumerate(xs): + key = (ys, xs_) + assert xs_.shape[-1] == 1, ( + f"The last dim of each xs should be 1, but xs[{k}] has shape " + f"{xs_.shape}" + ) + if key not in self.Js: + self.Js[key] = _Jacobian(ys, xs_, {0: grads_require[idx]}) + idx += 1 + Js_list.append(self.Js[key](i, j, retain_graph, create_graph)) + return Js_list + + def _clear(self): + """Clear cached Jacobians.""" + self.Js = {} + + +# Use high-order differentiation with singleton pattern for convenient +jacobian: Callable[ + [ + "paddle.Tensor", + Union["paddle.Tensor", List["paddle.Tensor"]], + int, + Optional[int], + Optional[bool], + bool, + ], + Union["paddle.Tensor", List["paddle.Tensor"]], +] = Jacobians() + + +class _Hessian: + """Compute Hessian matrix H: H[i][j] = d^2y / dx_i dx_j, where i,j = 0,..., dim_x-1. + + It is lazy evaluation, i.e., it only computes H[i][j] when needed. + + Args: + ys: Output Tensor of shape (batch_size, 1) or (batch_size, dim_y > 1). + xs: Input Tensor of shape (batch_size, dim_x). + component: If `y` has the shape (batch_size, dim_y > 1), then `y[:, component]` + is used to compute the Hessian. Do not use if `y` has the shape (batch_size, + 1). + grad_y: The gradient of `y` w.r.t. `xs`. Provide `grad_y` if known to avoid + duplicate computation. `grad_y` can be computed from ``Jacobian``. + """ + + def __init__( + self, + ys: "paddle.Tensor", + xs: "paddle.Tensor", + component: Optional[int] = None, + grad_y: Optional["paddle.Tensor"] = None, + ): + dim_y = ys.shape[1] + + if dim_y > 1: + if component is None: + raise ValueError( + f"component({component}) can not be None when dim_y({dim_y})>1." + ) + if component >= dim_y: + raise ValueError( + f"component({component}) should be smaller than dim_y({dim_y})." + ) + else: + if component is not None: + raise ValueError( + f"component{component} should be set to None when dim_y({dim_y})=1." + ) + component = 0 + + if grad_y is None: + # `create_graph` of first order(jacobian) should be `True` in _Hessian. + grad_y = jacobian( + ys, xs, i=component, j=None, retain_graph=None, create_graph=True + ) + self.H = _Jacobian(grad_y, xs) + + def __call__( + self, + i: int = 0, + j: int = 0, + retain_graph: Optional[bool] = None, + create_graph: bool = True, + ): + """Returns H[`i`][`j`].""" + return self.H(i, j, retain_graph, create_graph) + + +class Hessians: + r"""Compute multiple Hessians. + + $$ + \rm Hessian(ys, xs, component, i, j) = \dfrac{\partial ys_{component}}{\partial xs_i \partial xs_j} + $$ + + A new instance will be created for a new pair of (output, input). For the (output, + input) pair that has been computed before, it will reuse the previous instance, + rather than creating a new one. + """ + + def __init__(self): + self.Hs = {} + + def __call__( + self, + ys: "paddle.Tensor", + xs: "paddle.Tensor", + component: Optional[int] = None, + i: int = 0, + j: int = 0, + grad_y: Optional["paddle.Tensor"] = None, + retain_graph: Optional[bool] = None, + create_graph: bool = True, + ) -> "paddle.Tensor": + """Compute hessian matrix for given ys and xs. + + Args: + ys (paddle.Tensor): Output tensor. + xs (paddle.Tensor): Input tensor. + component (Optional[int]): If `y` has the shape (batch_size, dim_y > 1), then `y[:, component]` + is used to compute the Hessian. Do not use if `y` has the shape (batch_size, + 1). Defaults to None. + i (int, optional): I-th input variable. Defaults to 0. + j (int, optional): J-th input variable. Defaults to 0. + grad_y (Optional[paddle.Tensor]): The gradient of `y` w.r.t. `xs`. Provide `grad_y` if known to avoid + duplicate computation. Defaults to None. + retain_graph (Optional[bool]): Whether to retain the forward graph which + is used to calculate the gradient. When it is True, the graph would + be retained, in which way users can calculate backward twice for the + same graph. When it is False, the graph would be freed. Default None, + which means it is equal to `create_graph`. + create_graph (bool, optional): Whether to create the gradient graphs of + the computing process. When it is True, higher order derivatives are + supported to compute; when it is False, the gradient graphs of the + computing process would be discarded. Default False. + + Returns: + paddle.Tensor: Hessian matrix. + + Examples: + >>> import paddle + >>> import ppsci + >>> x = paddle.randn([4, 3]) + >>> x.stop_gradient = False + >>> y = (x * x).sin() + >>> dy_dxx = ppsci.autodiff.hessian(y, x, component=0) + >>> print(dy_dxx.shape) + [4, 1] + """ + key = (ys, xs, component) + if key not in self.Hs: + self.Hs[key] = _Hessian(ys, xs, component=component, grad_y=grad_y) + return self.Hs[key](i, j, retain_graph, create_graph) + + def _clear(self): + """Clear cached Hessians.""" + self.Hs = {} + + +# Use high-order differentiation with singleton pattern for convenient +hessian: Callable[ + [ + "paddle.Tensor", + "paddle.Tensor", + Optional[int], + int, + int, + Optional["paddle.Tensor"], + Optional[bool], + bool, + ], + "paddle.Tensor", +] = Hessians() + + +def clear(): + """Clear cached Jacobians and Hessians. + + Examples: + >>> import paddle + >>> import ppsci + >>> x = paddle.randn([4, 3]) + >>> x.stop_gradient = False + >>> y = (x * x).sin() + >>> dy_dxx = ppsci.autodiff.hessian(y, x, component=0) + >>> ppsci.autodiff.clear() + >>> print(ppsci.autodiff.hessian.Hs) + {} + """ + jacobian._clear() + hessian._clear() diff --git a/examples/smc_reac/ppsci/constraint/__init__.py b/examples/smc_reac/ppsci/constraint/__init__.py new file mode 100644 index 0000000000..9179439436 --- /dev/null +++ b/examples/smc_reac/ppsci/constraint/__init__.py @@ -0,0 +1,86 @@ +# 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. + +from __future__ import annotations + +import copy + +from ppsci.constraint.base import Constraint +from ppsci.constraint.boundary_constraint import BoundaryConstraint +from ppsci.constraint.initial_constraint import InitialConstraint +from ppsci.constraint.integral_constraint import IntegralConstraint +from ppsci.constraint.interior_constraint import InteriorConstraint +from ppsci.constraint.periodic_constraint import PeriodicConstraint +from ppsci.constraint.supervised_constraint import SupervisedConstraint +from ppsci.loss import build_loss +from ppsci.utils import logger +from ppsci.utils import misc + +__all__ = [ + "Constraint", + "BoundaryConstraint", + "InitialConstraint", + "IntegralConstraint", + "InteriorConstraint", + "PeriodicConstraint", + "SupervisedConstraint", +] + + +def build_constraint(cfg, equation_dict, geom_dict): + """Build constraint(s). + + Args: + cfg (List[DictConfig]): Constraint config list. + equation_dict (Dct[str, Equation]): Equation(s) in dict. + geom_dict (Dct[str, Geometry]): Geometry(ies) in dict. + + Returns: + Dict[str, constraint]: Constraint(s) in dict. + """ + if cfg is None: + return None + cfg = copy.deepcopy(cfg) + global_dataloader_cfg = cfg["dataloader"] + constraint_cfg = cfg["content"] + + constraint_dict = misc.PrettyOrderedDict() + for _item in constraint_cfg: + constraint_cls = next(iter(_item.keys())) + _constraint_cfg = _item[constraint_cls] + constraint_name = _constraint_cfg.get("name", constraint_cls) + + # select equation + if isinstance(_constraint_cfg["output_expr"], str): + equation_name = _constraint_cfg.pop("output_expr") + _constraint_cfg["output_expr"] = equation_dict[equation_name].equations + + # select geometry + geom_name = _constraint_cfg.pop("geom") + _constraint_cfg["geom"] = geom_dict[geom_name] + + # update complete dataloader config + local_dataloader_cfg = _constraint_cfg["dataloader"] + local_dataloader_cfg.update(global_dataloader_cfg) + + # build loss + _constraint_cfg["loss"] = build_loss(_constraint_cfg["loss"]) + + # instantiate constraint + _constraint_cfg["dataloader_cfg"] = _constraint_cfg.pop("dataloader") + constraint_dict[constraint_name] = eval(constraint_cls)(**_constraint_cfg) + + logger.debug(str(constraint_dict[constraint_name])) + + return constraint_dict diff --git a/examples/smc_reac/ppsci/constraint/base.py b/examples/smc_reac/ppsci/constraint/base.py new file mode 100644 index 0000000000..c3b4c8a122 --- /dev/null +++ b/examples/smc_reac/ppsci/constraint/base.py @@ -0,0 +1,62 @@ +# 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. + +from __future__ import annotations + +from typing import TYPE_CHECKING +from typing import Any +from typing import Dict + +from paddle import io + +from ppsci import data + +if TYPE_CHECKING: + from ppsci import loss + + +class Constraint: + """Base class for constraint. + + Args: + dataset (io.Dataset): Dataset. + dataloader_cfg (Dict[str, Any]): Dataloader config. + loss (loss.Loss): Loss functor. + name (str): Name of constraint. + """ + + def __init__( + self, + dataset: io.Dataset, + dataloader_cfg: Dict[str, Any], + loss: "loss.Loss", + name: str, + ): + self.data_loader = data.build_dataloader(dataset, dataloader_cfg) + self.data_iter = iter(self.data_loader) + self.loss = loss + self.name = name + + def __str__(self): + return ", ".join( + [ + self.__class__.__name__, + f"name = {self.name}", + f"input_keys = {self.input_keys}", + f"output_keys = {self.output_keys}", + f"output_expr = {self.output_expr}", + f"label_dict = {self.label_dict}", + f"loss = {self.loss}", + ] + ) diff --git a/examples/smc_reac/ppsci/constraint/boundary_constraint.py b/examples/smc_reac/ppsci/constraint/boundary_constraint.py new file mode 100644 index 0000000000..8ac30fcb41 --- /dev/null +++ b/examples/smc_reac/ppsci/constraint/boundary_constraint.py @@ -0,0 +1,163 @@ +# 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. + +from __future__ import annotations + +from typing import TYPE_CHECKING +from typing import Any +from typing import Callable +from typing import Dict +from typing import Optional +from typing import Union + +import numpy as np +import sympy +from typing_extensions import Literal + +from ppsci import geometry +from ppsci.constraint import base +from ppsci.data import dataset + +if TYPE_CHECKING: + from ppsci import loss + + +class BoundaryConstraint(base.Constraint): + """Class for boundary constraint. + + Args: + output_expr (Dict[str, Callable]): Function in dict for computing output. + e.g. {"u_mul_v": lambda out: out["u"] * out["v"]} means the model output u + will be multiplied by model output v and the result will be named "u_mul_v". + label_dict (Dict[str, Union[float, Callable]]): Function in dict for computing + label, which will be a reference value to participate in the loss calculation. + geom (geometry.Geometry): Geometry where data sampled from. + dataloader_cfg (Dict[str, Any]): Dataloader config. + loss (loss.Loss): Loss functor. + random (Literal["pseudo", "Halton", "LHS"], optional): Random method for sampling data in + geometry. Defaults to "pseudo". + criteria (Optional[Callable]): Criteria for refining specified boundaries. + Defaults to None. + evenly (bool, optional): Whether to use evenly distribution sampling. + Defaults to False. + weight_dict (Optional[Dict[str, Union[float, Callable]]]): Define the weight of each + constraint variable. Defaults to None. + name (str, optional): Name of constraint object. Defaults to "BC". + + Examples: + >>> import ppsci + >>> rect = ppsci.geometry.Rectangle((0, 0), (1, 1)) + >>> bc = ppsci.constraint.BoundaryConstraint( + ... {"u": lambda out: out["u"]}, + ... {"u": 0}, + ... rect, + ... { + ... "dataset": "IterableNamedArrayDataset", + ... "iters_per_epoch": 1, + ... "batch_size": 16, + ... }, + ... ppsci.loss.MSELoss("mean"), + ... name="BC", + ... ) # doctest: +SKIP + """ + + def __init__( + self, + output_expr: Dict[str, Callable], + label_dict: Dict[str, Union[float, Callable]], + geom: geometry.Geometry, + dataloader_cfg: Dict[str, Any], + loss: "loss.Loss", + random: Literal["pseudo", "Halton", "LHS"] = "pseudo", + criteria: Optional[Callable] = None, + evenly: bool = False, + weight_dict: Optional[Dict[str, Union[float, Callable]]] = None, + name: str = "BC", + ): + self.label_dict = label_dict + self.input_keys = geom.dim_keys + self.output_keys = tuple(label_dict.keys()) + self.output_expr = { + k: v for k, v in output_expr.items() if k in self.output_keys + } + + if isinstance(criteria, str): + criteria = eval(criteria) + + # prepare input + input = geom.sample_boundary( + dataloader_cfg["batch_size"] * dataloader_cfg["iters_per_epoch"], + random, + criteria, + evenly, + ) + if "area" in input: + input["area"] *= dataloader_cfg["iters_per_epoch"] + + # prepare label + label = {} + for key, value in label_dict.items(): + if isinstance(value, (int, float)): + label[key] = np.full_like(next(iter(input.values())), value) + elif isinstance(value, sympy.Basic): + func = sympy.lambdify( + sympy.symbols(geom.dim_keys), + value, + [{"amax": lambda xy, axis: np.maximum(xy[0], xy[1])}, "numpy"], + ) + label[key] = func( + **{k: v for k, v in input.items() if k in geom.dim_keys} + ) + elif callable(value): + func = value + label[key] = func(input) + if isinstance(label[key], (int, float)): + label[key] = np.full_like(next(iter(input.values())), label[key]) + else: + raise NotImplementedError(f"type of {type(value)} is invalid yet.") + + # prepare weight + weight = None + if weight_dict is not None: + weight = {key: np.ones_like(next(iter(label.values()))) for key in label} + for key, value in weight_dict.items(): + if isinstance(value, (int, float)): + weight[key] = np.full_like(next(iter(label.values())), value) + elif isinstance(value, sympy.Basic): + func = sympy.lambdify( + [sympy.Symbol(k) for k in geom.dim_keys], + value, + [{"amax": lambda xy, _: np.maximum(xy[0], xy[1])}, "numpy"], + ) + weight[key] = func(**{k: input[k] for k in geom.dim_keys}) + elif callable(value): + func = value + weight[key] = func(input) + if isinstance(weight[key], (int, float)): + weight[key] = np.full_like( + next(iter(input.values())), weight[key] + ) + else: + raise NotImplementedError(f"type of {type(value)} is invalid yet.") + + # wrap input, label, weight into a dataset + if isinstance(dataloader_cfg["dataset"], str): + dataloader_cfg["dataset"] = {"name": dataloader_cfg["dataset"]} + dataloader_cfg["dataset"].update( + {"input": input, "label": label, "weight": weight} + ) + _dataset = dataset.build_dataset(dataloader_cfg["dataset"]) + + # construct dataloader with dataset and dataloader_cfg + super().__init__(_dataset, dataloader_cfg, loss, name) diff --git a/examples/smc_reac/ppsci/constraint/initial_constraint.py b/examples/smc_reac/ppsci/constraint/initial_constraint.py new file mode 100644 index 0000000000..63ae320993 --- /dev/null +++ b/examples/smc_reac/ppsci/constraint/initial_constraint.py @@ -0,0 +1,172 @@ +# 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. + +from __future__ import annotations + +from typing import TYPE_CHECKING +from typing import Any +from typing import Callable +from typing import Dict +from typing import Optional +from typing import Union + +import numpy as np +import sympy +from typing_extensions import Literal + +from ppsci import geometry +from ppsci.constraint import base +from ppsci.data import dataset + +if TYPE_CHECKING: + from ppsci import loss + + +class InitialConstraint(base.Constraint): + """Class for initial interior constraint. + + Args: + output_expr (Dict[str, Callable]): Function in dict for computing output. + e.g. {"u_mul_v": lambda out: out["u"] * out["v"]} means the model output u + will be multiplied by model output v and the result will be named "u_mul_v". + label_dict (Dict[str, Union[float, Callable]]): Function in dict for computing + label, which will be a reference value to participate in the loss calculation. + geom (geometry.TimeXGeometry): Geometry where data sampled from. + dataloader_cfg (Dict[str, Any]): Dataloader config. + loss (loss.Loss): Loss functor. + random (Literal["pseudo", "Halton", "LHS"], optional): Random method for sampling data in + geometry. Defaults to "pseudo". + criteria (Optional[Callable]): Criteria for refining specified boundaries. + Defaults to None. + evenly (bool, optional): Whether to use evenly distribution sampling. + Defaults to False. + weight_dict (Optional[Dict[str, Callable]]): Define the weight of each + constraint variable. Defaults to None. + compute_sdf_derivatives (Optional[bool]): Whether compute derivatives for SDF. + Defaults to False. + name (str, optional): Name of constraint object. Defaults to "IC". + + Examples: + >>> import ppsci + >>> rect = ppsci.geometry.TimeXGeometry( + ... ppsci.geometry.TimeDomain(0, 1), + ... ppsci.geometry.Rectangle((0, 0), (1, 1)), + ... ) + >>> ic = ppsci.constraint.InitialConstraint( + ... {"u": lambda out: out["u"]}, + ... {"u": 0}, + ... rect, + ... { + ... "dataset": "IterableNamedArrayDataset", + ... "iters_per_epoch": 1, + ... "batch_size": 16, + ... }, + ... ppsci.loss.MSELoss("mean"), + ... name="IC", + ... ) # doctest: +SKIP + """ + + def __init__( + self, + output_expr: Dict[str, Callable], + label_dict: Dict[str, Union[float, Callable]], + geom: geometry.TimeXGeometry, + dataloader_cfg: Dict[str, Any], + loss: "loss.Loss", + random: Literal["pseudo", "Halton", "LHS"] = "pseudo", + criteria: Optional[Callable] = None, + evenly: bool = False, + weight_dict: Optional[Dict[str, Callable]] = None, + compute_sdf_derivatives: bool = False, + name: str = "IC", + ): + self.label_dict = label_dict + self.input_keys = geom.dim_keys + self.output_keys = tuple(label_dict.keys()) + self.output_expr = { + k: v for k, v in output_expr.items() if k in self.output_keys + } + + if isinstance(criteria, str): + criteria = eval(criteria) + + # prepare input + input = geom.sample_initial_interior( + dataloader_cfg["batch_size"] * dataloader_cfg["iters_per_epoch"], + random, + criteria, + evenly, + compute_sdf_derivatives, + ) + if "area" in input: + input["area"] *= dataloader_cfg["iters_per_epoch"] + + # prepare label + label = {} + for key, value in label_dict.items(): + if isinstance(value, (int, float)): + label[key] = np.full_like(next(iter(input.values())), value) + elif isinstance(value, sympy.Basic): + func = sympy.lambdify( + sympy.symbols(geom.dim_keys), + value, + [{"amax": lambda xy, _: np.maximum(xy[0], xy[1])}, "numpy"], + ) + label[key] = func( + **{k: v for k, v in input.items() if k in geom.dim_keys} + ) + elif callable(value): + func = value + label[key] = func(input) + if isinstance(label[key], (int, float)): + label[key] = np.full_like(next(iter(input.values())), label[key]) + else: + raise NotImplementedError(f"type of {type(value)} is invalid yet.") + + # prepare weight + weight = None + if weight_dict is not None: + weight = {key: np.ones_like(next(iter(label.values()))) for key in label} + for key, value in weight_dict.items(): + if isinstance(value, (int, float)): + weight[key] = np.full_like(next(iter(label.values())), value) + elif isinstance(value, sympy.Basic): + func = sympy.lambdify( + sympy.symbols(geom.dim_keys), + value, + [{"amax": lambda xy, _: np.maximum(xy[0], xy[1])}, "numpy"], + ) + weight[key] = func( + **{k: v for k, v in input.items() if k in geom.dim_keys} + ) + elif callable(value): + func = value + weight[key] = func(input) + if isinstance(weight[key], (int, float)): + weight[key] = np.full_like( + next(iter(input.values())), weight[key] + ) + else: + raise NotImplementedError(f"type of {type(value)} is invalid yet.") + + # wrap input, label, weight into a dataset + if isinstance(dataloader_cfg["dataset"], str): + dataloader_cfg["dataset"] = {"name": dataloader_cfg["dataset"]} + dataloader_cfg["dataset"].update( + {"input": input, "label": label, "weight": weight} + ) + _dataset = dataset.build_dataset(dataloader_cfg["dataset"]) + + # construct dataloader with dataset and dataloader_cfg + super().__init__(_dataset, dataloader_cfg, loss, name) diff --git a/examples/smc_reac/ppsci/constraint/integral_constraint.py b/examples/smc_reac/ppsci/constraint/integral_constraint.py new file mode 100644 index 0000000000..19f6f8a1fe --- /dev/null +++ b/examples/smc_reac/ppsci/constraint/integral_constraint.py @@ -0,0 +1,178 @@ +# 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. + +from __future__ import annotations + +from typing import TYPE_CHECKING +from typing import Any +from typing import Callable +from typing import Dict +from typing import List +from typing import Optional +from typing import Union + +import numpy as np +import paddle +import sympy +from typing_extensions import Literal + +from ppsci import geometry +from ppsci.constraint import base +from ppsci.data import dataset +from ppsci.utils import misc + +if TYPE_CHECKING: + from ppsci import loss + + +class IntegralConstraint(base.Constraint): + """Class for integral constraint. + + Args: + output_expr (Dict[str, Callable]): Function in dict for computing output. + e.g. {"u_mul_v": lambda out: out["u"] * out["v"]} means the model output u + will be multiplied by model output v and the result will be named "u_mul_v". + label_dict (Dict[str, Union[float, Callable]]): Function in dict for computing + label, which will be a reference value to participate in the loss calculation. + geom (geometry.Geometry): Geometry where data sampled from. + dataloader_cfg (Dict[str, Any]): Dataloader config. + loss (loss.Loss): Loss functor. + random (Literal["pseudo", "Halton", "LHS"], optional): Random method for sampling data in + geometry. Defaults to "pseudo". + criteria (Optional[Callable]): Criteria for refining specified boundaries. + Defaults to None. + weight_dict (Optional[Dict[str, Callable]]): Define the weight of each + constraint variable. Defaults to None. + name (str, optional): Name of constraint object. Defaults to "IgC". + + Examples: + >>> import ppsci + >>> rect = ppsci.geometry.Rectangle((0, 0), (1, 1)) + >>> igc = ppsci.constraint.IntegralConstraint( + ... {"u": lambda out: out["u"]}, + ... {"u": 0}, + ... rect, + ... { + ... "dataset": "IterableNamedArrayDataset", + ... "iters_per_epoch": 1, + ... "batch_size": 16, + ... "integral_batch_size": 8, + ... }, + ... ppsci.loss.MSELoss("mean"), + ... name="IgC", + ... ) # doctest: +SKIP + """ + + def __init__( + self, + output_expr: Dict[str, Callable], + label_dict: Dict[str, Union[float, Callable]], + geom: geometry.Geometry, + dataloader_cfg: Dict[str, Any], + loss: "loss.Loss", + random: Literal["pseudo", "Halton", "LHS"] = "pseudo", + criteria: Optional[Callable] = None, + weight_dict: Optional[Dict[str, Callable]] = None, + name: str = "IgC", + ): + self.label_dict = label_dict + self.input_keys = geom.dim_keys + self.output_keys = tuple(label_dict.keys()) + self.output_expr = { + k: v for k, v in output_expr.items() if k in self.output_keys + } + + if isinstance(criteria, str): + criteria = eval(criteria) + + # prepare input + input_list: List[Dict[str, np.ndarray]] = [] + for _ in range( + dataloader_cfg["batch_size"] * dataloader_cfg["iters_per_epoch"] + ): + input = geom.sample_boundary( + dataloader_cfg["integral_batch_size"], random, criteria + ) + input_list.append(input) + input = misc.stack_dict_list(input_list) + # shape of each input is [batch_size, integral_batch_size, ndim] + + # prepare label + # shape of each label is [batch_size, ndim] + label = {} + for key, value in label_dict.items(): + if isinstance(value, (int, float)): + label[key] = np.full( + (next(iter(input.values())).shape[0], 1), + value, + paddle.get_default_dtype(), + ) + elif isinstance(value, sympy.Basic): + func = sympy.lambdify( + sympy.symbols(geom.dim_keys), + value, + [{"amax": lambda xy, _: np.maximum(xy[0], xy[1])}, "numpy"], + ) + label[key] = func( + **{k: v for k, v in input.items() if k in geom.dim_keys} + ) + elif callable(value): + func = value + label[key] = func(input) + if isinstance(label[key], (int, float)): + label[key] = np.full( + (next(iter(input.values())).shape[0], 1), + label[key], + paddle.get_default_dtype(), + ) + else: + raise NotImplementedError(f"type of {type(value)} is invalid yet.") + + # prepare weight + # shape of each weight is [batch_size, ndim] + weight = None + if weight_dict is not None: + weight = {key: np.ones_like(next(iter(label.values()))) for key in label} + for key, value in weight_dict.items(): + if isinstance(value, (int, float)): + weight[key] = np.full_like(next(iter(label.values())), value) + elif isinstance(value, sympy.Basic): + func = sympy.lambdify( + sympy.symbols(geom.dim_keys), + value, + [{"amax": lambda xy, _: np.maximum(xy[0], xy[1])}, "numpy"], + ) + weight[key] = func( + **{k: v for k, v in input.items() if k in geom.dim_keys} + ) + elif callable(value): + func = value + weight[key] = func(input) + if isinstance(weight[key], (int, float)): + weight[key] = np.full_like( + next(iter(input.values())), weight[key] + ) + else: + raise NotImplementedError(f"type of {type(value)} is invalid yet.") + + # wrap input, label, weight into a dataset + if isinstance(dataloader_cfg["dataset"], str): + dataloader_cfg["dataset"] = {"name": dataloader_cfg["dataset"]} + dataloader_cfg["dataset"].update( + {"input": input, "label": label, "weight": weight} + ) + _dataset = dataset.build_dataset(dataloader_cfg["dataset"]) + + # construct dataloader with dataset and dataloader_cfg + super().__init__(_dataset, dataloader_cfg, loss, name) diff --git a/examples/smc_reac/ppsci/constraint/interior_constraint.py b/examples/smc_reac/ppsci/constraint/interior_constraint.py new file mode 100644 index 0000000000..3c1eb7ed3f --- /dev/null +++ b/examples/smc_reac/ppsci/constraint/interior_constraint.py @@ -0,0 +1,174 @@ +# 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. + +from __future__ import annotations + +from typing import TYPE_CHECKING +from typing import Any +from typing import Callable +from typing import Dict +from typing import Optional +from typing import Union + +import numpy as np +import sympy +from typing_extensions import Literal + +from ppsci import geometry +from ppsci.constraint import base +from ppsci.data import dataset + +if TYPE_CHECKING: + from ppsci import loss + + +class InteriorConstraint(base.Constraint): + """Class for interior constraint. + + Args: + output_expr (Dict[str, Callable]): Function in dict for computing output. + e.g. {"u_mul_v": lambda out: out["u"] * out["v"]} means the model output u + will be multiplied by model output v and the result will be named "u_mul_v". + label_dict (Dict[str, Union[float, Callable]]): Function in dict for computing + label, which will be a reference value to participate in the loss calculation. + geom (geometry.Geometry): Geometry where data sampled from. + dataloader_cfg (Dict[str, Any]): Dataloader config. + loss (loss.Loss): Loss functor. + random (Literal["pseudo", "Halton", "LHS"], optional): Random method for sampling data in + geometry. Defaults to "pseudo". + criteria (Optional[Callable]): Criteria for refining specified boundaries. + Defaults to None. + evenly (bool, optional): Whether to use evenly distribution sampling. + Defaults to False. + weight_dict (Optional[Dict[str, Union[Callable, float]]]): Define the + weight of each constraint variable. Defaults to None. + compute_sdf_derivatives (Optional[bool]): Whether compute derivatives for SDF. + Defaults to False. + name (str, optional): Name of constraint object. Defaults to "EQ". + + Examples: + >>> import ppsci + >>> rect = ppsci.geometry.Rectangle((0, 0), (1, 1)) + >>> pde_constraint = ppsci.constraint.InteriorConstraint( + ... {"u": lambda out: out["u"]}, + ... {"u": 0}, + ... rect, + ... { + ... "dataset": "IterableNamedArrayDataset", + ... "iters_per_epoch": 1, + ... "batch_size": 16, + ... }, + ... ppsci.loss.MSELoss("mean"), + ... name="EQ", + ... ) # doctest: +SKIP + """ + + def __init__( + self, + output_expr: Dict[str, Callable], + label_dict: Dict[str, Union[float, Callable]], + geom: geometry.Geometry, + dataloader_cfg: Dict[str, Any], + loss: "loss.Loss", + random: Literal["pseudo", "Halton", "LHS"] = "pseudo", + criteria: Optional[Callable] = None, + evenly: bool = False, + weight_dict: Optional[Dict[str, Union[Callable, float]]] = None, + compute_sdf_derivatives: bool = False, + name: str = "EQ", + ): + self.label_dict = label_dict + self.input_keys = geom.dim_keys + self.output_keys = tuple(label_dict.keys()) + self.output_expr = { + k: v for k, v in output_expr.items() if k in self.output_keys + } + + if isinstance(criteria, str): + criteria = eval(criteria) + + # prepare input + input = geom.sample_interior( + dataloader_cfg["batch_size"] * dataloader_cfg["iters_per_epoch"], + random, + criteria, + evenly, + compute_sdf_derivatives, + ) + if "area" in input: + input["area"] *= dataloader_cfg["iters_per_epoch"] + + # prepare label + label = {} + for key, value in label_dict.items(): + if isinstance(value, (int, float)): + label[key] = np.full_like(next(iter(input.values())), value) + elif isinstance(value, sympy.Basic): + func = sympy.lambdify( + sympy.symbols(geom.dim_keys), + value, + [{"amax": lambda xy, _: np.maximum(xy[0], xy[1])}, "numpy"], + ) + label[key] = func( + **{k: v for k, v in input.items() if k in geom.dim_keys} + ) + elif callable(value): + func = value + label[key] = func(input) + if isinstance(label[key], (int, float)): + label[key] = np.full_like(next(iter(input.values())), label[key]) + else: + raise NotImplementedError(f"type of {type(value)} is invalid yet.") + + # prepare weight + weight = None + if weight_dict is not None: + weight = {key: np.ones_like(next(iter(label.values()))) for key in label} + for key, value in weight_dict.items(): + if isinstance(value, str): + if value == "sdf": + weight[key] = input["sdf"] + else: + raise NotImplementedError(f"string {value} is invalid yet.") + elif isinstance(value, (int, float)): + weight[key] = np.full_like(next(iter(label.values())), float(value)) + elif isinstance(value, sympy.Basic): + func = sympy.lambdify( + sympy.symbols(geom.dim_keys), + value, + [{"amax": lambda xy, _: np.maximum(xy[0], xy[1])}, "numpy"], + ) + weight[key] = func( + **{k: v for k, v in input.items() if k in geom.dim_keys} + ) + elif callable(value): + func = value + weight[key] = func(input) + if isinstance(weight[key], (int, float)): + weight[key] = np.full_like( + next(iter(input.values())), weight[key] + ) + else: + raise NotImplementedError(f"type of {type(value)} is invalid yet.") + + # wrap input, label, weight into a dataset + if isinstance(dataloader_cfg["dataset"], str): + dataloader_cfg["dataset"] = {"name": dataloader_cfg["dataset"]} + dataloader_cfg["dataset"].update( + {"input": input, "label": label, "weight": weight} + ) + _dataset = dataset.build_dataset(dataloader_cfg["dataset"]) + + # construct dataloader with dataset and dataloader_cfg + super().__init__(_dataset, dataloader_cfg, loss, name) diff --git a/examples/smc_reac/ppsci/constraint/periodic_constraint.py b/examples/smc_reac/ppsci/constraint/periodic_constraint.py new file mode 100644 index 0000000000..cb5fc1a332 --- /dev/null +++ b/examples/smc_reac/ppsci/constraint/periodic_constraint.py @@ -0,0 +1,169 @@ +# 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. + +from __future__ import annotations + +from typing import TYPE_CHECKING +from typing import Any +from typing import Callable +from typing import Dict +from typing import Optional +from typing import Union + +import numpy as np +import paddle +import sympy +from typing_extensions import Literal + +from ppsci import geometry +from ppsci.constraint import base +from ppsci.data import dataset + +if TYPE_CHECKING: + from ppsci import loss + + +class PeriodicConstraint(base.Constraint): + """Class for periodic constraint. + + Args: + output_expr (Dict[str, Callable]): Function in dict for computing output. + e.g. {"u_mul_v": lambda out: out["u"] * out["v"]} means the model output u + will be multiplied by model output v and the result will be named "u_mul_v". + label_dict (Dict[str, Union[float, Callable]]): Function in dict for computing + label, which will be a reference value to participate in the loss calculation. + geom (geometry.Geometry): Geometry where data sampled from. + dataloader_cfg (Dict[str, Any]): Dataloader config. + periodic_key (str): Name of dimension which periodic constraint applied to. + loss (loss.Loss): Loss functor. + random (Literal["pseudo", "Halton", "LHS"], optional): Random method for sampling data in + geometry. Defaults to "pseudo". + criteria (Optional[Callable]): Criteria for refining specified boundaries. + Defaults to None. + evenly (bool, optional): Whether to use evenly distribution sampling. + Defaults to False. + weight_dict (Optional[Dict[str, Callable]]): Define the weight of each + constraint variable. Defaults to None. + name (str, optional): Name of constraint object. Defaults to "PeriodicBC". + """ + + def __init__( + self, + output_expr: Dict[str, Callable], + label_dict: Dict[str, Union[float, Callable]], + geom: geometry.Geometry, + periodic_key: str, + dataloader_cfg: Dict[str, Any], + loss: "loss.Loss", + random: Literal["pseudo", "Halton", "LHS"] = "pseudo", + criteria: Optional[Callable] = None, + evenly: bool = False, + weight_dict: Optional[Dict[str, Callable]] = None, + name: str = "PeriodicBC", + ): + self.input_keys = geom.dim_keys + self.output_keys = tuple(output_expr.keys()) + self.output_expr = { + k: v for k, v in output_expr.items() if k in self.output_keys + } + + if isinstance(criteria, str): + criteria = eval(criteria) + + if dataloader_cfg["batch_size"] % 2 > 0: + raise ValueError( + f"batch_size({dataloader_cfg['sampler']['batch_size']}) " + "should be positive and even when using PeriodicConstraint" + ) + if dataloader_cfg.get("shuffle", False): + raise ValueError( + f"shuffle({dataloader_cfg['sampler']['batch_size']}) " + "should be False when using PeriodicConstraint" + ) + + # prepare input + _bs_half = dataloader_cfg["batch_size"] // 2 + input = geom.sample_boundary( + _bs_half * dataloader_cfg["iters_per_epoch"], + random, + criteria, + evenly, + ) + if "area" in input: + input["area"] *= dataloader_cfg["iters_per_epoch"] + + input_periodic = geom.periodic_point( + input, + geom.geometry.dim_keys.index(periodic_key) + if isinstance(geom, geometry.TimeXGeometry) + else geom.dim_keys.index(periodic_key), + ) + # concatenate original data next to periodic data, i.e. + # [orignal1, periodic1, orignal2, periodic2, ..., orignalN, periodicN] + mixed_input = {} + for key in input: + mixed_input[key] = [] + for iter_id in range(dataloader_cfg["iters_per_epoch"]): + mixed_input[key].append( + input[key][iter_id * _bs_half : (iter_id + 1) * _bs_half] + ) + mixed_input[key].append( + input_periodic[key][iter_id * _bs_half : (iter_id + 1) * _bs_half] + ) + mixed_input[key] = np.vstack(mixed_input[key]) + + # prepare label, keep label the same shape as input_periodic + label = {} + for key, value in label_dict.items(): + # set all label's to zero for dummy data. + label[key] = np.full( + (next(iter(mixed_input.values())).shape[0], 1), + 0, + paddle.get_default_dtype(), + ) + + # # prepare weight, keep weight the same shape as input_periodic + weight = None + if weight_dict is not None: + weight = {key: np.ones_like(next(iter(label.values()))) for key in label} + for key, value in weight_dict.items(): + if isinstance(value, (int, float)): + weight[key] = np.full_like(next(iter(label.values())), value) + elif isinstance(value, sympy.Basic): + func = sympy.lambdify( + [sympy.Symbol(k) for k in geom.dim_keys], + value, + [{"amax": lambda xy, _: np.maximum(xy[0], xy[1])}, "numpy"], + ) + weight[key] = func(**{k: mixed_input[k] for k in geom.dim_keys}) + elif callable(value): + func = value + weight[key] = func(mixed_input) + if isinstance(weight[key], (int, float)): + weight[key] = np.full_like( + next(iter(mixed_input.values())), weight[key] + ) + else: + raise NotImplementedError(f"type of {type(value)} is invalid yet.") + + # wrap input, label, weight into a dataset + if isinstance(dataloader_cfg["dataset"], str): + dataloader_cfg["dataset"] = {"name": dataloader_cfg["dataset"]} + dataloader_cfg["dataset"].update( + {"input": mixed_input, "label": label, "weight": weight} + ) + _dataset = dataset.build_dataset(dataloader_cfg["dataset"]) + + # construct dataloader with dataset and dataloader_cfg + super().__init__(_dataset, dataloader_cfg, loss, name) diff --git a/examples/smc_reac/ppsci/constraint/supervised_constraint.py b/examples/smc_reac/ppsci/constraint/supervised_constraint.py new file mode 100644 index 0000000000..84b8816222 --- /dev/null +++ b/examples/smc_reac/ppsci/constraint/supervised_constraint.py @@ -0,0 +1,92 @@ +# 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. + +from __future__ import annotations + +from typing import TYPE_CHECKING +from typing import Any +from typing import Callable +from typing import Dict +from typing import Optional + +from ppsci.constraint import base +from ppsci.data import dataset + +if TYPE_CHECKING: + from ppsci import loss + + +class SupervisedConstraint(base.Constraint): + """Class for supervised constraint. + + Args: + dataloader_cfg (Dict[str, Any]): Dataloader config. + loss (loss.Loss): Loss functor. + output_expr (Optional[Dict[str, Callable]]): List of label expression. + Defaults to None. + name (str, optional): Name of constraint object. Defaults to "Sup". + + Examples: + >>> import ppsci + >>> bc_sup = ppsci.constraint.SupervisedConstraint( + ... { + ... "dataset": { + ... "name": "IterableCSVDataset", + ... "file_path": "/path/to/file.csv", + ... "input_keys": ("x", "y"), + ... "label_keys": ("u", "v"), + ... }, + ... }, + ... ppsci.loss.MSELoss("mean"), + ... name="bc_sup", + ... ) # doctest: +SKIP + """ + + def __init__( + self, + dataloader_cfg: Dict[str, Any], + loss: "loss.Loss", + output_expr: Optional[Dict[str, Callable]] = None, + name: str = "Sup", + ): + # build dataset + _dataset = dataset.build_dataset(dataloader_cfg["dataset"]) + + self.input_keys = _dataset.input_keys + self.output_keys = ( + tuple(output_expr.keys()) + if output_expr is not None + else _dataset.label_keys + ) + + self.output_expr = output_expr + if self.output_expr is None: + self.output_expr = { + key: (lambda out, k=key: out[k]) for key in self.output_keys + } + + # construct dataloader with dataset and dataloader_cfg + super().__init__(_dataset, dataloader_cfg, loss, name) + + def __str__(self): + return ", ".join( + [ + self.__class__.__name__, + f"name = {self.name}", + f"input_keys = {self.input_keys}", + f"output_keys = {self.output_keys}", + f"output_expr = {self.output_expr}", + f"loss = {self.loss}", + ] + ) diff --git a/examples/smc_reac/ppsci/data/__init__.py b/examples/smc_reac/ppsci/data/__init__.py new file mode 100644 index 0000000000..41ea77147d --- /dev/null +++ b/examples/smc_reac/ppsci/data/__init__.py @@ -0,0 +1,205 @@ +# 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. + +import copy +import random +from functools import partial +from typing import Callable +from typing import Optional + +import numpy as np +import paddle.distributed as dist +from paddle import device +from paddle import io + +from ppsci.data import dataloader +from ppsci.data import dataset +from ppsci.data import process +from ppsci.data.process import batch_transform +from ppsci.data.process import transform +from ppsci.utils import logger + +__all__ = [ + "dataset", + "process", + "dataloader", + "build_dataloader", + "transform", + "batch_transform", +] + + +def worker_init_fn(worker_id: int, num_workers: int, rank: int, base_seed: int) -> None: + """Callback function on each worker subprocess after seeding and before data loading. + + Args: + worker_id (int): Worker id in [0, num_workers - 1]. + num_workers (int): Number of subprocesses to use for data loading. + rank (int): Rank of process in distributed environment. If in non-distributed + environment, it is a constant number `0`. + base_seed (int): Base random seed. + """ + # The seed of each worker equals to: user_seed + num_worker * rank + worker_id + worker_seed = base_seed + num_workers * rank + worker_id + np.random.seed(worker_seed) + random.seed(worker_seed) + + +def build_dataloader(_dataset, cfg): + world_size = dist.get_world_size() + # just return IterableDataset as dataloader + if isinstance(_dataset, io.IterableDataset): + return _dataset + + cfg = copy.deepcopy(cfg) + + # build sampler + sampler_cfg = cfg.pop("sampler", None) + if sampler_cfg is not None: + batch_sampler_cls = sampler_cfg.pop("name") + + if batch_sampler_cls == "BatchSampler": + if world_size > 1: + batch_sampler_cls = "DistributedBatchSampler" + logger.warning( + f"Automatically use 'DistributedBatchSampler' instead of " + f"'BatchSampler' when world_size({world_size}) > 1." + ) + + sampler_cfg["batch_size"] = cfg["batch_size"] + batch_sampler = getattr(io, batch_sampler_cls)(_dataset, **sampler_cfg) + else: + batch_sampler_cls = "BatchSampler" + if world_size > 1: + batch_sampler_cls = "DistributedBatchSampler" + logger.warning( + f"Automatically use 'DistributedBatchSampler' instead of " + f"'BatchSampler' when world_size({world_size}) > 1." + ) + batch_sampler = getattr(io, batch_sampler_cls)( + _dataset, + batch_size=cfg["batch_size"], + shuffle=False, + drop_last=False, + ) + logger.message( + "'shuffle' and 'drop_last' are both set to False in default as sampler config is not specified." + ) + + # build collate_fn if specified + batch_transforms_cfg = cfg.pop("batch_transforms", None) + collate_fn: Optional[Callable] = cfg.pop("collate_fn", None) + if isinstance(batch_transforms_cfg, (list, tuple)): + collate_fn = batch_transform.build_batch_transforms( + batch_transforms_cfg, collate_fn + ) + + # build init function + _DEFAULT_NUM_WORKERS = 1 + _DEFAULT_SEED = 42 + init_fn = partial( + worker_init_fn, + num_workers=cfg.get("num_workers", _DEFAULT_NUM_WORKERS), + rank=dist.get_rank(), + base_seed=cfg.get("seed", _DEFAULT_SEED), + ) + + # build dataloader + if getattr(_dataset, "use_pgl", False): + # Use special dataloader from "Paddle Graph Learning" toolkit. + try: + from pgl.utils import data as pgl_data + except ModuleNotFoundError as e: + logger.error("Please install pgl with `pip install pgl`.") + raise ModuleNotFoundError(str(e)) + + if collate_fn is None: + collate_fn = batch_transform.default_collate_fn + dataloader_ = pgl_data.Dataloader( + dataset=_dataset, + batch_size=cfg["batch_size"], + drop_last=sampler_cfg.get("drop_last", False), + shuffle=sampler_cfg.get("shuffle", False), + num_workers=cfg.get("num_workers", _DEFAULT_NUM_WORKERS), + collate_fn=collate_fn, + ) + elif getattr(_dataset, "use_graph_grid_mesh", False): + # Use special dataloader `GridMeshAtmosphericDataset`. + + if collate_fn is None: + collate_fn = batch_transform.default_collate_fn + dataloader_ = io.DataLoader( + dataset=_dataset, + places=device.get_device(), + batch_sampler=batch_sampler, + collate_fn=collate_fn, + num_workers=cfg.get("num_workers", _DEFAULT_NUM_WORKERS), + use_shared_memory=cfg.get("use_shared_memory", False), + worker_init_fn=init_fn, + ) + else: + if ( + cfg.get("auto_collation", not getattr(_dataset, "batch_index", False)) + is False + ): + if "transforms" in cfg["dataset"] and "auto_collation" not in cfg: + logger.warning( + "'transforms' and batch indexing(auto_collation=False) are both " + "enabled. If you do want to apply transforms to the batch samples, " + "please explicitly set 'auto_collation' to False in dataloader_cfg;" + " otherwise, the 'transforms' will be retained, but batch indexing " + "will be disabled." + ) + else: + # 1. wrap batch_sampler again into BatchSampler for disabling auto collation, + # which can speed up the process of batch samples indexing from dataset. See + # details at: https://discuss.pytorch.org/t/efficiency-of-dataloader-and-collate-for-large-array-like-datasets/59569/8 + batch_sampler = io.BatchSampler(sampler=batch_sampler, batch_size=1) + if collate_fn is not None: + raise NotImplementedError( + "Detected collate_fn is not None for 'batch_transforms' might " + "be specified in 'dataloader_cfg', which is not supported yet " + "with 'auto_collation' is False at the same time" + ) + # 2. disable auto collation by given identity collate_fn which return the first + # (also the only) batch data in batch list, or there will be a redundant + # axis at the first dimension returned by dataloader. This step is necessary + # because paddle do not support 'sampler' as instantiation argument of 'io.DataLoader' + collate_fn = lambda batch: batch[0] # noqa: E731 + _DEFAULT_NUM_WORKERS = 0 + logger.info( + "Auto collation is disabled and set num_workers to " + f"{_DEFAULT_NUM_WORKERS} to speed up batch sampling." + ) + + dataloader_ = io.DataLoader( + dataset=_dataset, + places=device.get_device(), + batch_sampler=batch_sampler, + collate_fn=collate_fn, + num_workers=cfg.get("num_workers", _DEFAULT_NUM_WORKERS), + use_shared_memory=cfg.get("use_shared_memory", False), + worker_init_fn=init_fn, + # TODO: Do not enable 'persistent_workers' below for + # 'IndexError: pop from empty list ...' will be raised in certain cases + # persistent_workers=cfg.get("num_workers", _DEFAULT_NUM_WORKERS) > 0, + ) + + if len(dataloader_) == 0: + raise ValueError( + f"batch_size({sampler_cfg['batch_size']}) should not bigger than number of " + f"samples({len(_dataset)}) when drop_last is {sampler_cfg.get('drop_last', False)}." + ) + + return dataloader_ diff --git a/examples/smc_reac/ppsci/data/dataloader.py b/examples/smc_reac/ppsci/data/dataloader.py new file mode 100644 index 0000000000..4c01da873e --- /dev/null +++ b/examples/smc_reac/ppsci/data/dataloader.py @@ -0,0 +1,47 @@ +# 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. + +from __future__ import annotations + +from typing import Union + +from paddle import io + + +class InfiniteDataLoader: + """A wrapper for infinite dataloader. + + Args: + dataloader (Union[io.DataLoader, io.IterableDataset]): A finite and iterable loader or iterable dataset to be wrapped. + """ + + def __init__(self, dataloader: Union[io.DataLoader, io.IterableDataset]): + self.dataloader = dataloader + if isinstance(dataloader, io.DataLoader): + self.dataset = dataloader.dataset + elif isinstance(dataloader, io.IterableDataset): + self.dataset = dataloader + else: + raise TypeError( + f"dataloader should be io.DataLoader or io.IterableDataset, but got {type(dataloader)}" + ) + + def __iter__(self): + while True: + dataloader_iter = iter(self.dataloader) + for batch in dataloader_iter: + yield batch + + def __len__(self): + return len(self.dataloader) diff --git a/examples/smc_reac/ppsci/data/dataset/__init__.py b/examples/smc_reac/ppsci/data/dataset/__init__.py new file mode 100644 index 0000000000..9ece354700 --- /dev/null +++ b/examples/smc_reac/ppsci/data/dataset/__init__.py @@ -0,0 +1,118 @@ +# 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. + +import copy +from typing import TYPE_CHECKING + +from ppsci.data.dataset.airfoil_dataset import MeshAirfoilDataset +from ppsci.data.dataset.array_dataset import ChipHeatDataset +from ppsci.data.dataset.array_dataset import ContinuousNamedArrayDataset +from ppsci.data.dataset.array_dataset import IterableNamedArrayDataset +from ppsci.data.dataset.array_dataset import NamedArrayDataset +from ppsci.data.dataset.atmospheric_dataset import GridMeshAtmosphericDataset +from ppsci.data.dataset.cgcnn_dataset import CIFData as CGCNNDataset +from ppsci.data.dataset.csv_dataset import CSVDataset +from ppsci.data.dataset.csv_dataset import IterableCSVDataset +from ppsci.data.dataset.cylinder_dataset import MeshCylinderDataset +from ppsci.data.dataset.darcyflow_dataset import DarcyFlowDataset +from ppsci.data.dataset.dgmr_dataset import DGMRDataset +from ppsci.data.dataset.drivaernet_dataset import DrivAerNetDataset +from ppsci.data.dataset.drivaernetplusplus_dataset import DrivAerNetPlusPlusDataset +from ppsci.data.dataset.enso_dataset import ENSODataset +from ppsci.data.dataset.era5_dataset import ERA5Dataset +from ppsci.data.dataset.era5_dataset import ERA5SampledDataset +from ppsci.data.dataset.ext_moe_enso_dataset import ExtMoEENSODataset +from ppsci.data.dataset.fwi_dataset import FWIDataset +from ppsci.data.dataset.ifm_moe_dataset import IFMMoeDataset +from ppsci.data.dataset.mat_dataset import IterableMatDataset +from ppsci.data.dataset.mat_dataset import MatDataset +from ppsci.data.dataset.moflow_dataset import MOlFLOWDataset +from ppsci.data.dataset.mrms_dataset import MRMSDataset +from ppsci.data.dataset.mrms_dataset import MRMSSampledDataset +from ppsci.data.dataset.npz_dataset import IterableNPZDataset +from ppsci.data.dataset.npz_dataset import NPZDataset +from ppsci.data.dataset.pems_dataset import PEMSDataset +from ppsci.data.dataset.radar_dataset import RadarDataset +from ppsci.data.dataset.sevir_dataset import SEVIRDataset +from ppsci.data.dataset.spherical_swe_dataset import SphericalSWEDataset +from ppsci.data.dataset.trphysx_dataset import CylinderDataset +from ppsci.data.dataset.trphysx_dataset import LorenzDataset +from ppsci.data.dataset.trphysx_dataset import RosslerDataset +from ppsci.data.dataset.vtu_dataset import VtuDataset +from ppsci.data.process import transform +from ppsci.utils import logger + +if TYPE_CHECKING: + from paddle import io + +__all__ = [ + "IterableNamedArrayDataset", + "NamedArrayDataset", + "ContinuousNamedArrayDataset", + "ChipHeatDataset", + "CSVDataset", + "IterableCSVDataset", + "ERA5Dataset", + "ERA5SampledDataset", + "GridMeshAtmosphericDataset", + "IterableMatDataset", + "MatDataset", + "MRMSDataset", + "MRMSSampledDataset", + "IterableNPZDataset", + "NPZDataset", + "PEMSDataset", + "CylinderDataset", + "LorenzDataset", + "RadarDataset", + "RosslerDataset", + "VtuDataset", + "DGMRDataset", + "MeshAirfoilDataset", + "MeshCylinderDataset", + "DarcyFlowDataset", + "SphericalSWEDataset", + "ENSODataset", + "ExtMoEENSODataset", + "SEVIRDataset", + "MOlFLOWDataset", + "build_dataset", + "CGCNNDataset", + "FWIDataset", + "DrivAerNetDataset", + "DrivAerNetPlusPlusDataset", + "IFMMoeDataset", +] + + +def build_dataset(cfg) -> "io.Dataset": + """Build dataset + + Args: + cfg (List[DictConfig]): Dataset config list. + + Returns: + Dict[str, io.Dataset]: dataset. + """ + cfg = copy.deepcopy(cfg) + + dataset_cls = cfg.pop("name") + if "transforms" in cfg: + cfg["transforms"] = transform.build_transforms(cfg.pop("transforms")) + + dataset = eval(dataset_cls)(**cfg) + + logger.debug(str(dataset)) + + return dataset diff --git a/examples/smc_reac/ppsci/data/dataset/airfoil_dataset.py b/examples/smc_reac/ppsci/data/dataset/airfoil_dataset.py new file mode 100644 index 0000000000..2a249104f7 --- /dev/null +++ b/examples/smc_reac/ppsci/data/dataset/airfoil_dataset.py @@ -0,0 +1,241 @@ +# 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. + +from __future__ import annotations + +import os +import pickle +from os import path as osp +from typing import Dict +from typing import List +from typing import Tuple + +import numpy as np +import paddle +from paddle import io + +try: + import pgl +except ModuleNotFoundError: + pass + +SU2_SHAPE_IDS = { + "line": 3, + "triangle": 5, + "quad": 9, +} + + +# HACK: Simplify code below +def _get_mesh_graph( + mesh_filename: str, dtype: np.dtype = np.float32 +) -> Tuple[np.ndarray, np.ndarray, List[List[List[int]]], Dict[str, List[List[int]]]]: + def get_rhs(s: str) -> str: + return s.split("=")[-1] + + marker_dict = {} + with open(mesh_filename) as f: + for line in f: + if line.startswith("NPOIN"): + num_points = int(get_rhs(line)) + mesh_points = [ + [float(p) for p in f.readline().split()[:2]] + for _ in range(num_points) + ] + nodes = np.array(mesh_points, dtype=dtype) + elif line.startswith("NMARK"): + num_markers = int(get_rhs(line)) + for _ in range(num_markers): + line = f.readline() + assert line.startswith("MARKER_TAG") + marker_tag = get_rhs(line).strip() + num_elems = int(get_rhs(f.readline())) + marker_elems = [ + [int(e) for e in f.readline().split()[-2:]] + for _ in range(num_elems) + ] + marker_dict[marker_tag] = marker_elems + elif line.startswith("NELEM"): + edges = [] + triangles = [] + quads = [] + num_edges = int(get_rhs(line)) + for _ in range(num_edges): + elem = [int(p) for p in f.readline().split()] + if elem[0] == SU2_SHAPE_IDS["triangle"]: + n = 3 + triangles.append(elem[1 : 1 + n]) + elif elem[0] == SU2_SHAPE_IDS["quad"]: + n = 4 + quads.append(elem[1 : 1 + n]) + else: + raise NotImplementedError + elem = elem[1 : 1 + n] + edges += [[elem[i], elem[(i + 1) % n]] for i in range(n)] + edges = np.array(edges, dtype=np.compat.long).transpose() + elems = [triangles, quads] + return nodes, edges, elems, marker_dict + + +class MeshAirfoilDataset(io.Dataset): + """Dataset for `MeshAirfoil`. + + Args: + input_keys (Tuple[str, ...]): Name of input data. + label_keys (Tuple[str, ...]): Name of label data. + data_dir (str): Directory of MeshAirfoil data. + mesh_graph_path (str): Path of mesh graph. + transpose_edges (bool, optional): Whether transpose the edges array from (2, num_edges) to (num_edges, 2) for convenient of slicing. + + Examples: + >>> import ppsci + >>> dataset = ppsci.data.dataset.MeshAirfoilDataset( + ... "input_keys": ("input",), + ... "label_keys": ("output",), + ... "data_dir": "/path/to/MeshAirfoilDataset", + ... "mesh_graph_path": "/path/to/file.su2", + ... "transpose_edges": False, + ... ) # doctest: +SKIP + """ + + # Whether support batch indexing for speeding up fetching process. + batch_index: bool = False + + use_pgl: bool = True + + def __init__( + self, + input_keys: Tuple[str, ...], + label_keys: Tuple[str, ...], + data_dir: str, + mesh_graph_path: str, + transpose_edges: bool = False, + ): + self.input_keys = input_keys + self.label_keys = label_keys + self.data_dir = data_dir + self.file_list = os.listdir(self.data_dir) + self.len = len(self.file_list) + self.mesh_graph = _get_mesh_graph(mesh_graph_path) + + with open(osp.join(osp.dirname(self.data_dir), "train_max_min.pkl"), "rb") as f: + self.normalization_factors = pickle.load(f) + + self.nodes = self.mesh_graph[0] + self.edges = self.mesh_graph[1] + if transpose_edges: + self.edges = self.edges.transpose([1, 0]) + self.elems_list = self.mesh_graph[2] + self.marker_dict = self.mesh_graph[3] + self.node_markers = np.full([self.nodes.shape[0], 1], fill_value=-1) + for i, (marker_tag, marker_elems) in enumerate(self.marker_dict.items()): + for elem in marker_elems: + self.node_markers[elem[0]] = i + self.node_markers[elem[1]] = i + + self.raw_graphs = [self.get(i) for i in range(len(self))] + + def __len__(self): + return self.len + + def __getitem__(self, idx): + return ( + { + self.input_keys[0]: self.raw_graphs[idx], + }, + { + self.label_keys[0]: self.raw_graphs[idx], + }, + None, + ) + + def get(self, idx): + with open(osp.join(self.data_dir, self.file_list[idx]), "rb") as f: + fields = pickle.load(f) + fields = self._preprocess(fields) + aoa, reynolds, mach = self._get_params_from_name(self.file_list[idx]) + # aoa = aoa + mach_or_reynolds = mach if reynolds is None else reynolds + # mach_or_reynolds = mach_or_reynolds + norm_aoa = aoa / 10 + norm_mach_or_reynolds = ( + mach_or_reynolds if reynolds is None else (mach_or_reynolds - 1.5e6) / 1.5e6 + ) + + nodes = np.concatenate( + [ + self.nodes, + np.repeat(a=norm_aoa, repeats=self.nodes.shape[0])[:, np.newaxis], + np.repeat(a=norm_mach_or_reynolds, repeats=self.nodes.shape[0])[ + :, np.newaxis + ], + self.node_markers, + ], + axis=-1, + ).astype(paddle.get_default_dtype()) + + data = pgl.Graph( + num_nodes=nodes.shape[0], + edges=self.edges, + ) + data.x = nodes + data.y = fields + data.pos = self.nodes + data.edge_index = self.edges + + sender = data.x[data.edge_index[0]] + receiver = data.x[data.edge_index[1]] + relation_pos = sender[:, 0:2] - receiver[:, 0:2] + post = np.linalg.norm(relation_pos, ord=2, axis=1, keepdims=True).astype( + paddle.get_default_dtype() + ) + data.edge_attr = post + std_epsilon = [1e-8] + a = np.mean(data.edge_attr, axis=0) + b = data.edge_attr.std(axis=0) + b = np.maximum(b, std_epsilon).astype(paddle.get_default_dtype()) + data.edge_attr = (data.edge_attr - a) / b + data.aoa = aoa + data.norm_aoa = norm_aoa + data.mach_or_reynolds = mach_or_reynolds + data.norm_mach_or_reynolds = norm_mach_or_reynolds + return data + + def _preprocess(self, tensor_list, stack_output=True): + data_max, data_min = self.normalization_factors + normalized_tensors = [] + for i in range(len(tensor_list)): + normalized = (tensor_list[i] - data_min[i]) / ( + data_max[i] - data_min[i] + ) * 2 - 1 + normalized_tensors.append(normalized) + if stack_output: + normalized_tensors = np.stack(normalized_tensors, axis=1) + return normalized_tensors + + def _get_params_from_name(self, filename): + s = filename.rsplit(".", 1)[0].split("_") + aoa = np.array(s[s.index("aoa") + 1])[np.newaxis].astype( + paddle.get_default_dtype() + ) + reynolds = s[s.index("re") + 1] + reynolds = ( + np.array(reynolds)[np.newaxis].astype(paddle.get_default_dtype()) + if reynolds != "None" + else None + ) + mach = np.array(s[s.index("mach") + 1])[np.newaxis].astype( + paddle.get_default_dtype() + ) + return aoa, reynolds, mach diff --git a/examples/smc_reac/ppsci/data/dataset/array_dataset.py b/examples/smc_reac/ppsci/data/dataset/array_dataset.py new file mode 100644 index 0000000000..65ac0ca7c4 --- /dev/null +++ b/examples/smc_reac/ppsci/data/dataset/array_dataset.py @@ -0,0 +1,390 @@ +# 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. + +from __future__ import annotations + +from typing import Callable +from typing import Dict +from typing import Optional + +import numpy as np +import paddle +from paddle import distributed as dist +from paddle import io +from paddle import vision + +from ppsci.utils import logger + + +def _group_array_into_ranks( + data: Optional[np.ndarray], rank: int, world_size: int +) -> Optional[np.ndarray]: + """ + Group data into different ranks. For example, if data is [1, 2, 3, 4, 5, 6, 7, 8, 9] and + world_size is 3, then the result will be rank0: [1, 4, 7], rank1: [2, 5, 8], rank2: [3, 6, 9]. + + Args: + data (Optional[np.ndarray]): Data to be grouped, can be np.ndarray or None. + rank (int): Rank number. + world_size (int): Number of workers. + + Returns: + np.ndarray: Grouped data. + """ + if data is None: + # skip grouping if data is None + return None + + # check if data can be grouped evenly into different ranks + if len(data) < world_size: + raise ValueError( + f"Length of data to be grouped{len(data)} must be larger than world_size." + ) + if len(data) % world_size != 0: + raise ValueError( + f"Length of data to be grouped{len(data)} must be divisible by world_size." + ) + + return data[rank::world_size] + + +def _group_dict_into_ranks( + data_dict: Optional[Dict[str, Optional[np.ndarray]]], rank: int, world_size: int +) -> Optional[Dict[str, Optional[np.ndarray]]]: + """ + Group data dict into different ranks for each key-value pair. + + Args: + data_dict (Dict[str, Optional[np.ndarray]]): Data to be grouped, can be Dict[str, Optional[np.ndarray]] or None. + rank (int): Rank number. + world_size (int): Number of workers. + + Returns: + Optional[Dict[str, Optional[np.ndarray]]]: Grouped data dict. + """ + + if data_dict is None: + return data_dict + + return { + k: _group_array_into_ranks(v, rank, world_size) for k, v in data_dict.items() + } + + +class NamedArrayDataset(io.Dataset): + """Class for Named Array Dataset. + + Args: + input (Dict[str, np.ndarray]): Input dict. + label (Optional[Dict[str, np.ndarray]]): Label dict. Defaults to None. + weight (Optional[Dict[str, np.ndarray]]): Weight dict. Defaults to None. + transforms (Optional[vision.Compose]): Compose object contains sample wise + transform(s). Defaults to None. + + Examples: + >>> import ppsci + >>> input = {"x": np.random.randn(100, 1)} + >>> output = {"u": np.random.randn(100, 1)} + >>> weight = {"u": np.random.randn(100, 1)} + >>> dataset = ppsci.data.dataset.NamedArrayDataset(input, output, weight) + """ + + # Whether support batch indexing for speeding up fetching process. + batch_index: bool = True + + def __init__( + self, + input: Dict[str, np.ndarray], + label: Optional[Dict[str, np.ndarray]] = None, + weight: Optional[Dict[str, np.ndarray]] = None, + transforms: Optional[vision.Compose] = None, + ): + super().__init__() + self.input = input + self.label = {} if label is None else label + self.input_keys = tuple(input.keys()) + self.label_keys = tuple(self.label.keys()) + self.weight = {} if weight is None else weight + self.transforms = transforms + self._len = len(next(iter(input.values()))) + for key in input: + if key in self.label and len(input[key]) != len(self.label[key]): + logger.warning( + f"The length of input {key}({len(input[key])}) is not equal to " + f"the length of label {key}({len(self.label[key])})." + ) + + def __getitem__(self, idx): + input_item = {key: value[idx] for key, value in self.input.items()} + label_item = {key: value[idx] for key, value in self.label.items()} + weight_item = {key: value[idx] for key, value in self.weight.items()} + + if self.transforms is not None: + input_item, label_item, weight_item = self.transforms( + input_item, label_item, weight_item + ) + + return (input_item, label_item, weight_item) + + def __len__(self): + return self._len + + +class IterableNamedArrayDataset(io.IterableDataset): + """IterableNamedArrayDataset for full-data loading. + + Args: + input (Dict[str, np.ndarray]): Input dict. + label (Optional[Dict[str, np.ndarray]]): Label dict. Defaults to None. + weight (Optional[Dict[str, np.ndarray]]): Weight dict. Defaults to None. + transforms (Optional[vision.Compose]): Compose object contains sample wise + transform(s). Defaults to None. + + Examples: + >>> import ppsci + >>> input = {"x": np.random.randn(100, 1)} + >>> label = {"u": np.random.randn(100, 1)} + >>> weight = {"u": np.random.randn(100, 1)} + >>> dataset = ppsci.data.dataset.IterableNamedArrayDataset(input, label, weight) + """ + + # Whether support batch indexing for speeding up fetching process. + batch_index: bool = False + + def __init__( + self, + input: Dict[str, np.ndarray], + label: Optional[Dict[str, np.ndarray]] = None, + weight: Optional[Dict[str, np.ndarray]] = None, + transforms: Optional[vision.Compose] = None, + ): + super().__init__() + self.input = {key: paddle.to_tensor(value) for key, value in input.items()} + self.label = ( + {key: paddle.to_tensor(value) for key, value in label.items()} + if label is not None + else {} + ) + self.input_keys = tuple(input.keys()) + self.label_keys = tuple(self.label.keys()) + self.weight = ( + { + key: paddle.to_tensor(value, paddle.get_default_dtype()) + for key, value in weight.items() + } + if weight is not None + else None + ) + self._len = len(next(iter(self.input.values()))) + self.transforms = transforms + self.world_size_ = dist.get_world_size() + self.rank_ = dist.get_rank() + + @property + def num_samples(self): + """Number of samples within current dataset.""" + return self._len + + def __iter__(self): + if callable(self.transforms): + input_, label_, weight_ = self.transforms( + self.input, self.label, self.weight + ) + else: + input_, label_, weight_ = self.input, self.label, self.weight + + if self.world_size_ > 1: + input_ = _group_dict_into_ranks(input_, self.rank_, self.world_size_) + label_ = _group_dict_into_ranks(label_, self.rank_, self.world_size_) + weight_ = _group_dict_into_ranks(weight_, self.rank_, self.world_size_) + + yield input_, label_, weight_ + + def __len__(self): + return 1 + + +class ContinuousNamedArrayDataset(io.IterableDataset): + """ContinuousNamedArrayDataset for iterable sampling. + + Args: + input (Callable): Function generate input dict. + label (Callable): Function generate label dict. + weight (Optional[Callable]): Function generate weight dict. Defaults to None. + transforms (Optional[vision.Compose]): Compose object contains sample wise + transform(s). Defaults to None. + + Examples: + >>> import ppsci + >>> import numpy as np + >>> input = lambda : {"x": np.random.randn(100, 1)} + >>> label = lambda inp: {"u": np.random.randn(100, 1)} + >>> weight = lambda inp, label: {"u": 1 - (label["u"] ** 2)} + >>> dataset = ppsci.data.dataset.ContinuousNamedArrayDataset(input, label, weight) + >>> input_batch, label_batch, weight_batch = next(iter(dataset)) + >>> print(input_batch["x"].shape) + [100, 1] + >>> print(label_batch["u"].shape) + [100, 1] + >>> print(weight_batch["u"].shape) + [100, 1] + """ + + # Whether support batch indexing for speeding up fetching process. + batch_index: bool = False + + def __init__( + self, + input: Callable, + label: Callable, + weight: Optional[Callable] = None, + transforms: Optional[vision.Compose] = None, + ): + super().__init__() + self.input_fn = input + self.input_keys = tuple(self.input_fn().keys()) + + self.label_fn = label + input_ = self.input_fn() + self.label_keys = tuple(self.label_fn(input_).keys()) + + self.weight_fn = weight + self.transforms = transforms + self.world_size_ = dist.get_world_size() + self.rank_ = dist.get_rank() + + @property + def num_samples(self): + """Number of samples within current dataset.""" + raise NotImplementedError( + "ContinuousNamedArrayDataset has no fixed number of samples." + ) + + def __iter__(self): + def to_tensor_dict(_dict): + if _dict is None: + return None + return {k: paddle.to_tensor(v) for k, v in _dict.items()} + + while True: + input_batch = self.input_fn() + label_batch = self.label_fn(input_batch) + if callable(self.weight_fn): + weight_batch = self.weight_fn(input_batch, label_batch) + else: + weight_batch = None + + if callable(self.transforms): + input_batch, label_batch, weight_batch = self.transforms( + input_batch, label_batch, weight_batch + ) + + if self.world_size_ > 1: + input_batch = _group_dict_into_ranks( + input_batch, self.rank_, self.world_size_ + ) + label_batch = _group_dict_into_ranks( + label_batch, self.rank_, self.world_size_ + ) + weight_batch = _group_dict_into_ranks( + weight_batch, self.rank_, self.world_size_ + ) + + yield to_tensor_dict(input_batch), to_tensor_dict( + label_batch + ), to_tensor_dict(weight_batch) + + def __len__(self): + return 1 + + +class ChipHeatDataset(io.Dataset): + """ChipHeatDataset for data loading of multi-branch DeepONet model. + + Args: + input (Dict[str, np.ndarray]): Input dict. + label (Optional[Dict[str, np.ndarray]]): Label dict. Defaults to None. + index (tuple[str, ...]): Key of input dict. + data_type (str): One of key of input dict. + weight (Optional[Dict[str, np.ndarray]]): Weight dict. Defaults to None. + transforms (Optional[vision.Compose]): Compose object contains sample wise + transform(s). Defaults to None. + + Examples: + >>> import ppsci + >>> input = {"x": np.random.randn(100, 1)} + >>> label = {"u": np.random.randn(100, 1)} + >>> index = ('x', 'u', 'bc', 'bc_data') + >>> data_type = 'u' + >>> weight = {"u": np.random.randn(100, 1)} + >>> dataset = ppsci.data.dataset.ChipHeatDataset(input, label, index, data_type, weight) + """ + + def __init__( + self, + input: Dict[str, np.ndarray], + label: Dict[str, np.ndarray], + index: tuple[str, ...], + data_type: str, + weight: Optional[Dict[str, float]] = None, + transforms: Optional[vision.Compose] = None, + ): + super().__init__() + self.input = input + self.label = label + self.input_keys = tuple(input.keys()) + self.label_keys = tuple(label.keys()) + self.index = index + self.data_type = data_type + self.weight = {} if weight is None else weight + self.transforms = transforms + + def __getitem__(self, idx): + quotient = idx + index_ir = dict() + for i in self.index: + index_ir[i] = 0 + + for i in index_ir: + num = len(self.input[i]) + index_ir[i] = quotient % num + quotient = quotient // num + + input_item = {} + for key in self.input: + if key == "y": + input_item[key] = self.input[key][index_ir["x"]] + elif key == "u_one": + input_item[key] = self.input[key][ + len(self.input[self.data_type]) * index_ir["x"] + + index_ir[self.data_type] + ] + else: + input_item[key] = self.input[key][index_ir[key]] + + label_item = {key: value for key, value in self.label.items()} + weight_item = {key: value for key, value in self.weight.items()} + + if self.transforms is not None: + input_item, label_item, weight_item = self.transforms( + (input_item, label_item, weight_item) + ) + + return (input_item, label_item, weight_item) + + def __len__(self): + _len = 1 + for i in self.index: + _len *= len(self.input[i]) + return _len diff --git a/examples/smc_reac/ppsci/data/dataset/atmospheric_dataset.py b/examples/smc_reac/ppsci/data/dataset/atmospheric_dataset.py new file mode 100644 index 0000000000..c0e4c9ee6d --- /dev/null +++ b/examples/smc_reac/ppsci/data/dataset/atmospheric_dataset.py @@ -0,0 +1,1781 @@ +# Copyright (c) 2024 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. + +from typing import List +from typing import NamedTuple +from typing import Optional +from typing import Sequence +from typing import Tuple + +import numpy as np +import paddle +import pandas as pd +import scipy +from paddle import io + +try: + import trimesh + import xarray +except ModuleNotFoundError: + pass + +# https://www.ecmwf.int/en/forecasts/dataset/ecmwf-reanalysis-v5 +PRESSURE_LEVELS_ERA5_37 = ( + 1, + 2, + 3, + 5, + 7, + 10, + 20, + 30, + 50, + 70, + 100, + 125, + 150, + 175, + 200, + 225, + 250, + 300, + 350, + 400, + 450, + 500, + 550, + 600, + 650, + 700, + 750, + 775, + 800, + 825, + 850, + 875, + 900, + 925, + 950, + 975, + 1000, +) + +# https://www.ecmwf.int/en/forecasts/datasets/set-i +PRESSURE_LEVELS_HRES_25 = ( + 1, + 2, + 3, + 5, + 7, + 10, + 20, + 30, + 50, + 70, + 100, + 150, + 200, + 250, + 300, + 400, + 500, + 600, + 700, + 800, + 850, + 900, + 925, + 950, + 1000, +) + +# https://agupubs.onlinelibrary.wiley.com/doi/full/10.1029/2020MS002203 +PRESSURE_LEVELS_WEATHERBENCH_13 = ( + 50, + 100, + 150, + 200, + 250, + 300, + 400, + 500, + 600, + 700, + 850, + 925, + 1000, +) + +PRESSURE_LEVELS = { + 13: PRESSURE_LEVELS_WEATHERBENCH_13, + 25: PRESSURE_LEVELS_HRES_25, + 37: PRESSURE_LEVELS_ERA5_37, +} + + +TARGET_SURFACE_VARS = ( + "2m_temperature", + "mean_sea_level_pressure", + "10m_v_component_of_wind", + "10m_u_component_of_wind", + "total_precipitation_6hr", +) +TARGET_SURFACE_NO_PRECIP_VARS = ( + "2m_temperature", + "mean_sea_level_pressure", + "10m_v_component_of_wind", + "10m_u_component_of_wind", +) +TARGET_ATMOSPHERIC_VARS = ( + "temperature", + "geopotential", + "u_component_of_wind", + "v_component_of_wind", + "vertical_velocity", + "specific_humidity", +) +TARGET_ATMOSPHERIC_NO_W_VARS = ( + "temperature", + "geopotential", + "u_component_of_wind", + "v_component_of_wind", + "specific_humidity", +) +EXTERNAL_FORCING_VARS = ("toa_incident_solar_radiation",) +GENERATED_FORCING_VARS = ( + "year_progress_sin", + "year_progress_cos", + "day_progress_sin", + "day_progress_cos", +) +FORCING_VARS = EXTERNAL_FORCING_VARS + GENERATED_FORCING_VARS +STATIC_VARS = ( + "geopotential_at_surface", + "land_sea_mask", +) + +TASK_input_variables = ( + TARGET_SURFACE_VARS + TARGET_ATMOSPHERIC_VARS + FORCING_VARS + STATIC_VARS +) +TASK_target_variables = TARGET_SURFACE_VARS + TARGET_ATMOSPHERIC_VARS +TASK_forcing_variables = FORCING_VARS +TASK_pressure_levels = PRESSURE_LEVELS_ERA5_37 +TASK_input_duration = ("12h",) + +TASK_13_input_variables = ( + TARGET_SURFACE_VARS + TARGET_ATMOSPHERIC_VARS + FORCING_VARS + STATIC_VARS +) +TASK_13_target_variables = TARGET_SURFACE_VARS + TARGET_ATMOSPHERIC_VARS +TASK_13_forcing_variables = FORCING_VARS +TASK_13_pressure_levels = PRESSURE_LEVELS_WEATHERBENCH_13 +TASK_13_input_duration = ("12h",) + + +TASK_13_PRECIP_OUT_input_variables = ( + TARGET_SURFACE_NO_PRECIP_VARS + TARGET_ATMOSPHERIC_VARS + FORCING_VARS + STATIC_VARS +) +TASK_13_PRECIP_OUT_target_variables = TARGET_SURFACE_VARS + TARGET_ATMOSPHERIC_VARS +TASK_13_PRECIP_OUT_forcing_variables = FORCING_VARS +TASK_13_PRECIP_OUT_pressure_levels = PRESSURE_LEVELS_WEATHERBENCH_13 +TASK_13_PRECIP_OUT_input_duration = ("12h",) + +_SEC_PER_HOUR = 3600 +_HOUR_PER_DAY = 24 +SEC_PER_DAY = _SEC_PER_HOUR * _HOUR_PER_DAY +_AVG_DAY_PER_YEAR = 365.24219 +AVG_SEC_PER_YEAR = SEC_PER_DAY * _AVG_DAY_PER_YEAR + +DAY_PROGRESS = "day_progress" +YEAR_PROGRESS = "year_progress" + + +def stacked_to_dataset( + stacked_array: "xarray.Variable", + template_dataset: "xarray.Dataset", + preserved_dims: Tuple[str, ...] = ("batch", "lat", "lon"), +) -> "xarray.Dataset": + """The inverse of dataset_to_stacked. + + Requires a template dataset to demonstrate the variables/shapes/coordinates + required. + All variables must have preserved_dims dimensions. + + Args: + stacked_array: Data in BHWC layout, encoded the same as dataset_to_stacked would if it was asked to encode `template_dataset`. + template_dataset: A template Dataset (or other mapping of DataArrays) demonstrating the shape of output required (variables, shapes, coordinates etc). + preserved_dims: dimensions from the target_template that were not folded in the predictions channels. The preserved_dims need to be a subset of the dims of all the variables of template_dataset. + + Returns: + An xarray.Dataset (or other mapping of DataArrays) with the same shape and type as template_dataset. + """ + unstack_from_channels_sizes = {} + var_names = sorted(template_dataset.keys()) + for name in var_names: + template_var = template_dataset[name] + if not all(dim in template_var.dims for dim in preserved_dims): + raise ValueError( + f"stacked_to_dataset requires all Variables to have {preserved_dims} " + f"dimensions, but found only {template_var.dims}." + ) + unstack_from_channels_sizes[name] = { + dim: size + for dim, size in template_var.sizes.items() + if dim not in preserved_dims + } + + channels = { + name: np.prod(list(unstack_sizes.values()), dtype=np.int64) + for name, unstack_sizes in unstack_from_channels_sizes.items() + } + total_expected_channels = sum(channels.values()) + found_channels = stacked_array.sizes["channels"] + if total_expected_channels != found_channels: + raise ValueError( + f"Expected {total_expected_channels} channels but found " + f"{found_channels}, when trying to convert a stacked array of shape " + f"{stacked_array.sizes} to a dataset of shape {template_dataset}." + ) + + data_vars = {} + index = 0 + for name in var_names: + template_var = template_dataset[name] + var = stacked_array.isel({"channels": slice(index, index + channels[name])}) + index += channels[name] + var = var.unstack({"channels": unstack_from_channels_sizes[name]}) + var = var.transpose(*template_var.dims) + data_vars[name] = xarray.DataArray( + data=var, + coords=template_var.coords, + # This might not always be the same as the name it's keyed under; it + # will refer to the original variable name, whereas the key might be + # some alias e.g. temperature_850 under which it should be logged: + name=template_var.name, + ) + return type(template_dataset)( + data_vars + ) # pytype:disable=not-callable,wrong-arg-count + + +def get_graph_spatial_features( + *, + node_lat: np.ndarray, + node_lon: np.ndarray, + senders: np.ndarray, + receivers: np.ndarray, + add_node_positions: bool, + add_node_latitude: bool, + add_node_longitude: bool, + add_relative_positions: bool, + relative_longitude_local_coordinates: bool, + relative_latitude_local_coordinates: bool, + sine_cosine_encoding: bool = False, + encoding_num_freqs: int = 10, + encoding_multiplicative_factor: float = 1.2, +) -> Tuple[np.ndarray, np.ndarray]: + """Computes spatial features for the nodes. + + Args: + node_lat: Latitudes in the [-90, 90] interval of shape [num_nodes] + node_lon: Longitudes in the [0, 360] interval of shape [num_nodes] + senders: Sender indices of shape [num_edges] + receivers: Receiver indices of shape [num_edges] + add_node_positions: Add unit norm absolute positions. + add_node_latitude: Add a feature for latitude (cos(90 - lat)) + Note even if this is set to False, the model may be able to infer the longitude from relative features, unless `relative_latitude_local_coordinates` is also True, or if there is any bias on the relative edge sizes for different longitudes. + add_node_longitude: Add features for longitude (cos(lon), sin(lon)). + Note even if this is set to False, the model may be able to infer the longitude from relative features, unless `relative_longitude_local_coordinates` is also True, or if there is any bias on the relative edge sizes for different longitudes. + add_relative_positions: Whether to relative positions in R3 to the edges. + relative_longitude_local_coordinates: If True, relative positions are computed in a local space where the receiver is at 0 longitude. + relative_latitude_local_coordinates: If True, relative positions are computed in a local space where the receiver is at 0 latitude. + sine_cosine_encoding: If True, we will transform the node/edge features with sine and cosine functions, similar to NERF. + encoding_num_freqs: frequency parameter + encoding_multiplicative_factor: used for calculating the frequency. + + Returns: + Arrays of shape: [num_nodes, num_features] and [num_edges, num_features]. + with node and edge features. + """ + + num_nodes = node_lat.shape[0] + num_edges = senders.shape[0] + dtype = node_lat.dtype + node_phi, node_theta = lat_lon_deg_to_spherical(node_lat, node_lon) + + # Computing some node features. + node_features = [] + if add_node_positions: + # Already in [-1, 1.] range. + node_features.extend(spherical_to_cartesian(node_phi, node_theta)) + + if add_node_latitude: + # Using the cos of theta. + # From 1. (north pole) to -1 (south pole). + node_features.append(np.cos(node_theta)) + + if add_node_longitude: + # Using the cos and sin, which is already normalized. + node_features.append(np.cos(node_phi)) + node_features.append(np.sin(node_phi)) + + if not node_features: + node_features = np.zeros([num_nodes, 0], dtype=dtype) + else: + node_features = np.stack(node_features, axis=-1) + + # Computing some edge features. + edge_features = [] + + if add_relative_positions: + + relative_position = get_relative_position_in_receiver_local_coordinates( + node_phi=node_phi, + node_theta=node_theta, + senders=senders, + receivers=receivers, + latitude_local_coordinates=relative_latitude_local_coordinates, + longitude_local_coordinates=relative_longitude_local_coordinates, + ) + + # Note this is L2 distance in 3d space, rather than geodesic distance. + relative_edge_distances = np.linalg.norm( + relative_position, axis=-1, keepdims=True + ) + + # Normalize to the maximum edge distance. Note that we expect to always + # have an edge that goes in the opposite direction of any given edge + # so the distribution of relative positions should be symmetric around + # zero. So by scaling by the maximum length, we expect all relative + # positions to fall in the [-1., 1.] interval, and all relative distances + # to fall in the [0., 1.] interval. + max_edge_distance = relative_edge_distances.max() + edge_features.append(relative_edge_distances / max_edge_distance) + edge_features.append(relative_position / max_edge_distance) + + if not edge_features: + edge_features = np.zeros([num_edges, 0], dtype=dtype) + else: + edge_features = np.concatenate(edge_features, axis=-1) + + if sine_cosine_encoding: + + def sine_cosine_transform(x: np.ndarray) -> np.ndarray: + freqs = encoding_multiplicative_factor ** np.arange(encoding_num_freqs) + phases = freqs * x[..., None] + x_sin = np.sin(phases) + x_cos = np.cos(phases) + x_cat = np.concatenate([x_sin, x_cos], axis=-1) + return x_cat.reshape([x.shape[0], -1]) + + node_features = sine_cosine_transform(node_features) + edge_features = sine_cosine_transform(edge_features) + + return node_features, edge_features + + +def lat_lon_to_leading_axes(grid_xarray: "xarray.DataArray") -> "xarray.DataArray": + """Reorders xarray so lat/lon axes come first.""" + # leading + ["lat", "lon"] + trailing + # to + # ["lat", "lon"] + leading + trailing + return grid_xarray.transpose("lat", "lon", ...) + + +def restore_leading_axes(grid_xarray: "xarray.DataArray") -> "xarray.DataArray": + """Reorders xarray so batch/time/level axes come first (if present).""" + + # ["lat", "lon"] + [(batch,) (time,) (level,)] + trailing + # to + # [(batch,) (time,) (level,)] + ["lat", "lon"] + trailing + + input_dims = list(grid_xarray.dims) + output_dims = list(input_dims) + for leading_key in ["level", "time", "batch"]: # reverse order for insert + if leading_key in input_dims: + output_dims.remove(leading_key) + output_dims.insert(0, leading_key) + return grid_xarray.transpose(*output_dims) + + +def lat_lon_deg_to_spherical( + node_lat: np.ndarray, + node_lon: np.ndarray, +) -> Tuple[np.ndarray, np.ndarray]: + phi = np.deg2rad(node_lon) + theta = np.deg2rad(90 - node_lat) + return phi, theta + + +def spherical_to_lat_lon( + phi: np.ndarray, + theta: np.ndarray, +) -> Tuple[np.ndarray, np.ndarray]: + lon = np.mod(np.rad2deg(phi), 360) + lat = 90 - np.rad2deg(theta) + return lat, lon + + +def cartesian_to_spherical( + x: np.ndarray, + y: np.ndarray, + z: np.ndarray, +) -> Tuple[np.ndarray, np.ndarray]: + phi = np.arctan2(y, x) + with np.errstate(invalid="ignore"): # circumventing b/253179568 + theta = np.arccos(z) # Assuming unit radius. + return phi, theta + + +def spherical_to_cartesian( + phi: np.ndarray, theta: np.ndarray +) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: + # Assuming unit radius. + return (np.cos(phi) * np.sin(theta), np.sin(phi) * np.sin(theta), np.cos(theta)) + + +def get_relative_position_in_receiver_local_coordinates( + node_phi: np.ndarray, + node_theta: np.ndarray, + senders: np.ndarray, + receivers: np.ndarray, + latitude_local_coordinates: bool, + longitude_local_coordinates: bool, +) -> np.ndarray: + """Returns relative position features for the edges. + + The relative positions will be computed in a rotated space for a local + coordinate system as defined by the receiver. The relative positions are + simply obtained by subtracting sender position minues receiver position in + that local coordinate system after the rotation in R^3. + + Args: + node_phi: [num_nodes] with polar angles. + node_theta: [num_nodes] with azimuthal angles. + senders: [num_edges] with indices. + receivers: [num_edges] with indices. + latitude_local_coordinates: Whether to rotate edges such that in the positions are computed such that the receiver is always at latitude 0. + longitude_local_coordinates: Whether to rotate edges such that in the positions are computed such that the receiver is always at longitude 0. + + Returns: + Array of relative positions in R3 [num_edges, 3] + """ + + node_pos = np.stack(spherical_to_cartesian(node_phi, node_theta), axis=-1) + + # No rotation in this case. + if not (latitude_local_coordinates or longitude_local_coordinates): + return node_pos[senders] - node_pos[receivers] + + # Get rotation matrices for the local space space for every node. + rotation_matrices = get_rotation_matrices_to_local_coordinates( + reference_phi=node_phi, + reference_theta=node_theta, + rotate_latitude=latitude_local_coordinates, + rotate_longitude=longitude_local_coordinates, + ) + + # Each edge will be rotated according to the rotation matrix of its receiver + # node. + edge_rotation_matrices = rotation_matrices[receivers] + + # Rotate all nodes to the rotated space of the corresponding edge. + # Note for receivers we can also do the matmul first and the gather second: + # ``` + # receiver_pos_in_rotated_space = rotate_with_matrices( + # rotation_matrices, node_pos)[receivers] + # ``` + # which is more efficient, however, we do gather first to keep it more + # symmetric with the sender computation. + receiver_pos_in_rotated_space = rotate_with_matrices( + edge_rotation_matrices, node_pos[receivers] + ) + sender_pos_in_in_rotated_space = rotate_with_matrices( + edge_rotation_matrices, node_pos[senders] + ) + # Note, here, that because the rotated space is chosen according to the + # receiver, if: + # * latitude_local_coordinates = True: latitude for the receivers will be + # 0, that is the z coordinate will always be 0. + # * longitude_local_coordinates = True: longitude for the receivers will be + # 0, that is the y coordinate will be 0. + + # Now we can just subtract. + # Note we are rotating to a local coordinate system, where the y-z axes are + # parallel to a tangent plane to the sphere, but still remain in a 3d space. + # Note that if both `latitude_local_coordinates` and + # `longitude_local_coordinates` are True, and edges are short, + # then the difference in x coordinate between sender and receiver + # should be small, so we could consider dropping the new x coordinate if + # we wanted to the tangent plane, however in doing so + # we would lose information about the curvature of the mesh, which may be + # important for very coarse meshes. + return sender_pos_in_in_rotated_space - receiver_pos_in_rotated_space + + +def get_rotation_matrices_to_local_coordinates( + reference_phi: np.ndarray, + reference_theta: np.ndarray, + rotate_latitude: bool, + rotate_longitude: bool, +) -> np.ndarray: + """Returns a rotation matrix to rotate to a point based on a reference vector. + + The rotation matrix is build such that, a vector in the + same coordinate system at the reference point that points towards the pole + before the rotation, continues to point towards the pole after the rotation. + + Args: + reference_phi: [leading_axis] Polar angles of the reference. + reference_theta: [leading_axis] Azimuthal angles of the reference. + rotate_latitude: Whether to produce a rotation matrix that would rotate R^3 vectors to zero latitude. + rotate_longitude: Whether to produce a rotation matrix that would rotate R^3 vectors to zero longitude. + + Returns: + Matrices of shape [leading_axis] such that when applied to the reference + position with `rotate_with_matrices(rotation_matrices, reference_pos)` + + * phi goes to 0. if "rotate_longitude" is True. + + * theta goes to np.pi / 2 if "rotate_latitude" is True. + + The rotation consists of: + * rotate_latitude = False, rotate_longitude = True: + Latitude preserving rotation. + * rotate_latitude = True, rotate_longitude = True: + Latitude preserving rotation, followed by longitude preserving rotation. + * rotate_latitude = True, rotate_longitude = False: + Latitude preserving rotation, followed by longitude preserving rotation, and the inverse of the latitude preserving rotation. Note this is computationally different from rotating the longitude only and is. We do it like this, so the polar geodesic curve, continues to be aligned with one of the axis after the rotation. + """ + + if rotate_longitude and rotate_latitude: + + # We first rotate around the z axis "minus the azimuthal angle", to get the + # point with zero longitude + azimuthal_rotation = -reference_phi + + # One then we will do a polar rotation (which can be done along the y + # axis now that we are at longitude 0.), "minus the polar angle plus 2pi" + # to get the point with zero latitude. + polar_rotation = -reference_theta + np.pi / 2 + + return scipy.spatial.transform.Rotation.from_euler( + "zy", np.stack([azimuthal_rotation, polar_rotation], axis=1) + ).as_matrix() + elif rotate_longitude: + # Just like the previous case, but applying only the azimuthal rotation. + azimuthal_rotation = -reference_phi + return scipy.spatial.transform.Rotation.from_euler( + "z", -reference_phi + ).as_matrix() + elif rotate_latitude: + # Just like the first case, but after doing the polar rotation, undoing + # the azimuthal rotation. + azimuthal_rotation = -reference_phi + polar_rotation = -reference_theta + np.pi / 2 + + return scipy.spatial.transform.Rotation.from_euler( + "zyz", + np.stack([azimuthal_rotation, polar_rotation, -azimuthal_rotation], axis=1), + ).as_matrix() + else: + raise ValueError("At least one of longitude and latitude should be rotated.") + + +def rotate_with_matrices( + rotation_matrices: np.ndarray, positions: np.ndarray +) -> np.ndarray: + return np.einsum("bji,bi->bj", rotation_matrices, positions) + + +def get_bipartite_graph_spatial_features( + *, + senders_node_lat: np.ndarray, + senders_node_lon: np.ndarray, + senders: np.ndarray, + receivers_node_lat: np.ndarray, + receivers_node_lon: np.ndarray, + receivers: np.ndarray, + add_node_positions: bool, + add_node_latitude: bool, + add_node_longitude: bool, + add_relative_positions: bool, + edge_normalization_factor: Optional[float] = None, + relative_longitude_local_coordinates: bool, + relative_latitude_local_coordinates: bool, +) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: + """Computes spatial features for the nodes. + + This function is almost identical to `get_graph_spatial_features`. The only + difference is that sender nodes and receiver nodes can be in different arrays. + This is necessary to enable combination with typed Graph. + + Args: + senders_node_lat: Latitudes in the [-90, 90] interval of shape [num_sender_nodes] + senders_node_lon: Longitudes in the [0, 360] interval of shape [num_sender_nodes] + senders: Sender indices of shape [num_edges], indices in [0, num_sender_nodes) + receivers_node_lat: Latitudes in the [-90, 90] interval of shape [num_receiver_nodes] + receivers_node_lon: Longitudes in the [0, 360] interval of shape [num_receiver_nodes] + receivers: Receiver indices of shape [num_edges], indices in [0, num_receiver_nodes) + add_node_positions: Add unit norm absolute positions. + add_node_latitude: Add a feature for latitude (cos(90 - lat)). + Note even ifthis is set to False, the model may be able to infer the longitude from relative features, unless `relative_latitude_local_coordinates` is also True, or if there is any bias on the relative edge sizes for different longitudes. + add_node_longitude: Add features for longitude (cos(lon), sin(lon)). + Note even if this is set to False, the model may be able to infer the longitude from relative features, unless `relative_longitude_local_coordinates` is also True, or if there is any bias on the relative edge sizes for different longitudes. + add_relative_positions: Whether to relative positions in R3 to the edges. + edge_normalization_factor: Allows explicitly controlling edge normalization. If None, defaults to max edge length. This supports using pre-trained model weights with a different graph structure to what it was trained on. + relative_longitude_local_coordinates: If True, relative positions are computed in a local space where the receiver is at 0 longitude. + relative_latitude_local_coordinates: If True, relative positions are computed in a local space where the receiver is at 0 latitude. + + Returns: + Arrays of shape: [num_nodes, num_features] and [num_edges, num_features]. with node and edge features. + """ + + num_senders = senders_node_lat.shape[0] + num_receivers = receivers_node_lat.shape[0] + num_edges = senders.shape[0] + dtype = senders_node_lat.dtype + assert receivers_node_lat.dtype == dtype + senders_node_phi, senders_node_theta = lat_lon_deg_to_spherical( + senders_node_lat, senders_node_lon + ) + receivers_node_phi, receivers_node_theta = lat_lon_deg_to_spherical( + receivers_node_lat, receivers_node_lon + ) + + # Computing some node features. + senders_node_features = [] + receivers_node_features = [] + if add_node_positions: + # Already in [-1, 1.] range. + senders_node_features.extend( + spherical_to_cartesian(senders_node_phi, senders_node_theta) + ) + receivers_node_features.extend( + spherical_to_cartesian(receivers_node_phi, receivers_node_theta) + ) + + if add_node_latitude: + # Using the cos of theta. + # From 1. (north pole) to -1 (south pole). + senders_node_features.append(np.cos(senders_node_theta)) + receivers_node_features.append(np.cos(receivers_node_theta)) + + if add_node_longitude: + # Using the cos and sin, which is already normalized. + senders_node_features.append(np.cos(senders_node_phi)) + senders_node_features.append(np.sin(senders_node_phi)) + + receivers_node_features.append(np.cos(receivers_node_phi)) + receivers_node_features.append(np.sin(receivers_node_phi)) + + if not senders_node_features: + senders_node_features = np.zeros([num_senders, 0], dtype=dtype) + receivers_node_features = np.zeros([num_receivers, 0], dtype=dtype) + else: + senders_node_features = np.stack(senders_node_features, axis=-1) + receivers_node_features = np.stack(receivers_node_features, axis=-1) + + # Computing some edge features. + edge_features = [] + + if add_relative_positions: + + relative_position = ( + get_bipartite_relative_position_in_receiver_local_coordinates( + senders_node_phi=senders_node_phi, + senders_node_theta=senders_node_theta, + receivers_node_phi=receivers_node_phi, + receivers_node_theta=receivers_node_theta, + senders=senders, + receivers=receivers, + latitude_local_coordinates=relative_latitude_local_coordinates, + longitude_local_coordinates=relative_longitude_local_coordinates, + ) + ) + + # Note this is L2 distance in 3d space, rather than geodesic distance. + relative_edge_distances = np.linalg.norm( + relative_position, axis=-1, keepdims=True + ) + + if edge_normalization_factor is None: + # Normalize to the maximum edge distance. Note that we expect to always + # have an edge that goes in the opposite direction of any given edge + # so the distribution of relative positions should be symmetric around + # zero. So by scaling by the maximum length, we expect all relative + # positions to fall in the [-1., 1.] interval, and all relative distances + # to fall in the [0., 1.] interval. + edge_normalization_factor = relative_edge_distances.max() + + edge_features.append(relative_edge_distances / edge_normalization_factor) + edge_features.append(relative_position / edge_normalization_factor) + + if not edge_features: + edge_features = np.zeros([num_edges, 0], dtype=dtype) + else: + edge_features = np.concatenate(edge_features, axis=-1) + + return senders_node_features, receivers_node_features, edge_features + + +def get_bipartite_relative_position_in_receiver_local_coordinates( + senders_node_phi: np.ndarray, + senders_node_theta: np.ndarray, + senders: np.ndarray, + receivers_node_phi: np.ndarray, + receivers_node_theta: np.ndarray, + receivers: np.ndarray, + latitude_local_coordinates: bool, + longitude_local_coordinates: bool, +) -> np.ndarray: + """Returns relative position features for the edges. + + This function is equivalent to + `get_relative_position_in_receiver_local_coordinates`, but adapted to work + with bipartite typed graphs. + + The relative positions will be computed in a rotated space for a local + coordinate system as defined by the receiver. The relative positions are + simply obtained by subtracting sender position minues receiver position in + that local coordinate system after the rotation in R^3. + + Args: + senders_node_phi: [num_sender_nodes] with polar angles. + senders_node_theta: [num_sender_nodes] with azimuthal angles. + senders: [num_edges] with indices into sender nodes. + receivers_node_phi: [num_sender_nodes] with polar angles. + receivers_node_theta: [num_sender_nodes] with azimuthal angles. + receivers: [num_edges] with indices into receiver nodes. + latitude_local_coordinates: Whether to rotate edges such that in the positions are computed such that the receiver is always at latitude 0. + longitude_local_coordinates: Whether to rotate edges such that in the positions are computed such that the receiver is always at longitude 0. + + Returns: + Array of relative positions in R3 [num_edges, 3] + """ + + senders_node_pos = np.stack( + spherical_to_cartesian(senders_node_phi, senders_node_theta), axis=-1 + ) + + receivers_node_pos = np.stack( + spherical_to_cartesian(receivers_node_phi, receivers_node_theta), axis=-1 + ) + + # No rotation in this case. + if not (latitude_local_coordinates or longitude_local_coordinates): + return senders_node_pos[senders] - receivers_node_pos[receivers] + + # Get rotation matrices for the local space space for every receiver node. + receiver_rotation_matrices = get_rotation_matrices_to_local_coordinates( + reference_phi=receivers_node_phi, + reference_theta=receivers_node_theta, + rotate_latitude=latitude_local_coordinates, + rotate_longitude=longitude_local_coordinates, + ) + + # Each edge will be rotated according to the rotation matrix of its receiver + # node. + edge_rotation_matrices = receiver_rotation_matrices[receivers] + + # Rotate all nodes to the rotated space of the corresponding edge. + # Note for receivers we can also do the matmul first and the gather second: + # ``` + # receiver_pos_in_rotated_space = rotate_with_matrices( + # rotation_matrices, node_pos)[receivers] + # ``` + # which is more efficient, however, we do gather first to keep it more + # symmetric with the sender computation. + receiver_pos_in_rotated_space = rotate_with_matrices( + edge_rotation_matrices, receivers_node_pos[receivers] + ) + sender_pos_in_in_rotated_space = rotate_with_matrices( + edge_rotation_matrices, senders_node_pos[senders] + ) + # Note, here, that because the rotated space is chosen according to the + # receiver, if: + # * latitude_local_coordinates = True: latitude for the receivers will be + # 0, that is the z coordinate will always be 0. + # * longitude_local_coordinates = True: longitude for the receivers will be + # 0, that is the y coordinate will be 0. + + # Now we can just subtract. + # Note we are rotating to a local coordinate system, where the y-z axes are + # parallel to a tangent plane to the sphere, but still remain in a 3d space. + # Note that if both `latitude_local_coordinates` and + # `longitude_local_coordinates` are True, and edges are short, + # then the difference in x coordinate between sender and receiver + # should be small, so we could consider dropping the new x coordinate if + # we wanted to the tangent plane, however in doing so + # we would lose information about the curvature of the mesh, which may be + # important for very coarse meshes. + return sender_pos_in_in_rotated_space - receiver_pos_in_rotated_space + + +class GraphGridMesh: + """Graph datatype of GraphCast. + + Args: + mesh_size (int): size of mesh. + radius_query_fraction_edge_length (float): _description_ + mesh2grid_edge_normalization_factor (float): Normalization factor of edge in Mesh2Grid GNN. + resolution (float): resolution of atmospheric data. + mesh2mesh_src_index (np.array, optional): Index of Mesh2Mesh source node. Defaults to None. + mesh2mesh_dst_index (np.array, optional): Index of Mesh2Mesh destination node. Defaults to None. + grid2mesh_src_index (np.array, optional): Index of Grid2Mesh source node. Defaults to None. + grid2mesh_dst_index (np.array, optional): Index of Grid2Mesh destination node. + mesh2grid_src_index (np.array, optional): Index of Mesh2Grid source node. Defaults to None. + mesh2grid_dst_index (np.array, optional): Index of Mesh2Grid destination node. Defaults to None. + mesh_num_nodes (int, optional): Number of mesh nodes. Defaults to None. + grid_num_nodes (int, optional): Number of grid nodes. Defaults to None. + mesh_num_edges (int, optional): Number of mesh edges. Defaults to None. + grid2mesh_num_edges (int, optional): Number of edges in Grid2Mesh GNN. Defaults to None. + mesh2grid_num_edges (int, optional): Number of edges in Mesh2Grid GNN. Defaults to None. + grid_node_feat (np.array, optional): Feature of grid nodes. Defaults to None. + mesh_node_feat (np.array, optional): Feature of mehs nodes. Defaults to None. + mesh_edge_feat (np.array, optional): Feature of mesh edges. Defaults to None. + grid2mesh_edge_feat (np.array, optional): Feature of edges in Grid2Mesh GNN. Defaults to None. + mesh2grid_edge_feat (np.array, optional): Feature of edges in Mesh2Grid GNN. Defaults to None. + """ + + def __init__( + self, + mesh_size: int, + radius_query_fraction_edge_length: float, + mesh2grid_edge_normalization_factor: float, + resolution: float, + mesh2mesh_src_index: np.array = None, + mesh2mesh_dst_index: np.array = None, + grid2mesh_src_index: np.array = None, + grid2mesh_dst_index: np.array = None, + mesh2grid_src_index: np.array = None, + mesh2grid_dst_index: np.array = None, + mesh_num_nodes: int = None, + grid_num_nodes: int = None, + mesh_num_edges: int = None, + grid2mesh_num_edges: np.array = None, + mesh2grid_num_edges: np.array = None, + grid_node_feat: np.array = None, + mesh_node_feat: np.array = None, + mesh_edge_feat: np.array = None, + grid2mesh_edge_feat: np.array = None, + mesh2grid_edge_feat: np.array = None, + ): + self.meshes = get_hierarchy_of_triangular_meshes_for_sphere(mesh_size) + + all_input_vars = [ + mesh2mesh_src_index, + mesh2mesh_dst_index, + grid2mesh_src_index, + grid2mesh_dst_index, + mesh2grid_src_index, + mesh2grid_dst_index, + mesh_num_nodes, + grid_num_nodes, + mesh_num_edges, + grid2mesh_num_edges, + mesh2grid_num_edges, + grid_node_feat, + mesh_node_feat, + mesh_edge_feat, + grid2mesh_edge_feat, + mesh2grid_edge_feat, + ] + should_init = any(var is None for var in all_input_vars) + + if should_init: + self.query_radius = ( + self._get_max_edge_distance(self.finest_mesh) + * radius_query_fraction_edge_length + ) + self._mesh2grid_edge_normalization_factor = ( + mesh2grid_edge_normalization_factor + ) + self._spatial_features_kwargs = dict( + add_node_positions=False, + add_node_latitude=True, + add_node_longitude=True, + add_relative_positions=True, + relative_longitude_local_coordinates=True, + relative_latitude_local_coordinates=True, + ) + + self.init_mesh_properties() + self._init_grid_properties( + grid_lat=np.arange(-90.0, 90.0 + resolution, resolution), + grid_lon=np.arange(0.0, 360.0, resolution), + ) + self._grid2mesh_graph_structure = self._init_grid2mesh_graph() + self._mesh_graph_structure = self._init_mesh_graph() + self._mesh2grid_graph_structure = self._init_mesh2grid_graph() + else: + self.mesh2mesh_src_index = mesh2mesh_src_index + self.mesh2mesh_dst_index = mesh2mesh_dst_index + self.grid2mesh_src_index = grid2mesh_src_index + self.grid2mesh_dst_index = grid2mesh_dst_index + self.mesh2grid_src_index = mesh2grid_src_index + self.mesh2grid_dst_index = mesh2grid_dst_index + + self.mesh_num_nodes = mesh_num_nodes + self.grid_num_nodes = grid_num_nodes + + self.mesh_num_edges = mesh_num_edges + self.grid2mesh_num_edges = grid2mesh_num_edges + self.mesh2grid_num_edges = mesh2grid_num_edges + + self.grid_node_feat = grid_node_feat + self.mesh_node_feat = mesh_node_feat + self.mesh_edge_feat = mesh_edge_feat + self.grid2mesh_edge_feat = grid2mesh_edge_feat + self.mesh2grid_edge_feat = mesh2grid_edge_feat + + def update(self, name, value): + if hasattr(self, name): + setattr(self, name, value) + else: + raise ValueError + + def tensor(self): + self.mesh2mesh_src_index = paddle.to_tensor( + self.mesh2mesh_src_index, dtype=paddle.int64 + ) + + self.mesh2mesh_dst_index = paddle.to_tensor( + self.mesh2mesh_dst_index, dtype=paddle.int64 + ) + self.grid2mesh_src_index = paddle.to_tensor( + self.grid2mesh_src_index, dtype=paddle.int64 + ) + self.grid2mesh_dst_index = paddle.to_tensor( + self.grid2mesh_dst_index, dtype=paddle.int64 + ) + self.mesh2grid_src_index = paddle.to_tensor( + self.mesh2grid_src_index, dtype=paddle.int64 + ) + self.mesh2grid_dst_index = paddle.to_tensor( + self.mesh2grid_dst_index, dtype=paddle.int64 + ) + self.grid_node_feat = paddle.to_tensor( + self.grid_node_feat, dtype=paddle.get_default_dtype() + ) + self.mesh_node_feat = paddle.to_tensor( + self.mesh_node_feat, dtype=paddle.get_default_dtype() + ) + self.mesh_edge_feat = paddle.to_tensor( + self.mesh_edge_feat, dtype=paddle.get_default_dtype() + ) + self.grid2mesh_edge_feat = paddle.to_tensor( + self.grid2mesh_edge_feat, dtype=paddle.get_default_dtype() + ) + self.mesh2grid_edge_feat = paddle.to_tensor( + self.mesh2grid_edge_feat, dtype=paddle.get_default_dtype() + ) + return self + + @property + def finest_mesh(self): + return self.meshes[-1] + + def init_mesh_properties(self): + """Inits static properties that have to do with mesh nodes.""" + self.mesh_num_nodes = self.finest_mesh.vertices.shape[0] + mesh_phi, mesh_theta = cartesian_to_spherical( + self.finest_mesh.vertices[:, 0], + self.finest_mesh.vertices[:, 1], + self.finest_mesh.vertices[:, 2], + ) + (mesh_nodes_lat, mesh_nodes_lon) = spherical_to_lat_lon( + phi=mesh_phi, + theta=mesh_theta, + ) + # Convert to f32 to ensure the lat/lon features aren't in f64. + self._mesh_nodes_lat = mesh_nodes_lat.astype(np.float32) + self._mesh_nodes_lon = mesh_nodes_lon.astype(np.float32) + + def _init_grid_properties(self, grid_lat: np.ndarray, grid_lon: np.ndarray): + """Inits static properties that have to do with grid nodes.""" + self._grid_lat = grid_lat.astype(np.float32) + self._grid_lon = grid_lon.astype(np.float32) + # Initialized the counters. + self.grid_num_nodes = grid_lat.shape[0] * grid_lon.shape[0] + + # Initialize lat and lon for the grid. + grid_nodes_lon, grid_nodes_lat = np.meshgrid(grid_lon, grid_lat) + self._grid_nodes_lon = grid_nodes_lon.reshape([-1]).astype(np.float32) + self._grid_nodes_lat = grid_nodes_lat.reshape([-1]).astype(np.float32) + + def _init_grid2mesh_graph(self): + """Build Grid2Mesh graph.""" + + # Create some edges according to distance between mesh and grid nodes. + assert self._grid_lat is not None and self._grid_lon is not None + (grid_indices, mesh_indices) = radius_query_indices( + grid_latitude=self._grid_lat, + grid_longitude=self._grid_lon, + mesh=self.finest_mesh, + radius=self.query_radius, + ) + + # Edges sending info from grid to mesh. + senders = grid_indices + receivers = mesh_indices + + # Precompute structural node and edge features according to config options. + # Structural features are those that depend on the fixed values of the + # latitude and longitudes of the nodes. + ( + senders_node_features, + _, + edge_features, + ) = get_bipartite_graph_spatial_features( + senders_node_lat=self._grid_nodes_lat, + senders_node_lon=self._grid_nodes_lon, + receivers_node_lat=self._mesh_nodes_lat, + receivers_node_lon=self._mesh_nodes_lon, + senders=senders, + receivers=receivers, + edge_normalization_factor=None, + **self._spatial_features_kwargs, + ) + + self.grid_node_feat = np.expand_dims(senders_node_features, axis=1) + + self.grid2mesh_src_index = senders + self.grid2mesh_dst_index = receivers + self.grid2mesh_edge_feat = np.expand_dims(edge_features, axis=1) + self.grid2mesh_num_edges = len(edge_features) + + def _init_mesh_graph(self): + """Build Mesh graph.""" + merged_mesh = merge_meshes(self.meshes) + # Work simply on the mesh edges. + senders, receivers = faces_to_edges(merged_mesh.faces) + # Precompute structural node and edge features according to config options. + # Structural features are those that depend on the fixed values of the + # latitude and longitudes of the nodes. + assert self._mesh_nodes_lat is not None and self._mesh_nodes_lon is not None + node_features, edge_features = get_graph_spatial_features( + node_lat=self._mesh_nodes_lat, + node_lon=self._mesh_nodes_lon, + senders=senders, + receivers=receivers, + **self._spatial_features_kwargs, + ) + + self.mesh_node_feat = np.expand_dims(node_features, axis=1) + self.mesh2mesh_src_index = senders + self.mesh2mesh_dst_index = receivers + self.mesh_edge_feat = np.expand_dims(edge_features, axis=1) + self.mesh_num_edges = len(edge_features) + + def _init_mesh2grid_graph(self): + """Build Mesh2Grid graph.""" + + # Create some edges according to how the grid nodes are contained by + # mesh triangles. + (grid_indices, mesh_indices) = in_mesh_triangle_indices( + grid_latitude=self._grid_lat, + grid_longitude=self._grid_lon, + mesh=self.finest_mesh, + ) + + # Edges sending info from mesh to grid. + senders = mesh_indices + receivers = grid_indices + + # Precompute structural node and edge features according to config options. + assert self._mesh_nodes_lat is not None and self._mesh_nodes_lon is not None + (_, _, edge_features) = get_bipartite_graph_spatial_features( + senders_node_lat=self._mesh_nodes_lat, + senders_node_lon=self._mesh_nodes_lon, + receivers_node_lat=self._grid_nodes_lat, + receivers_node_lon=self._grid_nodes_lon, + senders=senders, + receivers=receivers, + edge_normalization_factor=self._mesh2grid_edge_normalization_factor, + **self._spatial_features_kwargs, + ) + + self.mesh2grid_src_index = senders + self.mesh2grid_dst_index = receivers + self.mesh2grid_edge_feat = np.expand_dims(edge_features, axis=1) + self.mesh2grid_num_edges = len(edge_features) + + @staticmethod + def _get_max_edge_distance(mesh): + senders, receivers = faces_to_edges(mesh.faces) + edge_distances = np.linalg.norm( + mesh.vertices[senders] - mesh.vertices[receivers], axis=-1 + ) + return edge_distances.max() + + def grid_node_outputs_to_prediction( + self, + grid_node_outputs: np.ndarray, + targets_template: "xarray.Dataset", + ) -> "xarray.Dataset": + """[num_grid_nodes, batch, num_outputs] -> xarray.""" + # numpy array with shape [lat_lon_node, batch, channels] + # to xarray `DataArray` (batch, lat, lon, channels) + assert self._grid_lat is not None and self._grid_lon is not None + grid_shape = (self._grid_lat.shape[0], self._grid_lon.shape[0]) + grid_outputs_lat_lon_leading = grid_node_outputs.reshape( + grid_shape + grid_node_outputs.shape[1:] + ) + dims = ("lat", "lon", "batch", "channels") + grid_xarray_lat_lon_leading = xarray.DataArray( + data=grid_outputs_lat_lon_leading, dims=dims + ) + grid_xarray = restore_leading_axes(grid_xarray_lat_lon_leading) + + # xarray `DataArray` (batch, lat, lon, channels) + # to xarray `Dataset` (batch, one time step, lat, lon, level, multiple vars) + return stacked_to_dataset(grid_xarray.variable, targets_template) + + +class TriangularMesh(NamedTuple): + vertices: np.ndarray + faces: np.ndarray + + +def merge_meshes(mesh_list: Sequence[TriangularMesh]) -> TriangularMesh: + for i in range(len(mesh_list) - 1): + mesh_i, mesh_ip1 = mesh_list[i], mesh_list[i + 1] + num_nodes_mesh_i = mesh_i.vertices.shape[0] + assert np.allclose(mesh_i.vertices, mesh_ip1.vertices[:num_nodes_mesh_i]) + + return TriangularMesh( + vertices=mesh_list[-1].vertices, + faces=np.concatenate([mesh.faces for mesh in mesh_list], axis=0), + ) + + +def get_icosahedron(): + phi = (1 + np.sqrt(5)) / 2 + product = [[1.0, phi], [1.0, -phi], [-1.0, phi], [-1.0, -phi]] + vertices = [] + for p in product: + c1 = p[0] + c2 = p[1] + vertices.append((c1, c2, 0.0)) + vertices.append((0.0, c1, c2)) + vertices.append((c2, 0.0, c1)) + + vertices = np.array(vertices, dtype=np.float32) + vertices /= np.linalg.norm([1.0, phi]) + + faces = [ + (0, 1, 2), + (0, 6, 1), + (8, 0, 2), + (8, 4, 0), + (3, 8, 2), + (3, 2, 7), + (7, 2, 1), + (0, 4, 6), + (4, 11, 6), + (6, 11, 5), + (1, 5, 7), + (4, 10, 11), + (4, 8, 10), + (10, 8, 3), + (10, 3, 9), + (11, 10, 9), + (11, 9, 5), + (5, 9, 7), + (9, 3, 7), + (1, 6, 5), + ] + + angle_between_faces = 2 * np.arcsin(phi / np.sqrt(3)) + rotation_angle = (np.pi - angle_between_faces) / 2 + rotation = scipy.spatial.transform.Rotation.from_euler( + seq="y", angles=rotation_angle + ) + rotation_matrix = rotation.as_matrix() + vertices = np.dot(vertices, rotation_matrix) + + return TriangularMesh( + vertices=vertices.astype(np.float32), faces=np.array(faces, dtype=np.int32) + ) + + +def get_hierarchy_of_triangular_meshes_for_sphere( + splits: int, +) -> List[TriangularMesh]: + current_mesh = get_icosahedron() + output_meshes = [current_mesh] + for _ in range(splits): + current_mesh = _two_split_unit_sphere_triangle_faces(current_mesh) + output_meshes.append(current_mesh) + return output_meshes + + +def _two_split_unit_sphere_triangle_faces( + triangular_mesh: TriangularMesh, +) -> TriangularMesh: + """Splits each triangular face into 4 triangles keeping the orientation.""" + new_vertices_builder = _ChildVerticesBuilder(triangular_mesh.vertices) + + new_faces = [] + for ind1, ind2, ind3 in triangular_mesh.faces: + ind12 = new_vertices_builder.get_new_child_vertex_index((ind1, ind2)) + ind23 = new_vertices_builder.get_new_child_vertex_index((ind2, ind3)) + ind31 = new_vertices_builder.get_new_child_vertex_index((ind3, ind1)) + new_faces.extend( + [ + [ind1, ind12, ind31], # 1 + [ind12, ind2, ind23], # 2 + [ind31, ind23, ind3], # 3 + [ind12, ind23, ind31], # 4 + ] + ) + return TriangularMesh( + vertices=new_vertices_builder.get_all_vertices(), + faces=np.array(new_faces, dtype=np.int32), + ) + + +class _ChildVerticesBuilder: + """Bookkeeping of new child vertices added to an existing set of vertices.""" + + def __init__(self, parent_vertices): + self._child_vertices_index_mapping = {} + self._parent_vertices = parent_vertices + # We start with all previous vertices. + self._all_vertices_list = list(parent_vertices) + + def _get_child_vertex_key(self, parent_vertex_indices): + return tuple(sorted(parent_vertex_indices)) + + def _create_child_vertex(self, parent_vertex_indices): + """Creates a new vertex.""" + # Position for new vertex is the middle point, between the parent points, + # projected to unit sphere. + child_vertex_position = self._parent_vertices[list(parent_vertex_indices)].mean( + 0 + ) + child_vertex_position /= np.linalg.norm(child_vertex_position) + + # Add the vertex to the output list. The index for this new vertex will + # match the length of the list before adding it. + child_vertex_key = self._get_child_vertex_key(parent_vertex_indices) + self._child_vertices_index_mapping[child_vertex_key] = len( + self._all_vertices_list + ) + self._all_vertices_list.append(child_vertex_position) + + def get_new_child_vertex_index(self, parent_vertex_indices): + """Returns index for a child vertex, creating it if necessary.""" + # Get the key to see if we already have a new vertex in the middle. + child_vertex_key = self._get_child_vertex_key(parent_vertex_indices) + if child_vertex_key not in self._child_vertices_index_mapping: + self._create_child_vertex(parent_vertex_indices) + return self._child_vertices_index_mapping[child_vertex_key] + + def get_all_vertices(self): + """Returns an array with old vertices.""" + return np.array(self._all_vertices_list) + + +def faces_to_edges(faces: np.ndarray): + """Transforms polygonal faces to sender and receiver indices. + + It does so by transforming every face into N_i edges. Such if the triangular + face has indices [0, 1, 2], three edges are added 0->1, 1->2, and 2->0. + + If all faces have consistent orientation, and the surface represented by the + faces is closed, then every edge in a polygon with a certain orientation + is also part of another polygon with the opposite orientation. In this + situation, the edges returned by the method are always bidirectional. + + Args: + faces: Integer array of shape [num_faces, 3]. Contains node indices adjacent to each face. + Returns: + Tuple with sender/receiver indices, each of shape [num_edges=num_faces*3]. + """ + + assert faces.ndim == 2 + assert faces.shape[-1] == 3 + senders = np.concatenate([faces[:, 0], faces[:, 1], faces[:, 2]]) + receivers = np.concatenate([faces[:, 1], faces[:, 2], faces[:, 0]]) + return senders, receivers + + +def _grid_lat_lon_to_coordinates( + grid_latitude: np.ndarray, grid_longitude: np.ndarray +) -> np.ndarray: + """Lat [num_lat] lon [num_lon] to 3d coordinates [num_lat, num_lon, 3].""" + # Convert to spherical coordinates phi and theta defined in the grid. + # Each [num_latitude_points, num_longitude_points] + phi_grid, theta_grid = np.meshgrid( + np.deg2rad(grid_longitude), np.deg2rad(90 - grid_latitude) + ) + + # [num_latitude_points, num_longitude_points, 3] + # Note this assumes unit radius, since for now we model the earth as a + # sphere of unit radius, and keep any vertical dimension as a regular grid. + return np.stack( + [ + np.cos(phi_grid) * np.sin(theta_grid), + np.sin(phi_grid) * np.sin(theta_grid), + np.cos(theta_grid), + ], + axis=-1, + ) + + +def radius_query_indices( + *, + grid_latitude: np.ndarray, + grid_longitude: np.ndarray, + mesh: TriangularMesh, + radius: float, +) -> Tuple[np.ndarray, np.ndarray]: + """Returns mesh-grid edge indices for radius query. + + Args: + grid_latitude: Latitude values for the grid [num_lat_points] + grid_longitude: Longitude values for the grid [num_lon_points] + mesh: Mesh object. + radius: Radius of connectivity in R3. for a sphere of unit radius. + + Returns: + tuple with `grid_indices` and `mesh_indices` indicating edges between the grid and the mesh such that the distances in a straight line (not geodesic) are smaller than or equal to `radius`. + grid_indices: Indices of shape [num_edges], that index into a + [num_lat_points, num_lon_points] grid, after flattening the leading axes. + mesh_indices: Indices of shape [num_edges], that index into mesh.vertices. + """ + + # [num_grid_points=num_lat_points * num_lon_points, 3] + grid_positions = _grid_lat_lon_to_coordinates( + grid_latitude, grid_longitude + ).reshape([-1, 3]) + + # [num_mesh_points, 3] + mesh_positions = mesh.vertices + kd_tree = scipy.spatial.cKDTree(mesh_positions) + + # [num_grid_points, num_mesh_points_per_grid_point] + # Note `num_mesh_points_per_grid_point` is not constant, so this is a list + # of arrays, rather than a 2d array. + query_indices = kd_tree.query_ball_point(x=grid_positions, r=radius) + + grid_edge_indices = [] + mesh_edge_indices = [] + for grid_index, mesh_neighbors in enumerate(query_indices): + grid_edge_indices.append(np.repeat(grid_index, len(mesh_neighbors))) + mesh_edge_indices.append(mesh_neighbors) + + # [num_edges] + grid_edge_indices = np.concatenate(grid_edge_indices, axis=0).astype(int) + mesh_edge_indices = np.concatenate(mesh_edge_indices, axis=0).astype(int) + + return grid_edge_indices, mesh_edge_indices + + +def in_mesh_triangle_indices( + *, grid_latitude: np.ndarray, grid_longitude: np.ndarray, mesh: TriangularMesh +) -> Tuple[np.ndarray, np.ndarray]: + """Returns mesh-grid edge indices for grid points contained in mesh triangles. + + Args: + grid_latitude: Latitude values for the grid [num_lat_points] + grid_longitude: Longitude values for the grid [num_lon_points] + mesh: Mesh object. + + Returns: + tuple with `grid_indices` and `mesh_indices` indicating edges between the grid and the mesh vertices of the triangle that contain each grid point. The number of edges is always num_lat_points * num_lon_points * 3 + grid_indices: Indices of shape [num_edges], that index into a [num_lat_points, num_lon_points] grid, after flattening the leading axes. + mesh_indices: Indices of shape [num_edges], that index into mesh.vertices. + """ + + # [num_grid_points=num_lat_points * num_lon_points, 3] + grid_positions = _grid_lat_lon_to_coordinates( + grid_latitude, grid_longitude + ).reshape([-1, 3]) + + mesh_trimesh = trimesh.Trimesh(vertices=mesh.vertices, faces=mesh.faces) + + # [num_grid_points] with mesh face indices for each grid point. + _, _, query_face_indices = trimesh.proximity.closest_point( + mesh_trimesh, grid_positions + ) + + # [num_grid_points, 3] with mesh node indices for each grid point. + mesh_edge_indices = mesh.faces[query_face_indices] + + # [num_grid_points, 3] with grid node indices, where every row simply contains + # the row (grid_point) index. + grid_indices = np.arange(grid_positions.shape[0]) + grid_edge_indices = np.tile(grid_indices.reshape([-1, 1]), [1, 3]) + + # Flatten to get a regular list. + # [num_edges=num_grid_points*3] + mesh_edge_indices = mesh_edge_indices.reshape([-1]) + grid_edge_indices = grid_edge_indices.reshape([-1]) + + return grid_edge_indices, mesh_edge_indices + + +def get_year_progress(seconds_since_epoch: np.ndarray) -> np.ndarray: + """Computes year progress for times in seconds. + Args: + seconds_since_epoch: Times in seconds since the "epoch" (the point at which UNIX time starts). + Returns: + Year progress normalized to be in the `[0, 1)` interval for each time point. + """ + # Start with the pure integer division, and then float at the very end. + # We will try to keep as much precision as possible. + years_since_epoch = ( + seconds_since_epoch / SEC_PER_DAY / np.float64(_AVG_DAY_PER_YEAR) + ) + # Note depending on how these ops are down, we may end up with a "weak_type" + # which can cause issues in subtle ways, and hard to track here. + # In any case, casting to float32 should get rid of the weak type. + # [0, 1.) Interval. + return np.mod(years_since_epoch, 1.0).astype(np.float32) + + +def get_day_progress( + seconds_since_epoch: np.ndarray, + longitude: np.ndarray, +) -> np.ndarray: + """Computes day progress for times in seconds at each longitude. + Args: + seconds_since_epoch: 1D array of times in seconds since the 'epoch' (the point at which UNIX time starts). + longitude: 1D array of longitudes at which day progress is computed. + Returns: + 2D array of day progress values normalized to be in the [0, 1) interval for each time point at each longitude. + """ + # [0.0, 1.0) Interval. + day_progress_greenwich = np.mod(seconds_since_epoch, SEC_PER_DAY) / SEC_PER_DAY + # Offset the day progress to the longitude of each point on Earth. + longitude_offsets = np.deg2rad(longitude) / (2 * np.pi) + day_progress = np.mod( + day_progress_greenwich[..., np.newaxis] + longitude_offsets, 1.0 + ) + return day_progress.astype(np.float32) + + +def datetime_features(seconds_since_epoch, longitude_offsets): + year_progress = get_year_progress(seconds_since_epoch) + day_progress = get_day_progress(seconds_since_epoch, longitude_offsets) + year_progress_phase = year_progress * (2 * np.pi) + day_progress_phase = day_progress * (2 * np.pi) + returned_data = { + "year_progress_sin": np.sin(year_progress_phase), + "year_progress_cos": np.cos(year_progress_phase), + "day_progress_sin": np.sin(day_progress_phase), + "day_progress_cos": np.cos(day_progress_phase), + } + return returned_data + + +def add_var_into_nc_dataset( + nc_dataset, + var_name, + var_value, + var_dims=( + "batch", + "time", + ), +): + new_var = nc_dataset.createVariable(var_name, "f8", var_dims) + new_var[:] = var_value + return nc_dataset + + +def extract_input_target_times( + dataset: "xarray.Dataset", + input_duration: str, + target_lead_times: str, +): + (target_lead_times, target_duration) = _process_target_lead_times_and_get_duration( + target_lead_times + ) + time = dataset.coords["time"] + dataset = dataset.assign_coords(time=time + target_duration - time[-1]) + + targets = dataset.sel({"time": target_lead_times}) + + input_duration = pd.Timedelta(input_duration) + zero = pd.Timedelta(0) + epsilon = pd.Timedelta(1, "ns") + inputs = dataset.sel({"time": slice(-input_duration + epsilon, zero)}) + return inputs, targets + + +def _process_target_lead_times_and_get_duration(target_lead_times: str): + """Returns the minimum duration for the target lead times.""" + if isinstance(target_lead_times, slice): + if target_lead_times.start is None: + target_lead_times = slice( + pd.Timedelta(1, "ns"), target_lead_times.stop, target_lead_times.step + ) + target_duration = pd.Timedelta(target_lead_times.stop) + else: + if not isinstance(target_lead_times, (list, tuple, set)): + target_lead_times = [target_lead_times] + + target_lead_times = [pd.Timedelta(x) for x in target_lead_times] + target_lead_times.sort() + target_duration = target_lead_times[-1] + return target_lead_times, target_duration + + +def variable_to_stacked( + variable: "xarray.Variable", + sizes: "xarray.core.utils.Frozen", + preserved_dims=("batch", "lat", "lon"), +) -> "xarray.Variable": + """Converts an xarray.Variable to preserved_dims + ("channels",). + + Any dimensions other than those included in preserved_dims get stacked into a final "channels" dimension. If any of the preserved_dims are missing then they are added, with the data broadcast/tiled to match the sizes specified in `sizes`. + + Args: + variable: An xarray.Variable. + sizes: Mapping including sizes for any dimensions which are not present in `variable` but are needed for the output. This may be needed for example for a static variable with only ("lat", "lon") dims, or if you want to encode just the latitude coordinates (a variable with dims ("lat",)). + preserved_dims: dimensions of variable to not be folded in channels. + + Returns: + An xarray.Variable with dimensions preserved_dims + ("channels",). + """ + stack_to_channels_dims = [d for d in variable.dims if d not in preserved_dims] + if stack_to_channels_dims: + variable = variable.stack(channels=stack_to_channels_dims) + dims = {dim: variable.sizes.get(dim) or sizes[dim] for dim in preserved_dims} + dims["channels"] = variable.sizes.get("channels", 1) + return variable.set_dims(dims) + + +def dataset_to_stacked( + dataset: "xarray.Dataset", + sizes=None, + preserved_dims=("batch", "lat", "lon"), +) -> "xarray.DataArray": + """Converts an xarray.Dataset to a single stacked array. + + This takes each consistuent data_var, converts it into BHWC layout + using `variable_to_stacked`, then concats them all along the channels axis. + + Args: + dataset: An xarray.Dataset. + sizes: Mapping including sizes for any dimensions which are not present in the `dataset` but are needed for the output. See variable_to_stacked. + preserved_dims: dimensions from the dataset that should not be folded in the predictions channels. + + Returns: + An xarray.DataArray with dimensions preserved_dims + ("channels",). Existing coordinates for preserved_dims axes will be preserved, however there will be no coordinates for "channels". + """ + data_vars = [ + variable_to_stacked( + dataset.variables[name], sizes or dataset.sizes, preserved_dims + ) + for name in sorted(dataset.data_vars.keys()) + ] + coords = { + dim: coord for dim, coord in dataset.coords.items() if dim in preserved_dims + } + return xarray.DataArray( + data=xarray.Variable.concat(data_vars, dim="channels"), coords=coords + ) + + +class GridMeshAtmosphericDataset(io.Dataset): + """This class is used to process ERA5 re-analyze data, and is used to generate the dataset generator supported by MindSpore. This class inherits the Data class. + + Args: + input_keys (Tuple[str, ...]): Name of input data. + label_keys (Tuple[str, ...]): Name of label data. + data_path: Path of atmospheric datafile. + mean_path: Path of mean datafile. + stddev_path: Path of standard deviation datafile. + stddev_diffs_path: Path of standard deviation different datafile. + type: Type of GraphCast network. + mesh_size: Size of mesh. + mesh2grid_edge_normalization_factor: Factor of normalization of edges in Mesh2Grid GNN. + radius_query_fraction_edge_length: Length of radius query fraction edges. + resolution: Resolution of atmospheric data. + + Examples: + >>> import ppsci + >>> dataset = ppsci.data.dataset.GridMeshAtmosphericDataset( + ... "input_keys": ("input",), + ... "label_keys": ("output",), + ... "data_path": "/path/to/file.nc", + ... "mean_path": "/path/to/file.nc", + ... "stddev_path": "/path/to/file.nc", + ... "stddev_diffs_path": "/path/to/file.nc", + ... "type": "graphcast_small", + ... "mesh_size": 5, + ... "mesh2grid_edge_normalization_factor": 0.06, + ... "radius_query_fraction_edge_length": 0.6180338738074472, + ... "resolution": 1, + ... ) # doctest: +SKIP + """ + + use_graph_grid_mesh: bool = True + + def __init__( + self, + input_keys: Tuple[str, ...], + label_keys: Tuple[str, ...], + data_path: str, + mean_path: str, + stddev_path: str, + stddev_diffs_path: str, + type: str, + mesh_size: int, + mesh2grid_edge_normalization_factor: float, + radius_query_fraction_edge_length: float, + resolution: float, + ): + super().__init__() + self.input_keys = input_keys + self.label_keys = label_keys + if type == "graphcast": + self.input_variables = TASK_input_variables + self.forcing_variables = TASK_forcing_variables + self.target_variables = TASK_target_variables + self.level_variables = PRESSURE_LEVELS[37] + elif type == "graphcast_small": + self.input_variables = TASK_13_input_variables + self.forcing_variables = TASK_13_forcing_variables + self.target_variables = TASK_13_target_variables + self.level_variables = PRESSURE_LEVELS[13] + elif type == "graphcast_operational": + self.input_variables = TASK_13_PRECIP_OUT_input_variables + self.forcing_variables = TASK_13_PRECIP_OUT_forcing_variables + self.target_variables = TASK_13_PRECIP_OUT_target_variables + self.level_variables = PRESSURE_LEVELS[13] + + nc_dataset = xarray.open_dataset(data_path) + + longitude_offsets = nc_dataset.coords["lon"].data + second_since_epoch = ( + nc_dataset.coords["datetime"].data.astype("datetime64[s]").astype(np.int64) + ) + datetime_feats = datetime_features(second_since_epoch, longitude_offsets) + nc_dataset.update( + { + "year_progress_sin": xarray.Variable( + ("batch", "time"), datetime_feats["year_progress_sin"] + ), + "year_progress_cos": xarray.Variable( + ("batch", "time"), datetime_feats["year_progress_cos"] + ), + "day_progress_sin": xarray.Variable( + ("batch", "time", "lon"), datetime_feats["day_progress_sin"] + ), + "day_progress_cos": xarray.Variable( + ("batch", "time", "lon"), datetime_feats["day_progress_cos"] + ), + } + ) + + inputs, targets = extract_input_target_times( + nc_dataset, input_duration="12h", target_lead_times="6h" + ) + + stddev_data = xarray.open_dataset(stddev_path).sel( + level=list(self.level_variables) + ) + stddev_diffs_data = xarray.open_dataset(stddev_diffs_path).sel( + level=list(self.level_variables) + ) + mean_data = xarray.open_dataset(mean_path).sel(level=list(self.level_variables)) + + missing_variables = set(self.target_variables) - set(self.input_variables) + exist_variables = set(self.target_variables) - missing_variables + targets_stddev = stddev_diffs_data[list(exist_variables)] + target_mean = inputs[list(exist_variables)].isel(time=-1) + if missing_variables: + targets_stddev.update({var: stddev_data[var] for var in missing_variables}) + target_mean.update( + {var: mean_data.variables[var] for var in missing_variables} + ) + + stacked_targets_stddev = dataset_to_stacked(targets_stddev, preserved_dims=()) + stacked_targets_mean = dataset_to_stacked(target_mean) + stacked_targets_mean = stacked_targets_mean.transpose("lat", "lon", ...) + + inputs = inputs[list(self.input_variables)] + forcings = targets[list(self.forcing_variables)] + targets = targets[list(self.target_variables)] + inputs = self.normalize(inputs, stddev_data, mean_data) + forcings = self.normalize(forcings, stddev_data, mean_data) + + self.targets_template = targets + + stacked_inputs = dataset_to_stacked(inputs) + stacked_forcings = dataset_to_stacked(forcings) + stacked_targets = dataset_to_stacked(targets) + stacked_inputs = xarray.concat( + [stacked_inputs, stacked_forcings], dim="channels" + ) + + stacked_inputs = stacked_inputs.transpose("lat", "lon", ...) + stacked_targets = stacked_targets.transpose("lat", "lon", ...) + + lat_dim, lon_dim, batch_dim, feat_dim = stacked_inputs.shape + stacked_inputs = stacked_inputs.data.reshape(lat_dim * lon_dim, batch_dim, -1) + stacked_targets = stacked_targets.data.reshape(lat_dim * lon_dim, batch_dim, -1) + self.stacked_targets_stddev = stacked_targets_stddev.data + self.stacked_targets_mean = stacked_targets_mean.data.reshape( + lat_dim * lon_dim, batch_dim, -1 + ) + + self.input_data = [] + self.target_data = [] + + graph = GraphGridMesh( + mesh_size=mesh_size, + radius_query_fraction_edge_length=radius_query_fraction_edge_length, + mesh2grid_edge_normalization_factor=mesh2grid_edge_normalization_factor, + resolution=resolution, + ) + + graph.grid_node_feat = np.concatenate( + [stacked_inputs, graph.grid_node_feat], axis=-1 + ) + mesh_node_feat = np.zeros([graph.mesh_num_nodes, batch_dim, feat_dim]) + graph.mesh_node_feat = np.concatenate( + [mesh_node_feat, graph.mesh_node_feat], axis=-1 + ) + + self.input_data.append(graph) + self.target_data.append(stacked_targets) + + def __len__(self): + return len(self.input_data) + + def __getitem__(self, idx): + return ( + { + self.input_keys[0]: self.input_data[idx], + }, + { + self.label_keys[0]: self.target_data[idx], + }, + None, + ) + + def normalize(self, inputs_data, stddev_data, mean_data): + for name in list(inputs_data.keys()): + inputs_data[name] = (inputs_data[name] - mean_data[name]) / stddev_data[ + name + ] + return inputs_data + + def denormalize(self, inputs_data): + return inputs_data * self.stacked_targets_stddev + self.stacked_targets_mean diff --git a/examples/smc_reac/ppsci/data/dataset/cgcnn_dataset.py b/examples/smc_reac/ppsci/data/dataset/cgcnn_dataset.py new file mode 100644 index 0000000000..0c04c5e319 --- /dev/null +++ b/examples/smc_reac/ppsci/data/dataset/cgcnn_dataset.py @@ -0,0 +1,312 @@ +# Copyright (c) 2024 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. + +import csv +import functools +import json +import os +import random +import warnings +from typing import Tuple + +import numpy as np +import paddle +from paddle import io + +try: + from pymatgen.core.structure import Structure +except ModuleNotFoundError: + pass + + +def collate_pool(dataset_list): + + """ + Collate a list of data and return a batch for predicting crystal properties. + + Args: + dataset_list (list): A list of tuples for each data point containing: + - atom_fea (paddle.Tensor): Shape (n_i, atom_fea_len). + - nbr_fea (paddle.Tensor): Shape (n_i, M, nbr_fea_len). + - nbr_fea_idx (paddle.Tensor): Shape (n_i, M). + - target (paddle.Tensor): Shape (1,). + - cif_id (str or int). + + Returns: + tuple: Contains the following: + - batch_atom_fea (paddle.Tensor): Shape (N, orig_atom_fea_len). Atom features from atom type. + - batch_nbr_fea (paddle.Tensor): Shape (N, M, nbr_fea_len). Bond features of each atom's M neighbors. + - batch_nbr_fea_idx (paddle.Tensor): Shape (N, M). Indices of M neighbors of each atom. + - crystal_atom_idx (list): List of paddle.Tensor of length N0. Mapping from the crystal idx to atom idx. + - target (paddle.Tensor): Shape (N, 1). Target value for prediction. + - batch_cif_ids (list): List of CIF IDs. + + Notes: + - N = sum(n_i); N0 = sum(i) + """ + batch_atom_fea, batch_nbr_fea, batch_nbr_fea_idx = [], [], [] + crystal_atom_idx, batch_target = [], [] + batch_cif_ids = [] + base_idx = 0 + for i, item in enumerate(dataset_list): + input: Tuple[np.ndarray, np.ndarray, np.ndarray] = item[0]["i"] + label = item[1]["l"] + id = item[2]["c"] + atom_fea, nbr_fea, nbr_fea_idx = input + target = label + cif_id = id + n_i = atom_fea.shape[0] # number of atoms for this crystal + batch_atom_fea.append(atom_fea) + batch_nbr_fea.append(nbr_fea) + batch_nbr_fea_idx.append(nbr_fea_idx + base_idx) + new_idx = np.arange(n_i, dtype="int64") + int(base_idx) + crystal_atom_idx.append(new_idx) + batch_target.append(target) + batch_cif_ids.append(cif_id) + base_idx += n_i + # Debugging: print shapes of the tensors to ensure they are consistent + # print("Shapes of batch_atom_fea:", [x.shape for x in batch_atom_fea]) + # print("Shapes of batch_nbr_fea:", [x.shape for x in batch_nbr_fea]) + # print("Shapes of batch_nbr_fea_idx:", [x.shape for x in batch_nbr_fea_idx]) + # Ensure all tensors in the lists have consistent shapes before concatenation + batch_atom_fea = np.concatenate(batch_atom_fea, axis=0) + batch_nbr_fea = np.concatenate(batch_nbr_fea, axis=0) + batch_nbr_fea_idx = np.concatenate(batch_nbr_fea_idx, axis=0) + return ( + { + "i": ( + np.array(batch_atom_fea, dtype="float32"), + np.array(batch_nbr_fea, dtype="float32"), + np.array(batch_nbr_fea_idx), + [np.array(crys_idx) for crys_idx in crystal_atom_idx], + ) + }, + {"l": np.array(np.stack(batch_target, axis=0))}, + {"c": batch_cif_ids}, + ) + + +class GaussianDistance(object): + """ + Expands the distance by Gaussian basis. + + Args: + dmin (float): Minimum interatomic distance. + dmax (float): Maximum interatomic distance. + step (float): Step size for the Gaussian filter. + """ + + def __init__(self, dmin, dmax, step, var=None): + assert dmin < dmax + assert dmax - dmin > step + self.filter = np.arange(dmin, dmax + step, step) + if var is None: + var = step + self.var = var + + def expand(self, distances): + """ + Apply Gaussian distance filter to a numpy distance array. + + Args: + distance (np.array): n-dimensional distance matrix of any shape. + + Returns: + np.array: Expanded distance matrix with the last dimension of length len(self.filter). + """ + + return np.exp( + -((distances[..., np.newaxis] - self.filter) ** 2) / self.var**2 + ) + + +class AtomInitializer(object): + """ + Base class for initializing the vector representation for atoms. + + !!! Use one AtomInitializer per dataset !!! + """ + + def __init__(self, atom_types): + self.atom_types = set(atom_types) + self._embedding = {} + + def get_atom_fea(self, atom_type): + assert atom_type in self.atom_types + return self._embedding[atom_type] + + def load_state_dict(self, state_dict): + self._embedding = state_dict + self.atom_types = set(self._embedding.keys()) + self._decodedict = { + idx: atom_type for atom_type, idx in self._embedding.items() + } + + def state_dict(self): + return self._embedding + + def decode(self, idx): + if not hasattr(self, "_decodedict"): + self._decodedict = { + idx: atom_type for atom_type, idx in self._embedding.items() + } + return self._decodedict[idx] + + +class AtomCustomJSONInitializer(AtomInitializer): + """ + Initialize atom feature vectors using a JSON file, which is a Python dictionary mapping from element number to a list representing the feature vector of the element. + + Args: + elem_embedding_file (str): The path to the .json file. + """ + + def __init__(self, elem_embedding_file): + with open(elem_embedding_file) as f: + elem_embedding = json.load(f) + elem_embedding = {int(key): value for key, value in elem_embedding.items()} + atom_types = set(elem_embedding.keys()) + super(AtomCustomJSONInitializer, self).__init__(atom_types) + for key, value in elem_embedding.items(): + self._embedding[key] = np.array(value, dtype=float) + + +class CIFData(io.Dataset): + """ + The CIFData dataset is a wrapper for a dataset where the crystal structures + are stored in the form of CIF files. The dataset should have the following + directory structure: + + root_dir + ├── id_prop.csv + ├── atom_init.json + ├── id0.cif + ├── id1.cif + ├── ... + + id_prop.csv: a CSV file with two columns. The first column recodes a + unique ID for each crystal, and the second column recodes the value of + target property. + + atom_init.json: a JSON file that stores the initialization vector for each element. + + ID.cif: a CIF file that recodes the crystal structure, where ID is the + unique ID for the crystal. + + Args + root_dir (str): The path to the root directory of the dataset + max_num_nbr (int): The maximum number of neighbors while constructing the crystal graph + radius (float): The cutoff radius for searching neighbors + dmin (float): The minimum distance for constructing GaussianDistance + step (float): The step size for constructing GaussianDistance + random_seed (int): Random seed for shuffling the dataset + + + Returns + atom_fea (paddle.Tensor): Shape (n_i, atom_fea_len) + nbr_fea (paddle.Tensor): Shape (n_i, M, nbr_fea_len) + nbr_fea_idx (paddle.Tensor): Shape (n_i, M) + target (paddle.Tensor): Shape (1, ) + cif_id (str or int) + + Examples: + >>> import ppsci + >>> dataset = ppsci.data.dataset.CGCNNDataset( + ... "file_path": "/path/to/CGCNNDataset", + ... "input_keys": "i", + ... "label_keys": "l", + ... "id_keys": "c", + ... ) # doctest: +SKIP + """ + + def __init__( + self, + root_dir: str, + input_keys: Tuple[str, ...], + label_keys: Tuple[str, ...], + id_keys: Tuple[str, ...], + max_num_nbr: int = 12, + radius: int = 8, + dmin: int = 0, + step: float = 0.2, + random_seed: int = 123, + ): + super().__init__() + self.input_keys = input_keys + self.label_keys = label_keys + self.id_keys = id_keys + self.root_dir = root_dir + self.max_num_nbr, self.radius = max_num_nbr, radius + assert os.path.exists(root_dir), "root_dir does not exist!" + id_prop_file = os.path.join(self.root_dir, "id_prop.csv") + assert os.path.exists(id_prop_file), "id_prop.csv does not exist!" + with open(id_prop_file) as f: + reader = csv.reader(f) + self.id_prop_data = [row for row in reader] + random.seed(random_seed) + random.shuffle(self.id_prop_data) + atom_init_file = os.path.join(self.root_dir, "atom_init.json") + assert os.path.exists(atom_init_file), f"{atom_init_file} does not exist!" + self.ari = AtomCustomJSONInitializer(atom_init_file) + self.gdf = GaussianDistance(dmin=dmin, dmax=self.radius, step=step) + self.raw_data = [self.get(i) for i in range(len(self))] + + def __len__(self): + return len(self.id_prop_data) + + @functools.lru_cache(maxsize=None) # Cache loaded structures + def __getitem__(self, idx): + return ( + {self.input_keys[0]: self.raw_data[idx][0]}, + {self.label_keys[0]: self.raw_data[idx][1]}, + {self.id_keys[0]: self.raw_data[idx][2]}, + ) + + def get(self, idx): + cif_id, target = self.id_prop_data[idx] + crystal = Structure.from_file(os.path.join(self.root_dir, cif_id + ".cif")) + atom_fea = np.vstack( + [ + self.ari.get_atom_fea(crystal[i].specie.number) + for i in range(len(crystal)) + ] + ) + atom_fea = paddle.Tensor(atom_fea) + all_nbrs = crystal.get_all_neighbors(self.radius, include_index=True) + all_nbrs = [sorted(nbrs, key=lambda x: x[1]) for nbrs in all_nbrs] + nbr_fea_idx, nbr_fea = [], [] + for nbr in all_nbrs: + if len(nbr) < self.max_num_nbr: + warnings.warn( + f"{cif_id} not find enough neighbors to build graph. " + "If it happens frequently, consider increase " + "radius." + ) + nbr_fea_idx.append( + list(map(lambda x: x[2], nbr)) + [0] * (self.max_num_nbr - len(nbr)) + ) + nbr_fea.append( + list(map(lambda x: x[1], nbr)) + + [self.radius + 1.0] * (self.max_num_nbr - len(nbr)) + ) + else: + nbr_fea_idx.append(list(map(lambda x: x[2], nbr[: self.max_num_nbr]))) + nbr_fea.append(list(map(lambda x: x[1], nbr[: self.max_num_nbr]))) + nbr_fea_idx, nbr_fea = np.array(nbr_fea_idx), np.array(nbr_fea) + nbr_fea = self.gdf.expand(nbr_fea) + atom_fea = np.array(atom_fea) + nbr_fea = np.array(nbr_fea) + nbr_fea_idx = np.array(nbr_fea_idx, dtype="int64") + target = np.array([float(target)], dtype="float32") + return (atom_fea, nbr_fea, nbr_fea_idx), target, cif_id diff --git a/examples/smc_reac/ppsci/data/dataset/csv_dataset.py b/examples/smc_reac/ppsci/data/dataset/csv_dataset.py new file mode 100644 index 0000000000..c14bb107da --- /dev/null +++ b/examples/smc_reac/ppsci/data/dataset/csv_dataset.py @@ -0,0 +1,287 @@ +# 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. + +from __future__ import annotations + +from typing import Callable +from typing import Dict +from typing import Optional +from typing import Tuple +from typing import Union + +import numpy as np +import paddle +from paddle import io +from paddle import vision + +from ppsci.utils import misc +from ppsci.utils import reader + + +class CSVDataset(io.Dataset): + """Dataset class for .csv file. + + Args: + file_path (str): CSV file path. + input_keys (Tuple[str, ...]): List of input keys. + label_keys (Tuple[str, ...]): List of label keys. + alias_dict (Optional[Dict[str, str]]): Dict of alias(es) for input and label keys. + i.e. {inner_key: outer_key}. Defaults to None. + weight_dict (Optional[Dict[str, Union[Callable, float]]]): Define the weight of + each constraint variable. Defaults to None. + timestamps (Optional[Tuple[float, ...]]): The number of repetitions of the data + in the time dimension. Defaults to None. + transforms (Optional[vision.Compose]): Compose object contains sample wise + transform(s). Defaults to None. + + Examples: + >>> import ppsci + >>> dataset = ppsci.data.dataset.CSVDataset( + ... "/path/to/file.csv", + ... ("x",), + ... ("u",), + ... ) # doctest: +SKIP + """ + + # Whether support batch indexing for speeding up fetching process. + batch_index: bool = True + + def __init__( + self, + file_path: str, + input_keys: Tuple[str, ...], + label_keys: Tuple[str, ...], + alias_dict: Optional[Dict[str, str]] = None, + weight_dict: Optional[Dict[str, Union[Callable, float]]] = None, + timestamps: Optional[Tuple[float, ...]] = None, + transforms: Optional[vision.Compose] = None, + ): + super().__init__() + self.input_keys = input_keys + self.label_keys = label_keys + + # read raw data from file + raw_data = reader.load_csv_file( + file_path, + input_keys + label_keys, + alias_dict, + ) + # filter raw data by given timestamps if specified + if timestamps is not None: + if "t" in raw_data: + # filter data according to given timestamps + raw_time_array = raw_data["t"] + mask = [] + for ti in timestamps: + mask.append(np.nonzero(np.isclose(raw_time_array, ti).flatten())[0]) + raw_data = misc.convert_to_array( + raw_data, self.input_keys + self.label_keys + ) + mask = np.concatenate(mask, 0) + raw_data = raw_data[mask] + raw_data = misc.convert_to_dict( + raw_data, self.input_keys + self.label_keys + ) + else: + # repeat data according to given timestamps + raw_data = misc.convert_to_array( + raw_data, self.input_keys + self.label_keys + ) + raw_data = misc.combine_array_with_time(raw_data, timestamps) + self.input_keys = ("t",) + tuple(self.input_keys) + raw_data = misc.convert_to_dict( + raw_data, self.input_keys + self.label_keys + ) + + # fetch input data + self.input = { + key: value for key, value in raw_data.items() if key in self.input_keys + } + # fetch label data + self.label = { + key: value for key, value in raw_data.items() if key in self.label_keys + } + + # prepare weights + self.weight = ( + {key: np.ones_like(next(iter(self.label.values()))) for key in self.label} + if weight_dict is not None + else {} + ) + if weight_dict is not None: + for key, value in weight_dict.items(): + if isinstance(value, (int, float)): + self.weight[key] = np.full_like( + next(iter(self.label.values())), value + ) + elif callable(value): + func = value + self.weight[key] = func(self.input) + if isinstance(self.weight[key], (int, float)): + self.weight[key] = np.full_like( + next(iter(self.label.values())), self.weight[key] + ) + else: + raise NotImplementedError(f"type of {type(value)} is invalid yet.") + + self.transforms = transforms + self._len = len(next(iter(self.input.values()))) + + def __getitem__(self, idx): + input_item = {key: value[idx] for key, value in self.input.items()} + label_item = {key: value[idx] for key, value in self.label.items()} + weight_item = {key: value[idx] for key, value in self.weight.items()} + + if self.transforms is not None: + input_item, label_item, weight_item = self.transforms( + input_item, label_item, weight_item + ) + + return (input_item, label_item, weight_item) + + def __len__(self): + return self._len + + +class IterableCSVDataset(io.IterableDataset): + """IterableCSVDataset for full-data loading. + + Args: + file_path (str): CSV file path. + input_keys (Tuple[str, ...]): List of input keys. + label_keys (Tuple[str, ...]): List of label keys. + alias_dict (Optional[Dict[str, str]]): Dict of alias(es) for input and label keys. + Defaults to None. + weight_dict (Optional[Dict[str, Union[Callable, float]]]): Define the weight of + each constraint variable. Defaults to None. + timestamps (Optional[Tuple[float, ...]]): The number of repetitions of the data + in the time dimension. Defaults to None. + transforms (Optional[vision.Compose]): Compose object contains sample wise + transform(s). Defaults to None. + + Examples: + >>> import ppsci + >>> dataset = ppsci.data.dataset.IterableCSVDataset( + ... "/path/to/file.csv" + ... ("x",), + ... ("u",), + ... ) # doctest: +SKIP + """ + + # Whether support batch indexing for speeding up fetching process. + batch_index: bool = False + + def __init__( + self, + file_path: str, + input_keys: Tuple[str, ...], + label_keys: Tuple[str, ...], + alias_dict: Optional[Dict[str, str]] = None, + weight_dict: Optional[Dict[str, Union[Callable, float]]] = None, + timestamps: Optional[Tuple[float, ...]] = None, + transforms: Optional[vision.Compose] = None, + ): + super().__init__() + self.input_keys = input_keys + self.label_keys = label_keys + + # read raw data from file + raw_data = reader.load_csv_file( + file_path, + input_keys + label_keys, + alias_dict, + ) + # filter raw data by given timestamps if specified + if timestamps is not None: + if "t" in raw_data: + # filter data according to given timestamps + raw_time_array = raw_data["t"] + mask = [] + for ti in timestamps: + mask.append(np.nonzero(np.isclose(raw_time_array, ti).flatten())[0]) + raw_data = misc.convert_to_array( + raw_data, self.input_keys + self.label_keys + ) + mask = np.concatenate(mask, 0) + raw_data = raw_data[mask] + raw_data = misc.convert_to_dict( + raw_data, self.input_keys + self.label_keys + ) + else: + # repeat data according to given timestamps + raw_data = misc.convert_to_array( + raw_data, self.input_keys + self.label_keys + ) + raw_data = misc.combine_array_with_time(raw_data, timestamps) + self.input_keys = ("t",) + tuple(self.input_keys) + raw_data = misc.convert_to_dict( + raw_data, self.input_keys + self.label_keys + ) + + # fetch input data + self.input = { + key: value for key, value in raw_data.items() if key in self.input_keys + } + # fetch label data + self.label = { + key: value for key, value in raw_data.items() if key in self.label_keys + } + + # prepare weights + self.weight = ( + {key: np.ones_like(next(iter(self.label.values()))) for key in self.label} + if weight_dict is not None + else {} + ) + if weight_dict is not None: + for key, value in weight_dict.items(): + if isinstance(value, (int, float)): + self.weight[key] = np.full_like( + next(iter(self.label.values())), value + ) + elif callable(value): + func = value + self.weight[key] = func(self.input) + if isinstance(self.weight[key], (int, float)): + self.weight[key] = np.full_like( + next(iter(self.label.values())), self.weight[key] + ) + else: + raise NotImplementedError(f"type of {type(value)} is invalid yet.") + + self.input = {key: paddle.to_tensor(value) for key, value in self.input.items()} + self.label = {key: paddle.to_tensor(value) for key, value in self.label.items()} + self.weight = { + key: paddle.to_tensor(value) for key, value in self.weight.items() + } + + self.transforms = transforms + self._len = len(next(iter(self.input.values()))) + + @property + def num_samples(self): + """Number of samples within current dataset.""" + return self._len + + def __iter__(self): + if callable(self.transforms): + input_, label_, weight_ = self.transforms( + self.input, self.label, self.weight + ) + yield input_, label_, weight_ + else: + yield self.input, self.label, self.weight + + def __len__(self): + return 1 diff --git a/examples/smc_reac/ppsci/data/dataset/cylinder_dataset.py b/examples/smc_reac/ppsci/data/dataset/cylinder_dataset.py new file mode 100644 index 0000000000..3a49b7d436 --- /dev/null +++ b/examples/smc_reac/ppsci/data/dataset/cylinder_dataset.py @@ -0,0 +1,215 @@ +# 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. + +from __future__ import annotations + +import os +from os import path as osp +from typing import Tuple + +import numpy as np +import paddle +from paddle import io + +from ppsci.data.dataset import airfoil_dataset + +try: + import pgl +except ModuleNotFoundError: + pass + +SU2_SHAPE_IDS = { + "line": 3, + "triangle": 5, + "quad": 9, +} + + +class MeshCylinderDataset(io.Dataset): + """Dataset for `MeshCylinder`. + + Args: + input_keys (Tuple[str, ...]): Name of input data. + label_keys (Tuple[str, ...]): Name of label data. + data_dir (str): Directory of MeshCylinder data. + mesh_graph_path (str): Path of mesh graph. + + Examples: + >>> import ppsci + >>> dataset = ppsci.data.dataset.MeshAirfoilDataset( + ... "input_keys": ("input",), + ... "label_keys": ("output",), + ... "data_dir": "/path/to/MeshAirfoilDataset", + ... "mesh_graph_path": "/path/to/file.su2", + ... ) # doctest: +SKIP + """ + + # Whether support batch indexing for speeding up fetching process. + batch_index: bool = False + + use_pgl: bool = True + + def __init__( + self, + input_keys: Tuple[str, ...], + label_keys: Tuple[str, ...], + data_dir: str, + mesh_graph_path: str, + ): + self.input_keys = input_keys + self.label_keys = label_keys + self.data_dir = data_dir + self.file_list = os.listdir(self.data_dir) + self.len = len(self.file_list) + self.mesh_graph = airfoil_dataset._get_mesh_graph(mesh_graph_path) + + self.normalization_factors = np.array( + [[978.6001, 48.9258, 24.8404], [-692.3159, -6.9950, -24.8572]], + dtype=paddle.get_default_dtype(), + ) + + self.nodes = self.mesh_graph[0] + self.meshnodes = self.mesh_graph[0] + self.edges = self.mesh_graph[1] + self.elems_list = self.mesh_graph[2] + self.marker_dict = self.mesh_graph[3] + self.bounder = [] + self.node_markers = np.full([self.nodes.shape[0], 1], fill_value=-1) + for i, (marker_tag, marker_elems) in enumerate(self.marker_dict.items()): + for elem in marker_elems: + self.node_markers[elem[0]] = i + self.node_markers[elem[1]] = i + + self.raw_graphs = [self.get(i) for i in range(len(self))] + + def __len__(self): + return self.len + + def __getitem__(self, idx): + return ( + { + self.input_keys[0]: self.raw_graphs[idx], + }, + { + self.label_keys[0]: self.raw_graphs[idx], + }, + None, + ) + + def get(self, idx): + with open(osp.join(self.data_dir, self.file_list[idx]), "r") as f: + field = [] + pos = [] + for line in f.read().splitlines()[1:]: + lines_pos = line.split(",")[1:3] + lines_field = line.split(",")[3:] + numbers_float = list(eval(i) for i in lines_pos) + array = np.array(numbers_float, paddle.get_default_dtype()) + pos.append(array) + numbers_float = list(eval(i) for i in lines_field) + array = np.array(numbers_float, paddle.get_default_dtype()) + field.append(array) + + field = np.stack(field, axis=0) + pos = np.stack(pos, axis=0) + indexlist = [] + for i in range(self.meshnodes.shape[0]): + b = self.meshnodes[i : (i + 1)] + b = np.squeeze(b) + index = np.nonzero( + np.sum((pos == b), axis=1, dtype=paddle.get_default_dtype()) + == pos.shape[1] + ) + indexlist.append(index) + indexlist = np.stack(indexlist, axis=0) + indexlist = np.squeeze(indexlist) + fields = field[indexlist] + velocity = self._get_params_from_name(self.file_list[idx]) + + norm_aoa = velocity / 40 + # add physics parameters to graph + nodes = np.concatenate( + [ + self.nodes, + np.repeat(a=norm_aoa, repeats=self.nodes.shape[0])[:, np.newaxis], + self.node_markers, + ], + axis=-1, + ).astype(paddle.get_default_dtype()) + + data = pgl.Graph( + num_nodes=nodes.shape[0], + edges=self.edges, + ) + data.x = nodes + data.y = fields + data.pos = self.nodes + data.edge_index = self.edges + data.velocity = velocity + + sender = data.x[data.edge_index[0]] + receiver = data.x[data.edge_index[1]] + relation_pos = sender[:, 0:2] - receiver[:, 0:2] + post = np.linalg.norm(relation_pos, ord=2, axis=1, keepdims=True).astype( + paddle.get_default_dtype() + ) + data.edge_attr = post + std_epsilon = [1e-8] + a = np.mean(data.edge_attr, axis=0) + b = data.edge_attr.std(axis=0) + b = np.maximum(b, std_epsilon).astype(paddle.get_default_dtype()) + data.edge_attr = (data.edge_attr - a) / b + a = np.mean(data.y, axis=0) + b = data.y.std(axis=0) + b = np.maximum(b, std_epsilon).astype(paddle.get_default_dtype()) + data.y = (data.y - a) / b + data.norm_max = a + data.norm_min = b + + # find the face of the boundary,our cylinder dataset come from fluent solver + with open(osp.join(osp.dirname(self.data_dir), "bounder"), "r") as f: + field = [] + pos = [] + for line in f.read().splitlines()[1:]: + lines_pos = line.split(",")[1:3] + lines_field = line.split(",")[3:] + numbers_float = list(eval(i) for i in lines_pos) + array = np.array(numbers_float, paddle.get_default_dtype()) + pos.append(array) + numbers_float = list(eval(i) for i in lines_field) + array = np.array(numbers_float, paddle.get_default_dtype()) + field.append(array) + + field = np.stack(field, axis=0) + pos = np.stack(pos, axis=0) + + indexlist = [] + for i in range(pos.shape[0]): + b = pos[i : (i + 1)] + b = np.squeeze(b) + index = np.nonzero( + np.sum((self.nodes == b), axis=1, dtype=paddle.get_default_dtype()) + == self.nodes.shape[1] + ) + indexlist.append(index) + + indexlist = np.stack(indexlist, axis=0) + indexlist = np.squeeze(indexlist) + self.bounder = indexlist + return data + + def _get_params_from_name(self, filename): + s = filename.rsplit(".", 1)[0] + reynolds = np.array(s[13:])[np.newaxis].astype(paddle.get_default_dtype()) + return reynolds diff --git a/examples/smc_reac/ppsci/data/dataset/darcyflow_dataset.py b/examples/smc_reac/ppsci/data/dataset/darcyflow_dataset.py new file mode 100644 index 0000000000..3e748eb785 --- /dev/null +++ b/examples/smc_reac/ppsci/data/dataset/darcyflow_dataset.py @@ -0,0 +1,296 @@ +from pathlib import Path +from typing import Dict +from typing import Optional +from typing import Tuple + +import numpy as np +import paddle +from paddle import io + + +# normalization, pointwise gaussian +class UnitGaussianNormalizer: + def __init__(self, x, eps=1e-7, reduce_dim=[0], verbose=True): + super().__init__() + n_samples, *shape = x.shape + self.sample_shape = shape + self.verbose = verbose + self.reduce_dim = reduce_dim + + # x could be in shape of ntrain*n or ntrain*T*n or ntrain*n*T + self.mean = paddle.mean(x, reduce_dim, keepdim=True).squeeze(0) + self.std = paddle.std(x, reduce_dim, keepdim=True).squeeze(0) + self.eps = eps + + if verbose: + print( + f"UnitGaussianNormalizer init on {n_samples}, reducing over {reduce_dim}, samples of shape {shape}." + ) + print(f" Mean and std of shape {self.mean.shape}, eps={eps}") + + def encode(self, x): + + x -= self.mean + x /= self.std + self.eps + return x + + def decode(self, x, sample_idx=None): + if sample_idx is None: + std = self.std + self.eps # n + mean = self.mean + else: + if len(self.mean.shape) == len(sample_idx[0].shape): + std = self.std[sample_idx] + self.eps # batch*n + mean = self.mean[sample_idx] + if len(self.mean.shape) > len(sample_idx[0].shape): + std = self.std[:, sample_idx] + self.eps # T*batch*n + mean = self.mean[:, sample_idx] + + # x is in shape of batch*n or T*batch*n + x *= std + x += mean + + return x + + +def get_grid_positional_encoding( + input_tensor, grid_boundaries=[[0, 1], [0, 1]], channel_dim=1 +): + """Appends grid positional encoding to an input tensor, concatenating as additional dimensions along the channels. + + Args: + input_tensor (paddle.Tensor): The input tensor. + grid_boundaries (list, optional): The boundaries of the grid. Defaults to [[0, 1], [0, 1]]. + channel_dim (int, optional): The location of unsqueeze. Defaults to 1. + """ + + shape = list(input_tensor.shape) + if len(shape) == 2: + height, width = shape + else: + _, height, width = shape + + xt = paddle.linspace(grid_boundaries[0][0], grid_boundaries[0][1], height + 1)[:-1] + yt = paddle.linspace(grid_boundaries[1][0], grid_boundaries[1][1], width + 1)[:-1] + + grid_x, grid_y = paddle.meshgrid(xt, yt, indexing="ij") + + if len(shape) == 2: + grid_x = grid_x.unsqueeze(channel_dim) + grid_y = grid_y.unsqueeze(channel_dim) + else: + grid_x = grid_x.unsqueeze(0).unsqueeze(channel_dim) + grid_y = grid_y.unsqueeze(0).unsqueeze(channel_dim) + + return grid_x, grid_y + + +def regular_grid(spatial_dims, grid_boundaries=[[0, 1], [0, 1]]): + """ + Appends grid positional encoding to an input tensor, concatenating as additional dimensions along the channels + """ + height, width = spatial_dims + + xt = paddle.linspace(grid_boundaries[0][0], grid_boundaries[0][1], height + 1)[:-1] + yt = paddle.linspace(grid_boundaries[1][0], grid_boundaries[1][1], width + 1)[:-1] + + grid_x, grid_y = paddle.meshgrid(xt, yt, indexing="ij") + + grid_x = grid_x.tile((1, 1)) + grid_y = grid_y.tile((1, 1)) + + return grid_x, grid_y + + +class PositionalEmbedding2D: + def __init__(self, grid_boundaries=[[0, 1], [0, 1]]): + self.grid_boundaries = grid_boundaries + self._grid = None + self._res = None + + def grid(self, spatial_dims, dtype): + """Grid generates 2D grid needed for pos encoding + and caches the grid associated with MRU resolution + + Args: + spatial_dims (tuple[int,...]): Sizes of spatial resolution. + dtype (str): Dtype to encode data. + + Returns: + paddle.Tensor: Output grids to concatenate + """ + # handle case of multiple train resolutions + if self._grid is None or self._res != spatial_dims: + grid_x, grid_y = regular_grid( + spatial_dims, grid_boundaries=self.grid_boundaries + ) + + grid_x = grid_x.astype(dtype).unsqueeze(0).unsqueeze(0) + grid_y = grid_y.astype(dtype).unsqueeze(0).unsqueeze(0) + self._grid = grid_x, grid_y + self._res = spatial_dims + + return self._grid + + def __call__(self, data): + if data.ndim == 3: + data = data.unsqueeze(0) + x, y = self.grid(data.shape[-2:], data.dtype) + out = paddle.concat( + (data, x.expand([1, -1, -1, -1]), y.expand([1, -1, -1, -1])), axis=1 + ) + return out.squeeze(0) + + +class DarcyFlowDataset(io.Dataset): + """Loads a small Darcy-Flow dataset + + Training contains 1000 samples in resolution 16x16. + Testing contains 100 samples at resolution 16x16 and + 50 samples at resolution 32x32. + + Args: + input_keys (Tuple[str, ...]): Input keys, such as ("input",). + label_keys (Tuple[str, ...]): Output keys, such as ("output",). + data_dir (str): The directory to load data from. + weight_dict (Optional[Dict[str, float]], optional): Define the weight of each constraint variable. Defaults to None. + test_resolutions (List[int,...]): The resolutions to test dataset. Default is [16, 32]. + grid_boundaries (List[int,...]): The boundaries of the grid. Default is [[0,1],[0,1]]. + positional_encoding (bool): Whether to use positional encoding. Default is True + encode_input (bool): Whether to encode the input. Default is False + encode_output (bool): Whether to encode the output. Default is True + encoding (str): The type of encoding. Default is 'channel-wise'. + channel_dim (int): The location of unsqueeze. Default is 1. + where to put the channel dimension. Defaults size is batch, channel, height, width + data_split (str): Wether to use training or test dataset. Default is 'train'. + """ + + def __init__( + self, + input_keys: Tuple[str, ...], + label_keys: Tuple[str, ...], + data_dir: str, + weight_dict: Optional[Dict[str, float]] = None, + test_resolutions: Tuple[int, ...] = [32], + train_resolution: int = 32, + grid_boundaries: Tuple[Tuple[int, ...], ...] = [[0, 1], [0, 1]], + positional_encoding: bool = True, + encode_input: bool = False, + encode_output: bool = True, + encoding: str = "channel-wise", + channel_dim: int = 1, + data_split: str = "train", + ): + super().__init__() + for res in test_resolutions: + if res not in [16, 32]: + raise ValueError( + f"Only 32 and 64 are supported for test resolution, but got {test_resolutions}" + ) + + self.input_keys = input_keys + self.label_keys = label_keys + self.data_dir = data_dir + self.weight_dict = {} if weight_dict is None else weight_dict + if weight_dict is not None: + self.weight_dict = {key: 1.0 for key in self.label_keys} + self.weight_dict.update(weight_dict) + + self.test_resolutions = test_resolutions + self.train_resolution = train_resolution + self.grid_boundaries = grid_boundaries + self.positional_encoding = positional_encoding + self.encode_input = encode_input + self.encode_output = encode_output + self.encoding = encoding + self.channel_dim = channel_dim + self.data_split = data_split + + # train path + path_train = ( + Path(self.data_dir) + .joinpath(f"darcy_train_{self.train_resolution}.npy") + .as_posix() + ) + self.x_train, self.y_train = self.read_data(path_train) + # test path + path_test_1 = ( + Path(self.data_dir) + .joinpath(f"darcy_test_{self.test_resolutions[0]}.npy") + .as_posix() + ) + self.x_test_1, self.y_test_1 = self.read_data(path_test_1) + path_test_2 = ( + Path(self.data_dir) + .joinpath(f"darcy_test_{self.test_resolutions[1]}.npy") + .as_posix() + ) + self.x_test_2, self.y_test_2 = self.read_data(path_test_2) + + # input encoder + if self.encode_input: + self.input_encoder = self.encode_data(self.x_train) + self.x_train = self.input_encoder.encode(self.x_train) + self.x_test_1 = self.input_encoder.encode(self.x_test_1) + self.x_test_2 = self.input_encoder.encode(self.x_test_2) + else: + self.input_encoder = None + # output encoder + if self.encode_output: + self.output_encoder = self.encode_data(self.y_train) + self.y_train = self.output_encoder.encode(self.y_train) + else: + self.output_encoder = None + + if positional_encoding: + self.transform_x = PositionalEmbedding2D(grid_boundaries) + + def read_data(self, path): + # load with numpy + data = np.load(path, allow_pickle=True).item() + x = ( + paddle.to_tensor(data["x"]) + .unsqueeze(self.channel_dim) + .astype("float32") + .clone() + ) + y = paddle.to_tensor(data["y"]).unsqueeze(self.channel_dim).clone() + del data + return x, y + + def encode_data(self, data): + if self.encoding == "channel-wise": + reduce_dims = list(range(data.ndim)) + elif self.encoding == "pixel-wise": + reduce_dims = [0] + input_encoder = UnitGaussianNormalizer(data, reduce_dim=reduce_dims) + return input_encoder + + def __len__(self): + if self.data_split == "train": + return self.x_train.shape[0] + elif self.data_split == "test_16x16": + return self.x_test_1.shape[0] + else: + return self.x_test_2.shape[0] + + def __getitem__(self, index): + if self.data_split == "train": + x = self.x_train[index] + y = self.y_train[index] + + elif self.data_split == "test_16x16": + x = self.x_test_1[index] + y = self.y_test_1[index] + else: + x = self.x_test_2[index] + y = self.y_test_2[index] + + if self.transform_x is not None: + x = self.transform_x(x) + + input_item = {self.input_keys[0]: x} + label_item = {self.label_keys[0]: y} + weight_item = self.weight_dict + + return input_item, label_item, weight_item diff --git a/examples/smc_reac/ppsci/data/dataset/dgmr_dataset.py b/examples/smc_reac/ppsci/data/dataset/dgmr_dataset.py new file mode 100644 index 0000000000..8490f679bc --- /dev/null +++ b/examples/smc_reac/ppsci/data/dataset/dgmr_dataset.py @@ -0,0 +1,95 @@ +# Copyright (c) 2024 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. + +from __future__ import annotations + +import importlib +from typing import Tuple + +import numpy as np +from numpy.random import default_rng +from paddle import io + + +class DGMRDataset(io.Dataset): + """ + Dataset class for DGMR (Deep Generative Model for Radar) model. + This open-sourced UK dataset has been mirrored to HuggingFace Datasets https://huggingface.co/datasets/openclimatefix/nimrod-uk-1km. + If the reader cannot load the dataset from Hugging Face, please manually download it and modify the dataset_path to the local path for loading. + + Args: + input_keys (Tuple[str, ...]): Input keys, such as ("input",). + label_keys (Tuple[str, ...]): Output keys, such as ("output",). + split (str, optional): The split of the dataset, "validation" or "train". Defaults to "validation". + num_input_frames (int, optional): Number of input frames. Defaults to 4. + num_target_frames (int, optional): Number of target frames. Defaults to 18. + dataset_path (str, optional): Path to the dataset. Defaults to "openclimatefix/nimrod-uk-1km". + + Examples: + >>> import ppsci + >>> dataset = ppsci.data.dataset.DGMRDataset(("input", ), ("output", )) # doctest: +SKIP + """ + + def __init__( + self, + input_keys: Tuple[str, ...], + label_keys: Tuple[str, ...], + split: str = "validation", + num_input_frames: int = 4, + num_target_frames: int = 18, + dataset_path: str = "openclimatefix/nimrod-uk-1km", + ): + super().__init__() + self.input_keys = input_keys + self.label_keys = label_keys + self.num_input_frames = num_input_frames + self.num_target_frames = num_target_frames + if not importlib.util.find_spec("datasets"): + raise ModuleNotFoundError( + "Please install datasets with `pip install datasets`" + " before exporting onnx model." + ) + import datasets + + self.reader = datasets.load_dataset( + dataset_path, "sample", split=split, streaming=True, trust_remote_code=True + ) + self.iter_reader = self.reader + + def __len__(self): + return 1000 + + def __getitem__(self, idx): + try: + row = next(self.iter_reader) + except Exception: + rng = default_rng(42) + self.iter_reader = iter( + self.reader.shuffle( + seed=rng.integers(low=0, high=100000), buffer_size=1000 + ) + ) + row = next(self.iter_reader) + radar_frames = row["radar_frames"] + input_frames = radar_frames[ + -self.num_target_frames - self.num_input_frames : -self.num_target_frames + ] + target_frames = radar_frames[-self.num_target_frames :] + input_item = { + self.input_keys[0]: np.moveaxis(input_frames, [0, 1, 2, 3], [0, 2, 3, 1]) + } + label_item = { + self.label_keys[0]: np.moveaxis(target_frames, [0, 1, 2, 3], [0, 2, 3, 1]) + } + return input_item, label_item diff --git a/examples/smc_reac/ppsci/data/dataset/drivaernet_dataset.py b/examples/smc_reac/ppsci/data/dataset/drivaernet_dataset.py new file mode 100644 index 0000000000..23634bebb2 --- /dev/null +++ b/examples/smc_reac/ppsci/data/dataset/drivaernet_dataset.py @@ -0,0 +1,316 @@ +# Copyright (c) 2024 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. + +""" +Created on Tue Dec 19 20:54:56 2023 + +@author: Mohamed Elrefaie, mohamed.elrefaie@mit.edu mohamed.elrefaie@tum.de + +This module is part of the research presented in the paper": +"DrivAerNet: A Parametric Car Dataset for Data-driven Aerodynamic Design and Graph-Based Drag Prediction". +""" +from __future__ import annotations + +import logging +import os +from typing import Callable +from typing import Dict +from typing import Optional +from typing import Tuple + +import numpy as np +import paddle +import pandas as pd + + +class DataAugmentation: + """ + Class encapsulating various data augmentation techniques for point clouds. + """ + + @staticmethod + def translate_pointcloud( + pointcloud: np.ndarray, + translation_range: Tuple[float, float] = (2.0 / 3.0, 3.0 / 2.0), + ) -> np.ndarray: + """ + Translates the pointcloud by a random factor within a given range. + + Args: + pointcloud: The input point cloud as a np.ndarray. + translation_range: A tuple specifying the range for translation factors. + + Returns: + Translated point cloud as a np.ndarray. + """ + xyz1 = np.random.uniform( + low=translation_range[0], high=translation_range[1], size=[3] + ) + xyz2 = np.random.uniform(low=-0.2, high=0.2, size=[3]) + translated_pointcloud = np.add(np.multiply(pointcloud, xyz1), xyz2).astype( + "float32" + ) + return paddle.to_tensor(data=translated_pointcloud, dtype="float32") + + @staticmethod + def jitter_pointcloud( + pointcloud: np.ndarray, sigma: float = 0.01, clip: float = 0.02 + ) -> np.ndarray: + """ + Adds Gaussian noise to the pointcloud. + + Args: + pointcloud: The input point cloud as a np.ndarray. + sigma: Standard deviation of the Gaussian noise. + clip: Maximum absolute value for noise. + + Returns: + Jittered point cloud as a np.ndarray. + """ + N, C = tuple(pointcloud.shape) + jittered_pointcloud = pointcloud + paddle.clip( + x=sigma * paddle.randn(shape=[N, C]), min=-clip, max=clip + ) + return jittered_pointcloud + + @staticmethod + def drop_points(pointcloud: np.ndarray, drop_rate: float = 0.1) -> np.ndarray: + """ + Randomly removes points from the point cloud based on the drop rate. + + Args: + pointcloud: The input point cloud as a np.ndarray. + drop_rate: The percentage of points to be randomly dropped. + + Returns: + The point cloud with points dropped as a np.ndarray. + """ + num_drop = int(drop_rate * pointcloud.shape[0]) + drop_indices = np.random.choice(pointcloud.shape[0], num_drop, replace=False) + keep_indices = np.setdiff1d(np.arange(pointcloud.shape[0]), drop_indices) + dropped_pointcloud = pointcloud[keep_indices, :] + return dropped_pointcloud + + +class DrivAerNetDataset(paddle.io.Dataset): + """ + Paddle Dataset class for the DrivAerNet dataset, handling loading, transforming, and augmenting 3D car models. + + This dataset is specifically designed for aerodynamic tasks, including training machine learning models + to predict aerodynamic coefficients such as drag coefficient (Cd) from 3D car models. + + Args: + input_keys (Tuple[str, ...]): Tuple specifying the keys for input features. + These keys correspond to the attributes of the dataset used as input to the model. + For example, "vertices" represents the 3D point cloud vertices of car models. + label_keys (Tuple[str, ...]): Tuple specifying the keys for ground-truth labels. + These keys correspond to the target values, such as aerodynamic coefficients like Cd. + Example: ("cd_value",) + weight_keys (Tuple[str, ...]): Tuple specifying the keys for optional sample weights. + These keys represent weighting factors that may be used to adjust loss computation + during model training. Useful for handling sample imbalance. + Example: ("weight_keys",) + subset_dir (str): Path to the directory containing subset information. + This directory typically contains files that divide the dataset into training, + validation, and test subsets using a list of model IDs. + ids_file (str): Path to the text file containing model IDs for the current subset. + Each line in the file corresponds to a unique model ID that defines which + models belong to the subset (e.g., training set or test set). + root_dir (str): Directory containing the STL files of 3D car models. + Each STL file is expected to represent a single car model and is named according + to the corresponding model ID. This is the primary data source. + csv_file (str): Path to the CSV file containing metadata for car models. + This file typically includes aerodynamic properties (e.g., drag coefficient) + and other descriptive attributes mapped to each model ID. + num_points (int): Fixed number of points to sample from each 3D model. + If a 3D model has more points than `num_points`, it will be randomly subsampled. + If it has fewer points, it will be zero-padded to reach the desired number. + transform (Optional[Callable]): An optional callable for applying data transformations. + This can include augmentations such as scaling, rotation, jittering, or other preprocessing + steps applied to the 3D point clouds before they are passed to the model. + pointcloud_exist (bool): Whether the point clouds are pre-processed and saved as `.pt` files. + If `True`, the dataset will directly load the pre-saved point clouds instead of generating them from STL files. + train_fractions (float): Fraction of the training data to use. Useful for experiments where only a portion of the data is needed. + mode (str): Mode of operation, either "train", "eval", or "test". Determines how the dataset behaves. + + Examples: + >>> import ppsci + >>> dataset = ppsci.data.dataset.DrivAerNetDataset( + ... input_keys=("vertices",), + ... label_keys=("cd_value",), + ... weight_keys=("weight_keys",), + ... subset_dir="/path/to/subset_dir", + ... ids_file="train_ids.txt", + ... root_dir="/path/to/DrivAerNetDataset", + ... csv_file="/path/to/aero_metadata.csv", + ... num_points=1024, + ... transform=None, + ... ) # doctest: +SKIP + """ + + def __init__( + self, + input_keys: Tuple[str, ...], + label_keys: Tuple[str, ...], + weight_keys: Tuple[str, ...], + subset_dir: str, + ids_file: str, + root_dir: str, + csv_file: str, + num_points: int, + transform: Optional[Callable] = None, + pointcloud_exist: bool = True, + train_fractions=1.0, + mode="eval", + ): + + super().__init__() + self.root_dir = root_dir + try: + self.data_frame = pd.read_csv(csv_file) + except Exception as e: + logging.error(f"Failed to load CSV file: {csv_file}. Error: {e}") + raise + self.input_keys = input_keys + self.label_keys = label_keys + self.weight_keys = weight_keys + self.subset_dir = subset_dir + self.ids_file = ids_file + self.transform = transform + self.num_points = num_points + self.pointcloud_exist = pointcloud_exist + self.mode = mode + self.train_fractions = train_fractions + self.augmentation = DataAugmentation() + self.cache = {} + + try: + with open(os.path.join(self.subset_dir, self.ids_file), "r") as file: + subset_ids = file.read().split() + except FileNotFoundError as e: + raise FileNotFoundError(f"Error loading subset file {self.ids_file}: {e}") + + self.subset_indices = self.data_frame[ + self.data_frame["Design"].isin(subset_ids) + ].index.tolist() + self.data_frame = self.data_frame.loc[self.subset_indices].reset_index( + drop=True + ) + + if self.mode == "train": + self.data_frame = self.data_frame.sample(frac=self.train_fractions) + else: + self.data_frame = self.data_frame + + def __len__(self) -> int: + """Returns the total number of samples in the dataset.""" + return len(self.data_frame) + + def _sample_or_pad_vertices( + self, vertices: paddle.Tensor, num_points: int + ) -> paddle.Tensor: + """ + Subsamples or pads the vertices of the model to a fixed number of points. + + Args: + vertices: The vertices of the 3D model as a paddle.Tensor. + num_points: The desired number of points for the model. + + Returns: + The vertices standardized to the specified number of points. + """ + num_vertices = vertices.shape[0] + if num_vertices > num_points: + indices = np.random.choice(num_vertices, num_points, replace=False) + vertices = vertices[indices] + elif num_vertices < num_points: + padding = paddle.zeros( + shape=(num_points - num_vertices, 3), dtype="float32" + ) + vertices = paddle.concat(x=(vertices, padding), axis=0) + return vertices + + def _load_point_cloud(self, design_id: str) -> Optional[paddle.Tensor]: + load_path = os.path.join(self.root_dir, f"{design_id}.paddle_tensor") + if os.path.exists(load_path) and os.path.getsize(load_path) > 0: + try: + vertices = paddle.load(path=str(load_path)) + num_vertices = vertices.shape[0] + + if num_vertices > self.num_points: + indices = np.random.choice( + num_vertices, self.num_points, replace=False + ) + vertices = vertices.numpy()[indices] + vertices = paddle.to_tensor(vertices) + + return vertices + except (EOFError, RuntimeError, ValueError) as e: + raise Exception( + f"Error loading point cloud from {load_path}: {e}" + ) from e + + def __getitem__( + self, idx: int, apply_augmentations: bool = True + ) -> Tuple[Dict[str, np.ndarray], Dict[str, np.ndarray], Dict[str, np.ndarray],]: + """ + Retrieves a sample and its corresponding label from the dataset, with an option to apply augmentations. + + Args: + idx (int): Index of the sample to retrieve. + apply_augmentations (bool, optional): Whether to apply data augmentations. Defaults to True. + + Tuple[Dict[str, np.ndarray], Dict[str, np.ndarray], Dict[str, np.ndarray]]: + A tuple containing three dictionaries: + - The first dictionary contains the input data (point cloud) under the key specified by `self.input_keys[0]`. + - The second dictionary contains the label (Cd value) under the key specified by `self.label_keys[0]`. + - The third dictionary contains the weight (default is 1) under the key specified by `self.weight_keys[0]`. + """ + if paddle.is_tensor(x=idx): + idx = idx.tolist() + + if idx in self.cache: + return self.cache[idx] + + row = self.data_frame.iloc[idx] + design_id = row["Design"] + cd_value = row["Average Cd"].reshape([-1]) + if self.pointcloud_exist: + try: + vertices = self._load_point_cloud(design_id) + if vertices is None: + raise ValueError( + f"Point cloud for design {design_id} is not found or corrupted." + ) + except Exception as e: + raise ValueError( + f"Failed to load point cloud for design {design_id}: {e}" + ) + if apply_augmentations: + vertices = self.augmentation.translate_pointcloud(vertices.numpy()) + vertices = self.augmentation.jitter_pointcloud(vertices) + if self.transform: + vertices = self.transform(vertices) + + self.cache[idx] = ( + {self.input_keys[0]: vertices}, + {self.label_keys[0]: cd_value}, + {self.weight_keys[0]: np.array(1, dtype=np.float32)}, + ) + + return ( + {self.input_keys[0]: vertices}, + {self.label_keys[0]: cd_value}, + {self.weight_keys[0]: np.array(1, dtype=np.float32)}, + ) diff --git a/examples/smc_reac/ppsci/data/dataset/drivaernetplusplus_dataset.py b/examples/smc_reac/ppsci/data/dataset/drivaernetplusplus_dataset.py new file mode 100644 index 0000000000..3194ece235 --- /dev/null +++ b/examples/smc_reac/ppsci/data/dataset/drivaernetplusplus_dataset.py @@ -0,0 +1,321 @@ +# Copyright (c) 2024 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. + +""" +@author: Mohamed Elrefaie, mohamed.elrefaie@mit.edu mohamed.elrefaie@tum.de + +This module is part of the research presented in the paper: +"DrivAerNet++: A Large-Scale Multimodal Car Dataset with Computational Fluid Dynamics Simulations and Deep Learning Benchmarks". + +The module defines two Paddle Datasets for loading and transforming 3D car models from the DrivAerNet++ dataset: +1. DrivAerNetPlusPlusDataset: Handles point cloud data, allowing loading, transforming, and augmenting 3D car models. +""" + + +from __future__ import annotations + +import logging +import os +from typing import Callable +from typing import Dict +from typing import Optional +from typing import Tuple + +import numpy as np +import paddle +import pandas as pd + + +class DataAugmentation: + """ + Class encapsulating various data augmentation techniques for point clouds. + """ + + @staticmethod + def translate_pointcloud( + pointcloud: np.ndarray, + translation_range: Tuple[float, float] = (2.0 / 3.0, 3.0 / 2.0), + ) -> np.ndarray: + """ + Translates the pointcloud by a random factor within a given range. + + Args: + pointcloud: The input point cloud as a np.ndarray. + translation_range: A tuple specifying the range for translation factors. + + Returns: + Translated point cloud as a np.ndarray. + """ + xyz1 = np.random.uniform( + low=translation_range[0], high=translation_range[1], size=[3] + ) + xyz2 = np.random.uniform(low=-0.2, high=0.2, size=[3]) + translated_pointcloud = np.add(np.multiply(pointcloud, xyz1), xyz2).astype( + "float32" + ) + return paddle.to_tensor(data=translated_pointcloud, dtype="float32") + + @staticmethod + def jitter_pointcloud( + pointcloud: np.ndarray, sigma: float = 0.01, clip: float = 0.02 + ) -> np.ndarray: + """ + Adds Gaussian noise to the pointcloud. + + Args: + pointcloud: The input point cloud as a np.ndarray. + sigma: Standard deviation of the Gaussian noise. + clip: Maximum absolute value for noise. + + Returns: + Jittered point cloud as a np.ndarray. + """ + N, C = tuple(pointcloud.shape) + jittered_pointcloud = pointcloud + paddle.clip( + x=sigma * paddle.randn(shape=[N, C]), min=-clip, max=clip + ) + return jittered_pointcloud + + @staticmethod + def drop_points(pointcloud: np.ndarray, drop_rate: float = 0.1) -> np.ndarray: + """ + Randomly removes points from the point cloud based on the drop rate. + + Args: + pointcloud: The input point cloud as a np.ndarray. + drop_rate: The percentage of points to be randomly dropped. + + Returns: + The point cloud with points dropped as a np.ndarray. + """ + num_drop = int(drop_rate * pointcloud.shape[0]) + drop_indices = np.random.choice(pointcloud.shape[0], num_drop, replace=False) + keep_indices = np.setdiff1d(np.arange(pointcloud.shape[0]), drop_indices) + dropped_pointcloud = pointcloud[keep_indices, :] + return dropped_pointcloud + + +class DrivAerNetPlusPlusDataset(paddle.io.Dataset): + """ + Paddle Dataset class for the DrivAerNet dataset, handling loading, transforming, and augmenting 3D car models. + + This dataset is designed for tasks involving aerodynamic simulations and deep learning models, + specifically for predicting aerodynamic coefficients (e.g., Cd values) from 3D car models. + + Args: + input_keys (Tuple[str, ...]): Tuple of strings specifying the input keys. + These keys correspond to the features extracted from the dataset, + typically the 3D vertices of car models. + Example: ("vertices",) + label_keys (Tuple[str, ...]): Tuple of strings specifying the label keys. + These keys correspond to the ground-truth labels, such as aerodynamic + coefficients (e.g., Cd values). + Example: ("cd_value",) + weight_keys (Tuple[str, ...]): Tuple of strings specifying the weight keys. + These keys represent optional weighting factors used during model training + to handle class imbalance or sample importance. + Example: ("weight_keys",) + subset_dir (str): Path to the directory containing subsets of the dataset. + This directory is used to divide the dataset into different subsets + (e.g., train, validation, test) based on provided IDs. + ids_file (str): Path to the file containing the list of IDs for the subset. + The file specifies which models belong to the current subset (e.g., training IDs). + root_dir (str): Root directory containing the 3D STL files of car models. + Each 3D model is expected to be stored in a file named according to its ID. + csv_file (str): Path to the CSV file containing metadata for the car models. + The CSV file includes information such as aerodynamic coefficients, + and may also map model IDs to specific attributes. + num_points (int): Number of points to sample or pad each 3D point cloud to. + If the model has more points than `num_points`, it will be subsampled. + If it has fewer points, zero-padding will be applied. + transform (Optional[Callable]): Optional transformation function applied to each sample. + This can include augmentations like scaling, rotation, or jittering. + pointcloud_exist (bool): Whether the point clouds are pre-processed and saved as `.pt` files. + If `True`, the dataset will directly load the pre-saved point clouds + instead of generating them from STL files. + + Examples: + >>> import ppsci + >>> dataset = ppsci.data.dataset.DrivAerNetPlusPlusDataset( + ... input_keys=("vertices",), + ... label_keys=("cd_value",), + ... weight_keys=("weight_keys",), + ... subset_dir="/path/to/subset_dir", + ... ids_file="train_ids.txt", + ... root_dir="/path/to/DrivAerNetPlusPlusDataset", + ... csv_file="/path/to/aero_metadata.csv", + ... num_points=1024, + ... transform=None, + ... ) # doctest: +SKIP + """ + + def __init__( + self, + input_keys: Tuple[str, ...], + label_keys: Tuple[str, ...], + weight_keys: Tuple[str, ...], + subset_dir: str, + ids_file: str, + root_dir: str, + csv_file: str, + num_points: int, + transform: Optional[Callable] = None, + pointcloud_exist: bool = True, + ): + super().__init__() + self.root_dir = root_dir + self.input_keys = input_keys + self.label_keys = label_keys + self.weight_keys = weight_keys + self.subset_dir = subset_dir + self.ids_file = ids_file + self.augmentation = DataAugmentation() + self.cache = {} + + try: + self.data_frame = pd.read_csv(csv_file) + except Exception as e: + logging.error(f"Failed to load CSV file: {csv_file}. Error: {e}") + raise + self.transform = transform + self.num_points = num_points + self.pointcloud_exist = pointcloud_exist + + try: + with open(os.path.join(self.subset_dir, self.ids_file), "r") as file: + subset_ids = file.read().split() + except FileNotFoundError as e: + raise FileNotFoundError(f"Error loading subset file {self.ids_file}: {e}") + + self.subset_indices = self.data_frame[ + self.data_frame["Design"].isin(subset_ids) + ].index.tolist() + self.data_frame = self.data_frame.loc[self.subset_indices].reset_index( + drop=True + ) + + def __len__(self) -> int: + """Returns the total number of samples in the dataset.""" + return len(self.data_frame) + + def min_max_normalize(self, data: np.ndarray) -> np.ndarray: + """ + Normalizes the data to the range [0, 1] based on min and max values. + """ + min_vals = data.min(axis=0, keepdim=True) + max_vals = data.max(axis=0, keepdim=True) + normalized_data = (data - min_vals) / (max_vals - min_vals) + return normalized_data + + def _sample_or_pad_vertices( + self, vertices: paddle.Tensor, num_points: int + ) -> paddle.Tensor: + """ + Subsamples or pads the vertices of the model to a fixed number of points. + + Args: + vertices: The vertices of the 3D model as a paddle.Tensor. + num_points: The desired number of points for the model. + + Returns: + The vertices standardized to the specified number of points. + """ + num_vertices = vertices.shape[0] + if num_vertices > num_points: + indices = np.random.choice(num_vertices, num_points, replace=False) + vertices = vertices[indices] + elif num_vertices < num_points: + padding = paddle.zeros( + shape=(num_points - num_vertices, 3), dtype="float32" + ) + vertices = paddle.concat(x=(vertices, padding), axis=0) + return vertices + + def _load_point_cloud(self, design_id: str): + load_path = os.path.join(self.root_dir, f"{design_id}.paddle_tensor") + if os.path.exists(load_path) and os.path.getsize(load_path) > 0: + try: + vertices = paddle.load(path=str(load_path)) + except (EOFError, RuntimeError, ValueError) as e: + raise Exception( + f"Error loading point cloud from {load_path}: {e}" + ) from e + num_vertices = vertices.shape[0] + + if num_vertices > self.num_points: + indices = np.random.choice(num_vertices, self.num_points, replace=False) + vertices = vertices.numpy()[indices] + + return vertices + + def __getitem__( + self, idx: int, apply_augmentations: bool = True + ) -> Tuple[Dict[str, np.ndarray], Dict[str, np.ndarray], Dict[str, np.ndarray]]: + """ + Retrieves a sample and its corresponding label from the dataset, with an option to apply augmentations. + + Args: + idx (int): Index of the sample to retrieve. + apply_augmentations (bool, optional): Whether to apply data augmentations. Defaults to True. + + Returns: + Tuple[Dict[str, np.ndarray], Dict[str, np.ndarray], Dict[str, np.ndarray]]: + A tuple containing three dictionaries: + - The first dictionary contains the input data (point cloud) under the key specified by `self.input_keys[0]`. + - The second dictionary contains the label (Cd value) under the key specified by `self.label_keys[0]`. + - The third dictionary contains the weight (default is 1) under the key specified by `self.weight_keys[0]`. + """ + if paddle.is_tensor(idx): + idx = idx.tolist() + + if idx in self.cache: + return self.cache[idx] + + row = self.data_frame.iloc[idx] + design_id = row["Design"] + cd_value = row["Average Cd"] + if self.pointcloud_exist: + try: + vertices = self._load_point_cloud(design_id) + if vertices is None: + raise ValueError( + f"Point cloud for design {design_id} is not found or corrupted." + ) + except Exception as e: + raise ValueError( + f"Failed to load point cloud for design {design_id}: {e}" + ) + + if apply_augmentations: + vertices = self.augmentation.translate_pointcloud(vertices.numpy()) + vertices = self.augmentation.jitter_pointcloud(vertices) + + if self.transform: + vertices = self.transform(vertices) + + vertices = self.min_max_normalize(vertices) + + cd_value = np.array(float(cd_value), dtype=np.float32).reshape([-1]) + + self.cache[idx] = ( + {self.input_keys[0]: vertices}, + {self.label_keys[0]: cd_value}, + {self.weight_keys[0]: np.array(1, dtype=np.float32)}, + ) + + return ( + {self.input_keys[0]: vertices}, + {self.label_keys[0]: cd_value}, + {self.weight_keys[0]: np.array(1, dtype=np.float32)}, + ) diff --git a/examples/smc_reac/ppsci/data/dataset/enso_dataset.py b/examples/smc_reac/ppsci/data/dataset/enso_dataset.py new file mode 100644 index 0000000000..601fcec413 --- /dev/null +++ b/examples/smc_reac/ppsci/data/dataset/enso_dataset.py @@ -0,0 +1,405 @@ +# 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. + +from __future__ import annotations + +import importlib +from pathlib import Path +from typing import Dict +from typing import Optional +from typing import Tuple + +import numpy as np +from paddle import io + +NINO_WINDOW_T = 3 # Nino index is the sliding average over sst, window size is 3 +CMIP6_SST_MAX = 10.198975563049316 +CMIP6_SST_MIN = -16.549121856689453 +CMIP5_SST_MAX = 8.991744995117188 +CMIP5_SST_MIN = -9.33076286315918 +CMIP6_NINO_MAX = 4.138188362121582 +CMIP6_NINO_MIN = -3.5832221508026123 +CMIP5_NINO_MAX = 3.8253555297851562 +CMIP5_NINO_MIN = -2.691682815551758 +SST_MAX = max(CMIP6_SST_MAX, CMIP5_SST_MAX) +SST_MIN = min(CMIP6_SST_MIN, CMIP5_SST_MIN) + + +def scale_sst(sst): + return (sst - SST_MIN) / (SST_MAX - SST_MIN) + + +def scale_back_sst(sst): + return (SST_MAX - SST_MIN) * sst + SST_MIN + + +def prepare_inputs_targets( + len_time, input_gap, input_length, pred_shift, pred_length, samples_gap +): + """Prepares the input and target indices for training. + + Args: + len_time (int): The total number of time steps in the dataset. + input_gap (int): Time gaps between two consecutive input frames. + input_length (int): The number of input frames. + pred_shift (int): The lead_time of the last target to be predicted. + pred_length (int): The number of frames to be predicted. + samples_gap (int): Stride of seq sampling. + """ + + if pred_shift < pred_length: + raise ValueError("pred_shift should be small than pred_length") + input_span = input_gap * (input_length - 1) + 1 + pred_gap = pred_shift // pred_length + input_ind = np.arange(0, input_span, input_gap) + target_ind = np.arange(0, pred_shift, pred_gap) + input_span + pred_gap - 1 + ind = np.concatenate([input_ind, target_ind]).reshape(1, input_length + pred_length) + max_n_sample = len_time - (input_span + pred_shift - 1) + ind = ind + np.arange(max_n_sample)[:, np.newaxis] @ np.ones( + (1, input_length + pred_length), dtype=int + ) + return ind[::samples_gap] + + +def fold(data, size=36, stride=12): + """Inverse of unfold/sliding window operation + only applicable to the case where the size of the sliding windows is n*stride + + Args: + data (tuple[int,...]): The input data.(N, size, *). + size (int, optional): The size of a single datum.The Defaults to 36. + stride (int, optional): The step.Defaults to 12. + + Returns: + outdata (np.array): (N_, *).N/size is the number/width of sliding blocks + """ + if size % stride != 0: + raise ValueError("size modulo stride should be zero") + times = size // stride + remain = (data.shape[0] - 1) % times + if remain > 0: + ls = list(data[::times]) + [data[-1, -(remain * stride) :]] + outdata = np.concatenate(ls, axis=0) # (36*(151//3+1)+remain*stride, *, 15) + else: + outdata = np.concatenate(data[::times], axis=0) # (36*(151/3+1), *, 15) + assert ( + outdata.shape[0] == size * ((data.shape[0] - 1) // times + 1) + remain * stride + ) + return outdata + + +def data_transform(data, num_years_per_model): + """The transform of the input data. + + Args: + data (Tuple[list,...]): The input data.Shape of (N, 36, *). + num_years_per_model (int): The number of years associated with each model.151/140. + """ + length = data.shape[0] + assert length % num_years_per_model == 0 + num_models = length // num_years_per_model + outdata = np.stack( + np.split(data, length / num_years_per_model, axis=0), axis=-1 + ) # (151, 36, *, 15) + # cmip6sst outdata.shape = (151, 36, 24, 48, 15) = (year, month, lat, lon, model) + # cmip5sst outdata.shape = (140, 36, 24, 48, 17) + # cmip6nino outdata.shape = (151, 36, 15) + # cmip5nino outdata.shape = (140, 36, 17) + outdata = fold(outdata, size=36, stride=12) + # cmip6sst outdata.shape = (1836, 24, 48, 15), 1836 == 151 * 12 + 24 + # cmip5sst outdata.shape = (1704, 24, 48, 17) + # cmip6nino outdata.shape = (1836, 15) + # cmip5nino outdata.shape = (1704, 17) + + # check output data + assert outdata.shape[-1] == num_models + assert not np.any(np.isnan(outdata)) + return outdata + + +def read_raw_data(ds_dir, out_dir=None): + """Read and process raw cmip data from CMIP_train.nc and CMIP_label.nc + + Args: + ds_dir (str): The path of the dataset. + out_dir (str): The path of output. Defaults to None. + """ + + import xarray as xr + + train_cmip = xr.open_dataset(Path(ds_dir) / "CMIP_train.nc").transpose( + "year", "month", "lat", "lon" + ) + label_cmip = xr.open_dataset(Path(ds_dir) / "CMIP_label.nc").transpose( + "year", "month" + ) + # train_cmip.sst.values.shape = (4645, 36, 24, 48) + + # select longitudes + lon = train_cmip.lon.values + lon = lon[np.logical_and(lon >= 95, lon <= 330)] + train_cmip = train_cmip.sel(lon=lon) + + cmip6sst = data_transform( + data=train_cmip.sst.values[:2265], num_years_per_model=151 + ) + cmip5sst = data_transform( + data=train_cmip.sst.values[2265:], num_years_per_model=140 + ) + cmip6nino = data_transform( + data=label_cmip.nino.values[:2265], num_years_per_model=151 + ) + cmip5nino = data_transform( + data=label_cmip.nino.values[2265:], num_years_per_model=140 + ) + + # cmip6sst.shape = (1836, 24, 48, 15) + # cmip5sst.shape = (1704, 24, 48, 17) + assert len(cmip6sst.shape) == 4 + assert len(cmip5sst.shape) == 4 + assert len(cmip6nino.shape) == 2 + assert len(cmip5nino.shape) == 2 + # store processed data for faster data access + if out_dir is not None: + ds_cmip6 = xr.Dataset( + { + "sst": (["month", "lat", "lon", "model"], cmip6sst), + "nino": (["month", "model"], cmip6nino), + }, + coords={ + "month": np.repeat( + np.arange(1, 13)[None], cmip6nino.shape[0] // 12, axis=0 + ).flatten(), + "lat": train_cmip.lat.values, + "lon": train_cmip.lon.values, + "model": np.arange(15) + 1, + }, + ) + ds_cmip6.to_netcdf(Path(out_dir) / "cmip6.nc") + ds_cmip5 = xr.Dataset( + { + "sst": (["month", "lat", "lon", "model"], cmip5sst), + "nino": (["month", "model"], cmip5nino), + }, + coords={ + "month": np.repeat( + np.arange(1, 13)[None], cmip5nino.shape[0] // 12, axis=0 + ).flatten(), + "lat": train_cmip.lat.values, + "lon": train_cmip.lon.values, + "model": np.arange(17) + 1, + }, + ) + ds_cmip5.to_netcdf(Path(out_dir) / "cmip5.nc") + train_cmip.close() + label_cmip.close() + return cmip6sst, cmip5sst, cmip6nino, cmip5nino + + +def cat_over_last_dim(data): + """Treat different models (15 from CMIP6, 17 from CMIP5) as batch_size + e.g., cmip6sst.shape = (178, 38, 24, 48, 15), converted_cmip6sst.shape = (2670, 38, 24, 48) + e.g., cmip5sst.shape = (165, 38, 24, 48, 15), converted_cmip6sst.shape = (2475, 38, 24, 48) + """ + + return np.concatenate(np.moveaxis(data, -1, 0), axis=0) + + +class ENSODataset(io.Dataset): + """The El Niño/Southern Oscillation dataset. + + Args: + input_keys (Tuple[str, ...]): Name of input keys, such as ("input",). + label_keys (Tuple[str, ...]): Name of label keys, such as ("output",). + data_dir (str): The directory of data. + weight_dict (Optional[Dict[str, Union[Callable, float]]]): Define the weight of each constraint variable. Defaults to None. + in_len (int, optional): The length of input data. Defaults to 12. + out_len (int, optional): The length of out data. Defaults to 26. + in_stride (int, optional): The stride of input data. Defaults to 1. + out_stride (int, optional): The stride of output data. Defaults to 1. + train_samples_gap (int, optional): The stride of sequence sampling during training. Defaults to 10. + e.g., samples_gap = 10, the first seq contains [0, 1, ..., T-1] frame indices, the second seq contains [10, 11, .., T+9] + eval_samples_gap (int, optional): The stride of sequence sampling during eval. Defaults to 11. + normalize_sst (bool, optional): Whether to use normalization. Defaults to True. + batch_size (int, optional): Batch size. Defaults to 1. + num_workers (int, optional): The num of workers. Defaults to 1. + training (str, optional): Training pathse. Defaults to "train". + """ + + # Whether support batch indexing for speeding up fetching process. + batch_index: bool = False + + def __init__( + self, + input_keys: Tuple[str, ...], + label_keys: Tuple[str, ...], + data_dir: str, + weight_dict: Optional[Dict[str, float]] = None, + in_len=12, + out_len=26, + in_stride=1, + out_stride=1, + train_samples_gap=10, + eval_samples_gap=11, + normalize_sst=True, + # datamodule_only + batch_size=1, + num_workers=1, + training="train", + ): + super(ENSODataset, self).__init__() + if importlib.util.find_spec("xarray") is None: + raise ModuleNotFoundError( + "To use RadarDataset, please install 'xarray' with: `pip install " + "xarray` first." + ) + self.input_keys = input_keys + self.label_keys = label_keys + self.data_dir = data_dir + self.weight_dict = {} if weight_dict is None else weight_dict + if weight_dict is not None: + self.weight_dict = {key: 1.0 for key in self.label_keys} + self.weight_dict.update(weight_dict) + + self.in_len = in_len + self.out_len = out_len + self.in_stride = in_stride + self.out_stride = out_stride + self.train_samples_gap = train_samples_gap + self.eval_samples_gap = eval_samples_gap + self.normalize_sst = normalize_sst + # datamodule_only + self.batch_size = batch_size + if num_workers != 1: + raise ValueError( + "Current implementation does not support `num_workers != 1`!" + ) + self.num_workers = num_workers + self.training = training + + # pre-data + cmip6sst, cmip5sst, cmip6nino, cmip5nino = read_raw_data(self.data_dir) + # TODO: more flexible train/val/test split + self.sst_train = [cmip6sst, cmip5sst[..., :-2]] + self.nino_train = [cmip6nino, cmip5nino[..., :-2]] + self.sst_eval = [cmip5sst[..., -2:-1]] + self.nino_eval = [cmip5nino[..., -2:-1]] + self.sst_test = [cmip5sst[..., -1:]] + self.nino_test = [cmip5nino[..., -1:]] + + self.sst, self.target_nino = self.create_data() + + def create_data( + self, + ): + if self.training == "train": + sst_cmip6 = self.sst_train[0] + nino_cmip6 = self.nino_train[0] + sst_cmip5 = self.sst_train[1] + nino_cmip5 = self.nino_train[1] + samples_gap = self.train_samples_gap + elif self.training == "eval": + sst_cmip6 = None + nino_cmip6 = None + sst_cmip5 = self.sst_eval[0] + nino_cmip5 = self.nino_eval[0] + samples_gap = self.eval_samples_gap + elif self.training == "test": + sst_cmip6 = None + nino_cmip6 = None + sst_cmip5 = self.sst_test[0] + nino_cmip5 = self.nino_test[0] + samples_gap = self.eval_samples_gap + + # cmip6 (N, *, 15) + # cmip5 (N, *, 17) + sst = [] + target_nino = [] + + nino_idx_slice = slice( + self.in_len, self.in_len + self.out_len - NINO_WINDOW_T + 1 + ) # e.g., 12:36 + if sst_cmip6 is not None: + assert len(sst_cmip6.shape) == 4 + assert len(nino_cmip6.shape) == 2 + idx_sst = prepare_inputs_targets( + len_time=sst_cmip6.shape[0], + input_length=self.in_len, + input_gap=self.in_stride, + pred_shift=self.out_len * self.out_stride, + pred_length=self.out_len, + samples_gap=samples_gap, + ) + + sst.append(cat_over_last_dim(sst_cmip6[idx_sst])) + target_nino.append( + cat_over_last_dim(nino_cmip6[idx_sst[:, nino_idx_slice]]) + ) + if sst_cmip5 is not None: + assert len(sst_cmip5.shape) == 4 + assert len(nino_cmip5.shape) == 2 + idx_sst = prepare_inputs_targets( + len_time=sst_cmip5.shape[0], + input_length=self.in_len, + input_gap=self.in_stride, + pred_shift=self.out_len * self.out_stride, + pred_length=self.out_len, + samples_gap=samples_gap, + ) + sst.append(cat_over_last_dim(sst_cmip5[idx_sst])) + target_nino.append( + cat_over_last_dim(nino_cmip5[idx_sst[:, nino_idx_slice]]) + ) + + # sst data containing both the input and target + self.sst = np.concatenate(sst, axis=0) # (N, in_len+out_len, lat, lon) + if self.normalize_sst: + self.sst = scale_sst(self.sst) + # nino data containing the target only + self.target_nino = np.concatenate( + target_nino, axis=0 + ) # (N, out_len+NINO_WINDOW_T-1) + assert self.sst.shape[0] == self.target_nino.shape[0] + assert self.sst.shape[1] == self.in_len + self.out_len + assert self.target_nino.shape[1] == self.out_len - NINO_WINDOW_T + 1 + return self.sst, self.target_nino + + def get_datashape(self): + return {"sst": self.sst.shape, "nino target": self.target_nino.shape} + + def __len__(self): + return self.sst.shape[0] + + def __getitem__(self, idx): + sst_data = self.sst[idx].astype("float32") + sst_data = sst_data[..., np.newaxis] + in_seq = sst_data[: self.in_len, ...] # ( in_len, lat, lon, 1) + target_seq = sst_data[self.in_len :, ...] # ( in_len, lat, lon, 1) + weight_item = self.weight_dict + + if self.training == "train": + input_item = {self.input_keys[0]: in_seq} + label_item = { + self.label_keys[0]: target_seq, + } + + return input_item, label_item, weight_item + else: + input_item = {self.input_keys[0]: in_seq} + label_item = { + self.label_keys[0]: target_seq, + self.label_keys[1]: self.target_nino[idx], + } + + return input_item, label_item, weight_item diff --git a/examples/smc_reac/ppsci/data/dataset/era5_dataset.py b/examples/smc_reac/ppsci/data/dataset/era5_dataset.py new file mode 100644 index 0000000000..75cf89c754 --- /dev/null +++ b/examples/smc_reac/ppsci/data/dataset/era5_dataset.py @@ -0,0 +1,249 @@ +# 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. + +from __future__ import annotations + +import glob +from typing import Dict +from typing import Optional +from typing import Tuple + +try: + import h5py +except ModuleNotFoundError: + pass + +import numpy as np +import paddle +from paddle import io +from paddle import vision + + +class ERA5Dataset(io.Dataset): + """Class for ERA5 dataset. + + Args: + file_path (str): Data set path. + input_keys (Tuple[str, ...]): Input keys, such as ("input",). + label_keys (Tuple[str, ...]): Output keys, such as ("output",). + precip_file_path (Optional[str]): Precipitation data set path. Defaults to None. + weight_dict (Optional[Dict[str, float]]): Weight dictionary. Defaults to None. + vars_channel (Optional[Tuple[int, ...]]): The variable channel index in ERA5 dataset. Defaults to None. + num_label_timestamps (int, optional): Number of timestamp of label. Defaults to 1. + transforms (Optional[vision.Compose]): Compose object contains sample wise + transform(s). Defaults to None. + training (bool, optional): Whether in train mode. Defaults to True. + stride (int, optional): Stride of sampling data. Defaults to 1. + + Examples: + >>> import ppsci + >>> dataset = ppsci.data.dataset.ERA5Dataset( + ... "file_path": "/path/to/ERA5Dataset", + ... "input_keys": ("input",), + ... "label_keys": ("output",), + ... ) # doctest: +SKIP + """ + + # Whether support batch indexing for speeding up fetching process. + batch_index: bool = False + + def __init__( + self, + file_path: str, + input_keys: Tuple[str, ...], + label_keys: Tuple[str, ...], + precip_file_path: Optional[str] = None, + weight_dict: Optional[Dict[str, float]] = None, + vars_channel: Optional[Tuple[int, ...]] = None, + num_label_timestamps: int = 1, + transforms: Optional[vision.Compose] = None, + training: bool = True, + stride: int = 1, + ): + super().__init__() + self.file_path = file_path + self.input_keys = input_keys + self.label_keys = label_keys + self.precip_file_path = precip_file_path + + self.weight_dict = {} if weight_dict is None else weight_dict + if weight_dict is not None: + self.weight_dict = {key: 1.0 for key in self.label_keys} + self.weight_dict.update(weight_dict) + + self.vars_channel = list(range(20)) if vars_channel is None else vars_channel + self.num_label_timestamps = num_label_timestamps + self.transforms = transforms + self.training = training + self.stride = stride + + self.files = self.read_data(file_path) + self.n_years = len(self.files) + self.num_samples_per_year = self.files[0].shape[0] + self.num_samples = self.n_years * self.num_samples_per_year + if self.precip_file_path is not None: + self.precip_files = self.read_data(precip_file_path, "tp") + + def read_data(self, path: str, var="fields"): + paths = [path] if path.endswith(".h5") else glob.glob(path + "/*.h5") + paths.sort() + files = [] + for path_ in paths: + _file = h5py.File(path_, "r") + files.append(_file[var]) + return files + + def __len__(self): + return self.num_samples // self.stride + + def __getitem__(self, global_idx): + global_idx *= self.stride + year_idx = global_idx // self.num_samples_per_year + local_idx = global_idx % self.num_samples_per_year + step = 0 if local_idx >= self.num_samples_per_year - 1 else 1 + + if self.num_label_timestamps > 1: + if local_idx >= self.num_samples_per_year - self.num_label_timestamps: + local_idx = self.num_samples_per_year - self.num_label_timestamps - 1 + + input_file = self.files[year_idx] + label_file = ( + self.precip_files[year_idx] + if self.precip_file_path is not None + else input_file + ) + if self.precip_file_path is not None and year_idx == 0 and self.training: + # first year has 2 missing samples in precip (they are first two time points) + lim = self.num_samples_per_year - 2 + local_idx = local_idx % lim + step = 0 if local_idx >= lim - 1 else 1 + input_idx = local_idx + 2 + label_idx = local_idx + step + else: + input_idx, label_idx = local_idx, local_idx + step + + input_item = {self.input_keys[0]: input_file[input_idx, self.vars_channel]} + + label_item = {} + for i in range(self.num_label_timestamps): + if self.precip_file_path is not None: + label_item[self.label_keys[i]] = np.expand_dims( + label_file[label_idx + i], 0 + ) + else: + label_item[self.label_keys[i]] = label_file[ + label_idx + i, self.vars_channel + ] + + weight_shape = [1] * len(next(iter(label_item.values())).shape) + weight_item = { + key: np.full(weight_shape, value, paddle.get_default_dtype()) + for key, value in self.weight_dict.items() + } + + if self.transforms is not None: + input_item, label_item, weight_item = self.transforms( + input_item, label_item, weight_item + ) + + return input_item, label_item, weight_item + + +class ERA5SampledDataset(io.Dataset): + """Class for ERA5 sampled dataset. + + Args: + file_path (str): Data set path. + input_keys (Tuple[str, ...]): Input keys, such as ("input",). + label_keys (Tuple[str, ...]): Output keys, such as ("output",). + weight_dict (Optional[Dict[str, float]]): Weight dictionary. Defaults to None. + transforms (Optional[vision.Compose]): Compose object contains sample wise + transform(s). Defaults to None. + + Examples: + >>> import ppsci + >>> dataset = ppsci.data.dataset.ERA5SampledDataset( + ... "file_path": "/path/to/ERA5SampledDataset", + ... "input_keys": ("input",), + ... "label_keys": ("output",), + ... ) # doctest: +SKIP + >>> # get the length of the dataset + >>> dataset_size = len(dataset) # doctest: +SKIP + >>> # get the first sample of the data + >>> first_sample = dataset[0] # doctest: +SKIP + >>> print("First sample:", first_sample) # doctest: +SKIP + """ + + def __init__( + self, + file_path: str, + input_keys: Tuple[str, ...], + label_keys: Tuple[str, ...], + weight_dict: Optional[Dict[str, float]] = None, + transforms: Optional[vision.Compose] = None, + ): + super().__init__() + self.file_path = file_path + self.input_keys = input_keys + self.label_keys = label_keys + + self.weight_dict = {} if weight_dict is None else weight_dict + if weight_dict is not None: + self.weight_dict = {key: 1.0 for key in self.label_keys} + self.weight_dict.update(weight_dict) + + self.transforms = transforms + + self.files = self.read_data(file_path) + self.num_samples = len(self.files) + + def read_data(self, path: str): + paths = glob.glob(path + "/*.h5") + paths.sort() + files = [] + for _path in paths: + _file = h5py.File(_path, "r") + files.append(_file) + return files + + def __len__(self): + return self.num_samples + + def __getitem__(self, global_idx): + _file = self.files[global_idx] + + input_item = {} + for key in _file["input_dict"]: + input_item[key] = np.asarray( + _file["input_dict"][key], paddle.get_default_dtype() + ) + + label_item = {} + for key in _file["label_dict"]: + label_item[key] = np.asarray( + _file["label_dict"][key], paddle.get_default_dtype() + ) + + weight_shape = [1] * len(next(iter(label_item.values())).shape) + weight_item = { + key: np.full(weight_shape, value, paddle.get_default_dtype()) + for key, value in self.weight_dict.items() + } + + if self.transforms is not None: + input_item, label_item, weight_item = self.transforms( + input_item, label_item, weight_item + ) + + return input_item, label_item, weight_item diff --git a/examples/smc_reac/ppsci/data/dataset/ext_moe_enso_dataset.py b/examples/smc_reac/ppsci/data/dataset/ext_moe_enso_dataset.py new file mode 100644 index 0000000000..5286b4bfe2 --- /dev/null +++ b/examples/smc_reac/ppsci/data/dataset/ext_moe_enso_dataset.py @@ -0,0 +1,406 @@ +# 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. + +from __future__ import annotations + +import importlib +from pathlib import Path +from typing import Dict +from typing import Optional +from typing import Tuple + +import numpy as np +from paddle import io + +NINO_WINDOW_T = 3 # Nino index is the sliding average over sst, window size is 3 +CMIP6_SST_MAX = 10.198975563049316 +CMIP6_SST_MIN = -16.549121856689453 +CMIP5_SST_MAX = 8.991744995117188 +CMIP5_SST_MIN = -9.33076286315918 +CMIP6_NINO_MAX = 4.138188362121582 +CMIP6_NINO_MIN = -3.5832221508026123 +CMIP5_NINO_MAX = 3.8253555297851562 +CMIP5_NINO_MIN = -2.691682815551758 +SST_MAX = max(CMIP6_SST_MAX, CMIP5_SST_MAX) +SST_MIN = min(CMIP6_SST_MIN, CMIP5_SST_MIN) + + +def scale_sst(sst): + return (sst - SST_MIN) / (SST_MAX - SST_MIN) + + +def scale_back_sst(sst): + return (SST_MAX - SST_MIN) * sst + SST_MIN + + +def prepare_inputs_targets( + len_time, input_gap, input_length, pred_shift, pred_length, samples_gap +): + """Prepares the input and target indices for training. + + Args: + len_time (int): The total number of time steps in the dataset. + input_gap (int): Time gaps between two consecutive input frames. + input_length (int): The number of input frames. + pred_shift (int): The lead_time of the last target to be predicted. + pred_length (int): The number of frames to be predicted. + samples_gap (int): Stride of seq sampling. + """ + + if pred_shift < pred_length: + raise ValueError("Pred_shift should be small than pred_length") + input_span = input_gap * (input_length - 1) + 1 + pred_gap = pred_shift // pred_length + input_ind = np.arange(0, input_span, input_gap) + target_ind = np.arange(0, pred_shift, pred_gap) + input_span + pred_gap - 1 + ind = np.concatenate([input_ind, target_ind]).reshape(1, input_length + pred_length) + max_n_sample = len_time - (input_span + pred_shift - 1) + ind = ind + np.arange(max_n_sample)[:, np.newaxis] @ np.ones( + (1, input_length + pred_length), dtype=int + ) + return ind[::samples_gap] + + +def fold(data, size=36, stride=12): + """inverse of unfold/sliding window operation + only applicable to the case where the size of the sliding windows is n*stride + + Args: + data (tuple[int,...]): The input data.(N, size, *). + size (int, optional): The size of a single datum.The Defaults to 36. + stride (int, optional): The step.Defaults to 12. + + Returns: + outdata (np.ndarray): (N_, *).N/size is the number/width of sliding blocks + """ + + if size % stride != 0: + raise ValueError("size modulo stride should be zero") + times = size // stride + remain = (data.shape[0] - 1) % times + if remain > 0: + ls = list(data[::times]) + [data[-1, -(remain * stride) :]] + outdata = np.concatenate(ls, axis=0) # (36*(151//3+1)+remain*stride, *, 15) + else: + outdata = np.concatenate(data[::times], axis=0) # (36*(151/3+1), *, 15) + assert ( + outdata.shape[0] == size * ((data.shape[0] - 1) // times + 1) + remain * stride + ) + return outdata + + +def data_transform(data, num_years_per_model): + """The transform of the input data. + + Args: + data (Tuple[list,...]): The input data.Shape of (N, 36, *). + num_years_per_model (int): The number of years associated with each model.151/140. + """ + + length = data.shape[0] + assert length % num_years_per_model == 0 + num_models = length // num_years_per_model + outdata = np.stack( + np.split(data, length / num_years_per_model, axis=0), axis=-1 + ) # (151, 36, *, 15) + # cmip6sst outdata.shape = (151, 36, 24, 48, 15) = (year, month, lat, lon, model) + # cmip5sst outdata.shape = (140, 36, 24, 48, 17) + # cmip6nino outdata.shape = (151, 36, 15) + # cmip5nino outdata.shape = (140, 36, 17) + outdata = fold(outdata, size=36, stride=12) + # cmip6sst outdata.shape = (1836, 24, 48, 15), 1836 == 151 * 12 + 24 + # cmip5sst outdata.shape = (1704, 24, 48, 17) + # cmip6nino outdata.shape = (1836, 15) + # cmip5nino outdata.shape = (1704, 17) + + # check output data + assert outdata.shape[-1] == num_models + assert not np.any(np.isnan(outdata)) + return outdata + + +def read_raw_data(ds_dir, out_dir=None): + """read and process raw cmip data from CMIP_train.nc and CMIP_label.nc + + Args: + ds_dir (str): The path of the dataset. + out_dir (str): The path of output. Defaults to None. + """ + import xarray as xr + + train_cmip = xr.open_dataset( + Path(ds_dir) / "CMIP_train.nc", engine="h5netcdf" + ).transpose("year", "month", "lat", "lon") + label_cmip = xr.open_dataset( + Path(ds_dir) / "CMIP_label.nc", engine="h5netcdf" + ).transpose("year", "month") + # train_cmip.sst.values.shape = (4645, 36, 24, 48) + + # select longitudes + lon = train_cmip.lon.values + lon = lon[np.logical_and(lon >= 95, lon <= 330)] + train_cmip = train_cmip.sel(lon=lon) + + cmip6sst = data_transform( + data=train_cmip.sst.values[:2265], num_years_per_model=151 + ) + cmip5sst = data_transform( + data=train_cmip.sst.values[2265:], num_years_per_model=140 + ) + cmip6nino = data_transform( + data=label_cmip.nino.values[:2265], num_years_per_model=151 + ) + cmip5nino = data_transform( + data=label_cmip.nino.values[2265:], num_years_per_model=140 + ) + + # cmip6sst.shape = (1836, 24, 48, 15) + # cmip5sst.shape = (1704, 24, 48, 17) + assert len(cmip6sst.shape) == 4 + assert len(cmip5sst.shape) == 4 + assert len(cmip6nino.shape) == 2 + assert len(cmip5nino.shape) == 2 + # store processed data for faster data access + if out_dir is not None: + ds_cmip6 = xr.Dataset( + { + "sst": (["month", "lat", "lon", "model"], cmip6sst), + "nino": (["month", "model"], cmip6nino), + }, + coords={ + "month": np.repeat( + np.arange(1, 13)[None], cmip6nino.shape[0] // 12, axis=0 + ).flatten(), + "lat": train_cmip.lat.values, + "lon": train_cmip.lon.values, + "model": np.arange(15) + 1, + }, + ) + ds_cmip6.to_netcdf(Path(out_dir) / "cmip6.nc") + ds_cmip5 = xr.Dataset( + { + "sst": (["month", "lat", "lon", "model"], cmip5sst), + "nino": (["month", "model"], cmip5nino), + }, + coords={ + "month": np.repeat( + np.arange(1, 13)[None], cmip5nino.shape[0] // 12, axis=0 + ).flatten(), + "lat": train_cmip.lat.values, + "lon": train_cmip.lon.values, + "model": np.arange(17) + 1, + }, + ) + ds_cmip5.to_netcdf(Path(out_dir) / "cmip5.nc") + train_cmip.close() + label_cmip.close() + return cmip6sst, cmip5sst, cmip6nino, cmip5nino + + +def cat_over_last_dim(data): + """treat different models (15 from CMIP6, 17 from CMIP5) as batch_size + e.g., cmip6sst.shape = (178, 38, 24, 48, 15), converted_cmip6sst.shape = (2670, 38, 24, 48) + e.g., cmip5sst.shape = (165, 38, 24, 48, 15), converted_cmip6sst.shape = (2475, 38, 24, 48) + """ + + return np.concatenate(np.moveaxis(data, -1, 0), axis=0) + + +class ExtMoEENSODataset(io.Dataset): + """The El Niño/Southern Oscillation dataset. + + Args: + input_keys (Tuple[str, ...]): Name of input keys, such as ("input",). + label_keys (Tuple[str, ...]): Name of label keys, such as ("output",). + data_dir (str): The directory of data. + weight_dict (Optional[Dict[str, Union[Callable, float]]]): Define the weight of each constraint variable. Defaults to None. + in_len (int, optional): The length of input data. Defaults to 12. + out_len (int, optional): The length of out data. Defaults to 26. + in_stride (int, optional): The stride of input data. Defaults to 1. + out_stride (int, optional): The stride of output data. Defaults to 1. + train_samples_gap (int, optional): The stride of sequence sampling during training. Defaults to 10. + e.g., samples_gap = 10, the first seq contains [0, 1, ..., T-1] frame indices, the second seq contains [10, 11, .., T+9] + eval_samples_gap (int, optional): The stride of sequence sampling during eval. Defaults to 11. + normalize_sst (bool, optional): Whether to use normalization. Defaults to True. + batch_size (int, optional): Batch size. Defaults to 1. + num_workers (int, optional): The num of workers. Defaults to 1. + training (str, optional): Training pathse. Defaults to "train". + """ + + # Whether support batch indexing for speeding up fetching process. + batch_index: bool = False + + def __init__( + self, + input_keys: Tuple[str, ...], + label_keys: Tuple[str, ...], + data_dir: str, + weight_dict: Optional[Dict[str, float]] = None, + in_len: int = 12, + out_len: int = 26, + in_stride: int = 1, + out_stride: int = 1, + train_samples_gap: int = 10, + eval_samples_gap: int = 11, + normalize_sst: bool = True, + batch_size: int = 1, + num_workers: int = 1, + training: str = "train", + ): + super(ExtMoEENSODataset, self).__init__() + if importlib.util.find_spec("xarray") is None: + raise ModuleNotFoundError( + "To use RadarDataset, please install 'xarray' with: `pip install " + "xarray` first." + ) + self.input_keys = input_keys + self.label_keys = label_keys + self.data_dir = data_dir + self.weight_dict = {} if weight_dict is None else weight_dict + if weight_dict is not None: + self.weight_dict = {key: 1.0 for key in self.label_keys} + self.weight_dict.update(weight_dict) + + self.in_len = in_len + self.out_len = out_len + self.in_stride = in_stride + self.out_stride = out_stride + self.train_samples_gap = train_samples_gap + self.eval_samples_gap = eval_samples_gap + self.normalize_sst = normalize_sst + # datamodule_only + self.batch_size = batch_size + if num_workers != 1: + raise ValueError( + "Current implementation does not support `num_workers != 1`!" + ) + self.num_workers = num_workers + self.training = training + + # pre-data + cmip6sst, cmip5sst, cmip6nino, cmip5nino = read_raw_data(self.data_dir) + # TODO: more flexible train/val/test split + self.sst_train = [cmip6sst, cmip5sst[..., :-2]] + self.nino_train = [cmip6nino, cmip5nino[..., :-2]] + self.sst_eval = [cmip5sst[..., -2:-1]] + self.nino_eval = [cmip5nino[..., -2:-1]] + self.sst_test = [cmip5sst[..., -1:]] + self.nino_test = [cmip5nino[..., -1:]] + + self.sst, self.target_nino = self.create_data() + + def create_data( + self, + ): + if self.training == "train": + sst_cmip6 = self.sst_train[0] + nino_cmip6 = self.nino_train[0] + sst_cmip5 = self.sst_train[1] + nino_cmip5 = self.nino_train[1] + samples_gap = self.train_samples_gap + elif self.training == "eval": + sst_cmip6 = None + nino_cmip6 = None + sst_cmip5 = self.sst_eval[0] + nino_cmip5 = self.nino_eval[0] + samples_gap = self.eval_samples_gap + elif self.training == "test": + sst_cmip6 = None + nino_cmip6 = None + sst_cmip5 = self.sst_test[0] + nino_cmip5 = self.nino_test[0] + samples_gap = self.eval_samples_gap + + # cmip6 (N, *, 15) + # cmip5 (N, *, 17) + sst = [] + target_nino = [] + + nino_idx_slice = slice( + self.in_len, self.in_len + self.out_len - NINO_WINDOW_T + 1 + ) # e.g., 12:36 + if sst_cmip6 is not None: + assert len(sst_cmip6.shape) == 4 + assert len(nino_cmip6.shape) == 2 + idx_sst = prepare_inputs_targets( + len_time=sst_cmip6.shape[0], + input_length=self.in_len, + input_gap=self.in_stride, + pred_shift=self.out_len * self.out_stride, + pred_length=self.out_len, + samples_gap=samples_gap, + ) + + sst.append(cat_over_last_dim(sst_cmip6[idx_sst])) + target_nino.append( + cat_over_last_dim(nino_cmip6[idx_sst[:, nino_idx_slice]]) + ) + if sst_cmip5 is not None: + assert len(sst_cmip5.shape) == 4 + assert len(nino_cmip5.shape) == 2 + idx_sst = prepare_inputs_targets( + len_time=sst_cmip5.shape[0], + input_length=self.in_len, + input_gap=self.in_stride, + pred_shift=self.out_len * self.out_stride, + pred_length=self.out_len, + samples_gap=samples_gap, + ) + sst.append(cat_over_last_dim(sst_cmip5[idx_sst])) + target_nino.append( + cat_over_last_dim(nino_cmip5[idx_sst[:, nino_idx_slice]]) + ) + + # sst data containing both the input and target + self.sst = np.concatenate(sst, axis=0) # (N, in_len+out_len, lat, lon) + if self.normalize_sst: + self.sst = scale_sst(self.sst) + # nino data containing the target only + self.target_nino = np.concatenate( + target_nino, axis=0 + ) # (N, out_len+NINO_WINDOW_T-1) + assert self.sst.shape[0] == self.target_nino.shape[0] + assert self.sst.shape[1] == self.in_len + self.out_len + assert self.target_nino.shape[1] == self.out_len - NINO_WINDOW_T + 1 + + return self.sst, self.target_nino + + def get_datashape(self): + return {"sst": self.sst.shape, "nino target": self.target_nino.shape} + + def __len__(self): + return self.sst.shape[0] + + def __getitem__(self, idx): + sst_data = self.sst[idx].astype("float32") + sst_data = sst_data[..., np.newaxis] + in_seq = sst_data[: self.in_len, ...] # ( in_len, lat, lon, 1) + target_seq = sst_data[self.in_len :, ...] # ( in_len, lat, lon, 1) + weight_item = self.weight_dict + + if self.training == "train": + input_item = {self.input_keys[0]: in_seq, "sst_target": target_seq} + label_item = { + self.label_keys[0]: target_seq, + } + + return input_item, label_item, weight_item + else: + input_item = {self.input_keys[0]: in_seq, "sst_target": target_seq} + label_item = { + self.label_keys[0]: target_seq, + self.label_keys[1]: self.target_nino[idx], + } + + return input_item, label_item, weight_item diff --git a/examples/smc_reac/ppsci/data/dataset/fwi_dataset.py b/examples/smc_reac/ppsci/data/dataset/fwi_dataset.py new file mode 100644 index 0000000000..7eaef93187 --- /dev/null +++ b/examples/smc_reac/ppsci/data/dataset/fwi_dataset.py @@ -0,0 +1,103 @@ +import os +from typing import Dict +from typing import Optional +from typing import Tuple + +import numpy as np +from paddle import io +from paddle import vision + + +class FWIDataset(io.Dataset): + """Datasets for full waveform inversion tasks. + For convenience, in this class, a batch refers to a npy file instead of the batch used during training. + + Args: + input_keys (Tuple[str, ...]): List of input keys. + label_keys (Tuple[str, ...]): List of label keys. + weight: Define the weight dict for loss function. + anno: Path to annotation file. + preload: Whether to load the whole dataset into memory. + sample_ratio: Downsample ratio for seismic data. + file_size: Number of samples in each npy file. + transform_data: Transformation applied to data. + transform_label: Transformation applied to label. + + Examples: + >>> import ppsci + >>> dataset = ppsci.data.dataset.FWIDataset(("input", ), ("label", ), "path/to/anno_file") # doctest: +SKIP + """ + + def __init__( + self, + input_keys: Tuple[str, ...], + label_keys: Tuple[str, ...], + anno: str, + weight: Optional[Dict[str, np.ndarray]] = None, + preload: bool = True, + sample_ratio: int = 1, + file_size: int = 500, + transform_data: Optional[vision.Compose] = None, + transform_label: Optional[vision.Compose] = None, + ): + super().__init__() + self.input_keys = input_keys + self.label_keys = label_keys + self.weight = {} if weight is None else weight + if not os.path.exists(anno): + print(f"Annotation file {anno} does not exists") + self.preload = preload + self.sample_ratio = sample_ratio + self.file_size = file_size + self.transform_data = transform_data + self.transform_label = transform_label + with open(anno, "r") as f: + self.batches = f.readlines() + if preload: + self.data_list, self.label_list = [], [] + for batch in self.batches: + data, label = self.load_every(batch) + self.data_list.append(data) + if label is not None: + self.label_list.append(label) + + def load_every(self, batch): + batch = batch.split("\t") + data_path = batch[0] if len(batch) > 1 else batch[0][:-1] + data = np.load(data_path)[:, :, :: self.sample_ratio, :] + data = data.astype("float32") + if len(batch) > 1: + label_path = batch[1][:-1] + label = np.load(label_path) + label = label.astype("float32") + else: + label = None + + return data, label + + def __getitem__(self, idx): + batch_idx, sample_idx = idx // self.file_size, idx % self.file_size + if self.preload: + data = self.data_list[batch_idx][sample_idx] + label = ( + self.label_list[batch_idx][sample_idx] + if len(self.label_list) != 0 + else None + ) + else: + data, label = self.load_every(self.batches[batch_idx]) + data = data[sample_idx] + label = label[sample_idx] if label is not None else None + if self.transform_data: + data = self.transform_data(data) + if self.transform_label and label is not None: + label = self.transform_label(label) + + input_item = {self.input_keys[0]: data} + label_item = {self.label_keys[0]: label if label is not None else np.array([])} + weight_item = self.weight + + return input_item, label_item, weight_item + + def __len__(self): + return len(self.batches) * self.file_size diff --git a/examples/smc_reac/ppsci/data/dataset/ifm_moe_dataset.py b/examples/smc_reac/ppsci/data/dataset/ifm_moe_dataset.py new file mode 100644 index 0000000000..43cd50f419 --- /dev/null +++ b/examples/smc_reac/ppsci/data/dataset/ifm_moe_dataset.py @@ -0,0 +1,462 @@ +# Copyright (c) 2024 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. + +from __future__ import annotations + +import copy +import os +from typing import Tuple + +import numpy as np +import paddle +import pandas as pd +from paddle import io + +tasks_dic = { + "freesolv": ["activity"], + "esol": ["activity"], + "lipop": ["activity"], + "bace": ["activity"], + "bbbp": ["activity"], + "hiv": ["activity"], + "clintox": ["FDA_APPROVED", "CT_TOX"], + "sider": [ + "SIDER1", + "SIDER2", + "SIDER3", + "SIDER4", + "SIDER5", + "SIDER6", + "SIDER7", + "SIDER8", + "SIDER9", + "SIDER10", + "SIDER11", + "SIDER12", + "SIDER13", + "SIDER14", + "SIDER15", + "SIDER16", + "SIDER17", + "SIDER18", + "SIDER19", + "SIDER20", + "SIDER21", + "SIDER22", + "SIDER23", + "SIDER24", + "SIDER25", + "SIDER26", + "SIDER27", + ], + "tox21": [ + "NR-AR", + "NR-AR-LBD", + "NR-AhR", + "NR-Aromatase", + "NR-ER", + "NR-ER-LBD", + "NR-PPAR-gamma", + "SR-ARE", + "SR-ATAD5", + "SR-HSE", + "SR-MMP", + "SR-p53", + ], + "muv": [ + "MUV-466", + "MUV-548", + "MUV-600", + "MUV-644", + "MUV-652", + "MUV-689", + "MUV-692", + "MUV-712", + "MUV-713", + "MUV-733", + "MUV-737", + "MUV-810", + "MUV-832", + "MUV-846", + "MUV-852", + "MUV-858", + "MUV-859", + ], + "toxcast": [ + "ACEA_T47D_80hr_Negative", + "ACEA_T47D_80hr_Positive", + "APR_HepG2_CellCycleArrest_24h_dn", + "APR_HepG2_CellCycleArrest_72h_dn", + "APR_HepG2_CellLoss_24h_dn", + "APR_HepG2_CellLoss_72h_dn", + "APR_HepG2_MicrotubuleCSK_72h_up", + "APR_HepG2_MitoMass_24h_dn", + "APR_HepG2_MitoMass_72h_dn", + "APR_HepG2_MitoMembPot_24h_dn", + "APR_HepG2_MitoMembPot_72h_dn", + "APR_HepG2_MitoticArrest_24h_up", + "APR_HepG2_MitoticArrest_72h_up", + "APR_HepG2_OxidativeStress_24h_up", + "APR_HepG2_OxidativeStress_72h_up", + "APR_HepG2_StressKinase_72h_up", + "APR_HepG2_p53Act_24h_up", + "APR_HepG2_p53Act_72h_up", + "ATG_AP_1_CIS_up", + "ATG_Ahr_CIS_up", + "ATG_BRE_CIS_up", + "ATG_CMV_CIS_up", + "ATG_CRE_CIS_up", + "ATG_DR4_LXR_CIS_dn", + "ATG_DR5_CIS_up", + "ATG_EGR_CIS_up", + "ATG_ERE_CIS_up", + "ATG_ERa_TRANS_up", + "ATG_E_Box_CIS_dn", + "ATG_HIF1a_CIS_up", + "ATG_HSE_CIS_up", + "ATG_IR1_CIS_dn", + "ATG_ISRE_CIS_dn", + "ATG_MRE_CIS_up", + "ATG_NRF2_ARE_CIS_up", + "ATG_Oct_MLP_CIS_up", + "ATG_PBREM_CIS_up", + "ATG_PPARg_TRANS_up", + "ATG_PPRE_CIS_up", + "ATG_PXRE_CIS_dn", + "ATG_PXRE_CIS_up", + "ATG_PXR_TRANS_up", + "ATG_Pax6_CIS_up", + "ATG_RORE_CIS_up", + "ATG_RXRb_TRANS_up", + "ATG_SREBP_CIS_up", + "ATG_Sp1_CIS_up", + "ATG_TCF_b_cat_CIS_dn", + "ATG_VDRE_CIS_up", + "ATG_Xbp1_CIS_up", + "ATG_p53_CIS_dn", + "BSK_3C_Eselectin_down", + "BSK_3C_HLADR_down", + "BSK_3C_ICAM1_down", + "BSK_3C_IL8_down", + "BSK_3C_MCP1_down", + "BSK_3C_MIG_down", + "BSK_3C_Proliferation_down", + "BSK_3C_SRB_down", + "BSK_3C_Thrombomodulin_up", + "BSK_3C_TissueFactor_down", + "BSK_3C_VCAM1_down", + "BSK_3C_Vis_down", + "BSK_3C_uPAR_down", + "BSK_4H_Eotaxin3_down", + "BSK_4H_MCP1_down", + "BSK_4H_Pselectin_down", + "BSK_4H_SRB_down", + "BSK_4H_VCAM1_down", + "BSK_4H_VEGFRII_down", + "BSK_4H_uPAR_down", + "BSK_BE3C_HLADR_down", + "BSK_BE3C_IL1a_down", + "BSK_BE3C_IP10_down", + "BSK_BE3C_MIG_down", + "BSK_BE3C_MMP1_down", + "BSK_BE3C_MMP1_up", + "BSK_BE3C_PAI1_down", + "BSK_BE3C_SRB_down", + "BSK_BE3C_TGFb1_down", + "BSK_BE3C_tPA_down", + "BSK_BE3C_uPAR_down", + "BSK_BE3C_uPA_down", + "BSK_CASM3C_HLADR_down", + "BSK_CASM3C_IL6_down", + "BSK_CASM3C_IL8_down", + "BSK_CASM3C_LDLR_down", + "BSK_CASM3C_MCP1_down", + "BSK_CASM3C_MCSF_down", + "BSK_CASM3C_MIG_down", + "BSK_CASM3C_Proliferation_down", + "BSK_CASM3C_SAA_down", + "BSK_CASM3C_SRB_down", + "BSK_CASM3C_Thrombomodulin_up", + "BSK_CASM3C_TissueFactor_down", + "BSK_CASM3C_VCAM1_down", + "BSK_CASM3C_uPAR_down", + "BSK_KF3CT_ICAM1_down", + "BSK_KF3CT_IL1a_down", + "BSK_KF3CT_IP10_down", + "BSK_KF3CT_MCP1_down", + "BSK_KF3CT_MMP9_down", + "BSK_KF3CT_SRB_down", + "BSK_KF3CT_TGFb1_down", + "BSK_KF3CT_TIMP2_down", + "BSK_KF3CT_uPA_down", + "BSK_LPS_CD40_down", + "BSK_LPS_Eselectin_down", + "BSK_LPS_IL1a_down", + "BSK_LPS_IL8_down", + "BSK_LPS_MCP1_down", + "BSK_LPS_MCSF_down", + "BSK_LPS_PGE2_down", + "BSK_LPS_SRB_down", + "BSK_LPS_TNFa_down", + "BSK_LPS_TissueFactor_down", + "BSK_LPS_VCAM1_down", + "BSK_SAg_CD38_down", + "BSK_SAg_CD40_down", + "BSK_SAg_CD69_down", + "BSK_SAg_Eselectin_down", + "BSK_SAg_IL8_down", + "BSK_SAg_MCP1_down", + "BSK_SAg_MIG_down", + "BSK_SAg_PBMCCytotoxicity_down", + "BSK_SAg_Proliferation_down", + "BSK_SAg_SRB_down", + "BSK_hDFCGF_CollagenIII_down", + "BSK_hDFCGF_IL8_down", + "BSK_hDFCGF_IP10_down", + "BSK_hDFCGF_MCSF_down", + "BSK_hDFCGF_MIG_down", + "BSK_hDFCGF_MMP1_down", + "BSK_hDFCGF_PAI1_down", + "BSK_hDFCGF_Proliferation_down", + "BSK_hDFCGF_SRB_down", + "BSK_hDFCGF_TIMP1_down", + "BSK_hDFCGF_VCAM1_down", + "CEETOX_H295R_11DCORT_dn", + "CEETOX_H295R_ANDR_dn", + "CEETOX_H295R_CORTISOL_dn", + "CEETOX_H295R_ESTRONE_dn", + "CEETOX_H295R_ESTRONE_up", + "NHEERL_ZF_144hpf_TERATOSCORE_up", + "NVS_NR_bER", + "NVS_NR_hER", + "NVS_NR_hPPARg", + "NVS_NR_hPXR", + "NVS_NR_mERa", + "OT_AR_ARSRC1_0960", + "OT_ER_ERaERb_0480", + "OT_ER_ERaERb_1440", + "OT_ER_ERbERb_0480", + "OT_ER_ERbERb_1440", + "OT_ERa_EREGFP_0120", + "OT_FXR_FXRSRC1_0480", + "OT_NURR1_NURR1RXRa_0480", + "TOX21_ARE_BLA_agonist_ratio", + "TOX21_AR_BLA_Antagonist_ratio", + "TOX21_AR_LUC_MDAKB2_Antagonist", + "TOX21_AR_LUC_MDAKB2_Antagonist2", + "TOX21_AhR_LUC_Agonist", + "TOX21_Aromatase_Inhibition", + "TOX21_ERa_BLA_Antagonist_ratio", + "TOX21_ERa_LUC_BG1_Agonist", + "TOX21_FXR_BLA_antagonist_ratio", + "TOX21_MMP_ratio_down", + "TOX21_TR_LUC_GH3_Antagonist", + "TOX21_p53_BLA_p1_ratio", + "TOX21_p53_BLA_p2_ch2", + "TOX21_p53_BLA_p2_ratio", + "TOX21_p53_BLA_p2_viability", + "TOX21_p53_BLA_p3_ratio", + "TOX21_p53_BLA_p4_ratio", + "TOX21_p53_BLA_p5_ratio", + "Tanguay_ZF_120hpf_AXIS_up", + "Tanguay_ZF_120hpf_ActivityScore", + "Tanguay_ZF_120hpf_JAW_up", + "Tanguay_ZF_120hpf_MORT_up", + "Tanguay_ZF_120hpf_PE_up", + "Tanguay_ZF_120hpf_SNOU_up", + "Tanguay_ZF_120hpf_YSE_up", + ], +} + + +def standardize(col): + return (col - np.mean(col)) / np.std(col) + + +def get_pos_weight(Ys): + Ys = paddle.to_tensor(np.nan_to_num(Ys), dtype=paddle.float32) + num_pos = paddle.sum(Ys, axis=0) + num_indices = paddle.to_tensor(len(Ys)) + return (num_indices - num_pos) / num_pos + + +class IFMMoeDataset(io.Dataset): + """Dataset for `IFMMoe`. + + Args: + input_keys (Tuple[str, ...]): Name of input data. + label_keys (Tuple[str, ...]): Name of label data. + data_dir (str): Directory of IFMMoe data. + data_label (str): IFMMoe data label in tox21/esol/freesolv/lipop... + data_mode (str): train/val/test mode data. + + Examples: + >>> import ppsci + >>> dataset = ppsci.data.dataset.IFMMoeDataset( + ... "input_keys": ("input",), + ... "label_keys": ("output",), + ... "data_dir": "/path/to/IFMMoeDataset", + ... "data_label": "tox21", + ... "data_mode": "train", + ... ) # doctest: +SKIP + """ + + # Whether support batch indexing for speeding up fetching process. + batch_index: bool = False + use_pgl: bool = False + + def __init__( + self, + input_keys: Tuple[str, ...], + label_keys: Tuple[str, ...], + data_dir: str, + data_label: str, + data_mode: str, + ): + self.input_keys = input_keys + self.label_keys = label_keys + + self.data_label = data_label + self.data_dir = data_dir + self.data_mode = data_mode + + if data_label == "esol" or data_label == "freesolv" or data_label == "lipop": + self.task_type = "reg" + self.reg = True + # metric = "rmse" + else: + self.task_type = "cla" + self.reg = False + # metric = "roc_auc" + + self.task_dict = tasks_dic + + self.Xs = None + self.Ys = None + self.mask = None + self.process_data() + + def process_data(self): + file_name = os.path.join(self.data_dir, self.data_label + "_moe_pubsubfp.csv") + # preprocess data + dataset_all = pd.read_csv(file_name) + if self.data_label == "freesolv": + dataset_all.drop(columns=["vsa_pol", "h_emd", "a_donacc"], inplace=True) + elif self.data_label == "esol": + dataset_all.drop(columns=["logS", "h_logS", "SlogP"], inplace=True) + else: + dataset_all.drop(columns=["SlogP", "h_logD", "logS"], inplace=True) + tasks = tasks_dic[self.data_label] + cols = copy.deepcopy(tasks) + cols.extend(dataset_all.columns[len(tasks) + 1 :]) + dataset = dataset_all[cols] + x_cols = dataset_all.columns[len(tasks) + 1 :] + # remove the features with na + if self.data_label != "hiv": + rm_cols1 = ( + dataset[x_cols] + .isnull() + .any()[dataset[x_cols].isnull().any() == True] # noqa: E712 + .index + ) + dataset.drop(columns=rm_cols1, inplace=True) + else: + rm_indx1 = ( + dataset[x_cols] + .isnull() + .T.any()[dataset[x_cols].isnull().T.any() == True] # noqa: E712 + .index + ) + dataset.drop(index=rm_indx1, inplace=True) + x_cols = dataset.columns.drop(tasks) + + # Removing features with low variance + # threshold = 0.05 + data_fea_var = dataset[x_cols].var() + del_fea1 = list(data_fea_var[data_fea_var <= 0.05].index) + dataset.drop(columns=del_fea1, inplace=True) + x_cols = dataset.columns.drop(tasks) + + # pair correlations + # threshold = 0.95 + data_fea_corr = dataset[x_cols].corr() + del_fea2_col = [] + del_fea2_ind = [] + length = data_fea_corr.shape[1] + for i in range(length): + for j in range(i + 1, length): + if abs(data_fea_corr.iloc[i, j]) >= 0.95: + del_fea2_col.append(data_fea_corr.columns[i]) + del_fea2_ind.append(data_fea_corr.index[j]) + dataset.drop(columns=del_fea2_ind, inplace=True) + # standardize the features + cols_ = dataset.columns[len(tasks) + 1 :] + # print('the retained features for %s is %d' % (args.task, len(cols_))) + dataset[cols_] = dataset[cols_].apply(standardize, axis=0) + + dataseta = pd.read_csv( + os.path.join( + self.data_dir, "dataset_used_for_modeling", self.data_label + ".csv" + ) + ) + data_tr = dataset[dataseta.group == "train"] + data_va = dataset[dataseta.group == "valid"] + data_te = dataset[dataseta.group == "test"] + + # training set + data_tr_y = data_tr[tasks].values.reshape(-1, len(tasks)) + data_tr_x = data_tr.iloc[:, len(tasks) :].values # 249 + # data_tr_x = data_tr.iloc[:, len(tasks):].values + # test set + data_te_y = data_te[tasks].values.reshape(-1, len(tasks)) + data_te_x = data_te.iloc[:, len(tasks) :].values + # data_te_x = data_te.iloc[:, len(tasks):].values + + # validation set + data_va_y = data_va[tasks].values.reshape(-1, len(tasks)) + data_va_x = data_va.iloc[:, len(tasks) :].values + # data_va_x = data_va.iloc[:, len(tasks):].values + + # dataloader + # train_dataset = MyDataset(data_tr_x, data_tr_y) + # validation_dataset = MyDataset(data_va_x, data_va_y) + # test_dataset = MyDataset(data_te_x, data_te_y) + if self.data_mode == "train": + Xs, Ys = data_tr_x, data_tr_y + elif self.data_mode == "val": + Xs, Ys = data_va_x, data_va_y + elif self.data_mode == "test": + Xs, Ys = data_te_x, data_te_y + if not self.reg: + self.pos_weights = get_pos_weight(dataset[tasks].values) + + self.data_tr_x = data_tr_x + self.Xs = Xs + self.Ys = np.nan_to_num(Ys) + self.mask = ~np.isnan(Ys) * 1.0 + + def __len__(self): + return len(self.Ys) + + def __getitem__(self, idx): + return ( + { + self.input_keys[0]: paddle.to_tensor(self.Xs[idx], dtype="float32"), + }, + { + self.label_keys[0]: paddle.to_tensor(self.Ys[idx], dtype="float32"), + self.label_keys[1]: paddle.to_tensor(self.mask[idx], dtype="float32"), + }, + {}, + ) diff --git a/examples/smc_reac/ppsci/data/dataset/mat_dataset.py b/examples/smc_reac/ppsci/data/dataset/mat_dataset.py new file mode 100644 index 0000000000..609e35aeaa --- /dev/null +++ b/examples/smc_reac/ppsci/data/dataset/mat_dataset.py @@ -0,0 +1,287 @@ +# 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. + +from __future__ import annotations + +from typing import Callable +from typing import Dict +from typing import Optional +from typing import Tuple +from typing import Union + +import numpy as np +import paddle +from paddle import io +from paddle import vision + +from ppsci.utils import misc +from ppsci.utils import reader + + +class MatDataset(io.Dataset): + """Dataset class for .mat file. + + Args: + file_path (str): Mat file path. + input_keys (Tuple[str, ...]): List of input keys. + label_keys (Tuple[str, ...], optional): List of label keys. Defaults to (). + alias_dict (Optional[Dict[str, str]]): Dict of alias(es) for input and label keys. + i.e. {inner_key: outer_key}. Defaults to None. + weight_dict (Optional[Dict[str, Union[Callable, float]]]): Define the weight of + each constraint variable. Defaults to None. + timestamps (Optional[Tuple[float, ...]]): The number of repetitions of the data + in the time dimension. Defaults to None. + transforms (Optional[vision.Compose]): Compose object contains sample wise + transform(s). Defaults to None. + + Examples: + >>> import ppsci + >>> dataset = ppsci.data.dataset.MatDataset( + ... "/path/to/file.mat" + ... ("x",), + ... ("u",), + ... ) # doctest: +SKIP + """ + + # Whether support batch indexing for speeding up fetching process. + batch_index: bool = True + + def __init__( + self, + file_path: str, + input_keys: Tuple[str, ...], + label_keys: Tuple[str, ...] = (), + alias_dict: Optional[Dict[str, str]] = None, + weight_dict: Optional[Dict[str, Union[Callable, float]]] = None, + timestamps: Optional[Tuple[float, ...]] = None, + transforms: Optional[vision.Compose] = None, + ): + super().__init__() + self.input_keys = input_keys + self.label_keys = label_keys + + # read raw data from file + raw_data = reader.load_mat_file( + file_path, + input_keys + label_keys, + alias_dict, + ) + # filter raw data by given timestamps if specified + if timestamps is not None: + if "t" in raw_data: + # filter data according to given timestamps + raw_time_array = raw_data["t"] + mask = [] + for ti in timestamps: + mask.append(np.nonzero(np.isclose(raw_time_array, ti).flatten())[0]) + raw_data = misc.convert_to_array( + raw_data, self.input_keys + self.label_keys + ) + mask = np.concatenate(mask, 0) + raw_data = raw_data[mask] + raw_data = misc.convert_to_dict( + raw_data, self.input_keys + self.label_keys + ) + else: + # repeat data according to given timestamps + raw_data = misc.convert_to_array( + raw_data, self.input_keys + self.label_keys + ) + raw_data = misc.combine_array_with_time(raw_data, timestamps) + self.input_keys = ("t",) + tuple(self.input_keys) + raw_data = misc.convert_to_dict( + raw_data, self.input_keys + self.label_keys + ) + + # fetch input data + self.input = { + key: value for key, value in raw_data.items() if key in self.input_keys + } + # fetch label data + self.label = { + key: value for key, value in raw_data.items() if key in self.label_keys + } + + # prepare weights + self.weight = ( + {key: np.ones_like(next(iter(self.label.values()))) for key in self.label} + if weight_dict is not None + else {} + ) + if weight_dict is not None: + for key, value in weight_dict.items(): + if isinstance(value, (int, float)): + self.weight[key] = np.full_like( + next(iter(self.label.values())), value + ) + elif callable(value): + func = value + self.weight[key] = func(self.input) + if isinstance(self.weight[key], (int, float)): + self.weight[key] = np.full_like( + next(iter(self.label.values())), self.weight[key] + ) + else: + raise NotImplementedError(f"type of {type(value)} is invalid yet.") + + self.transforms = transforms + self._len = len(next(iter(self.input.values()))) + + def __getitem__(self, idx): + input_item = {key: value[idx] for key, value in self.input.items()} + label_item = {key: value[idx] for key, value in self.label.items()} + weight_item = {key: value[idx] for key, value in self.weight.items()} + + if self.transforms is not None: + input_item, label_item, weight_item = self.transforms( + input_item, label_item, weight_item + ) + + return (input_item, label_item, weight_item) + + def __len__(self): + return self._len + + +class IterableMatDataset(io.IterableDataset): + """IterableMatDataset for full-data loading. + + Args: + file_path (str): Mat file path. + input_keys (Tuple[str, ...]): List of input keys. + label_keys (Tuple[str, ...], optional): List of label keys. Defaults to (). + alias_dict (Optional[Dict[str, str]]): Dict of alias(es) for input and label keys. + i.e. {inner_key: outer_key}. Defaults to None. + weight_dict (Optional[Dict[str, Union[Callable, float]]]): Define the weight of + each constraint variable. Defaults to None. + timestamps (Optional[Tuple[float, ...]]): The number of repetitions of the data + in the time dimension. Defaults to None. + transforms (Optional[vision.Compose]): Compose object contains sample wise + transform(s). Defaults to None. + + Examples: + >>> import ppsci + >>> dataset = ppsci.data.dataset.IterableMatDataset( + ... "/path/to/file.mat" + ... ("x",), + ... ("u",), + ... ) # doctest: +SKIP + """ + + # Whether support batch indexing for speeding up fetching process. + batch_index: bool = False + + def __init__( + self, + file_path: str, + input_keys: Tuple[str, ...], + label_keys: Tuple[str, ...] = (), + alias_dict: Optional[Dict[str, str]] = None, + weight_dict: Optional[Dict[str, Union[Callable, float]]] = None, + timestamps: Optional[Tuple[float, ...]] = None, + transforms: Optional[vision.Compose] = None, + ): + super().__init__() + self.input_keys = input_keys + self.label_keys = label_keys + + # read raw data from file + raw_data = reader.load_mat_file( + file_path, + input_keys + label_keys, + alias_dict, + ) + # filter raw data by given timestamps if specified + if timestamps is not None: + if "t" in raw_data: + # filter data according to given timestamps + raw_time_array = raw_data["t"] + mask = [] + for ti in timestamps: + mask.append(np.nonzero(np.isclose(raw_time_array, ti).flatten())[0]) + raw_data = misc.convert_to_array( + raw_data, self.input_keys + self.label_keys + ) + mask = np.concatenate(mask, 0) + raw_data = raw_data[mask] + raw_data = misc.convert_to_dict( + raw_data, self.input_keys + self.label_keys + ) + else: + # repeat data according to given timestamps + raw_data = misc.convert_to_array( + raw_data, self.input_keys + self.label_keys + ) + raw_data = misc.combine_array_with_time(raw_data, timestamps) + self.input_keys = ("t",) + tuple(self.input_keys) + raw_data = misc.convert_to_dict( + raw_data, self.input_keys + self.label_keys + ) + + # fetch input data + self.input = { + key: value for key, value in raw_data.items() if key in self.input_keys + } + # fetch label data + self.label = { + key: value for key, value in raw_data.items() if key in self.label_keys + } + + # prepare weights + self.weight = ( + {key: np.ones_like(next(iter(self.label.values()))) for key in self.label} + if weight_dict is not None + else {} + ) + if weight_dict is not None: + for key, value in weight_dict.items(): + if isinstance(value, (int, float)): + self.weight[key] = np.full_like( + next(iter(self.label.values())), value + ) + elif callable(value): + func = value + self.weight[key] = func(self.input) + if isinstance(self.weight[key], (int, float)): + self.weight[key] = np.full_like( + next(iter(self.label.values())), self.weight[key] + ) + else: + raise NotImplementedError(f"type of {type(value)} is invalid yet.") + + self.input = {key: paddle.to_tensor(value) for key, value in self.input.items()} + self.label = {key: paddle.to_tensor(value) for key, value in self.label.items()} + self.weight = { + key: paddle.to_tensor(value) for key, value in self.weight.items() + } + + self.transforms = transforms + self._len = len(next(iter(self.input.values()))) + + @property + def num_samples(self): + """Number of samples within current dataset.""" + return self._len + + def __iter__(self): + if callable(self.transforms): + input_, label_, weight_ = self.transforms( + self.input, self.label, self.weight + ) + yield input_, label_, weight_ + else: + yield self.input, self.label, self.weight + + def __len__(self): + return 1 diff --git a/examples/smc_reac/ppsci/data/dataset/moflow_dataset.py b/examples/smc_reac/ppsci/data/dataset/moflow_dataset.py new file mode 100644 index 0000000000..627365c1dc --- /dev/null +++ b/examples/smc_reac/ppsci/data/dataset/moflow_dataset.py @@ -0,0 +1,437 @@ +# Copyright (c) 2024 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. + +# Copyright 2020 Chengxi Zang + +from __future__ import annotations + +import os +from typing import Callable +from typing import Dict +from typing import List +from typing import Optional +from typing import Tuple + +import numpy as np +import pandas as pd +from paddle import io +from tqdm import tqdm + +from ppsci.utils import logger + +try: + from rdkit import Chem + from rdkit.Chem import rdmolops +except ModuleNotFoundError: + pass + + +class MolGraph: + """ + Args: + max_atoms (int): Max number of atoms for each molecule, if the + number of atoms is more than this value, this data is simply + ignored. + Setting negative value indicates no limit for max atoms. + out_size (int): It specifies the size of array returned by + `get_input_features`. + If the number of atoms in the molecule is less than this value, + the returned arrays is padded to have fixed size. + Setting negative value indicates do not pad returned array. + add_Hs (bool): If True, implicit Hs are added. + kekulize (bool): If True, Kekulizes the molecule. + """ + + def __init__(self, max_atoms=-1, out_size=-1, add_Hs=False, kekulize=False): + super(MolGraph, self).__init__() + self.add_Hs = add_Hs + self.kekulize = kekulize + if max_atoms >= 0 and out_size >= 0 and max_atoms > out_size: + raise ValueError( + f"max_atoms {max_atoms} must be less or equal to out_size {out_size}" + ) + self.max_atoms = max_atoms + self.out_size = out_size + + def get_input_features(self, mol): + """ + get input features + Args: + mol (Mol): mol instance + + Returns: + (tuple): (`atom`, `adj`) + + """ + self.type_check_num_atoms(mol, self.max_atoms) + atom_array = self.construct_atomic_number_array(mol, out_size=self.out_size) + adj_array = self.construct_discrete_edge_matrix(mol, out_size=self.out_size) + return atom_array, adj_array + + def prepare_smiles_and_mol(self, mol): + """Prepare `smiles` and `mol` used in following preprocessing. + This method is called before `get_input_features` is called, by parser + class. + This method may be overridden to support custom `smile`/`mol` extraction + Args: + mol (mol): mol instance + + Returns (tuple): (`smiles`, `mol`) + """ + canonical_smiles = Chem.MolToSmiles(mol, isomericSmiles=False, canonical=True) + mol = Chem.MolFromSmiles(canonical_smiles) + if self.add_Hs: + mol = Chem.AddHs(mol) + if self.kekulize: + Chem.Kekulize(mol) + return canonical_smiles, mol + + def get_label(self, mol, label_names=None): + """Extracts label information from a molecule. + This method extracts properties whose keys are + specified by ``label_names`` from a molecule ``mol`` + and returns these values as a list. + The order of the values is same as that of ``label_names``. + If the molecule does not have a + property with some label, this function fills the corresponding + index of the returned list with ``None``. + + Args: + mol (rdkit.Chem.Mol): molecule whose features to be extracted + label_names (None or iterable): list of label names. + + Returns: + list of str: label information. Its length is equal to + that of ``label_names``. If ``label_names`` is ``None``, + this function returns an empty list. + + """ + if label_names is None: + return [] + label_list = [] + for label_name in label_names: + if mol.HasProp(label_name): + label_list.append(mol.GetProp(label_name)) + else: + label_list.append(None) + return label_list + + def type_check_num_atoms(self, mol, num_max_atoms=-1): + """Check number of atoms in `mol` does not exceed `num_max_atoms` + If number of atoms in `mol` exceeds the number `num_max_atoms`, it will + raise `MolGraphError` exception. + + Args: + mol (Mol): + num_max_atoms (int): If negative value is set, not check number of + atoms. + + """ + num_atoms = mol.GetNumAtoms() + if num_max_atoms >= 0 and num_atoms > num_max_atoms: + raise MolGraphError( + f"Number of atoms in mol {num_atoms} exceeds num_max_atoms {num_max_atoms}" + ) + + def construct_atomic_number_array(self, mol, out_size=-1): + """Returns atomic numbers of atoms consisting a molecule. + + Args: + mol (rdkit.Chem.Mol): Input molecule. + out_size (int): The size of returned array. + If this option is negative, it does not take any effect. + Otherwise, it must be larger than the number of atoms + in the input molecules. In that case, the tail of + the array is padded with zeros. + + Returns: + numpy.ndarray: an array consisting of atomic numbers + of atoms in the molecule. + """ + atom_list = [a.GetAtomicNum() for a in mol.GetAtoms()] + n_atom = len(atom_list) + if out_size < 0: + return np.array(atom_list, dtype=np.int32) + elif out_size >= n_atom: + atom_array = np.zeros(out_size, dtype=np.int32) + atom_array[:n_atom] = np.array(atom_list, dtype=np.int32) + return atom_array + else: + raise ValueError( + f"`out_size` (={out_size}) must be negative or larger than or equal to " + f"the number of atoms in the input molecules (={n_atom})." + ) + + def construct_adj_matrix(self, mol, out_size=-1, self_connection=True): + """Returns the adjacent matrix of the given molecule. + + This function returns the adjacent matrix of the given molecule. + Contrary to the specification of + :func:`rdkit.Chem.rdmolops.GetAdjacencyMatrix`, + The diagonal entries of the returned matrix are all-one. + + Args: + mol (rdkit.Chem.Mol): Input molecule. + out_size (int): The size of the returned matrix. + If this option is negative, it does not take any effect. + Otherwise, it must be larger than the number of atoms + in the input molecules. In that case, the adjacent + matrix is expanded and zeros are padded to right + columns and bottom rows. + self_connection (bool): Add self connection or not. + If True, diagonal element of adjacency matrix is filled with 1. + + Returns: + adj_array (numpy.ndarray): The adjacent matrix of the input molecule. + It is 2-dimensional array with shape (atoms1, atoms2), where + atoms1 & atoms2 represent from and to of the edge respectively. + If ``out_size`` is non-negative, the returned + its size is equal to that value. Otherwise, + it is equal to the number of atoms in the the molecule. + """ + adj = rdmolops.GetAdjacencyMatrix(mol) + s0, s1 = tuple(adj.shape) + if s0 != s1: + raise ValueError( + f"The adjacent matrix of the input moleculehas an invalid shape: ({s0}, " + f"{s1}). It must be square." + ) + if self_connection: + adj = adj + np.eye(s0) + if out_size < 0: + adj_array = adj.astype(np.float32) + elif out_size >= s0: + adj_array = np.zeros((out_size, out_size), dtype=np.float32) + adj_array[:s0, :s1] = adj + else: + raise ValueError( + f"`out_size` (={out_size}) must be negative or larger than or equal to " + f"the number of atoms in the input molecules (={s0})." + ) + return adj_array + + def construct_discrete_edge_matrix(self, mol, out_size=-1): + """Returns the edge-type dependent adjacency matrix of the given molecule. + + Args: + mol (rdkit.Chem.Mol): Input molecule. + out_size (int): The size of the returned matrix. + If this option is negative, it does not take any effect. + Otherwise, it must be larger than the number of atoms + in the input molecules. In that case, the adjacent + matrix is expanded and zeros are padded to right + columns and bottom rows. + + Returns: + adj_array (numpy.ndarray): The adjacent matrix of the input molecule. + It is 3-dimensional array with shape (edge_type, atoms1, atoms2), + where edge_type represents the bond type, + atoms1 & atoms2 represent from and to of the edge respectively. + If ``out_size`` is non-negative, its size is equal to that value. + Otherwise, it is equal to the number of atoms in the the molecule. + """ + if mol is None: + raise MolGraphError("mol is None") + N = mol.GetNumAtoms() + if out_size < 0: + size = N + elif out_size >= N: + size = out_size + else: + raise ValueError( + f"out_size {out_size} is smaller than number of atoms in mol {N}" + ) + adjs = np.zeros((4, size, size), dtype=np.float32) + bond_type_to_channel = { + Chem.BondType.SINGLE: 0, + Chem.BondType.DOUBLE: 1, + Chem.BondType.TRIPLE: 2, + Chem.BondType.AROMATIC: 3, + } + for bond in mol.GetBonds(): + bond_type = bond.GetBondType() + ch = bond_type_to_channel[bond_type] + i = bond.GetBeginAtomIdx() + j = bond.GetEndAtomIdx() + adjs[ch, i, j] = 1.0 + adjs[ch, j, i] = 1.0 + return adjs + + +class MolGraphError(Exception): + pass + + +class MOlFLOWDataset(io.Dataset): + """Class for moflow qm9 and zinc250k Dataset of a tuple of datasets. + + It combines multiple datasets into one dataset. Each example is represented + by a tuple whose ``i``-th item corresponds to the i-th dataset. + And each ``i``-th dataset is expected to be an instance of numpy.ndarray. + + Args: + file_path (str): Data set path. + data_name (str): Data name, "qm9" or "zinc250k" + valid_idx (List[int, ...]): Data for validate + mode (str): "train" or "eval", output Data + input_keys (Tuple[str, ...]): Input keys, such as ("nodes","edges",). + label_keys (Tuple[str, ...]): labels (str or list or None) . + smiles_col (str): smiles column + weight_dict (Optional[Dict[str, Union[Callable, float]]]): Define the weight of each constraint variable. Defaults to None. + transform_fn: An optional function applied to an item bofre returning + """ + + # Whether support batch indexing for speeding up fetching process. + batch_index: bool = True + + def __init__( + self, + file_path: str, + data_name: str, + valid_idx: List[int, ...], + mode: str, + input_keys: Tuple[str, ...], + label_keys: Tuple[str, ...], + smiles_col: str, + weight_dict: Optional[Dict[str, float]] = None, + transform_fn: Optional[Callable] = None, + ): + super().__init__() + self.file_path = file_path + self.data_name = data_name + self.input_keys = input_keys + self.label_keys = label_keys + self.smiles_col = smiles_col + self.weight_dict = weight_dict + + if data_name == "qm9": + max_atoms = 9 + elif data_name == "zinc250k": + max_atoms = 38 + + self.molgraph = MolGraph(out_size=max_atoms, kekulize=True) + self.logger = logger + # read and deal data from file + inputs, labels = self.load_csv_file(file_path, data_name + ".csv") + train_idx = [t for t in range(len(inputs[0])) if t not in valid_idx] + self.train_idx = train_idx + # data train or test + if mode == "train": + inputs = [ + np.array(list(io.Subset(dataset=in_put, indices=train_idx))) + for in_put in inputs + ] + labels = np.array(list(io.Subset(dataset=labels, indices=train_idx))) + elif mode == "eval": + inputs = [ + np.array(list(io.Subset(dataset=in_put, indices=valid_idx))) + for in_put in inputs + ] + labels = np.array(list(io.Subset(dataset=labels, indices=valid_idx))) + + # fetch input data + self.input = {key: inputs[i] for i, key in enumerate(self.input_keys)} + # fetch label data + self.label = {"label": labels} + + self.logger.message( + f"Dataload finished. MODE {mode}, " + f"inputs {len(next(iter(self.input.values())))}, " + f"labelS {len(next(iter(self.label.values())))}" + ) + + self._length = len(next(iter(self.input.values()))) + self.transform = transform_fn + + def __getitem__(self, index: int): + input_item = {key: value[index] for key, value in self.input.items()} + label_item = {key: value[index] for key, value in self.label.items()} + + if self.transform: + input_item, label_item = self.transform_func(input_item, label_item) + + return (input_item, label_item, {}) + + def __len__(self): + return self._length + + def load_csv_file(self, path: str, name: str): + """Parse DataFrame using `MolGraph` and prepare a dataset instance + Labels are extracted from `labels` columns and input features are + extracted from smiles information in `smiles` column. + """ + file = os.path.join(path, name) + df = pd.read_csv(file, index_col=0) + all_nodes = [] + all_edges = [] + # inputs = [] + + total_count = df.shape[0] + fail_count = 0 + success_count = 0 + if isinstance(self.molgraph, MolGraph): + for smiles in tqdm(df[self.smiles_col], total=df.shape[0]): + try: + mol = Chem.MolFromSmiles(smiles) + if mol is None: + fail_count += 1 + continue + canonical_smiles, mol = self.molgraph.prepare_smiles_and_mol(mol) + nodes, edges = self.molgraph.get_input_features(mol) + + except MolGraphError as e: + fail_count += 1 + self.logger.warning(f"parse(), type: {type(e).__name__}, {e.args}") + continue + except Exception as e: + self.logger.warning(f"parse(), type: {type(e).__name__}, {e.args}") + fail_count += 1 + continue + # raw_data = misc.convert_to_dict(np.array([nodes, edges]), self.input_keys) + + all_nodes.append(nodes) + all_edges.append(edges) + # inputs.append(raw_data) + + success_count += 1 + + labels = np.array( + [*(df[label_col].values for label_col in self.label_keys)] + ).T + result = [np.array(all_nodes), np.array(all_edges)], labels + self.logger.message( + f"Preprocess finished. FAIL {fail_count}, " + f"SUCCESS {success_count}, TOTAL {total_count}" + ) + else: + raise NotImplementedError + + return result + + def transform_func(self, data_dict, label_dict): + items = [] + length = len(next(iter(data_dict.values()))) + for idx in range(length): + input_item = [value[idx] for key, value in data_dict.items()] + label_item = [value[idx] for key, value in label_dict.items()] + item = input_item + label_item + if self.transform: + item = self.transform(item) + items.append(item) + items = np.array(items, dtype=object).T + + data_dict = {key: np.stack(items[i], axis=0) for i, key in enumerate(data_dict)} + label_dict = {key: np.vstack(item[2]) for key in label_dict} + + return data_dict, label_dict diff --git a/examples/smc_reac/ppsci/data/dataset/mrms_dataset.py b/examples/smc_reac/ppsci/data/dataset/mrms_dataset.py new file mode 100644 index 0000000000..bee3337f7e --- /dev/null +++ b/examples/smc_reac/ppsci/data/dataset/mrms_dataset.py @@ -0,0 +1,251 @@ +# 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. + +from __future__ import annotations + +import glob +import os.path as osp +from datetime import datetime +from datetime import timedelta +from typing import Dict +from typing import List +from typing import Optional +from typing import Tuple + +try: + import h5py +except ModuleNotFoundError: + pass +import numpy as np +import paddle +from paddle import io +from paddle import vision + + +class MRMSDataset(io.Dataset): + """Class for MRMS dataset. MRMS day's data is stored in a .h5 file. Each file includes keys "date"/"time_interval"/"dataset". + + Args: + file_path (str): Dataset path. + input_keys (Tuple[str, ...]): Input keys, usually there is only one, such as ("input",). + label_keys (Tuple[str, ...]): Output keys, usually there is only one, such as ("output",). + weight_dict (Optional[Dict[str, float]]): Weight dictionary. Defaults to None. + date_period (Tuple[str,...], optional): Dates of data. Scale is [start_date, end_date] with format "%Y%m%d". Defaults to ("20230101","20230101"). + num_input_timestamps (int, optional): Number of timestamp of input. Defaults to 1. + num_label_timestamps (int, optional): Number of timestamp of label. Defaults to 1. + stride (int, optional): Stride of sampling data. Defaults to 1. + transforms (Optional[vision.Compose]): Composed transform functor(s). Defaults to None. + + Examples: + >>> import ppsci + >>> dataset = ppsci.data.dataset.MRMSDataset( + ... "file_path": "/path/to/MRMSDataset", + ... "input_keys": ("input",), + ... "label_keys": ("output",), + ... "date_period": ("20230101","20230131"), + ... "num_input_timestamps": 9, + ... "num_label_timestamps": 20, + ... "transforms": transform, + ... "stride": 1, + ... ) # doctest: +SKIP + """ + + # Whether support batch indexing for speeding up fetching process. + batch_index: bool = False + + def __init__( + self, + file_path: str, + input_keys: Tuple[str, ...], + label_keys: Tuple[str, ...], + weight_dict: Optional[Dict[str, float]] = None, + date_period: Tuple[str, ...] = ("20230101", "20230101"), + num_input_timestamps: int = 1, + num_label_timestamps: int = 1, + stride: int = 1, + transforms: Optional[vision.Compose] = None, + ): + super().__init__() + self.file_path = file_path + self.input_keys = input_keys + self.label_keys = label_keys + + self.weight_dict = {} if weight_dict is None else weight_dict + if weight_dict is not None: + self.weight_dict = {key: 1.0 for key in self.label_keys} + self.weight_dict.update(weight_dict) + + self.date_list = self._get_date_strs(date_period) + self.num_input_timestamps = num_input_timestamps + self.num_label_timestamps = num_label_timestamps + self.stride = stride + self.transforms = transforms + + self.files = self._read_data(file_path) + self.num_samples_per_day = self.files[0].shape[0] + self.num_samples = self.num_samples_per_day * len(self.date_list) + + def _get_date_strs(self, date_period: Tuple[str, ...]) -> List: + """Get a string list of all dates within given period. + + Args: + date_period (Tuple[str,...]): Dates of data. Scale is [start_date, end_date] with format "%Y%m%d". + """ + start_time = datetime.strptime(date_period[0], "%Y%m%d") + end_time = datetime.strptime(date_period[1], "%Y%m%d") + results = [] + current_time = start_time + while current_time <= end_time: + date_str = current_time.strftime("%Y%m%d") + results.append(date_str) + current_time += timedelta(days=1) + return results + + def _read_data(self, path: str): + if path.endswith(".h5"): + paths = [path] + else: + paths = [ + _path + for _path in glob.glob(osp.join(path, "*.h5")) + if _path.split(".h5")[0].split("_")[-1] in self.date_list + ] + assert len(paths) == len( + self.date_list + ), f"Data of {len(self.date_list)} days wanted but only {len(paths)} days be found" + paths.sort() + + files = [h5py.File(_path, "r")["dataset"] for _path in paths] + return files + + def __len__(self): + return ( + self.num_samples // self.stride + - self.num_input_timestamps + - self.num_label_timestamps + + 1 + ) + + def __getitem__(self, global_idx): + global_idx *= self.stride + _samples = np.empty( + ( + self.num_input_timestamps + self.num_label_timestamps, + *self.files[0].shape[1:], + ), + dtype=paddle.get_default_dtype(), + ) + for idx in range(self.num_input_timestamps + self.num_label_timestamps): + sample_idx = global_idx + idx * self.stride + day_idx = sample_idx // self.num_samples_per_day + local_idx = sample_idx % self.num_samples_per_day + _samples[idx] = self.files[day_idx][local_idx] + + input_item = {self.input_keys[0]: _samples[: self.num_input_timestamps]} + label_item = {self.label_keys[0]: _samples[self.num_input_timestamps :]} + + weight_shape = [1] * len(next(iter(label_item.values())).shape) + weight_item = { + key: np.full(weight_shape, value, paddle.get_default_dtype()) + for key, value in self.weight_dict.items() + } + + if self.transforms is not None: + input_item, label_item, weight_item = self.transforms( + input_item, label_item, weight_item + ) + + return input_item, label_item, weight_item + + +class MRMSSampledDataset(io.Dataset): + """Class for MRMS sampled dataset. MRMS one sample's data is stored in a .h5 file. Each file includes keys "date"/"time_interval"/"dataset". + The class just return data by input_item and values of label_item are empty for all label_keys. + + Args: + file_path (str): Dataset path. + input_keys (Tuple[str, ...]): Input keys, such as ("input",). + label_keys (Tuple[str, ...]): Output keys, such as ("output",). + weight_dict (Optional[Dict[str, float]]): Weight dictionary. Defaults to None. + num_total_timestamps (int, optional): Number of timestamp of input+label. Defaults to 1. + transforms (Optional[vision.Compose]): Composed transform functor(s). Defaults to None. + + Examples: + >>> import ppsci + >>> dataset = ppsci.data.dataset.MRMSSampledDataset( + ... "file_path": "/path/to/MRMSSampledDataset", + ... "input_keys": ("input",), + ... "label_keys": ("output",), + ... "num_total_timestamps": 29, + ... ) # doctest: +SKIP + >>> # get the length of the dataset + >>> dataset_size = len(dataset) # doctest: +SKIP + >>> # get the first sample of the data + >>> first_sample = dataset[0] # doctest: +SKIP + >>> print("First sample:", first_sample) # doctest: +SKIP + """ + + def __init__( + self, + file_path: str, + input_keys: Tuple[str, ...], + label_keys: Tuple[str, ...], + weight_dict: Optional[Dict[str, float]] = None, + num_total_timestamps: int = 1, + transforms: Optional[vision.Compose] = None, + ): + super().__init__() + self.file_path = file_path + self.input_keys = input_keys + self.label_keys = label_keys + + self.weight_dict = {} if weight_dict is None else weight_dict + if weight_dict is not None: + self.weight_dict = {key: 1.0 for key in self.label_keys} + self.weight_dict.update(weight_dict) + + self.num_total_timestamps = num_total_timestamps + self.transforms = transforms + + self.files = self._read_data(file_path) + self.num_samples = len(self.files) + + def _read_data(self, path: str): + paths = glob.glob(osp.join(path, "*.h5")) + paths.sort() + files = [h5py.File(_path, "r")["dataset"] for _path in paths] + return files + + def __len__(self): + return self.num_samples - self.num_total_timestamps + 1 + + def __getitem__(self, global_idx): + _samples = [] + for idx in range(global_idx, global_idx + self.num_total_timestamps): + _samples.append(np.expand_dims(self.files[idx], axis=0)) + + input_item = { + self.input_keys[0]: np.concatenate(_samples, axis=0).astype( + paddle.get_default_dtype() + ) + } + label_item = {} + weight_item = {} + + if self.transforms is not None: + input_item, label_item, weight_item = self.transforms( + input_item, label_item, weight_item + ) + + return input_item, label_item, weight_item diff --git a/examples/smc_reac/ppsci/data/dataset/npz_dataset.py b/examples/smc_reac/ppsci/data/dataset/npz_dataset.py new file mode 100644 index 0000000000..76d737d021 --- /dev/null +++ b/examples/smc_reac/ppsci/data/dataset/npz_dataset.py @@ -0,0 +1,279 @@ +# 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. + +from __future__ import annotations + +from typing import Callable +from typing import Dict +from typing import Optional +from typing import Tuple +from typing import Union + +import numpy as np +import paddle +from paddle import io +from paddle import vision + +from ppsci.utils import misc +from ppsci.utils import reader + + +class NPZDataset(io.Dataset): + """Dataset class for .npz file. + + Args: + file_path (str): Npz file path. + input_keys (Tuple[str, ...]): List of input keys. + label_keys (Tuple[str, ...], optional): List of label keys. Defaults to (). + alias_dict (Optional[Dict[str, str]]): Dict of alias(es) for input and label keys. + i.e. {inner_key: outer_key}. Defaults to None. + weight_dict (Optional[Dict[str, Union[Callable, float]]]): Define the weight of + each constraint variable. Defaults to None. + timestamps (Optional[Tuple[float, ...]]): The number of repetitions of the data + in the time dimension. Defaults to None. + transforms (Optional[vision.Compose]): Compose object contains sample wise + transform(s). Defaults to None. + + Examples: + >>> import ppsci + >>> dataset = ppsci.data.dataset.NPZDataset( + ... "/path/to/file.npz" + ... ("x",), + ... ("u",), + ... ) # doctest: +SKIP + """ + + # Whether support batch indexing for speeding up fetching process. + batch_index: bool = True + + def __init__( + self, + file_path: str, + input_keys: Tuple[str, ...], + label_keys: Tuple[str, ...] = (), + alias_dict: Optional[Dict[str, str]] = None, + weight_dict: Optional[Dict[str, Union[Callable, float]]] = None, + timestamps: Optional[Tuple[float, ...]] = None, + transforms: Optional[vision.Compose] = None, + ): + super().__init__() + self.input_keys = input_keys + self.label_keys = label_keys + + # read raw data from file + raw_data = reader.load_npz_file( + file_path, + input_keys + label_keys, + alias_dict, + ) + # filter raw data by given timestamps if specified + if timestamps is not None: + if "t" in raw_data: + # filter data according to given timestamps + raw_time_array = raw_data["t"] + mask = [] + for ti in timestamps: + mask.append(np.nonzero(np.isclose(raw_time_array, ti).flatten())[0]) + raw_data = misc.convert_to_array( + raw_data, self.input_keys + self.label_keys + ) + mask = np.concatenate(mask, 0) + raw_data = raw_data[mask] + raw_data = misc.convert_to_dict( + raw_data, self.input_keys + self.label_keys + ) + else: + # repeat data according to given timestamps + raw_data = misc.convert_to_array( + raw_data, self.input_keys + self.label_keys + ) + raw_data = misc.combine_array_with_time(raw_data, timestamps) + self.input_keys = ("t",) + tuple(self.input_keys) + raw_data = misc.convert_to_dict( + raw_data, self.input_keys + self.label_keys + ) + + # fetch input data + self.input = { + key: value for key, value in raw_data.items() if key in self.input_keys + } + # fetch label data + self.label = { + key: value for key, value in raw_data.items() if key in self.label_keys + } + + # prepare weights + self.weight = {} + if weight_dict is not None: + for key, value in weight_dict.items(): + if isinstance(value, (int, float)): + self.weight[key] = np.full_like( + next(iter(self.label.values())), value + ) + elif callable(value): + func = value + self.weight[key] = func(self.input) + if isinstance(self.weight[key], (int, float)): + self.weight[key] = np.full_like( + next(iter(self.label.values())), self.weight[key] + ) + else: + raise NotImplementedError(f"type of {type(value)} is invalid yet.") + + self.transforms = transforms + self._len = len(next(iter(self.input.values()))) + + def __getitem__(self, idx): + input_item = {key: value[idx] for key, value in self.input.items()} + label_item = {key: value[idx] for key, value in self.label.items()} + weight_item = {key: value[idx] for key, value in self.weight.items()} + + if self.transforms is not None: + input_item, label_item, weight_item = self.transforms( + input_item, label_item, weight_item + ) + + return (input_item, label_item, weight_item) + + def __len__(self): + return self._len + + +class IterableNPZDataset(io.IterableDataset): + """IterableNPZDataset for full-data loading. + + Args: + file_path (str): Npz file path. + input_keys (Tuple[str, ...]): List of input keys. + label_keys (Tuple[str, ...], optional): List of label keys. Defaults to (). + alias_dict (Optional[Dict[str, str]]): Dict of alias(es) for input and label keys. + i.e. {inner_key: outer_key}. Defaults to None. + weight_dict (Optional[Dict[str, Union[Callable, float]]]): Define the weight of + each constraint variable. Defaults to None. + timestamps (Optional[Tuple[float, ...]]): The number of repetitions of the data + in the time dimension. Defaults to None. + transforms (Optional[vision.Compose]): Compose object contains sample wise + transform(s). Defaults to None. + + Examples: + >>> import ppsci + >>> dataset = ppsci.data.dataset.IterableNPZDataset( + ... "/path/to/file.npz" + ... ("x",), + ... ("u",), + ... ) # doctest: +SKIP + """ + + # Whether support batch indexing for speeding up fetching process. + batch_index: bool = False + + def __init__( + self, + file_path: str, + input_keys: Tuple[str, ...], + label_keys: Tuple[str, ...] = (), + alias_dict: Optional[Dict[str, str]] = None, + weight_dict: Optional[Dict[str, Union[Callable, float]]] = None, + timestamps: Optional[Tuple[float, ...]] = None, + transforms: Optional[vision.Compose] = None, + ): + super().__init__() + self.input_keys = input_keys + self.label_keys = label_keys + + # read raw data from file + raw_data = reader.load_npz_file( + file_path, + input_keys + label_keys, + alias_dict, + ) + # filter raw data by given timestamps if specified + if timestamps is not None: + if "t" in raw_data: + # filter data according to given timestamps + raw_time_array = raw_data["t"] + mask = [] + for ti in timestamps: + mask.append(np.nonzero(np.isclose(raw_time_array, ti).flatten())[0]) + raw_data = misc.convert_to_array( + raw_data, self.input_keys + self.label_keys + ) + mask = np.concatenate(mask, 0) + raw_data = raw_data[mask] + raw_data = misc.convert_to_dict( + raw_data, self.input_keys + self.label_keys + ) + else: + # repeat data according to given timestamps + raw_data = misc.convert_to_array( + raw_data, self.input_keys + self.label_keys + ) + raw_data = misc.combine_array_with_time(raw_data, timestamps) + self.input_keys = ("t",) + tuple(self.input_keys) + raw_data = misc.convert_to_dict( + raw_data, self.input_keys + self.label_keys + ) + + # fetch input data + self.input = { + key: value for key, value in raw_data.items() if key in self.input_keys + } + # fetch label data + self.label = { + key: value for key, value in raw_data.items() if key in self.label_keys + } + + # prepare weights + self.weight = {} + if weight_dict is not None: + for key, value in weight_dict.items(): + if isinstance(value, (int, float)): + self.weight[key] = np.full_like( + next(iter(self.label.values())), value + ) + elif callable(value): + func = value + self.weight[key] = func(self.input) + if isinstance(self.weight[key], (int, float)): + self.weight[key] = np.full_like( + next(iter(self.label.values())), self.weight[key] + ) + else: + raise NotImplementedError(f"type of {type(value)} is invalid yet.") + + self.input = {key: paddle.to_tensor(value) for key, value in self.input.items()} + self.label = {key: paddle.to_tensor(value) for key, value in self.label.items()} + self.weight = { + key: paddle.to_tensor(value) for key, value in self.weight.items() + } + + self.transforms = transforms + self._len = len(next(iter(self.input.values()))) + + @property + def num_samples(self): + """Number of samples within current dataset.""" + return self._len + + def __iter__(self): + if callable(self.transforms): + input_, label_, weight_ = self.transforms( + self.input, self.label, self.weight + ) + yield input_, label_, weight_ + else: + yield self.input, self.label, self.weight + + def __len__(self): + return 1 diff --git a/examples/smc_reac/ppsci/data/dataset/pems_dataset.py b/examples/smc_reac/ppsci/data/dataset/pems_dataset.py new file mode 100644 index 0000000000..1e3dde0a3d --- /dev/null +++ b/examples/smc_reac/ppsci/data/dataset/pems_dataset.py @@ -0,0 +1,151 @@ +import os +from typing import Dict +from typing import Optional +from typing import Tuple + +import numpy as np +import pandas as pd +from paddle.io import Dataset +from paddle.vision.transforms import Compose + + +class StandardScaler: + def __init__(self, mean, std): + self.mean = mean + self.std = std + + def transform(self, data): + return (data - self.mean) / self.std + + def inverse_transform(self, data): + return (data * self.std) + self.mean + + +def add_window_horizon(data, in_step=12, out_step=12): + length = len(data) + end_index = length - out_step - in_step + X = [] + Y = [] + for i in range(end_index + 1): + X.append(data[i : i + in_step]) + Y.append(data[i + in_step : i + in_step + out_step]) + return X, Y + + +def get_edge_index(file_path, bi=True, reduce="mean"): + TYPE_DICT = {0: np.int64, 1: np.int64, 2: np.float32} + df = pd.read_csv( + os.path.join(file_path, "dist.csv"), + skiprows=1, + header=None, + sep=",", + dtype=TYPE_DICT, + ) + + edge_index = df.loc[:, [0, 1]].values.T + edge_attr = df.loc[:, 2].values + + if bi: + re_edge_index = np.concatenate((edge_index[1:, :], edge_index[:1, :]), axis=0) + edge_index = np.concatenate((edge_index, re_edge_index), axis=-1) + edge_attr = np.concatenate((edge_attr, edge_attr), axis=0) + + num = np.max(edge_index) + 1 + adj = np.zeros((num, num), dtype=np.float32) + + if reduce == "sum": + adj[edge_index[0], edge_index[1]] = 1.0 + elif reduce == "mean": + adj[edge_index[0], edge_index[1]] = 1.0 + adj = adj / adj.sum(axis=-1) + else: + raise ValueError + + return edge_index, edge_attr, adj + + +class PEMSDataset(Dataset): + """Dataset class for PEMSD4 and PEMSD8 dataset. + + Args: + file_path (str): Dataset root path. + split (str): Dataset split label. + input_keys (Tuple[str, ...]): A tuple of input keys. + label_keys (Tuple[str, ...]): A tuple of label keys. + weight_dict (Optional[Dict[str, float]]): Define the weight of each constraint variable. Defaults to None. + transforms (Optional[Compose]): Compose object contains sample wise transform(s). Defaults to None. + norm_input (bool): Whether to normalize the input. Defaults to True. + norm_label (bool): Whether to normalize the output. Defaults to False. + input_len (int): The input timesteps. Defaults to 12. + label_len (int): The output timesteps. Defaults to 12. + + Examples: + >>> import ppsci + >>> dataset = ppsci.data.dataset.PEMSDataset( + ... "./Data/PEMSD4", + ... "train", + ... ("input",), + ... ("label",), + ... ) # doctest: +SKIP + """ + + def __init__( + self, + file_path: str, + split: str, + input_keys: Tuple[str, ...], + label_keys: Tuple[str, ...], + weight_dict: Optional[Dict[str, float]] = None, + transforms: Optional[Compose] = None, + norm_input: bool = True, + norm_label: bool = False, + input_len: int = 12, + label_len: int = 12, + ): + super().__init__() + + self.input_keys = input_keys + self.label_keys = label_keys + self.weight_dict = weight_dict + + self.transforms = transforms + self.norm_input = norm_input + self.norm_label = norm_label + + data = np.load(os.path.join(file_path, f"{split}.npy")).astype(np.float32) + + self.mean = np.load(os.path.join(file_path, "mean.npy")).astype(np.float32) + self.std = np.load(os.path.join(file_path, "std.npy")).astype(np.float32) + self.scaler = StandardScaler(self.mean, self.std) + + X, Y = add_window_horizon(data, input_len, label_len) + if norm_input: + X = self.scaler.transform(X) + if norm_label: + Y = self.scaler.transform(Y) + + self._len = X.shape[0] + + self.input = {input_keys[0]: X} + self.label = {label_keys[0]: Y} + + if weight_dict is not None: + self.weight_dict = {key: np.array(1.0) for key in self.label_keys} + self.weight_dict.update(weight_dict) + else: + self.weight = {} + + def __getitem__(self, idx): + input_item = {key: value[idx] for key, value in self.input.items()} + label_item = {key: value[idx] for key, value in self.label.items()} + weight_item = {key: value[idx] for key, value in self.weight.items()} + + if self.transforms is not None: + input_item, label_item, weight_item = self.transforms( + input_item, label_item, weight_item + ) + + return (input_item, label_item, weight_item) + + def __len__(self): + return self._len diff --git a/examples/smc_reac/ppsci/data/dataset/radar_dataset.py b/examples/smc_reac/ppsci/data/dataset/radar_dataset.py new file mode 100644 index 0000000000..e484558455 --- /dev/null +++ b/examples/smc_reac/ppsci/data/dataset/radar_dataset.py @@ -0,0 +1,146 @@ +# 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. + +from __future__ import annotations + +import os +from typing import Dict +from typing import Optional +from typing import Tuple + +try: + import cv2 +except ModuleNotFoundError: + pass + +import importlib + +import numpy as np +import paddle +from paddle import io + + +class RadarDataset(io.Dataset): + """Class for Radar dataset. + + Args: + input_keys (Tuple[str, ...]): Input keys, such as ("input",). + label_keys (Tuple[str, ...]): Output keys, such as ("output",). + image_width (int): Image width. + image_height (int): Image height. + total_length (int): Total length. + dataset_path (str): Dataset path. + data_type (str): Input and output data type. Defaults to paddle.get_default_dtype(). + weight_dict (Optional[Dict[str, float]]): Weight dictionary. Defaults to None. + + Examples: + >>> import ppsci + >>> dataset = ppsci.data.dataset.RadarDataset( + ... "input_keys": ("input",), + ... "label_keys": ("output",), + ... "image_width": 512, + ... "image_height": 512, + ... "total_length": 29, + ... "dataset_path": "datasets/mrms/figure", + ... "data_type": paddle.get_default_dtype(), + ... ) # doctest: +SKIP + """ + + # Whether support batch indexing for speeding up fetching process. + batch_index: bool = False + + def __init__( + self, + input_keys: Tuple[str, ...], + label_keys: Tuple[str, ...], + image_width: int, + image_height: int, + total_length: int, + dataset_path: str, + data_type: str = paddle.get_default_dtype(), + weight_dict: Optional[Dict[str, float]] = None, + ): + super().__init__() + if importlib.util.find_spec("cv2") is None: + raise ModuleNotFoundError( + "To use RadarDataset, please install 'opencv-python' with: `pip install " + "opencv-python` first." + ) + self.input_keys = input_keys + self.label_keys = label_keys + self.img_width = image_width + self.img_height = image_height + self.length = total_length + self.dataset_path = dataset_path + self.data_type = data_type + + self.weight_dict = {} if weight_dict is None else weight_dict + if weight_dict is not None: + self.weight_dict = {key: 1.0 for key in self.label_keys} + self.weight_dict.update(weight_dict) + + self.case_list = [] + name_list = os.listdir(self.dataset_path) + name_list.sort() + for name in name_list: + case = [] + for i in range(29): + case.append( + self.dataset_path + + "/" + + name + + "/" + + name + + "-" + + str(i).zfill(2) + + ".png" + ) + self.case_list.append(case) + + def _load(self, index): + data = [] + for img_path in self.case_list[index]: + img = cv2.imread(img_path, 2) + data.append(np.expand_dims(img, axis=0)) + data = np.concatenate(data, axis=0).astype(self.data_type) / 10.0 - 3.0 + assert data.shape[1] <= 1024 and data.shape[2] <= 1024 + return data + + def __getitem__(self, index): + data = self._load(index)[-self.length :].copy() + mask = np.ones_like(data) + mask[data < 0] = 0 + data[data < 0] = 0 + data = np.clip(data, 0, 128) + vid = np.zeros( + (self.length, self.img_height, self.img_width, 2), dtype=self.data_type + ) + vid[..., 0] = data + vid[..., 1] = mask + + input_item = {self.input_keys[0]: vid} + label_item = {} + weight_item = {} + for key in self.label_keys: + label_item[key] = np.asarray([], paddle.get_default_dtype()) + if len(label_item) > 0: + weight_shape = [1] * len(next(iter(label_item.values())).shape) + weight_item = { + key: np.full(weight_shape, value, paddle.get_default_dtype()) + for key, value in self.weight_dict.items() + } + return input_item, label_item, weight_item + + def __len__(self): + return len(self.case_list) diff --git a/examples/smc_reac/ppsci/data/dataset/sevir_dataset.py b/examples/smc_reac/ppsci/data/dataset/sevir_dataset.py new file mode 100644 index 0000000000..42ae274c2b --- /dev/null +++ b/examples/smc_reac/ppsci/data/dataset/sevir_dataset.py @@ -0,0 +1,814 @@ +import datetime +import os +from copy import deepcopy +from typing import Dict +from typing import Optional +from typing import Sequence +from typing import Tuple +from typing import Union + +try: + import h5py +except ModuleNotFoundError: + pass +import numpy as np +import paddle +import paddle.nn.functional as F +import pandas as pd +from paddle import io + +# SEVIR Dataset constants +SEVIR_DATA_TYPES = ["vis", "ir069", "ir107", "vil", "lght"] +SEVIR_RAW_DTYPES = { + "vis": np.int16, + "ir069": np.int16, + "ir107": np.int16, + "vil": np.uint8, + "lght": np.int16, +} +LIGHTING_FRAME_TIMES = np.arange(-120.0, 125.0, 5) * 60 +SEVIR_DATA_SHAPE = { + "lght": (48, 48), +} +PREPROCESS_SCALE_SEVIR = { + "vis": 1, # Not utilized in original paper + "ir069": 1 / 1174.68, + "ir107": 1 / 2562.43, + "vil": 1 / 47.54, + "lght": 1 / 0.60517, +} +PREPROCESS_OFFSET_SEVIR = { + "vis": 0, # Not utilized in original paper + "ir069": 3683.58, + "ir107": 1552.80, + "vil": -33.44, + "lght": -0.02990, +} +PREPROCESS_SCALE_01 = { + "vis": 1, + "ir069": 1, + "ir107": 1, + "vil": 1 / 255, # currently the only one implemented + "lght": 1, +} +PREPROCESS_OFFSET_01 = { + "vis": 0, + "ir069": 0, + "ir107": 0, + "vil": 0, # currently the only one implemented + "lght": 0, +} + + +def change_layout_np(data, in_layout="NHWT", out_layout="NHWT", ret_contiguous=False): + # first convert to 'NHWT' + if in_layout == "NHWT": + pass + elif in_layout == "NTHW": + data = np.transpose(data, axes=(0, 2, 3, 1)) + elif in_layout == "NWHT": + data = np.transpose(data, axes=(0, 2, 1, 3)) + elif in_layout == "NTCHW": + data = data[:, :, 0, :, :] + data = np.transpose(data, axes=(0, 2, 3, 1)) + elif in_layout == "NTHWC": + data = data[:, :, :, :, 0] + data = np.transpose(data, axes=(0, 2, 3, 1)) + elif in_layout == "NTWHC": + data = data[:, :, :, :, 0] + data = np.transpose(data, axes=(0, 3, 2, 1)) + elif in_layout == "TNHW": + data = np.transpose(data, axes=(1, 2, 3, 0)) + elif in_layout == "TNCHW": + data = data[:, :, 0, :, :] + data = np.transpose(data, axes=(1, 2, 3, 0)) + else: + raise NotImplementedError(f"{in_layout} is invalid.") + + if out_layout == "NHWT": + pass + elif out_layout == "NTHW": + data = np.transpose(data, axes=(0, 3, 1, 2)) + elif out_layout == "NWHT": + data = np.transpose(data, axes=(0, 2, 1, 3)) + elif out_layout == "NTCHW": + data = np.transpose(data, axes=(0, 3, 1, 2)) + data = np.expand_dims(data, axis=2) + elif out_layout == "NTHWC": + data = np.transpose(data, axes=(0, 3, 1, 2)) + data = np.expand_dims(data, axis=-1) + elif out_layout == "NTWHC": + data = np.transpose(data, axes=(0, 3, 2, 1)) + data = np.expand_dims(data, axis=-1) + elif out_layout == "TNHW": + data = np.transpose(data, axes=(3, 0, 1, 2)) + elif out_layout == "TNCHW": + data = np.transpose(data, axes=(3, 0, 1, 2)) + data = np.expand_dims(data, axis=2) + else: + raise NotImplementedError(f"{out_layout} is invalid.") + if ret_contiguous: + data = data.ascontiguousarray() + return data + + +def change_layout_paddle( + data, in_layout="NHWT", out_layout="NHWT", ret_contiguous=False +): + # first convert to 'NHWT' + if in_layout == "NHWT": + pass + elif in_layout == "NTHW": + data = data.transpose(perm=[0, 2, 3, 1]) + elif in_layout == "NTCHW": + data = data[:, :, 0, :, :] + data = data.transpose(perm=[0, 2, 3, 1]) + elif in_layout == "NTHWC": + data = data[:, :, :, :, 0] + data = data.transpose(perm=[0, 2, 3, 1]) + elif in_layout == "TNHW": + data = data.transpose(perm=[1, 2, 3, 0]) + elif in_layout == "TNCHW": + data = data[:, :, 0, :, :] + data = data.transpose(perm=[1, 2, 3, 0]) + else: + raise NotImplementedError(f"{in_layout} is invalid.") + + if out_layout == "NHWT": + pass + elif out_layout == "NTHW": + data = data.transpose(perm=[0, 3, 1, 2]) + elif out_layout == "NTCHW": + data = data.transpose(perm=[0, 3, 1, 2]) + data = paddle.unsqueeze(data, axis=2) + elif out_layout == "NTHWC": + data = data.transpose(perm=[0, 3, 1, 2]) + data = paddle.unsqueeze(data, axis=-1) + elif out_layout == "TNHW": + data = data.transpose(perm=[3, 0, 1, 2]) + elif out_layout == "TNCHW": + data = data.transpose(perm=[3, 0, 1, 2]) + data = paddle.unsqueeze(data, axis=2) + else: + raise NotImplementedError(f"{out_layout} is invalid.") + return data + + +def path_splitall(path): + allparts = [] + while 1: + parts = os.path.split(path) + if parts[0] == path: # sentinel for absolute paths + allparts.insert(0, parts[0]) + break + elif parts[1] == path: # sentinel for relative paths + allparts.insert(0, parts[1]) + break + else: + path = parts[0] + allparts.insert(0, parts[1]) + return allparts + + +class SEVIRDataset(io.Dataset): + """The Storm EVent ImagRy dataset. + + Args: + input_keys (Tuple[str, ...]): Name of input keys, such as ("input",). + label_keys (Tuple[str, ...]): Name of label keys, such as ("output",). + data_dir (str): The path of the dataset. + weight_dict (Optional[Dict[str, Union[Callable, float]]]): Define the weight of each constraint variable. Defaults to None. + data_types (Sequence[str], optional): A subset of SEVIR_DATA_TYPES. Defaults to [ "vil", ]. + seq_len (int, optional): The length of the data sequences. Should be smaller than the max length raw_seq_len. Defaults to 49. + raw_seq_len (int, optional): The length of the raw data sequences. Defaults to 49. + sample_mode (str, optional): The mode of sampling, eg.'random' or 'sequent'. Defaults to "sequent". + stride (int, optional): Useful when sample_mode == 'sequent' + stride must not be smaller than out_len to prevent data leakage in testing. Defaults to 12. + batch_size (int, optional): The batch size. Defaults to 1. + layout (str, optional): Consists of batch_size 'N', seq_len 'T', channel 'C', height 'H', width 'W' + The layout of sampled data. Raw data layout is 'NHWT'. + valid layout: 'NHWT', 'NTHW', 'NTCHW', 'TNHW', 'TNCHW'. Defaults to "NHWT". + in_len (int, optional): The length of input data. Defaults to 13. + out_len (int, optional): The length of output data. Defaults to 12. + num_shard (int, optional): Split the whole dataset into num_shard parts for distributed training. Defaults to 1. + rank (int, optional): Rank of the current process within num_shard. Defaults to 0. + split_mode (str, optional): If 'ceil', all `num_shard` dataloaders have the same length = ceil(total_len / num_shard). + Different dataloaders may have some duplicated data batches, if the total size of datasets is not divided by num_shard. + if 'floor', all `num_shard` dataloaders have the same length = floor(total_len / num_shard). + The last several data batches may be wasted, if the total size of datasets is not divided by num_shard. + if 'uneven', the last datasets has larger length when the total length is not divided by num_shard. + The uneven split leads to synchronization error in dist.all_reduce() or dist.barrier(). + See related issue: https://github.com/pytorch/pytorch/issues/33148 + Notice: this also affects the behavior of `self.use_up`. Defaults to "uneven". + start_date (datetime.datetime, optional): Start time of SEVIR samples to generate. Defaults to None. + end_date (datetime.datetime, optional): End time of SEVIR samples to generate. Defaults to None. + datetime_filter (function, optional): Mask function applied to time_utc column of catalog (return true to keep the row). + Pass function of the form lambda t : COND(t) + Example: lambda t: np.logical_and(t.dt.hour>=13,t.dt.hour<=21) # Generate only day-time events. Defaults to None. + catalog_filter (function, optional): Function or None or 'default' + Mask function applied to entire catalog dataframe (return true to keep row). + Pass function of the form lambda catalog: COND(catalog) + Example: lambda c: [s[0]=='S' for s in c.id] # Generate only the 'S' events + shuffle (bool, optional): If True, data samples are shuffled before each epoch. Defaults to False. + shuffle_seed (int, optional): Seed to use for shuffling. Defaults to 1. + output_type (np.dtype, optional): The type of generated tensors. Defaults to np.float32. + preprocess (bool, optional): If True, self.preprocess_data_dict(data_dict) is called before each sample generated. Defaults to True. + rescale_method (str, optional): The method of rescale. Defaults to "01". + downsample_dict (Dict[str, Sequence[int]], optional): Downsample_dict.keys() == data_types. + downsample_dict[key] is a Sequence of (t_factor, h_factor, w_factor),representing the downsampling factors of all dimensions. Defaults to None. + verbose (bool, optional): Verbose when opening raw data files. Defaults to False. + training (str, optional): Training pathse. Defaults to "train". + """ + + # Whether support batch indexing for speeding up fetching process. + batch_index: bool = False + + def __init__( + self, + input_keys: Tuple[str, ...], + label_keys: Tuple[str, ...], + data_dir: str, + weight_dict: Optional[Dict[str, float]] = None, + data_types: Sequence[str] = [ + "vil", + ], + seq_len: int = 49, + raw_seq_len: int = 49, + sample_mode: str = "sequent", + stride: int = 12, + batch_size: int = 1, + layout: str = "NHWT", + in_len: int = 13, + out_len: int = 12, + num_shard: int = 1, + rank: int = 0, + split_mode: str = "uneven", + start_date: datetime.datetime = None, + end_date: datetime.datetime = None, + datetime_filter=None, + catalog_filter="default", + shuffle: bool = False, + shuffle_seed: int = 1, + output_type=np.float32, + preprocess: bool = True, + rescale_method: str = "01", + downsample_dict: Dict[str, Sequence[int]] = None, + verbose: bool = False, + training="train", + ): + super(SEVIRDataset, self).__init__() + self.input_keys = input_keys + self.label_keys = label_keys + self.data_dir = data_dir + self.weight_dict = {} if weight_dict is None else weight_dict + if weight_dict is not None: + self.weight_dict = {key: 1.0 for key in self.label_keys} + self.weight_dict.update(weight_dict) + + # sevir + SEVIR_ROOT_DIR = os.path.join(self.data_dir, "sevir") + sevir_catalog = os.path.join(SEVIR_ROOT_DIR, "CATALOG.csv") + sevir_data_dir = os.path.join(SEVIR_ROOT_DIR, "data") + # sevir-lr + # SEVIR_ROOT_DIR = os.path.join(self.data_dir, "sevir_lr") + # SEVIR_CATALOG = os.path.join(SEVIR_ROOT_DIR, "CATALOG.csv") + # SEVIR_DATA_DIR = os.path.join(SEVIR_ROOT_DIR, "data") + + if data_types is None: + data_types = SEVIR_DATA_TYPES + else: + assert set(data_types).issubset(SEVIR_DATA_TYPES) + + # configs which should not be modified + self._dtypes = SEVIR_RAW_DTYPES + self.lght_frame_times = LIGHTING_FRAME_TIMES + self.data_shape = SEVIR_DATA_SHAPE + + self.raw_seq_len = raw_seq_len + self.seq_len = seq_len + + if seq_len > raw_seq_len: + raise ValueError("seq_len must be small than raw_seq_len") + + if sample_mode not in ["random", "sequent"]: + raise ValueError("sample_mode must be 'random' or 'sequent'.") + + self.sample_mode = sample_mode + self.stride = stride + self.batch_size = batch_size + valid_layout = ("NHWT", "NTHW", "NTCHW", "NTHWC", "TNHW", "TNCHW") + if layout not in valid_layout: + raise ValueError( + f"Invalid layout = {layout}! Must be one of {valid_layout}." + ) + self.layout = layout + self.in_len = in_len + self.out_len = out_len + + self.num_shard = num_shard + self.rank = rank + valid_split_mode = ("ceil", "floor", "uneven") + if split_mode not in valid_split_mode: + raise ValueError( + f"Invalid split_mode: {split_mode}! Must be one of {valid_split_mode}." + ) + self.split_mode = split_mode + self._samples = None + self._hdf_files = {} + self.data_types = data_types + if isinstance(sevir_catalog, str): + self.catalog = pd.read_csv( + sevir_catalog, parse_dates=["time_utc"], low_memory=False + ) + else: + self.catalog = sevir_catalog + self.sevir_data_dir = sevir_data_dir + self.datetime_filter = datetime_filter + self.catalog_filter = catalog_filter + self.start_date = start_date + self.end_date = end_date + # train val test split + self.start_date = ( + datetime.datetime(*start_date) if start_date is not None else None + ) + self.end_date = datetime.datetime(*end_date) if end_date is not None else None + + self.shuffle = shuffle + self.shuffle_seed = int(shuffle_seed) + self.output_type = output_type + self.preprocess = preprocess + self.downsample_dict = downsample_dict + self.rescale_method = rescale_method + self.verbose = verbose + + if self.start_date is not None: + self.catalog = self.catalog[self.catalog.time_utc > self.start_date] + if self.end_date is not None: + self.catalog = self.catalog[self.catalog.time_utc <= self.end_date] + if self.datetime_filter: + self.catalog = self.catalog[self.datetime_filter(self.catalog.time_utc)] + + if self.catalog_filter is not None: + if self.catalog_filter == "default": + self.catalog_filter = lambda c: c.pct_missing == 0 + self.catalog = self.catalog[self.catalog_filter(self.catalog)] + + self._compute_samples() + self._open_files(verbose=self.verbose) + + def _compute_samples(self): + """ + Computes the list of samples in catalog to be used. This sets self._samples + """ + # locate all events containing colocated data_types + imgt = self.data_types + imgts = set(imgt) + filtcat = self.catalog[ + np.logical_or.reduce([self.catalog.img_type == i for i in imgt]) + ] + # remove rows missing one or more requested img_types + filtcat = filtcat.groupby("id").filter( + lambda x: imgts.issubset(set(x["img_type"])) + ) + # If there are repeated IDs, remove them (this is a bug in SEVIR) + # TODO: is it necessary to keep one of them instead of deleting them all + filtcat = filtcat.groupby("id").filter(lambda x: x.shape[0] == len(imgt)) + self._samples = filtcat.groupby("id").apply( + lambda df: self._df_to_series(df, imgt) + ) + if self.shuffle: + self.shuffle_samples() + + def shuffle_samples(self): + self._samples = self._samples.sample(frac=1, random_state=self.shuffle_seed) + + def _df_to_series(self, df, imgt): + d = {} + df = df.set_index("img_type") + for i in imgt: + s = df.loc[i] + idx = s.file_index if i != "lght" else s.id + d.update({f"{i}_filename": [s.file_name], f"{i}_index": [idx]}) + + return pd.DataFrame(d) + + def _open_files(self, verbose=True): + """ + Opens HDF files + """ + imgt = self.data_types + hdf_filenames = [] + for t in imgt: + hdf_filenames += list(np.unique(self._samples[f"{t}_filename"].values)) + self._hdf_files = {} + for f in hdf_filenames: + if verbose: + print("Opening HDF5 file for reading", f) + self._hdf_files[f] = h5py.File(self.sevir_data_dir + "/" + f, "r") + + def close(self): + """ + Closes all open file handles + """ + for f in self._hdf_files: + self._hdf_files[f].close() + self._hdf_files = {} + + @property + def num_seq_per_event(self): + return 1 + (self.raw_seq_len - self.seq_len) // self.stride + + @property + def total_num_seq(self): + """ + The total number of sequences within each shard. + Notice that it is not the product of `self.num_seq_per_event` and `self.total_num_event`. + """ + return int(self.num_seq_per_event * self.num_event) + + @property + def total_num_event(self): + """ + The total number of events in the whole dataset, before split into different shards. + """ + return int(self._samples.shape[0]) + + @property + def start_event_idx(self): + """ + The event idx used in certain rank should satisfy event_idx >= start_event_idx + """ + return self.total_num_event // self.num_shard * self.rank + + @property + def end_event_idx(self): + """ + The event idx used in certain rank should satisfy event_idx < end_event_idx + + """ + if self.split_mode == "ceil": + _last_start_event_idx = ( + self.total_num_event // self.num_shard * (self.num_shard - 1) + ) + _num_event = self.total_num_event - _last_start_event_idx + return self.start_event_idx + _num_event + elif self.split_mode == "floor": + return self.total_num_event // self.num_shard * (self.rank + 1) + else: # self.split_mode == 'uneven': + if self.rank == self.num_shard - 1: # the last process + return self.total_num_event + else: + return self.total_num_event // self.num_shard * (self.rank + 1) + + @property + def num_event(self): + """ + The number of events split into each rank + """ + return self.end_event_idx - self.start_event_idx + + def __len__(self): + """ + Used only when self.sample_mode == 'sequent' + """ + return self.total_num_seq // self.batch_size + + def _read_data(self, row, data): + """ + Iteratively read data into data dict. Finally data[imgt] gets shape (batch_size, height, width, raw_seq_len). + + Args: + row (Dict,optional): A series with fields IMGTYPE_filename, IMGTYPE_index, IMGTYPE_time_index. + data (Dict,optional): , data[imgt] is a data tensor with shape = (tmp_batch_size, height, width, raw_seq_len). + + Returns: + data (np.array): Updated data. Updated shape = (tmp_batch_size + 1, height, width, raw_seq_len). + """ + + imgtyps = np.unique([x.split("_")[0] for x in list(row.keys())]) + for t in imgtyps: + fname = row[f"{t}_filename"] + idx = row[f"{t}_index"] + t_slice = slice(0, None) + # Need to bin lght counts into grid + if t == "lght": + lght_data = self._hdf_files[fname][idx][:] + data_i = self._lght_to_grid(lght_data, t_slice) + else: + data_i = self._hdf_files[fname][t][idx : idx + 1, :, :, t_slice] + data[t] = ( + np.concatenate((data[t], data_i), axis=0) if (t in data) else data_i + ) + return data + + def _lght_to_grid(self, data, t_slice=slice(0, None)): + """ + Converts Nx5 lightning data matrix into a 2D grid of pixel counts + """ + # out_size = (48,48,len(self.lght_frame_times)-1) if isinstance(t_slice,(slice,)) else (48,48) + out_size = ( + (*self.data_shape["lght"], len(self.lght_frame_times)) + if t_slice.stop is None + else (*self.data_shape["lght"], 1) + ) + if data.shape[0] == 0: + return np.zeros((1,) + out_size, dtype=np.float32) + + # filter out points outside the grid + x, y = data[:, 3], data[:, 4] + m = np.logical_and.reduce([x >= 0, x < out_size[0], y >= 0, y < out_size[1]]) + data = data[m, :] + if data.shape[0] == 0: + return np.zeros((1,) + out_size, dtype=np.float32) + + # Filter/separate times + t = data[:, 0] + if t_slice.stop is not None: # select only one time bin + if t_slice.stop > 0: + if t_slice.stop < len(self.lght_frame_times): + tm = np.logical_and( + t >= self.lght_frame_times[t_slice.stop - 1], + t < self.lght_frame_times[t_slice.stop], + ) + else: + tm = t >= self.lght_frame_times[-1] + else: # special case: frame 0 uses lght from frame 1 + tm = np.logical_and( + t >= self.lght_frame_times[0], t < self.lght_frame_times[1] + ) + # tm=np.logical_and( (t>=FRAME_TIMES[t_slice],t self.end_event_idx: + pad_size = event_idx_slice_end - self.end_event_idx + event_idx_slice_end = self.end_event_idx + pd_batch = self._samples.iloc[event_idx:event_idx_slice_end] + data = {} + for index, row in pd_batch.iterrows(): + data = self._read_data(row, data) + if pad_size > 0: + event_batch = [] + for t in self.data_types: + pad_shape = [ + pad_size, + ] + list(data[t].shape[1:]) + data_pad = np.concatenate( + ( + data[t].astype(self.output_type), + np.zeros(pad_shape, dtype=self.output_type), + ), + axis=0, + ) + event_batch.append(data_pad) + else: + event_batch = [data[t].astype(self.output_type) for t in self.data_types] + return event_batch + + def __iter__(self): + return self + + @staticmethod + def preprocess_data_dict( + data_dict, data_types=None, layout="NHWT", rescale="01" + ) -> Dict[str, Union[np.ndarray, paddle.Tensor]]: + """The preprocess of data dict. + + Args: + data_dict (Dict[str, Union[np.ndarray, paddle.Tensor]]): The dict of data. + data_types (Sequence[str]) : The data types that we want to rescale. This mainly excludes "mask" from preprocessing. + layout (str) : consists of batch_size 'N', seq_len 'T', channel 'C', height 'H', width 'W'. + rescale (str): + 'sevir': use the offsets and scale factors in original implementation. + '01': scale all values to range 0 to 1, currently only supports 'vil'. + + Returns: + data_dict (Dict[str, Union[np.ndarray, paddle.Tensor]]): preprocessed data. + """ + + if rescale == "sevir": + scale_dict = PREPROCESS_SCALE_SEVIR + offset_dict = PREPROCESS_OFFSET_SEVIR + elif rescale == "01": + scale_dict = PREPROCESS_SCALE_01 + offset_dict = PREPROCESS_OFFSET_01 + else: + raise ValueError(f"Invalid rescale option: {rescale}.") + if data_types is None: + data_types = data_dict.keys() + for key, data in data_dict.items(): + if key in data_types: + if isinstance(data, np.ndarray): + data = scale_dict[key] * ( + data.astype(np.float32) + offset_dict[key] + ) + data = change_layout_np( + data=data, in_layout="NHWT", out_layout=layout + ) + elif isinstance(data, paddle.Tensor): + data = scale_dict[key] * (data.astype("float32") + offset_dict[key]) + data = change_layout_paddle( + data=data, in_layout="NHWT", out_layout=layout + ) + data_dict[key] = data + return data_dict + + @staticmethod + def process_data_dict_back(data_dict, data_types=None, rescale="01"): + if rescale == "sevir": + scale_dict = PREPROCESS_SCALE_SEVIR + offset_dict = PREPROCESS_OFFSET_SEVIR + elif rescale == "01": + scale_dict = PREPROCESS_SCALE_01 + offset_dict = PREPROCESS_OFFSET_01 + else: + raise ValueError(f"Invalid rescale option: {rescale}.") + if data_types is None: + data_types = data_dict.keys() + for key in data_types: + data = data_dict[key] + data = data.astype("float32") / scale_dict[key] - offset_dict[key] + data_dict[key] = data + return data_dict + + @staticmethod + def data_dict_to_tensor(data_dict, data_types=None): + """ + Convert each element in data_dict to paddle.Tensor (copy without grad). + """ + ret_dict = {} + if data_types is None: + data_types = data_dict.keys() + for key, data in data_dict.items(): + if key in data_types: + if isinstance(data, paddle.Tensor): + ret_dict[key] = data.detach().clone() + elif isinstance(data, np.ndarray): + ret_dict[key] = paddle.to_tensor(data) + else: + raise ValueError( + f"Invalid data type: {type(data)}. Should be paddle.Tensor or np.ndarray" + ) + else: # key == "mask" + ret_dict[key] = data + return ret_dict + + @staticmethod + def downsample_data_dict( + data_dict, data_types=None, factors_dict=None, layout="NHWT" + ) -> Dict[str, paddle.Tensor]: + """The downsample of data. + + Args: + data_dict (Dict[str, Union[np.array, paddle.Tensor]]): The dict of data. + factors_dict (Optional[Dict[str, Sequence[int]]]):each element `factors` is + a Sequence of int, representing (t_factor, h_factor, w_factor). + + Returns: + downsampled_data_dict (Dict[str, paddle.Tensor]): Modify on a deep copy of + data_dict instead of directly modifying the original data_dict. + """ + + if factors_dict is None: + factors_dict = {} + if data_types is None: + data_types = data_dict.keys() + downsampled_data_dict = SEVIRDataset.data_dict_to_tensor( + data_dict=data_dict, data_types=data_types + ) # make a copy + for key, data in data_dict.items(): + factors = factors_dict.get(key, None) + if factors is not None: + downsampled_data_dict[key] = change_layout_paddle( + data=downsampled_data_dict[key], in_layout=layout, out_layout="NTHW" + ) + # downsample t dimension + t_slice = [ + slice(None, None), + ] * 4 + t_slice[1] = slice(None, None, factors[0]) + downsampled_data_dict[key] = downsampled_data_dict[key][tuple(t_slice)] + # downsample spatial dimensions + downsampled_data_dict[key] = F.avg_pool2d( + input=downsampled_data_dict[key], + kernel_size=(factors[1], factors[2]), + ) + + downsampled_data_dict[key] = change_layout_paddle( + data=downsampled_data_dict[key], in_layout="NTHW", out_layout=layout + ) + + return downsampled_data_dict + + def layout_to_in_out_slice( + self, + ): + t_axis = self.layout.find("T") + num_axes = len(self.layout) + in_slice = [ + slice(None, None), + ] * num_axes + out_slice = deepcopy(in_slice) + in_slice[t_axis] = slice(None, self.in_len) + if self.out_len is None: + out_slice[t_axis] = slice(self.in_len, None) + else: + out_slice[t_axis] = slice(self.in_len, self.in_len + self.out_len) + return in_slice, out_slice + + def __getitem__(self, index): + event_idx = (index * self.batch_size) // self.num_seq_per_event + seq_idx = (index * self.batch_size) % self.num_seq_per_event + num_sampled = 0 + sampled_idx_list = [] # list of (event_idx, seq_idx) records + while num_sampled < self.batch_size: + sampled_idx_list.append({"event_idx": event_idx, "seq_idx": seq_idx}) + seq_idx += 1 + if seq_idx >= self.num_seq_per_event: + event_idx += 1 + seq_idx = 0 + num_sampled += 1 + + start_event_idx = sampled_idx_list[0]["event_idx"] + event_batch_size = sampled_idx_list[-1]["event_idx"] - start_event_idx + 1 + + event_batch = self._load_event_batch( + event_idx=start_event_idx, event_batch_size=event_batch_size + ) + ret_dict = {} + for sampled_idx in sampled_idx_list: + batch_slice = [ + sampled_idx["event_idx"] - start_event_idx, + ] # use [] to keepdim + seq_slice = slice( + sampled_idx["seq_idx"] * self.stride, + sampled_idx["seq_idx"] * self.stride + self.seq_len, + ) + for imgt_idx, imgt in enumerate(self.data_types): + sampled_seq = event_batch[imgt_idx][batch_slice, :, :, seq_slice] + if imgt in ret_dict: + ret_dict[imgt] = np.concatenate( + (ret_dict[imgt], sampled_seq), axis=0 + ) + else: + ret_dict.update({imgt: sampled_seq}) + + ret_dict = self.data_dict_to_tensor( + data_dict=ret_dict, data_types=self.data_types + ) + if self.preprocess: + ret_dict = self.preprocess_data_dict( + data_dict=ret_dict, + data_types=self.data_types, + layout=self.layout, + rescale=self.rescale_method, + ) + + if self.downsample_dict is not None: + ret_dict = self.downsample_data_dict( + data_dict=ret_dict, + data_types=self.data_types, + factors_dict=self.downsample_dict, + layout=self.layout, + ) + in_slice, out_slice = self.layout_to_in_out_slice() + data_seq = ret_dict["vil"] + if isinstance(data_seq, paddle.Tensor): + data_seq = data_seq.numpy() + x = data_seq[in_slice[0], in_slice[1], in_slice[2], in_slice[3], in_slice[4]] + y = data_seq[ + out_slice[0], out_slice[1], out_slice[2], out_slice[3], out_slice[4] + ] + + weight_item = self.weight_dict + input_item = {self.input_keys[0]: x} + label_item = { + self.label_keys[0]: y, + } + + return input_item, label_item, weight_item diff --git a/examples/smc_reac/ppsci/data/dataset/spherical_swe_dataset.py b/examples/smc_reac/ppsci/data/dataset/spherical_swe_dataset.py new file mode 100644 index 0000000000..68e29e7883 --- /dev/null +++ b/examples/smc_reac/ppsci/data/dataset/spherical_swe_dataset.py @@ -0,0 +1,104 @@ +from pathlib import Path +from typing import Dict +from typing import Optional +from typing import Tuple + +import numpy as np +from paddle import io + + +class SphericalSWEDataset(io.Dataset): + """Loads a Spherical Shallow Water equations dataset + + Training contains 200 samples in resolution 32x64. + Testing contains 50 samples at resolution 32x64 and 50 samples at resolution 64x128. + + Args: + input_keys (Tuple[str, ...]): Input keys, such as ("input",). + label_keys (Tuple[str, ...]): Output keys, such as ("output",). + data_dir (str): The directory to load data from. + weight_dict (Optional[Dict[str, float]], optional): Define the weight of each constraint variable. + Defaults to None. + test_resolutions (Tuple[str, ...], optional): The resolutions to test dataset. Defaults to ["34x64", "64x128"]. + train_resolution (str, optional): The resolutions to train dataset. Defaults to "34x64". + data_split (str, optional): Specify the dataset split, either 'train' , 'test_32x64',or 'test_64x128'. + Defaults to "train". + """ + + def __init__( + self, + input_keys: Tuple[str, ...], + label_keys: Tuple[str, ...], + data_dir: str, + weight_dict: Optional[Dict[str, float]] = None, + test_resolutions: Tuple[str, ...] = ["34x64", "64x128"], + train_resolution: str = "34x64", + data_split: str = "train", + ): + super().__init__() + self.input_keys = input_keys + self.label_keys = label_keys + self.data_dir = data_dir + self.weight_dict = {} if weight_dict is None else weight_dict + if weight_dict is not None: + self.weight_dict = {key: 1.0 for key in self.label_keys} + self.weight_dict.update(weight_dict) + + self.test_resolutions = test_resolutions + self.train_resolution = train_resolution + self.data_split = data_split + + # train path + path_train = ( + Path(self.data_dir) + .joinpath(f"train_SWE_{self.train_resolution}.npy") + .as_posix() + ) + self.x_train, self.y_train = self.read_data(path_train) + # test path + path_test_1 = ( + Path(self.data_dir) + .joinpath(f"test_SWE_{self.test_resolutions[0]}.npy") + .as_posix() + ) + self.x_test_1, self.y_test_1 = self.read_data(path_test_1) + path_test_2 = ( + Path(self.data_dir) + .joinpath(f"test_SWE_{self.test_resolutions[1]}.npy") + .as_posix() + ) + self.x_test_2, self.y_test_2 = self.read_data(path_test_2) + + def read_data(self, path): + # load with numpy + data = np.load(path, allow_pickle=True).item() + x = data["x"].astype("float32") + y = data["y"].astype("float32") + del data + return x, y + + def __len__(self): + if self.data_split == "train": + return self.x_train.shape[0] + elif self.data_split == "test_32x64": + return self.x_test_1.shape[0] + else: + return self.x_test_2.shape[0] + + def __getitem__(self, index): + if self.data_split == "train": + x = self.x_train[index] + y = self.y_train[index] + + elif self.data_split == "test_32x64": + x = self.x_test_1[index] + y = self.y_test_1[index] + else: + x = self.x_test_2[index] + y = self.y_test_2[index] + + input_item = {self.input_keys[0]: x} + label_item = {self.label_keys[0]: y} + weight_item = self.weight_dict + + return input_item, label_item, weight_item diff --git a/examples/smc_reac/ppsci/data/dataset/trphysx_dataset.py b/examples/smc_reac/ppsci/data/dataset/trphysx_dataset.py new file mode 100644 index 0000000000..3160951530 --- /dev/null +++ b/examples/smc_reac/ppsci/data/dataset/trphysx_dataset.py @@ -0,0 +1,326 @@ +# 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. + +""" +Code below is heavily based on [transformer-physx](https://github.com/zabaras/transformer-physx) +""" + +from __future__ import annotations + +import os +from typing import Dict +from typing import Optional +from typing import Tuple + +try: + import h5py +except ModuleNotFoundError: + pass +import numpy as np +import paddle +from paddle import io + +from ppsci.arch import base + + +class LorenzDataset(io.Dataset): + """Dataset for training Lorenz model. + + Args: + file_path (str): Data set path. + input_keys (Tuple[str, ...]): Input keys, such as ("states",). + label_keys (Tuple[str, ...]): Output keys, such as ("pred_states", "recover_states"). + block_size (int): Data block size. + stride (int): Data stride. + ndata (Optional[int]): Number of data series to use. Defaults to None. + weight_dict (Optional[Dict[str, float]]): Weight dictionary. Defaults to None. + embedding_model (Optional[base.Arch]): Embedding model. Defaults to None. + + Examples: + >>> import ppsci + >>> dataset = ppsci.data.dataset.LorenzDataset( + ... "file_path": "/path/to/LorenzDataset", + ... "input_keys": ("x",), + ... "label_keys": ("v",), + ... "block_size": 32, + ... "stride": 16, + ... ) # doctest: +SKIP + """ + + # Whether support batch indexing for speeding up fetching process. + batch_index: bool = False + + def __init__( + self, + file_path: str, + input_keys: Tuple[str, ...], + label_keys: Tuple[str, ...], + block_size: int, + stride: int, + ndata: Optional[int] = None, + weight_dict: Optional[Dict[str, float]] = None, + embedding_model: Optional[base.Arch] = None, + ): + super().__init__() + if not os.path.exists(file_path): + raise FileNotFoundError( + f"file_path({file_path}) not exists. Please download dataset first. " + "Training: https://paddle-org.bj.bcebos.com/paddlescience/datasets/transformer_physx/lorenz_training_rk.hdf5. " + "Valid: https://paddle-org.bj.bcebos.com/paddlescience/datasets/transformer_physx/lorenz_valid_rk.hdf5." + ) + + self.file_path = file_path + self.input_keys = input_keys + self.label_keys = label_keys + + self.block_size = block_size + self.stride = stride + self.ndata = ndata + self.weight_dict = {key: 1.0 for key in self.label_keys} + if weight_dict is not None: + self.weight_dict.update(weight_dict) + + self.data = self.read_data(file_path, block_size, stride) + self.embedding_model = embedding_model + if embedding_model is None: + self.embedding_data = None + else: + embedding_model.eval() + with paddle.no_grad(): + data_tensor = paddle.to_tensor(self.data) + embedding_data_tensor = embedding_model.encoder(data_tensor) + self.embedding_data = embedding_data_tensor.numpy() + + def read_data(self, file_path: str, block_size: int, stride: int): + data = [] + with h5py.File(file_path, "r") as f: + data_num = 0 + for key in f.keys(): + data_series = np.asarray(f[key], dtype=paddle.get_default_dtype()) + for i in range(0, data_series.shape[0] - block_size + 1, stride): + data.append(data_series[i : i + block_size]) + data_num += 1 + if self.ndata is not None and data_num >= self.ndata: + break + return np.asarray(data) + + def __len__(self): + return len(self.data) + + def __getitem__(self, idx): + # when embedding data is None + if self.embedding_data is None: + data_item = self.data[idx] + input_item = {self.input_keys[0]: data_item} + label_item = { + self.label_keys[0]: data_item[1:, :], + self.label_keys[1]: data_item, + } + else: + data_item = self.embedding_data[idx] + input_item = {self.input_keys[0]: data_item[:-1, :]} + label_item = {self.label_keys[0]: data_item[1:, :]} + if len(self.label_keys) == 2: + label_item[self.label_keys[1]] = self.data[idx][1:, :] + + weight_shape = [1] * len(data_item.shape) + weight_item = { + key: np.full(weight_shape, value, paddle.get_default_dtype()) + for key, value in self.weight_dict.items() + } + return (input_item, label_item, weight_item) + + +class RosslerDataset(LorenzDataset): + """Dataset for training Rossler model. + + Args: + file_path (str): Data set path. + input_keys (Tuple[str, ...]): Input keys, such as ("states",). + label_keys (Tuple[str, ...]): Output keys, such as ("pred_states", "recover_states"). + block_size (int): Data block size. + stride (int): Data stride. + ndata (Optional[int]): Number of data series to use. Defaults to None. + weight_dict (Optional[Dict[str, float]]): Weight dictionary. Defaults to None. + embedding_model (Optional[base.Arch]): Embedding model. Defaults to None. + + Examples: + >>> import ppsci + >>> dataset = ppsci.data.dataset.RosslerDataset( + ... "file_path": "/path/to/RosslerDataset", + ... "input_keys": ("x",), + ... "label_keys": ("v",), + ... "block_size": 32, + ... "stride": 16, + ... ) # doctest: +SKIP + """ + + # Whether support batch indexing for speeding up fetching process. + batch_index: bool = False + + def __init__( + self, + file_path: str, + input_keys: Tuple[str, ...], + label_keys: Tuple[str, ...], + block_size: int, + stride: int, + ndata: Optional[int] = None, + weight_dict: Optional[Dict[str, float]] = None, + embedding_model: Optional[base.Arch] = None, + ): + if not os.path.exists(file_path): + raise FileNotFoundError( + f"file_path({file_path}) not exists. Please download dataset first. " + "Training: https://paddle-org.bj.bcebos.com/paddlescience/datasets/transformer_physx/rossler_training.hdf5. " + "Valid: https://paddle-org.bj.bcebos.com/paddlescience/datasets/transformer_physx/rossler_valid.hdf5." + ) + super().__init__( + file_path, + input_keys, + label_keys, + block_size, + stride, + ndata, + weight_dict, + embedding_model, + ) + + +class CylinderDataset(io.Dataset): + """Dataset for training Cylinder model. + + Args: + file_path (str): Data set path. + input_keys (Tuple[str, ...]): Input keys, such as ("states","visc"). + label_keys (Tuple[str, ...]): Output keys, such as ("pred_states", "recover_states"). + block_size (int): Data block size. + stride (int): Data stride. + ndata (Optional[int]): Number of data series to use. Defaults to None. + weight_dict (Optional[Dict[str, float]]): Weight dictionary. Defaults to None. + embedding_model (Optional[base.Arch]): Embedding model. Defaults to None. + embedding_batch_size (int, optional): The batch size of embedding model. Defaults to 64. + + Examples: + >>> import ppsci + >>> dataset = ppsci.data.dataset.CylinderDataset( + ... "file_path": "/path/to/CylinderDataset", + ... "input_keys": ("x",), + ... "label_keys": ("v",), + ... "block_size": 32, + ... "stride": 16, + ... ) # doctest: +SKIP + """ + + # Whether support batch indexing for speeding up fetching process. + batch_index: bool = False + + def __init__( + self, + file_path: str, + input_keys: Tuple[str, ...], + label_keys: Tuple[str, ...], + block_size: int, + stride: int, + ndata: Optional[int] = None, + weight_dict: Optional[Dict[str, float]] = None, + embedding_model: Optional[base.Arch] = None, + embedding_batch_size: int = 64, + ): + if not os.path.exists(file_path): + raise FileNotFoundError( + f"file_path({file_path}) not exists. Please download dataset first. " + "Training: https://paddle-org.bj.bcebos.com/paddlescience/datasets/transformer_physx/cylinder_training.hdf5. " + "Valid: https://paddle-org.bj.bcebos.com/paddlescience/datasets/transformer_physx/cylinder_valid.hdf5." + ) + super().__init__() + self.file_path = file_path + self.input_keys = input_keys + self.label_keys = label_keys + + self.block_size = block_size + self.stride = stride + self.ndata = ndata + self.weight_dict = {key: 1.0 for key in self.label_keys} + if weight_dict is not None: + self.weight_dict.update(weight_dict) + + self.data, self.visc = self.read_data(file_path, block_size, stride) + self.embedding_model = embedding_model + if embedding_model is None: + self.embedding_data = None + else: + embedding_model.eval() + with paddle.no_grad(): + data_tensor = paddle.to_tensor(self.data) + visc_tensor = paddle.to_tensor(self.visc) + embedding_data = [] + for i in range(0, len(data_tensor), embedding_batch_size): + start, end = i, min(i + embedding_batch_size, len(data_tensor)) + embedding_data_batch = embedding_model.encoder( + data_tensor[start:end], visc_tensor[start:end] + ) + embedding_data.append(embedding_data_batch.numpy()) + self.embedding_data = np.concatenate(embedding_data) + + def read_data(self, file_path: str, block_size: int, stride: int): + data = [] + visc = [] + with h5py.File(file_path, "r") as f: + data_num = 0 + for key in f.keys(): + visc0 = 2.0 / float(key) + ux = np.asarray(f[key + "/ux"], dtype=paddle.get_default_dtype()) + uy = np.asarray(f[key + "/uy"], dtype=paddle.get_default_dtype()) + p = np.asarray(f[key + "/p"], dtype=paddle.get_default_dtype()) + data_series = np.stack([ux, uy, p], axis=1) + + for i in range(0, data_series.shape[0] - block_size + 1, stride): + data.append(data_series[i : i + block_size]) + visc.append([visc0]) + + data_num += 1 + if self.ndata is not None and data_num >= self.ndata: + break + + data = np.asarray(data) + visc = np.asarray(visc, dtype=paddle.get_default_dtype()) + return data, visc + + def __len__(self): + return len(self.data) + + def __getitem__(self, i): + if self.embedding_data is None: + data_item = self.data[i] + input_item = { + self.input_keys[0]: data_item, + self.input_keys[1]: self.visc[i], + } + label_item = { + self.label_keys[0]: data_item[1:], + self.label_keys[1]: data_item, + } + else: + data_item = self.embedding_data[i] + input_item = {self.input_keys[0]: data_item[:-1, :]} + label_item = {self.label_keys[0]: data_item[1:, :]} + if len(self.label_keys) == 2: + label_item[self.label_keys[1]] = data_item[1:, :] + weight_shape = [1] * len(data_item.shape) + weight_item = { + key: np.full(weight_shape, value, paddle.get_default_dtype()) + for key, value in self.weight_dict.items() + } + return (input_item, label_item, weight_item) diff --git a/examples/smc_reac/ppsci/data/dataset/vtu_dataset.py b/examples/smc_reac/ppsci/data/dataset/vtu_dataset.py new file mode 100644 index 0000000000..fb0c9201b7 --- /dev/null +++ b/examples/smc_reac/ppsci/data/dataset/vtu_dataset.py @@ -0,0 +1,106 @@ +# 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. + +from __future__ import annotations + +from typing import Dict +from typing import Optional +from typing import Tuple + +import numpy as np +from paddle import io +from paddle import vision + +from ppsci.utils import reader + + +class VtuDataset(io.Dataset): + """Dataset class for .vtu file. + + Args: + file_path (str): *.vtu file path. + input_keys (Optional[Tuple[str, ...]]): Tuple of input keys. Defaults to None. + label_keys (Optional[Tuple[str, ...]]): Tuple of label keys. Defaults to None. + time_step (Optional[int]): Time step with unit second. Defaults to None. + time_index (Optional[Tuple[int, ...]]): Time index tuple in increasing order. + labels (Optional[Dict[str, float]]): Temporary variable for [load_vtk_with_time_file]. + transforms (vision.Compose, optional): Compose object contains sample wise. + transform(s). + + Examples: + >>> from ppsci.data.dataset import VtuDataset + + >>> dataset = VtuDataset(file_path='example.vtu') # doctest: +SKIP + + >>> # get the length of the dataset + >>> dataset_size = len(dataset) # doctest: +SKIP + >>> # get the first sample of the data + >>> first_sample = dataset[0] # doctest: +SKIP + >>> print("First sample:", first_sample) # doctest: +SKIP + """ + + # Whether support batch indexing for speeding up fetching process. + batch_index: bool = True + + def __init__( + self, + file_path: str, + input_keys: Optional[Tuple[str, ...]] = None, + label_keys: Optional[Tuple[str, ...]] = None, + time_step: Optional[int] = None, + time_index: Optional[Tuple[int, ...]] = None, + labels: Optional[Dict[str, float]] = None, + transforms: Optional[vision.Compose] = None, + ): + super().__init__() + + # load data from file + if time_step is not None and time_index is not None: + _input, _label = reader.load_vtk_file( + file_path, time_step, time_index, input_keys, label_keys + ) + _label = {key: _label[key] for key in label_keys} + elif time_step is None and time_index is None: + _input = reader.load_vtk_with_time_file(file_path) + _label = {} + for key, value in labels.items(): + if isinstance(value, (int, float)): + _label[key] = np.full_like( + next(iter(_input.values())), value, "float32" + ) + else: + _label[key] = value + else: + raise ValueError( + "Error, read vtu with time_step and time_index, or neither" + ) + + # transform + _input = transforms(_input) + _label = transforms(_label) + + self.input = _input + self.label = _label + self.input_keys = input_keys + self.label_keys = label_keys + self.transforms = transforms + self.num_samples = len(next(iter(self.input.values()))) + + def __getitem__(self, idx): + input_item = {key: value[idx] for key, value in self.input.items()} + label_item = {key: value[idx] for key, value in self.label.items()} + return (input_item, label_item, {}) + + def __len__(self): + return self.num_samples diff --git a/examples/smc_reac/ppsci/data/process/__init__.py b/examples/smc_reac/ppsci/data/process/__init__.py new file mode 100644 index 0000000000..f46c8dd9cf --- /dev/null +++ b/examples/smc_reac/ppsci/data/process/__init__.py @@ -0,0 +1,21 @@ +# 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. + +from ppsci.data.process import batch_transform +from ppsci.data.process import transform + +__all__ = [ + "batch_transform", + "transform", +] diff --git a/examples/smc_reac/ppsci/data/process/batch_transform/__init__.py b/examples/smc_reac/ppsci/data/process/batch_transform/__init__.py new file mode 100644 index 0000000000..9e98f39264 --- /dev/null +++ b/examples/smc_reac/ppsci/data/process/batch_transform/__init__.py @@ -0,0 +1,135 @@ +# 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. + +import copy +import numbers +from collections.abc import Mapping +from collections.abc import Sequence +from typing import Any +from typing import Callable +from typing import List +from typing import Optional + +import numpy as np +import paddle + +from ppsci.data.process import transform +from ppsci.data.process.batch_transform.preprocess import FunctionalBatchTransform + +try: + import pgl +except ModuleNotFoundError: + pass + + +__all__ = [ + "build_batch_transforms", + "default_collate_fn", + "FunctionalBatchTransform", +] + + +def default_collate_fn(batch: List[Any]) -> Any: + """Default_collate_fn for paddle dataloader. + + NOTE: This `default_collate_fn` is different from official `default_collate_fn` + which specially adapt case where sample is `None` and `pgl.Graph`. + + ref: https://github.com/PaddlePaddle/Paddle/blob/develop/python/paddle/io/dataloader/collate.py#L25 + + Args: + batch (List[Any]): Batch of samples to be collated. + + Returns: + Any: Collated batch data. + """ + sample = batch[0] + if sample is None: + return None + elif isinstance(sample, np.ndarray): + batch = np.stack(batch, axis=0) + return batch + elif isinstance(sample, (paddle.Tensor, paddle.framework.core.eager.Tensor)): + return paddle.stack(batch, axis=0) + elif isinstance(sample, numbers.Number): + batch = np.array(batch) + return batch + elif isinstance(sample, (str, bytes)): + return batch + elif isinstance(sample, Mapping): + return {key: default_collate_fn([d[key] for d in batch]) for key in sample} + elif isinstance(sample, Sequence): + sample_fields_num = len(sample) + if not all(len(sample) == sample_fields_num for sample in iter(batch)): + raise RuntimeError("Fields number not same among samples in a batch") + return [default_collate_fn(fields) for fields in zip(*batch)] + elif str(type(sample)) == "": + # use str(type()) instead of isinstance() in case of pgl is not installed. + graph = pgl.Graph(num_nodes=sample.num_nodes, edges=sample.edges) + graph.x = np.concatenate([g.x for g in batch]) + graph.y = np.concatenate([g.y for g in batch]) + graph.edge_index = np.concatenate([g.edge_index for g in batch], axis=1) + + graph.edge_attr = np.concatenate([g.edge_attr for g in batch]) + graph.pos = np.concatenate([g.pos for g in batch]) + if hasattr(sample, "aoa"): + graph.aoa = np.concatenate([g.aoa for g in batch]) + if hasattr(sample, "mach_or_reynolds"): + graph.mach_or_reynolds = np.concatenate([g.mach_or_reynolds for g in batch]) + graph.tensor() + graph.shape = [len(batch)] + return graph + elif ( + str(type(sample)) + == "" + ): + graph = sample + graph.tensor() + graph.shape = [1] + return graph + raise TypeError( + "batch data can only contains: paddle.Tensor, numpy.ndarray, " + f"dict, list, number, None, pgl.Graph, GraphGridMesh, but got {type(sample)}" + ) + + +def build_transforms(cfg): + if not cfg: + return transform.Compose([]) + cfg = copy.deepcopy(cfg) + + transform_list = [] + for _item in cfg: + transform_cls = next(iter(_item.keys())) + transform_cfg = _item[transform_cls] + transform_obj = eval(transform_cls)(**transform_cfg) + transform_list.append(transform_obj) + + return transform.Compose(transform_list) + + +def build_batch_transforms(cfg, collate_fn: Optional[Callable]): + cfg = copy.deepcopy(cfg) + batch_transforms: Callable[[List[Any]], List[Any]] = build_transforms(cfg) + if collate_fn is None: + collate_fn = default_collate_fn + + def collate_fn_batch_transforms(batch: List[Any]): + # apply batch transform on separate samples + batch = batch_transforms(batch) + + # then collate separate samples into batched data + return collate_fn(batch) + + return collate_fn_batch_transforms diff --git a/examples/smc_reac/ppsci/data/process/batch_transform/preprocess.py b/examples/smc_reac/ppsci/data/process/batch_transform/preprocess.py new file mode 100644 index 0000000000..62ca5d3be1 --- /dev/null +++ b/examples/smc_reac/ppsci/data/process/batch_transform/preprocess.py @@ -0,0 +1,74 @@ +# 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. + +from __future__ import annotations + +from typing import Any +from typing import Callable +from typing import Dict +from typing import List +from typing import Optional +from typing import Tuple + +import numpy as np + + +class FunctionalBatchTransform: + """Functional data transform class, which allows to use custom data transform function from given transform_func for special cases. + + Args: + transform_func (Callable): Function of batch data transform. + + Examples: + >>> import ppsci + >>> from typing import Tuple, Dict, Optional + >>> def batch_transform_func( + ... data_list: List[ + ... Tuple[Dict[str, np.ndarray], Dict[str, np.ndarray], Optional[Dict[str, np.ndarray]]] + ... ], + ... ) -> List[Tuple[Dict[str, np.ndarray], Dict[str, np.ndarray], Optional[Dict[str, np.ndarray]]]]: + ... input_dicts, label_dicts, weight_dicts = zip(*data_list) + ... + ... for input_dict in input_dicts: + ... for key in input_dict: + ... input_dict[key] = input_dict[key] * 2 + ... + ... for label_dict in label_dicts: + ... for key in label_dict: + ... label_dict[key] = label_dict[key] + 1.0 + ... + ... return list(zip(input_dicts, label_dicts, weight_dicts)) + ... + >>> # Create a FunctionalBatchTransform object with the batch_transform_func function + >>> transform = ppsci.data.batch_transform.FunctionalBatchTransform(batch_transform_func) + >>> # Define some sample data, labels, and weights + >>> data = [({'x': 1}, {'y': 2}, None), ({'x': 11}, {'y': 22}, None)] + >>> transformed_data = transform(data) + >>> for tuple in transformed_data: + ... print(tuple) + ({'x': 2}, {'y': 3.0}, None) + ({'x': 22}, {'y': 23.0}, None) + """ + + def __init__( + self, + transform_func: Callable[[List[Any]], List[Any]], + ): + self.transform_func = transform_func + + def __call__( + self, + data_list: List[Tuple[Optional[Dict[str, np.ndarray]], ...]], + ) -> List[Tuple[Optional[Dict[str, np.ndarray]], ...]]: + return self.transform_func(data_list) diff --git a/examples/smc_reac/ppsci/data/process/transform/__init__.py b/examples/smc_reac/ppsci/data/process/transform/__init__.py new file mode 100644 index 0000000000..f5a4baa287 --- /dev/null +++ b/examples/smc_reac/ppsci/data/process/transform/__init__.py @@ -0,0 +1,72 @@ +# 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. + +import copy +import traceback +from typing import Any +from typing import Tuple + +from paddle import vision + +from ppsci.data.process.transform.preprocess import CropData +from ppsci.data.process.transform.preprocess import FunctionalTransform +from ppsci.data.process.transform.preprocess import Log1p +from ppsci.data.process.transform.preprocess import Normalize +from ppsci.data.process.transform.preprocess import Scale +from ppsci.data.process.transform.preprocess import SqueezeData +from ppsci.data.process.transform.preprocess import Translate + +__all__ = [ + "CropData", + "FunctionalTransform", + "Log1p", + "Normalize", + "Scale", + "SqueezeData", + "Translate", + "build_transforms", +] + + +class Compose(vision.Compose): + """Custom Compose for multiple items in given data.""" + + def __call__(self, *data: Tuple[Any, ...]): + for f in self.transforms: + try: + # NOTE: This is different from vision.Compose to allow receive multiple data items + data = f(*data) + except Exception as e: + stack_info = traceback.format_exc() + print( + f"fail to perform transform [{f}] with error: " + f"{e} and stack:\n{str(stack_info)}" + ) + raise e + return data + + +def build_transforms(cfg): + if not cfg: + return Compose([]) + cfg = copy.deepcopy(cfg) + + transform_list = [] + for _item in cfg: + transform_cls = next(iter(_item.keys())) + transform_cfg = _item[transform_cls] + transform = eval(transform_cls)(**transform_cfg) + transform_list.append(transform) + + return Compose(transform_list) diff --git a/examples/smc_reac/ppsci/data/process/transform/preprocess.py b/examples/smc_reac/ppsci/data/process/transform/preprocess.py new file mode 100644 index 0000000000..48f0fa1222 --- /dev/null +++ b/examples/smc_reac/ppsci/data/process/transform/preprocess.py @@ -0,0 +1,331 @@ +# 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. + +from __future__ import annotations + +from typing import Callable +from typing import Dict +from typing import Tuple +from typing import Union + +import numpy as np + + +class Translate: + """Translate class. + + Args: + offset (Dict[str, float]): Shift the input data according to the variable name + and coefficient specified in offset. + + Examples: + >>> import ppsci + >>> import numpy as np + + >>> input_dict = {"x": np.array([5.0, 10.0]), "y": np.array([20.0, 40.0])} + >>> label_dict = {"x": np.array([1.0, 2.0]), "y": np.array([3.0, 4.0])} + >>> weight_dict = {"x": np.array([10.0, 20.0]), "y": np.array([30.0, 40.0])} + + >>> translate = ppsci.data.transform.Translate({"x": 1.0, "y": -1.0}) + >>> translated_input_dict, translated_label_dict, translated_weight_dict = translate(input_dict, label_dict, weight_dict) + + >>> print(translated_input_dict) + {'x': array([ 6., 11.]), 'y': array([19., 39.])} + >>> print(translated_label_dict) + {'x': array([1., 2.]), 'y': array([3., 4.])} + >>> print(translated_weight_dict) + {'x': array([10., 20.]), 'y': array([30., 40.])} + """ + + def __init__(self, offset: Dict[str, float]): + self.offset = offset + + def __call__(self, input_dict, label_dict, weight_dict): + input_dict_copy = {**input_dict} + for key in self.offset: + if key in input_dict: + input_dict_copy[key] += self.offset[key] + return input_dict_copy, label_dict, weight_dict + + +class Scale: + """Scale class for data transformation. + + Args: + scale (Dict[str, float]): Scale the input data according to the variable name + and coefficient specified in scale. + + Examples: + >>> import ppsci + >>> translate = ppsci.data.transform.Scale({"x": 1.5, "y": 2.0}) + >>> input_dict = {"x": 10, "y": 20} + >>> label_dict = {"x": 100, "y": 200} + >>> weight_dict = {"x": 1000, "y": 2000} + >>> input_dict_scaled, label_dict_scaled, weight_dict_scaled = translate(input_dict, label_dict, weight_dict) + >>> print(input_dict_scaled) + {'x': 15.0, 'y': 40.0} + >>> print(label_dict_scaled) + {'x': 100, 'y': 200} + >>> print(weight_dict_scaled) + {'x': 1000, 'y': 2000} + """ + + def __init__(self, scale: Dict[str, float]): + self.scale = scale + + def __call__(self, input_dict, label_dict, weight_dict): + input_dict_copy = {**input_dict} + for key in self.scale: + if key in input_dict: + input_dict_copy[key] *= self.scale[key] + return input_dict_copy, label_dict, weight_dict + + +class Normalize: + """Normalize data class. + + NOTE: This transform will modify the input data dict inplace. + + Args: + mean (Union[np.ndarray, Tuple[float, ...]]): Mean of training dataset. + std (Union[np.ndarray, Tuple[float, ...]]): Standard Deviation of training dataset. + apply_keys (Tuple[str, ...], optional): Which data is the normalization method applied to. Defaults to ("input", "label"). + + Examples: + >>> import ppsci + >>> normalize = ppsci.data.transform.Normalize((0.0, 0.0, 0.0), (1.0, 1.0, 1.0)) + >>> input_item = {"data": np.array([1.0, 2.0, 3.0])} + >>> label_item = {"data": np.array([4.0, 5.0, 6.0])} + >>> weight_item = np.array([0.1, 0.2, 0.3]) + >>> normalized_item = normalize(input_item, label_item, weight_item) + >>> print(normalized_item) + ({'data': array([1., 2., 3.])}, {'data': array([4., 5., 6.])}, array([0.1, 0.2, 0.3])) + """ + + def __init__( + self, + mean: Union[np.ndarray, Tuple[float, ...]], + std: Union[np.ndarray, Tuple[float, ...]], + apply_keys: Tuple[str, ...] = ("input", "label"), + ): + if len(apply_keys) == 0 or len(set(apply_keys) | {"input", "label"}) > 2: + raise ValueError( + f"apply_keys should be a non empty subset of ('input', 'label'), but got {apply_keys}" + ) + self.mean = mean + self.std = std + self.apply_keys = apply_keys + + def __call__(self, input_item, label_item, weight_item): + if "input" in self.apply_keys: + for key, value in input_item.items(): + input_item[key] = (value - self.mean) / self.std + if "label" in self.apply_keys: + for key, value in label_item.items(): + label_item[key] = (value - self.mean) / self.std + return input_item, label_item, weight_item + + +class Log1p: + """Calculates the natural logarithm of one plus the data, element-wise. + + NOTE: This transform will modify the input data dict inplace. + + Args: + scale (float, optional): Scale data. Defaults to 1.0. + apply_keys (Tuple[str, ...], optional): Which data is the log1p method applied to. Defaults to ("input", "label"). + + Examples: + >>> import ppsci + >>> log1p = ppsci.data.transform.Log1p(1e-5) + >>> input_item = {"data": np.array([1.0, 2.0, 3.0])} + >>> label_item = {"data": np.array([4.0, 5.0, 6.0])} + >>> weight_item = np.array([0.1, 0.2, 0.3]) + >>> input_item_transformed, label_item_transformed, weight_item_transformed = log1p(input_item, label_item, weight_item) + >>> print(input_item_transformed) + {'data': array([11.51293546, 12.20607765, 12.61154109])} + >>> print(label_item_transformed) + {'data': array([12.89922233, 13.12236538, 13.3046866 ])} + >>> print(weight_item_transformed) + [0.1 0.2 0.3] + """ + + def __init__( + self, + scale: float = 1.0, + apply_keys: Tuple[str, ...] = ("input", "label"), + ): + if len(apply_keys) == 0 or len(set(apply_keys) | {"input", "label"}) > 2: + raise ValueError( + f"apply_keys should be a non empty subset of ('input', 'label'), but got {apply_keys}" + ) + self.scale = scale + self.apply_keys = apply_keys + + def __call__(self, input_item, label_item, weight_item): + if "input" in self.apply_keys: + for key, value in input_item.items(): + input_item[key] = np.log1p(value / self.scale) + if "label" in self.apply_keys: + for key, value in label_item.items(): + label_item[key] = np.log1p(value / self.scale) + return input_item, label_item, weight_item + + +class CropData: + """Crop data class. + + This class is used to crop data based on a specified bounding box. + + NOTE: This transform will modify the input data dict inplace. + + Args: + xmin (Tuple[int, ...]): Bottom left corner point, [x0, y0]. + xmax (Tuple[int, ...]): Top right corner point, [x1, y1]. + apply_keys (Tuple[str, ...], optional): Which data is the crop method applied to. Defaults to ("input", "label"). + + Examples: + >>> import ppsci + >>> import numpy as np + >>> crop_data = ppsci.data.transform.CropData((0, 0), (256, 512)) + >>> input_item = {"input": np.zeros((3, 720, 1440))} + >>> label_item = {"label": np.zeros((3, 720, 1440))} + >>> weight_item = {"weight": np.ones((3, 720, 1440))} + >>> input_item, label_item, weight_item = crop_data(input_item, label_item, weight_item) + >>> print(input_item["input"].shape) + (3, 256, 512) + >>> print(label_item["label"].shape) + (3, 256, 512) + """ + + def __init__( + self, + xmin: Tuple[int, ...], + xmax: Tuple[int, ...], + apply_keys: Tuple[str, ...] = ("input", "label"), + ): + if len(apply_keys) == 0 or len(set(apply_keys) | {"input", "label"}) > 2: + raise ValueError( + f"apply_keys should be a non empty subset of ('input', 'label'), but got {apply_keys}" + ) + self.xmin = xmin + self.xmax = xmax + self.apply_keys = apply_keys + + def __call__(self, input_item, label_item, weight_item): + if "input" in self.apply_keys: + for key, value in input_item.items(): + input_item[key] = value[ + :, self.xmin[0] : self.xmax[0], self.xmin[1] : self.xmax[1] + ] + if "label" in self.apply_keys: + for key, value in label_item.items(): + label_item[key] = value[ + :, self.xmin[0] : self.xmax[0], self.xmin[1] : self.xmax[1] + ] + return input_item, label_item, weight_item + + +class SqueezeData: + """Squeeze data class. + + NOTE: This transform will modify the input data dict inplace. + + Args: + apply_keys (Tuple[str, ...], optional): Which data is the squeeze method applied to. Defaults to ("input", "label"). + + Examples: + >>> import ppsci + >>> import numpy as np + >>> squeeze_data = ppsci.data.transform.SqueezeData() + >>> input_data = {"input": np.random.rand(10, 224, 224)} + >>> label_data = {"label": np.random.rand(10, 224, 224)} + >>> weight_data = {"weight": np.random.rand(10, 224, 224)} + >>> input_data_squeezed, label_data_squeezed, weight_data_squeezed = squeeze_data(input_data, label_data, weight_data) + """ + + def __init__(self, apply_keys: Tuple[str, ...] = ("input", "label")): + if len(apply_keys) == 0 or len(set(apply_keys) | {"input", "label"}) > 2: + raise ValueError( + f"apply_keys should be a non empty subset of ('input', 'label'), but got {apply_keys}" + ) + self.apply_keys = apply_keys + + def __call__(self, input_item, label_item, weight_item): + if "input" in self.apply_keys: + for key, value in input_item.items(): + if value.ndim == 4: + B, C, H, W = value.shape + input_item[key] = value.reshape((B * C, H, W)) + if value.ndim != 3: + raise ValueError( + f"Only support squeeze data to ndim=3 now, but got ndim={value.ndim}" + ) + if "label" in self.apply_keys: + for key, value in label_item.items(): + if value.ndim == 4: + B, C, H, W = value.shape + label_item[key] = value.reshape((B * C, H, W)) + if value.ndim != 3: + raise ValueError( + f"Only support squeeze data to ndim=3 now, but got ndim={value.ndim}" + ) + return input_item, label_item, weight_item + + +class FunctionalTransform: + """Functional data transform class, which allows to use custom data transform function from given transform_func for special cases. + + Args: + transform_func (Callable): Function of data transform. + + Examples: + >>> # This is the transform_func function. It takes three dictionaries as input: data_dict, label_dict, and weight_dict. + >>> # The function will perform some transformations on the data in data_dict, convert all labels in label_dict to uppercase, + >>> # and modify the weights in weight_dict by dividing each weight by 10. + >>> # Finally, it returns the transformed data, labels, and weights as a tuple. + >>> import ppsci + >>> def transform_func(data_dict, label_dict, weight_dict): + ... for key in data_dict: + ... data_dict[key] = data_dict[key] * 2 + ... for key in label_dict: + ... label_dict[key] = label_dict[key] + 1.0 + ... for key in weight_dict: + ... weight_dict[key] = weight_dict[key] / 10 + ... return data_dict, label_dict, weight_dict + >>> transform = ppsci.data.transform.FunctionalTransform(transform_func) + >>> # Define some sample data, labels, and weights + >>> data = {'feature1': np.array([1, 2, 3]), 'feature2': np.array([4, 5, 6])} + >>> label = {'class': 0.0, 'instance': 0.1} + >>> weight = {'weight1': 0.5, 'weight2': 0.5} + >>> # Apply the transform function to the data, labels, and weights using the FunctionalTransform instance + >>> transformed_data = transform(data, label, weight) + >>> print(transformed_data) + ({'feature1': array([2, 4, 6]), 'feature2': array([ 8, 10, 12])}, {'class': 1.0, 'instance': 1.1}, {'weight1': 0.05, 'weight2': 0.05}) + """ + + def __init__( + self, + transform_func: Callable, + ): + self.transform_func = transform_func + + def __call__( + self, *data: Tuple[Dict[str, np.ndarray], ...] + ) -> Tuple[Dict[str, np.ndarray], ...]: + data_dict, label_dict, weight_dict = data + data_dict_copy = {**data_dict} + label_dict_copy = {**label_dict} + weight_dict_copy = {**weight_dict} if weight_dict is not None else {} + return self.transform_func(data_dict_copy, label_dict_copy, weight_dict_copy) diff --git a/examples/smc_reac/ppsci/equation/__init__.py b/examples/smc_reac/ppsci/equation/__init__.py new file mode 100644 index 0000000000..bcffef4060 --- /dev/null +++ b/examples/smc_reac/ppsci/equation/__init__.py @@ -0,0 +1,76 @@ +# 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. + +import copy + +from ppsci.equation.fpde import FractionalPoisson +from ppsci.equation.ide import Volterra +from ppsci.equation.pde import DETACH_FUNC_NAME +from ppsci.equation.pde import NLSMB +from ppsci.equation.pde import PDE +from ppsci.equation.pde import AllenCahn +from ppsci.equation.pde import Biharmonic +from ppsci.equation.pde import HeatExchanger +from ppsci.equation.pde import Helmholtz +from ppsci.equation.pde import Laplace +from ppsci.equation.pde import LinearElasticity +from ppsci.equation.pde import NavierStokes +from ppsci.equation.pde import NormalDotVec +from ppsci.equation.pde import Poisson +from ppsci.equation.pde import Vibration +from ppsci.utils import logger +from ppsci.utils import misc + +__all__ = [ + "PDE", + "DETACH_FUNC_NAME", + "AllenCahn", + "Biharmonic", + "HeatExchanger", + "Helmholtz", + "Laplace", + "LinearElasticity", + "NavierStokes", + "NormalDotVec", + "Poisson", + "Vibration", + "Volterra", + "NLSMB", + "FractionalPoisson", + "build_equation", +] + + +def build_equation(cfg): + """Build equation(s) + + Args: + cfg (List[DictConfig]): Equation(s) config list. + + Returns: + Dict[str, Equation]: Equation(s) in dict. + """ + if cfg is None: + return None + cfg = copy.deepcopy(cfg) + eq_dict = misc.PrettyOrderedDict() + for _item in cfg: + eq_cls = next(iter(_item.keys())) + eq_cfg = _item[eq_cls] + eq_name = eq_cfg.pop("name", eq_cls) + eq_dict[eq_name] = eval(eq_cls)(**eq_cfg) + + logger.debug(str(eq_dict[eq_name])) + + return eq_dict diff --git a/examples/smc_reac/ppsci/equation/fpde/__init__.py b/examples/smc_reac/ppsci/equation/fpde/__init__.py new file mode 100644 index 0000000000..3e74ec56c7 --- /dev/null +++ b/examples/smc_reac/ppsci/equation/fpde/__init__.py @@ -0,0 +1,19 @@ +# 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. + +from ppsci.equation.fpde.fractional_poisson import FractionalPoisson + +__all__ = [ + "FractionalPoisson", +] diff --git a/examples/smc_reac/ppsci/equation/fpde/fractional_poisson.py b/examples/smc_reac/ppsci/equation/fpde/fractional_poisson.py new file mode 100644 index 0000000000..01b6fc929f --- /dev/null +++ b/examples/smc_reac/ppsci/equation/fpde/fractional_poisson.py @@ -0,0 +1,196 @@ +# 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. + +from __future__ import annotations + +import math +from typing import Tuple + +import numpy as np +import paddle +from paddle import sparse +from scipy import special + +from ppsci import geometry +from ppsci.equation.pde import PDE +from ppsci.utils import misc + + +class FractionalPoisson(PDE): + """(TODO)Docstring of this class will be refined in the future. + + Args: + alpha (float): Alpha. + geom (geometry.Geometry): Computation geometry. + resolution (Tuple[int, ...]): Resolution. + + Examples: + >>> import ppsci + >>> geom_disk = ppsci.geometry.Disk([0, 0], 1) + >>> ALPHA = 0.5 + >>> fpde = ppsci.equation.FractionalPoisson(ALPHA, geom_disk, [8, 100]) + """ + + dtype = paddle.get_default_dtype() + + def __init__( + self, alpha: float, geom: geometry.Geometry, resolution: Tuple[int, ...] + ): + super().__init__() + self.alpha = alpha + self.geom = geom + self.resolution = resolution + self._w_init = self._init_weights() + + def compute_fpde_func(out): + x = paddle.concat((out["x"], out["y"]), axis=1) + y = out["u"] + indices, values, shape = self.int_mat + int_mat = sparse.sparse_coo_tensor( + [[p[0] for p in indices], [p[1] for p in indices]], + values, + shape, + stop_gradient=False, + ) + lhs = sparse.matmul(int_mat, y) + lhs = lhs[:, 0] + lhs *= ( + special.gamma((1 - self.alpha) / 2) + * special.gamma((2 + self.alpha) / 2) + / (2 * np.pi**1.5) + ) + x = x[: paddle.numel(lhs)] + rhs = ( + 2**self.alpha + * special.gamma(2 + self.alpha / 2) + * special.gamma(1 + self.alpha / 2) + * (1 - (1 + self.alpha / 2) * paddle.sum(x**2, axis=1)) + ) + res = lhs - rhs + return res + + self.add_equation("fpde", compute_fpde_func) + + def _init_weights(self): + n = self._dynamic_dist2npts(self.geom.diam) + 1 + w = [1.0] + for j in range(1, n): + w.append(w[-1] * (j - 1 - self.alpha) / j) + return np.array(w, dtype=self.dtype) + + def get_x(self, x_f): + if hasattr(self, "train_x"): + return self.train_x + + self.x0 = x_f + if np.any(self.geom.on_boundary(self.x0)): + raise ValueError("x0 contains boundary points.") + + if self.geom.ndim == 1: + dirns, dirn_w = [-1, 1], [1, 1] + elif self.geom.ndim == 2: + gauss_x, gauss_w = np.polynomial.legendre.leggauss(self.resolution[0]) + gauss_x, gauss_w = gauss_x.astype(self.dtype), gauss_w.astype(self.dtype) + thetas = np.pi * gauss_x + np.pi + dirns = np.vstack((np.cos(thetas), np.sin(thetas))).T + dirn_w = np.pi * gauss_w + elif self.geom.ndim == 3: + gauss_x, gauss_w = np.polynomial.legendre.leggauss(max(self.resolution[:2])) + gauss_x, gauss_w = gauss_x.astype(self.dtype), gauss_w.astype(self.dtype) + thetas = (np.pi * gauss_x[: self.resolution[0]] + np.pi) / 2 + phis = np.pi * gauss_x[: self.resolution[1]] + np.pi + dirns, dirn_w = [], [] + for i in range(self.resolution[0]): + for j in range(self.resolution[1]): + dirns.append( + [ + np.sin(thetas[i]) * np.cos(phis[j]), + np.sin(thetas[i]) * np.sin(phis[j]), + np.cos(thetas[i]), + ] + ) + dirn_w.append(gauss_w[i] * gauss_w[j] * np.sin(thetas[i])) + dirn_w = np.pi**2 / 2 * np.array(dirn_w) + + x, self.w = [], [] + for x0i in self.x0: + xi = list( + map( + lambda dirn: self.background_points( + x0i, dirn, self._dynamic_dist2npts, 0 + ), + dirns, + ) + ) + wi = list( + map( + lambda i: dirn_w[i] + * np.linalg.norm(xi[i][1] - xi[i][0]) ** (-self.alpha) + * self.get_weight(len(xi[i]) - 1), + range(len(dirns)), + ) + ) + # first order + # xi, wi = zip(self.modify_first_order(xij, wij) for xij, wij in zip(xi, wi)) + xi, wi = zip(*map(self.modify_first_order, xi, wi)) + # second order + # xi, wi = zip(*map(self.modify_second_order, xi, wi)) + # third order + # xi, wi = zip(*map(self.modify_third_order, xi, wi)) + x.append(np.vstack(xi)) + self.w.append(np.hstack(wi)) + self.x = np.vstack([self.x0] + x) + self.int_mat = self._get_int_matrix(self.x0) + self.train_x = misc.convert_to_dict(self.x, ("x", "y")) + return self.train_x + + def get_weight(self, n): + return self._w_init[: n + 1] + + def background_points(self, x, dirn, dist2npt, shift): + dirn = dirn / np.linalg.norm(dirn) + dx = self.distance2boundary_unitdirn(x, -dirn) + n = max(dist2npt(dx), 1) + h = dx / n + pts = x - np.arange(-shift, n - shift + 1, dtype=self.dtype)[:, None] * h * dirn + return pts + + def distance2boundary_unitdirn(self, x, dirn): + # https://en.wikipedia.org/wiki/Line%E2%80%93sphere_intersection + xc = x - self.geom.center + xc = xc + ad = np.dot(xc, dirn) + return ( + -ad + (ad**2 - np.sum(xc * xc, axis=-1) + self.geom.radius**2) ** 0.5 + ).astype(self.dtype) + + def modify_first_order(self, x, w): + x = np.vstack(([2 * x[0] - x[1]], x[:-1])) + if not self.geom.is_inside(x[0:1])[0]: + return x[1:], w[1:] + return x, w + + def _dynamic_dist2npts(self, dx): + return int(math.ceil(self.resolution[-1] * dx)) + + def _get_int_matrix(self, x: np.ndarray) -> np.ndarray: + dense_shape = (x.shape[0], self.x.shape[0]) + indices, values = [], [] + beg = x.shape[0] + for i in range(x.shape[0]): + for _ in range(self.w[i].shape[0]): + indices.append([i, beg]) + beg += 1 + values = np.hstack((values, self.w[i])) + return indices, values.astype(self.dtype), dense_shape diff --git a/examples/smc_reac/ppsci/equation/ide/__init__.py b/examples/smc_reac/ppsci/equation/ide/__init__.py new file mode 100644 index 0000000000..4d4cab56cb --- /dev/null +++ b/examples/smc_reac/ppsci/equation/ide/__init__.py @@ -0,0 +1,19 @@ +# 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. + +from ppsci.equation.ide.volterra import Volterra + +__all__ = [ + "Volterra", +] diff --git a/examples/smc_reac/ppsci/equation/ide/volterra.py b/examples/smc_reac/ppsci/equation/ide/volterra.py new file mode 100644 index 0000000000..77fb6f3173 --- /dev/null +++ b/examples/smc_reac/ppsci/equation/ide/volterra.py @@ -0,0 +1,127 @@ +# 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. + +from __future__ import annotations + +from typing import Callable + +import numpy as np +import paddle + +from ppsci.equation.pde import PDE + + +class Volterra(PDE): + r"""A second kind of volterra integral equation with Gaussian quadrature algorithm. + + $$ + x(t) - f(t)=\int_a^t K(t, s) x(s) d s + $$ + + [Volterra integral equation](https://en.wikipedia.org/wiki/Volterra_integral_equation) + + [Gaussian quadrature](https://en.wikipedia.org/wiki/Gaussian_quadrature#Change_of_interval) + + Args: + bound (float): Lower bound `a` for Volterra integral equation. + num_points (int): Sampled points in integral interval. + quad_deg (int): Number of quadrature. + kernel_func (Callable): Kernel func `K(t,s)`. + func (Callable): `x(t) - f(t)` in Volterra integral equation. + + Examples: + >>> import ppsci + >>> import numpy as np + >>> vol_eq = ppsci.equation.Volterra( + ... 0, 12, 20, lambda t, s: np.exp(s - t), lambda out: out["u"], + ... ) + """ + + dtype = paddle.get_default_dtype() + + def __init__( + self, + bound: float, + num_points: int, + quad_deg: int, + kernel_func: Callable, + func: Callable, + ): + super().__init__() + self.bound = bound + self.num_points = num_points + self.quad_deg = quad_deg + self.kernel_func = kernel_func + self.func = func + + self.quad_x, self.quad_w = np.polynomial.legendre.leggauss(quad_deg) + self.quad_x = self.quad_x.astype(Volterra.dtype).reshape([-1, 1]) # [Q, 1] + self.quad_x = paddle.to_tensor(self.quad_x) # [Q, 1] + + self.quad_w = self.quad_w.astype(Volterra.dtype) # [Q, ] + + def compute_volterra_func(out): + x, u = out["x"], out["u"] + lhs = self.func(out) + + int_mat = paddle.to_tensor(self._get_int_matrix(x), stop_gradient=False) + rhs = paddle.mm(int_mat, u) # (N, 1) + + volterra = lhs[: len(rhs)] - rhs + return volterra + + self.add_equation("volterra", compute_volterra_func) + + def get_quad_points(self, t: paddle.Tensor) -> paddle.Tensor: + """Scale and transform quad_x from [-1, 1] to range [a, b]. + + reference: https://en.wikipedia.org/wiki/Gaussian_quadrature#Change_of_interval + + Args: + t (paddle.Tensor): Tensor array of upper bounds 't' for integral. + + Returns: + paddle.Tensor: Transformed points in desired range with shape of [N, Q]. + """ + a, b = self.bound, t + return ((b - a) / 2) @ self.quad_x.T + (b + a) / 2 + + def _get_quad_weights(self, t: float) -> np.ndarray: + """Scale weights to range according to given t and lower bound of integral. + + reference: https://en.wikipedia.org/wiki/Gaussian_quadrature#Change_of_interval + + Args: + t (float): Array of upper bound 't' for integral. + + Returns: + np.ndarray: Transformed weights in desired range with shape of [Q, ]. + """ + a, b = self.bound, t + return (b - a) / 2 * self.quad_w + + def _get_int_matrix(self, x: np.ndarray) -> np.ndarray: + int_mat = np.zeros( + (self.num_points, self.num_points + (self.num_points * self.quad_deg)), + dtype=Volterra.dtype, + ) + for i in range(self.num_points): + xi = float(x[i]) + beg = self.num_points + self.quad_deg * i + end = self.num_points + self.quad_deg * (i + 1) + K = np.ravel( + self.kernel_func(np.full((self.quad_deg, 1), xi), x[beg:end].numpy()) + ) + int_mat[i, beg:end] = self._get_quad_weights(xi) * K + return int_mat diff --git a/examples/smc_reac/ppsci/equation/pde/__init__.py b/examples/smc_reac/ppsci/equation/pde/__init__.py new file mode 100644 index 0000000000..0dbcea2a8f --- /dev/null +++ b/examples/smc_reac/ppsci/equation/pde/__init__.py @@ -0,0 +1,43 @@ +# 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. + +from ppsci.equation.pde.allen_cahn import AllenCahn +from ppsci.equation.pde.base import DETACH_FUNC_NAME +from ppsci.equation.pde.base import PDE +from ppsci.equation.pde.biharmonic import Biharmonic +from ppsci.equation.pde.heat_exchanger import HeatExchanger +from ppsci.equation.pde.helmholtz import Helmholtz +from ppsci.equation.pde.laplace import Laplace +from ppsci.equation.pde.linear_elasticity import LinearElasticity +from ppsci.equation.pde.navier_stokes import NavierStokes +from ppsci.equation.pde.nls_m_b import NLSMB +from ppsci.equation.pde.normal_dot_vec import NormalDotVec +from ppsci.equation.pde.poisson import Poisson +from ppsci.equation.pde.viv import Vibration + +__all__ = [ + "PDE", + "DETACH_FUNC_NAME", + "AllenCahn", + "Biharmonic", + "HeatExchanger", + "Helmholtz", + "Laplace", + "LinearElasticity", + "NavierStokes", + "NLSMB", + "NormalDotVec", + "Poisson", + "Vibration", +] diff --git a/examples/smc_reac/ppsci/equation/pde/allen_cahn.py b/examples/smc_reac/ppsci/equation/pde/allen_cahn.py new file mode 100644 index 0000000000..44e0ec899f --- /dev/null +++ b/examples/smc_reac/ppsci/equation/pde/allen_cahn.py @@ -0,0 +1,64 @@ +# Copyright (c) 2024 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. + +from __future__ import annotations + +from typing import Optional +from typing import Tuple + +from ppsci.autodiff import jacobian +from ppsci.equation.pde import base + + +class AllenCahn(base.PDE): + r"""Class for Allen-Cahn equation. + + $$ + \dfrac{\partial u}{\partial t} - \epsilon^2 \Delta u + 5u^3 - 5u = 0 + $$ + + Args: + eps (float): Represents the characteristicscale of interfacial width, + influencing the thickness and dynamics of phase boundaries. + detach_keys (Optional[Tuple[str, ...]]): Keys used for detach during computing. + Defaults to None. + + Examples: + >>> import ppsci + >>> pde = ppsci.equation.AllenCahn(eps=0.01) + """ + + def __init__( + self, + eps: float, + detach_keys: Optional[Tuple[str, ...]] = None, + ): + super().__init__() + self.detach_keys = detach_keys + self.eps = eps + # t, x = self.create_symbols("t x") + # invars = (t, x, ) + # u = self.create_function("u", invars) + # allen_cahn = u.diff(t) + 5 * u**3 - 5 * u - 0.0001 * u.diff(x, 2) + + # TODO: Pow(u,3) seems cause slightly larger L2 error than multiply(u*u*u) + def allen_cahn(out): + t, x = out["t"], out["x"] + u = out["u"] + u__t, u__x = jacobian(u, [t, x]) + u__x__x = jacobian(u__x, x) + + return u__t - (self.eps**2) * u__x__x + 5 * u * u * u - 5 * u + + self.add_equation("allen_cahn", allen_cahn) diff --git a/examples/smc_reac/ppsci/equation/pde/base.py b/examples/smc_reac/ppsci/equation/pde/base.py new file mode 100644 index 0000000000..41f54b3861 --- /dev/null +++ b/examples/smc_reac/ppsci/equation/pde/base.py @@ -0,0 +1,243 @@ +# 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. + +from __future__ import annotations + +from typing import Callable +from typing import Dict +from typing import List +from typing import Optional +from typing import Tuple +from typing import Union + +import paddle +import sympy as sp +from paddle import nn + +DETACH_FUNC_NAME = "detach" + + +class PDE: + """Base class for Partial Differential Equation.""" + + def __init__(self): + super().__init__() + self.equations: Dict[str, Union[Callable, sp.Basic]] = {} + # for PDE which has learnable parameter(s) + self.learnable_parameters = nn.ParameterList() + + self.detach_keys: Optional[Tuple[str, ...]] = None + + @staticmethod + def create_symbols( + symbol_str: str, + ) -> Union[sp.Symbol, Tuple[sp.Symbol, ...]]: + """Create symbolic variables. + + Args: + symbol_str (str): String contains symbols, such as "x", "x y z". + + Returns: + Union[sympy.Symbol, Tuple[sympy.Symbol, ...]]: Created symbol(s). + + Examples: + >>> import ppsci + >>> pde = ppsci.equation.PDE() + >>> symbol_x = pde.create_symbols('x') + >>> symbols_xyz = pde.create_symbols('x y z') + >>> print(symbol_x) + x + >>> print(symbols_xyz) + (x, y, z) + """ + return sp.symbols(symbol_str) + + def create_function(self, name: str, invars: Tuple[sp.Symbol, ...]) -> sp.Function: + """Create named function depending on given invars. + + Args: + name (str): Function name. such as "u", "v", and "f". + invars (Tuple[sympy.Symbol, ...]): List of independent variable of function. + + Returns: + sympy.Function: Named sympy function. + + Examples: + >>> import ppsci + >>> pde = ppsci.equation.PDE() + >>> x, y, z = pde.create_symbols('x y z') + >>> u = pde.create_function('u', (x, y)) + >>> f = pde.create_function('f', (x, y, z)) + >>> print(u) + u(x, y) + >>> print(f) + f(x, y, z) + """ + expr = sp.Function(name)(*invars) + + return expr + + def _apply_detach(self): + """ + Wrap detached sub_expr into detach(sub_expr) to prevent gradient back-propagation, only for those items specified in self.detach_keys. + + NOTE: This function is expected to be called after self.equations is ready in PDE.__init__. + + Examples: + >>> import ppsci + >>> ns = ppsci.equation.NavierStokes(1.0, 1.0, 2, False) + >>> print(ns) + NavierStokes + continuity: Derivative(u(x, y), x) + Derivative(v(x, y), y) + momentum_x: u(x, y)*Derivative(u(x, y), x) + v(x, y)*Derivative(u(x, y), y) + 1.0*Derivative(p(x, y), x) - 1.0*Derivative(u(x, y), (x, 2)) - 1.0*Derivative(u(x, y), (y, 2)) + momentum_y: u(x, y)*Derivative(v(x, y), x) + v(x, y)*Derivative(v(x, y), y) + 1.0*Derivative(p(x, y), y) - 1.0*Derivative(v(x, y), (x, 2)) - 1.0*Derivative(v(x, y), (y, 2)) + >>> detach_keys = ("u", "v__y") + >>> ns = ppsci.equation.NavierStokes(1.0, 1.0, 2, False, detach_keys=detach_keys) + >>> print(ns) + NavierStokes + continuity: detach(Derivative(v(x, y), y)) + Derivative(u(x, y), x) + momentum_x: detach(u(x, y))*Derivative(u(x, y), x) + v(x, y)*Derivative(u(x, y), y) + 1.0*Derivative(p(x, y), x) - 1.0*Derivative(u(x, y), (x, 2)) - 1.0*Derivative(u(x, y), (y, 2)) + momentum_y: detach(u(x, y))*Derivative(v(x, y), x) + detach(Derivative(v(x, y), y))*v(x, y) + 1.0*Derivative(p(x, y), y) - 1.0*Derivative(v(x, y), (x, 2)) - 1.0*Derivative(v(x, y), (y, 2)) + """ + if self.detach_keys is None: + return + + from copy import deepcopy + + from sympy.core.traversal import postorder_traversal + + from ppsci.utils.symbolic import _cvt_to_key + + for name, expr in self.equations.items(): + if not isinstance(expr, sp.Basic): + continue + # only process sympy expression + expr_ = deepcopy(expr) + for item in postorder_traversal(expr): + if _cvt_to_key(item) in self.detach_keys: + # inplace all related sub_expr into detach(sub_expr) + expr_ = expr_.replace(item, sp.Function(DETACH_FUNC_NAME)(item)) + + # remove all detach wrapper for more-than-once wrapped items to prevent duplicated wrapping + expr_ = expr_.replace( + sp.Function(DETACH_FUNC_NAME)( + sp.Function(DETACH_FUNC_NAME)(item) + ), + sp.Function(DETACH_FUNC_NAME)(item), + ) + + # remove unccessary detach wrapping for the first arg of Derivative + for item_ in list(postorder_traversal(expr_)): + if isinstance(item_, sp.Derivative): + if item_.args[0].name == DETACH_FUNC_NAME: + expr_ = expr_.replace( + item_, + sp.Derivative( + item_.args[0].args[0], *item_.args[1:] + ), + ) + + self.equations[name] = expr_ + + def add_equation(self, name: str, equation: Callable): + """Add an equation. + + Args: + name (str): Name of equation + equation (Callable): Computation function for equation. + + Examples: + >>> import ppsci + >>> import sympy + >>> pde = ppsci.equation.PDE() + >>> x, y = pde.create_symbols('x y') + >>> u = x**2 + y**2 + >>> equation = sympy.diff(u, x) + sympy.diff(u, y) + >>> pde.add_equation('linear_pde', equation) + >>> print(pde) + PDE + linear_pde: 2*x + 2*y + """ + self.equations.update({name: equation}) + + def parameters(self) -> List[paddle.Tensor]: + """Return learnable parameters contained in PDE. + + Returns: + List[Tensor]: A list of learnable parameters. + + Examples: + >>> import ppsci + >>> pde = ppsci.equation.Vibration(2, -4, 0) + >>> print(pde.parameters()) + [Parameter containing: + Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=False, + -4.), Parameter containing: + Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=False, + 0.)] + """ + return self.learnable_parameters.parameters() + + def state_dict(self) -> Dict[str, paddle.Tensor]: + """Return named learnable parameters in dict. + + Returns: + Dict[str, Tensor]: A dict of states(str) and learnable parameters(Tensor). + + Examples: + >>> import ppsci + >>> pde = ppsci.equation.Vibration(2, -4, 0) + >>> print(pde.state_dict()) + OrderedDict([('0', Parameter containing: + Tensor(shape=[], dtype=float64, place=Place(gpu:0), stop_gradient=False, + -4.)), ('1', Parameter containing: + Tensor(shape=[], dtype=float64, place=Place(gpu:0), stop_gradient=False, + 0.))]) + """ + return self.learnable_parameters.state_dict() + + def set_state_dict( + self, state_dict: Dict[str, paddle.Tensor] + ) -> Tuple[List[str], List[str]]: + """Set state dict from dict. + + Args: + state_dict (Dict[str, paddle.Tensor]): The state dict to be set. + + Returns: + Tuple[List[str], List[str]]: List of missing_keys and unexpected_keys. + Expected to be two empty tuples mostly. + + Examples: + >>> import paddle + >>> import ppsci + >>> paddle.set_default_dtype("float64") + >>> pde = ppsci.equation.Vibration(2, -4, 0) + >>> state = pde.state_dict() + >>> state['0'] = paddle.to_tensor(-3.1) + >>> pde.set_state_dict(state) + ([], []) + >>> print(state) + OrderedDict([('0', Tensor(shape=[], dtype=float64, place=Place(gpu:0), stop_gradient=True, + -3.10000000)), ('1', Parameter containing: + Tensor(shape=[], dtype=float64, place=Place(gpu:0), stop_gradient=False, + 0.))]) + """ + return self.learnable_parameters.set_state_dict(state_dict) + + def __str__(self): + return "\n".join( + [self.__class__.__name__] + + [f" {name}: {eq}" for name, eq in self.equations.items()] + ) diff --git a/examples/smc_reac/ppsci/equation/pde/biharmonic.py b/examples/smc_reac/ppsci/equation/pde/biharmonic.py new file mode 100644 index 0000000000..933888ac60 --- /dev/null +++ b/examples/smc_reac/ppsci/equation/pde/biharmonic.py @@ -0,0 +1,74 @@ +# 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. + +from __future__ import annotations + +from typing import Optional +from typing import Tuple +from typing import Union + +import sympy + +from ppsci.equation.pde import base + + +class Biharmonic(base.PDE): + r"""Class for biharmonic equation with supporting special load. + + $$ + \nabla^4 \varphi = \dfrac{q}{D} + $$ + + Args: + dim (int): Dimension of equation. + q (Union[float, str, sympy.Basic]): Load. + D (Union[float, str]): Rigidity. + detach_keys (Optional[Tuple[str, ...]]): Keys used for detach during computing. + Defaults to None. + + Examples: + >>> import ppsci + >>> pde = ppsci.equation.Biharmonic(2, -1.0, 1.0) + """ + + def __init__( + self, + dim: int, + q: Union[float, str, sympy.Basic], + D: Union[float, str], + detach_keys: Optional[Tuple[str, ...]] = None, + ): + super().__init__() + self.detach_keys = detach_keys + + invars = self.create_symbols("x y z")[:dim] + u = self.create_function("u", invars) + + if isinstance(q, str): + q = self.create_function("q", invars) + if isinstance(D, str): + D = self.create_function("D", invars) + + self.dim = dim + self.q = q + self.D = D + + biharmonic = -self.q / self.D + for invar_i in invars: + for invar_j in invars: + biharmonic += u.diff(invar_i, 2).diff(invar_j, 2) + + self.add_equation("biharmonic", biharmonic) + + self._apply_detach() diff --git a/examples/smc_reac/ppsci/equation/pde/heat_exchanger.py b/examples/smc_reac/ppsci/equation/pde/heat_exchanger.py new file mode 100644 index 0000000000..c2e0107ff3 --- /dev/null +++ b/examples/smc_reac/ppsci/equation/pde/heat_exchanger.py @@ -0,0 +1,94 @@ +# 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. + +from __future__ import annotations + +from typing import Union + +from ppsci.equation.pde import base + + +class HeatExchanger(base.PDE): + r"""Class for heat exchanger equation. + + $$ + \begin{aligned} + & L\left(\frac{q_m c_p}{v}\right)_{\mathrm{c}} \frac{\partial T_{\mathrm{c}}}{\partial \tau}-L\left(q_m c_p\right)_{\mathrm{c}} \frac{\partial T_{\mathrm{c}}}{\partial x}=\left(\eta_{\mathrm{o}} \alpha A\right)_{\mathrm{c}}\left(T_{\mathrm{w}}-T_{\mathrm{c}}\right), \\ + & L\left(\frac{q_m c_p}{v}\right)_{\mathrm{h}} \frac{\partial T_{\mathrm{h}}}{\partial \tau}+L\left(q_m c_p\right)_{\mathrm{h}} \frac{\partial T_{\mathrm{h}}}{\partial x}=\left(\eta_{\mathrm{o}} \alpha A\right)_{\mathrm{h}}\left(T_{\mathrm{w}}-T_{\mathrm{h}}\right), \\ + & \left(M c_p\right)_{\mathrm{w}} \frac{\partial T_{\mathrm{w}}}{\partial \tau}=\left(\eta_{\mathrm{o}} \alpha A\right)_{\mathrm{h}}\left(T_{\mathrm{h}}-T_{\mathrm{w}}\right)+\left(\eta_{\mathrm{o}} \alpha A\right)_{\mathrm{c}}\left(T_{\mathrm{c}}-T_{\mathrm{w}}\right). + \end{aligned} + $$ + + where: + + - $T$ is temperature, + - $q_m$ is mass flow rate, + - $c_p$ represents specific heat capacity, + - $v$ denotes flow velocity, + - $L$ stands for flow length, + - $\eta_{\mathrm{o}}$ signifies fin surface efficiency, + - $\alpha$ stands for heat transfer coefficient, + - $A$ indicates heat transfer area, + - $M$ represents the mass of the heat transfer structure, + - $\tau$ correspond to time, + - $x$ correspond flow direction, + - Subscripts $\mathrm{h}$, $\mathrm{c}$, and $\mathrm{w}$ denote the hot fluid side, cold fluid side, and heat transfer wall, respectively. + + Args: + alpha_h: $\frac{(\eta_o\alpha A)_h}{L(c_p)_h}$ + alpha_c: $\frac{(\eta_o\alpha A)_c}{L(c_p)_c}$ + v_h: $v_h$ + v_c: $v_c$ + w_h: $\frac{(\eta_o\alpha A)_h}{M(c_p)_w}$ + w_c: $\frac{(\eta_o\alpha A)_c}{M(c_p)_w}$ + + Examples: + >>> import ppsci + >>> pde = ppsci.equation.HeatExchanger(1.0,1.0,1.0,1.0,1.0,1.0) + """ + + def __init__( + self, + alpha_h: Union[float, str], + alpha_c: Union[float, str], + v_h: Union[float, str], + v_c: Union[float, str], + w_h: Union[float, str], + w_c: Union[float, str], + ): + super().__init__() + x, t, qm_h, qm_c = self.create_symbols("x t qm_h qm_c") + + T_h = self.create_function("T_h", (x, t, qm_h)) + T_c = self.create_function("T_c", (x, t, qm_c)) + T_w = self.create_function("T_w", (x, t)) + + T_h_x = T_h.diff(x) + T_h_t = T_h.diff(t) + T_c_x = T_c.diff(x) + T_c_t = T_c.diff(t) + T_w_t = T_w.diff(t) + + beta_h = (alpha_h * v_h) / qm_h + beta_c = (alpha_c * v_c) / qm_c + + heat_boundary = T_h_t + v_h * T_h_x - beta_h * (T_w - T_h) + cold_boundary = T_c_t - v_c * T_c_x - beta_c * (T_w - T_c) + wall = T_w_t - w_h * (T_h - T_w) - w_c * (T_c - T_w) + + self.add_equation("heat_boundary", heat_boundary) + self.add_equation("cold_boundary", cold_boundary) + self.add_equation("wall", wall) + + self._apply_detach() diff --git a/examples/smc_reac/ppsci/equation/pde/helmholtz.py b/examples/smc_reac/ppsci/equation/pde/helmholtz.py new file mode 100644 index 0000000000..e71fdbe983 --- /dev/null +++ b/examples/smc_reac/ppsci/equation/pde/helmholtz.py @@ -0,0 +1,119 @@ +# 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. + +from __future__ import annotations + +from typing import Callable +from typing import Dict +from typing import Optional +from typing import Tuple + +import paddle + +from ppsci.equation.pde import base + + +def hvp_revrev(f: Callable, primals: Tuple[paddle.Tensor, ...]) -> paddle.Tensor: + """Compute the Hessian vector product of f with respect to primals using + double backward trick in reverse mode AD. + + Args: + f (Callable): Function to compute HVP. + primals (Tuple[paddle.Tensor, ...]): Input tensors. + + Returns: + paddle.Tensor: Hessian vector product of f with respect to primals. + """ + # TODO: Merge this option into ppsci.autodiff.ad + g = lambda primals: paddle.incubate.autograd.jvp(f, primals)[1] + tangents_out = paddle.incubate.autograd.jvp(g, primals)[1] + return tangents_out[0] + + +class Helmholtz(base.PDE): + r"""Class for helmholtz equation. + + $$ + \nabla^2 u + k^2 u = f + $$ + + $$ + \text{where } f \text{ is the source term}. + $$ + + Args: + dim (int): Dimension of equation. + k (float): The wave number, which is a parameter that affects the frequency of the solution. + detach_keys (Optional[Tuple[str, ...]]): Keys used for detach during computing. + Defaults to None. + + Examples: + >>> import ppsci + >>> model = ppsci.arch.MLP(("x", "y"), ("u",), 2, 32) + >>> pde = ppsci.equation.Helmholtz(2, -1.0, model) + """ + + def __init__( + self, + dim: int, + k: float, + model: paddle.nn.Layer, + detach_keys: Optional[Tuple[str, ...]] = None, + ): + super().__init__() + self.dim = dim + self.k = k + self.detach_keys = detach_keys + + invars = self.create_symbols("x y z")[:dim] + + # TODO: This is a hack, should be simplified in the future + self.model = model + + def helmholtz(data_dict: Dict[str, paddle.Tensor]) -> paddle.Tensor: + xs = tuple(data_dict[invar.name] for invar in invars) + + # TODO: Hard code here, for hvp_revrev requires tuple input(s) but not dict + if self.dim == 1: + u__x__x = hvp_revrev(lambda x_: self.model.forward_tensor(x_), (xs[0],)) + out = (self.k**2) * data_dict["u"] + u__x__x + elif self.dim == 2: + u__x__x = hvp_revrev( + lambda x_: self.model.forward_tensor(x_, xs[1]), (xs[0],) + ) + u__y__y = hvp_revrev( + lambda y_: self.model.forward_tensor(xs[0], y_), (xs[1],) + ) + out = (self.k**2) * data_dict["u"] + u__x__x + u__y__y + elif self.dim == 3: + u__x__x = hvp_revrev( + lambda x_: self.model.forward_tensor(x_, xs[1], xs[2]), (xs[0],) + ) + u__y__y = hvp_revrev( + lambda y_: self.model.forward_tensor(xs[0], y_, xs[2]), (xs[1],) + ) + u__z__z = hvp_revrev( + lambda z_: self.model.forward_tensor(xs[0], xs[1], z_), (xs[2],) + ) + out = (self.k**2) * data_dict["u"] + u__x__x + u__y__y + u__z__z + else: + raise NotImplementedError( + f"dim should be less or equal to 3, but got {self.dim}." + ) + + return out + + self.add_equation("helmholtz", helmholtz) + + self._apply_detach() diff --git a/examples/smc_reac/ppsci/equation/pde/laplace.py b/examples/smc_reac/ppsci/equation/pde/laplace.py new file mode 100644 index 0000000000..b99d7c8d9a --- /dev/null +++ b/examples/smc_reac/ppsci/equation/pde/laplace.py @@ -0,0 +1,55 @@ +# 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. + +from __future__ import annotations + +from typing import Optional +from typing import Tuple + +from ppsci.equation.pde import base + + +class Laplace(base.PDE): + r"""Class for laplace equation. + + $$ + \nabla^2 \varphi = 0 + $$ + + Args: + dim (int): Dimension of equation. + detach_keys (Optional[Tuple[str, ...]]): Keys used for detach during computing. + Defaults to None. + + Examples: + >>> import ppsci + >>> pde = ppsci.equation.Laplace(2) + """ + + def __init__(self, dim: int, detach_keys: Optional[Tuple[str, ...]] = None): + super().__init__() + self.detach_keys = detach_keys + + invars = self.create_symbols("x y z")[:dim] + u = self.create_function("u", invars) + + self.dim = dim + + laplace = 0 + for invar in invars: + laplace += u.diff(invar, 2) + + self.add_equation("laplace", laplace) + + self._apply_detach() diff --git a/examples/smc_reac/ppsci/equation/pde/linear_elasticity.py b/examples/smc_reac/ppsci/equation/pde/linear_elasticity.py new file mode 100644 index 0000000000..289d924899 --- /dev/null +++ b/examples/smc_reac/ppsci/equation/pde/linear_elasticity.py @@ -0,0 +1,184 @@ +# 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. + +from __future__ import annotations + +from typing import Optional +from typing import Tuple +from typing import Union + +import sympy as sp + +from ppsci.equation.pde import base + + +class LinearElasticity(base.PDE): + r"""Linear elasticity equations. + + Use either (E, nu) or (lambda_, mu) to define the material properties. + + $$ + \begin{cases} + stress\_disp_{xx} = \lambda(\dfrac{\partial u}{\partial x} + \dfrac{\partial v}{\partial y} + \dfrac{\partial w}{\partial z}) + 2\mu \dfrac{\partial u}{\partial x} - \sigma_{xx} \\ + stress\_disp_{yy} = \lambda(\dfrac{\partial u}{\partial x} + \dfrac{\partial v}{\partial y} + \dfrac{\partial w}{\partial z}) + 2\mu \dfrac{\partial v}{\partial y} - \sigma_{yy} \\ + stress\_disp_{zz} = \lambda(\dfrac{\partial u}{\partial x} + \dfrac{\partial v}{\partial y} + \dfrac{\partial w}{\partial z}) + 2\mu \dfrac{\partial w}{\partial z} - \sigma_{zz} \\ + stress\_disp_{xy} = \mu(\dfrac{\partial u}{\partial y} + \dfrac{\partial v}{\partial x}) - \sigma_{xy} \\ + stress\_disp_{xz} = \mu(\dfrac{\partial u}{\partial z} + \dfrac{\partial w}{\partial x}) - \sigma_{xz} \\ + stress\_disp_{yz} = \mu(\dfrac{\partial v}{\partial z} + \dfrac{\partial w}{\partial y}) - \sigma_{yz} \\ + equilibrium_{x} = \rho \dfrac{\partial^2 u}{\partial t^2} - (\dfrac{\partial \sigma_{xx}}{\partial x} + \dfrac{\partial \sigma_{xy}}{\partial y} + \dfrac{\partial \sigma_{xz}}{\partial z}) \\ + equilibrium_{y} = \rho \dfrac{\partial^2 u}{\partial t^2} - (\dfrac{\partial \sigma_{xy}}{\partial x} + \dfrac{\partial \sigma_{yy}}{\partial y} + \dfrac{\partial \sigma_{yz}}{\partial z}) \\ + equilibrium_{z} = \rho \dfrac{\partial^2 u}{\partial t^2} - (\dfrac{\partial \sigma_{xz}}{\partial x} + \dfrac{\partial \sigma_{yz}}{\partial y} + \dfrac{\partial \sigma_{zz}}{\partial z}) \\ + \end{cases} + $$ + + Args: + E (Optional[Union[float, str]]): The Young's modulus. Defaults to None. + nu (Optional[Union[float, str]]): The Poisson's ratio. Defaults to None. + lambda_ (Optional[Union[float, str]]): Lamé's first parameter. Defaults to None. + mu (Optional[Union[float, str]]): Lamé's second parameter (shear modulus). Defaults to None. + rho (Union[float, str], optional): Mass density. Defaults to 1. + dim (int, optional): Dimension of the linear elasticity (2 or 3). Defaults to 3. + time (bool, optional): Whether contains time data. Defaults to False. + detach_keys (Optional[Tuple[str, ...]]): Keys used for detach during computing. + Defaults to None. + + Examples: + >>> import ppsci + >>> pde = ppsci.equation.LinearElasticity( + ... E=None, nu=None, lambda_=1e4, mu=100, dim=3 + ... ) + """ + + def __init__( + self, + E: Optional[Union[float, str]] = None, + nu: Optional[Union[float, str]] = None, + lambda_: Optional[Union[float, str]] = None, + mu: Optional[Union[float, str]] = None, + rho: Union[float, str] = 1, + dim: int = 3, + time: bool = False, + detach_keys: Optional[Tuple[str, ...]] = None, + ): + super().__init__() + self.detach_keys = detach_keys + self.dim = dim + self.time = time + + t, x, y, z = self.create_symbols("t x y z") + normal_x, normal_y, normal_z = self.create_symbols("normal_x normal_y normal_z") + invars = (x, y) + if time: + invars = (t,) + invars + if self.dim == 3: + invars += (z,) + + u = self.create_function("u", invars) + v = self.create_function("v", invars) + w = self.create_function("w", invars) if dim == 3 else sp.Number(0) + + sigma_xx = self.create_function("sigma_xx", invars) + sigma_yy = self.create_function("sigma_yy", invars) + sigma_xy = self.create_function("sigma_xy", invars) + sigma_zz = ( + self.create_function("sigma_zz", invars) if dim == 3 else sp.Number(0) + ) + sigma_xz = ( + self.create_function("sigma_xz", invars) if dim == 3 else sp.Number(0) + ) + sigma_yz = ( + self.create_function("sigma_yz", invars) if dim == 3 else sp.Number(0) + ) + + # compute lambda and mu + if lambda_ is None: + if isinstance(nu, str): + nu = self.create_function(nu, invars) + if isinstance(E, str): + E = self.create_function(E, invars) + lambda_ = nu * E / ((1 + nu) * (1 - 2 * nu)) + mu = E / (2 * (1 + nu)) + else: + if isinstance(lambda_, str): + lambda_ = self.create_function(lambda_, invars) + if isinstance(mu, str): + mu = self.create_function(mu, invars) + + if isinstance(rho, str): + rho = self.create_function(rho, invars) + + self.E = E + self.nu = nu + self.lambda_ = lambda_ + self.mu = mu + self.rho = rho + + # compute stress equations + stress_disp_xx = ( + lambda_ * (u.diff(x) + v.diff(y) + w.diff(z)) + + 2 * mu * u.diff(x) + - sigma_xx + ) + stress_disp_yy = ( + lambda_ * (u.diff(x) + v.diff(y) + w.diff(z)) + + 2 * mu * v.diff(y) + - sigma_yy + ) + stress_disp_zz = ( + lambda_ * (u.diff(x) + v.diff(y) + w.diff(z)) + + 2 * mu * w.diff(z) + - sigma_zz + ) + stress_disp_xy = mu * (u.diff(y) + v.diff(x)) - sigma_xy + stress_disp_xz = mu * (u.diff(z) + w.diff(x)) - sigma_xz + stress_disp_yz = mu * (v.diff(z) + w.diff(y)) - sigma_yz + + # compute equilibrium equations + equilibrium_x = rho * ((u.diff(t)).diff(t)) - ( + sigma_xx.diff(x) + sigma_xy.diff(y) + sigma_xz.diff(z) + ) + equilibrium_y = rho * ((v.diff(t)).diff(t)) - ( + sigma_xy.diff(x) + sigma_yy.diff(y) + sigma_yz.diff(z) + ) + equilibrium_z = rho * ((w.diff(t)).diff(t)) - ( + sigma_xz.diff(x) + sigma_yz.diff(y) + sigma_zz.diff(z) + ) + + # compute traction equations + traction_x = normal_x * sigma_xx + normal_y * sigma_xy + normal_z * sigma_xz + traction_y = normal_x * sigma_xy + normal_y * sigma_yy + normal_z * sigma_yz + traction_z = normal_x * sigma_xz + normal_y * sigma_yz + normal_z * sigma_zz + + # add stress equations + self.add_equation("stress_disp_xx", stress_disp_xx) + self.add_equation("stress_disp_yy", stress_disp_yy) + self.add_equation("stress_disp_xy", stress_disp_xy) + if self.dim == 3: + self.add_equation("stress_disp_zz", stress_disp_zz) + self.add_equation("stress_disp_xz", stress_disp_xz) + self.add_equation("stress_disp_yz", stress_disp_yz) + + # add equilibrium equations + self.add_equation("equilibrium_x", equilibrium_x) + self.add_equation("equilibrium_y", equilibrium_y) + if self.dim == 3: + self.add_equation("equilibrium_z", equilibrium_z) + + # add traction equations + self.add_equation("traction_x", traction_x) + self.add_equation("traction_y", traction_y) + if self.dim == 3: + self.add_equation("traction_z", traction_z) + + self._apply_detach() diff --git a/examples/smc_reac/ppsci/equation/pde/navier_stokes.py b/examples/smc_reac/ppsci/equation/pde/navier_stokes.py new file mode 100644 index 0000000000..c0d3d193a2 --- /dev/null +++ b/examples/smc_reac/ppsci/equation/pde/navier_stokes.py @@ -0,0 +1,151 @@ +# 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. + +from __future__ import annotations + +from typing import Optional +from typing import Tuple +from typing import Union + +import sympy as sp +from sympy.parsing import sympy_parser as sp_parser + +from ppsci.equation.pde import base + + +class NavierStokes(base.PDE): + r"""Class for navier-stokes equation. + + $$ + \begin{cases} + \dfrac{\partial u}{\partial x} + \dfrac{\partial v}{\partial y} + \dfrac{\partial w}{\partial z} = 0 \\ + \dfrac{\partial u}{\partial t} + u\dfrac{\partial u}{\partial x} + v\dfrac{\partial u}{\partial y} + w\dfrac{\partial u}{\partial z} = + - \dfrac{1}{\rho}\dfrac{\partial p}{\partial x} + + \nu( + \dfrac{\partial ^2 u}{\partial x ^2} + + \dfrac{\partial ^2 u}{\partial y ^2} + + \dfrac{\partial ^2 u}{\partial z ^2} + ) \\ + \dfrac{\partial v}{\partial t} + u\dfrac{\partial v}{\partial x} + v\dfrac{\partial v}{\partial y} + w\dfrac{\partial v}{\partial z} = + - \dfrac{1}{\rho}\dfrac{\partial p}{\partial y} + + \nu( + \dfrac{\partial ^2 v}{\partial x ^2} + + \dfrac{\partial ^2 v}{\partial y ^2} + + \dfrac{\partial ^2 v}{\partial z ^2} + ) \\ + \dfrac{\partial w}{\partial t} + u\dfrac{\partial w}{\partial x} + v\dfrac{\partial w}{\partial y} + w\dfrac{\partial w}{\partial z} = + - \dfrac{1}{\rho}\dfrac{\partial p}{\partial z} + + \nu( + \dfrac{\partial ^2 w}{\partial x ^2} + + \dfrac{\partial ^2 w}{\partial y ^2} + + \dfrac{\partial ^2 w}{\partial z ^2} + ) \\ + \end{cases} + $$ + + Args: + nu (Union[float, str]): Dynamic viscosity. + rho (Union[float, str]): Density. + dim (int): Dimension of equation. + time (bool): Whether the equation is time-dependent. + detach_keys (Optional[Tuple[str, ...]]): Keys used for detach during computing. + Defaults to None. + + Examples: + >>> import ppsci + >>> pde = ppsci.equation.NavierStokes(0.1, 1.0, 3, False) + """ + + def __init__( + self, + nu: Union[float, str], + rho: Union[float, str], + dim: int, + time: bool, + detach_keys: Optional[Tuple[str, ...]] = None, + ): + super().__init__() + self.detach_keys = detach_keys + self.dim = dim + self.time = time + + t, x, y, z = self.create_symbols("t x y z") + invars = (x, y) + if time: + invars = (t,) + invars + if dim == 3: + invars += (z,) + + if isinstance(nu, str): + nu = sp_parser.parse_expr(nu) + if isinstance(nu, sp.Symbol): + invars += (nu,) + + if isinstance(rho, str): + rho = sp_parser.parse_expr(rho) + if isinstance(rho, sp.Symbol): + invars += (rho,) + + self.nu = nu + self.rho = rho + + u = self.create_function("u", invars) + v = self.create_function("v", invars) + w = self.create_function("w", invars) if dim == 3 else sp.Number(0) + p = self.create_function("p", invars) + + continuity = u.diff(x) + v.diff(y) + w.diff(z) + momentum_x = ( + u.diff(t) + + u * u.diff(x) + + v * u.diff(y) + + w * u.diff(z) + - ( + (nu * u.diff(x)).diff(x) + + (nu * u.diff(y)).diff(y) + + (nu * u.diff(z)).diff(z) + ) + + 1 / rho * p.diff(x) + ) + momentum_y = ( + v.diff(t) + + u * v.diff(x) + + v * v.diff(y) + + w * v.diff(z) + - ( + (nu * v.diff(x)).diff(x) + + (nu * v.diff(y)).diff(y) + + (nu * v.diff(z)).diff(z) + ) + + 1 / rho * p.diff(y) + ) + momentum_z = ( + w.diff(t) + + u * w.diff(x) + + v * w.diff(y) + + w * w.diff(z) + - ( + (nu * w.diff(x)).diff(x) + + (nu * w.diff(y)).diff(y) + + (nu * w.diff(z)).diff(z) + ) + + 1 / rho * p.diff(z) + ) + self.add_equation("continuity", continuity) + self.add_equation("momentum_x", momentum_x) + self.add_equation("momentum_y", momentum_y) + if self.dim == 3: + self.add_equation("momentum_z", momentum_z) + + self._apply_detach() diff --git a/examples/smc_reac/ppsci/equation/pde/nls_m_b.py b/examples/smc_reac/ppsci/equation/pde/nls_m_b.py new file mode 100644 index 0000000000..3db2984268 --- /dev/null +++ b/examples/smc_reac/ppsci/equation/pde/nls_m_b.py @@ -0,0 +1,101 @@ +# 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. + +from __future__ import annotations + +from typing import Optional +from typing import Tuple +from typing import Union + +from ppsci.equation.pde import base + + +class NLSMB(base.PDE): + r"""Class for nonlinear Schrodinger-Maxwell-Bloch equation. + + $$ + \begin{cases} + \dfrac{\partial E}{\partial x} = i \alpha_1 \dfrac{\partial^2 E}{\partial t ^2} - i \alpha_2 |E|^2 E+2 p \\ + \dfrac{\partial p}{\partial t} = 2 i \omega_0 p+2 E \eta \\ + \dfrac{\partial \eta}{\partial t} = -(E p^* + E^* p) + \end{cases} + $$ + + Args: + alpha_1 (Union[float, str]): Group velocity dispersion. + alpha_2 (Union[float, str]): Kerr nonlinearity. + omega_0 (Union[float, str]): The offset of resonance frequency. + time (bool): Whether the equation is time-dependent. + detach_keys (Optional[Tuple[str, ...]]): Keys used for detach during computing. + Defaults to None. + + Examples: + >>> import ppsci + >>> pde = ppsci.equation.NLSMB(0.5, -1.0, 0.5, True) + """ + + def __init__( + self, + alpha_1: Union[float, str], + alpha_2: Union[float, str], + omega_0: Union[float, str], + time: bool, + detach_keys: Optional[Tuple[str, ...]] = None, + ): + super().__init__() + self.detach_keys = detach_keys + self.time = time + + t, x = self.create_symbols("t x") + invars = (x,) + if time: + invars = (t,) + invars + + self.alpha_1 = alpha_1 + self.alpha_2 = alpha_2 + self.omega_0 = omega_0 + + Eu = self.create_function("Eu", invars) + Ev = self.create_function("Ev", invars) + pu = self.create_function("pu", invars) + pv = self.create_function("pv", invars) + eta = self.create_function("eta", invars) + + pu_t = pu.diff(t) + pv_t = pv.diff(t) + eta_t = eta.diff(t) + + Eu_x = Eu.diff(x) + Ev_x = Ev.diff(x) + + Eu_tt = Eu.diff(t).diff(t) + Ev_tt = Ev.diff(t).diff(t) + + Schrodinger_1 = ( + alpha_1 * Eu_tt - alpha_2 * Eu * (Eu**2 + Ev**2) + 2 * pv - Ev_x + ) + Schrodinger_2 = ( + alpha_1 * Ev_tt - alpha_2 * Ev * (Eu**2 + Ev**2) - 2 * pu + Eu_x + ) + Maxwell_1 = 2 * Ev * eta - pv_t + 2 * pu * omega_0 + Maxwell_2 = -2 * Eu * eta + pu_t + 2 * pv * omega_0 + Bloch = 2 * pv * Ev + 2 * pu * Eu + eta_t + + self.add_equation("Schrodinger_1", Schrodinger_1) + self.add_equation("Schrodinger_2", Schrodinger_2) + self.add_equation("Maxwell_1", Maxwell_1) + self.add_equation("Maxwell_2", Maxwell_2) + self.add_equation("Bloch", Bloch) + + self._apply_detach() diff --git a/examples/smc_reac/ppsci/equation/pde/normal_dot_vec.py b/examples/smc_reac/ppsci/equation/pde/normal_dot_vec.py new file mode 100644 index 0000000000..a6f3942eeb --- /dev/null +++ b/examples/smc_reac/ppsci/equation/pde/normal_dot_vec.py @@ -0,0 +1,59 @@ +# 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. + +from __future__ import annotations + +from typing import Optional +from typing import Tuple + +from ppsci.equation.pde import base + + +class NormalDotVec(base.PDE): + r"""Normal Dot Vector. + + $$ + \mathbf{n} \cdot \mathbf{v} = 0 + $$ + + Args: + vec_keys (Tuple[str, ...]): Keys for vectors, such as ("u", "v", "w") for + velocity vector. + detach_keys (Optional[Tuple[str, ...]]): Keys used for detach during computing. + Defaults to None. + + Examples: + >>> import ppsci + >>> pde = ppsci.equation.NormalDotVec(("u", "v", "w")) + """ + + def __init__( + self, vec_keys: Tuple[str, ...], detach_keys: Optional[Tuple[str, ...]] = None + ): + super().__init__() + self.detach_keys = detach_keys + if not vec_keys: + raise ValueError(f"len(vec_keys)({len(vec_keys)}) should be larger than 0.") + + self.vec_keys = vec_keys + vec_vars = self.create_symbols(" ".join(vec_keys)) + normals = self.create_symbols("normal_x normal_y normal_z") + + normal_dot_vec = 0 + for (normal, vec) in zip(normals, vec_vars): + normal_dot_vec += normal * vec + + self.add_equation("normal_dot_vec", normal_dot_vec) + + self._apply_detach() diff --git a/examples/smc_reac/ppsci/equation/pde/poisson.py b/examples/smc_reac/ppsci/equation/pde/poisson.py new file mode 100644 index 0000000000..4f9551a23a --- /dev/null +++ b/examples/smc_reac/ppsci/equation/pde/poisson.py @@ -0,0 +1,53 @@ +# 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. + +from __future__ import annotations + +from typing import Optional +from typing import Tuple + +from ppsci.equation.pde import base + + +class Poisson(base.PDE): + r"""Class for poisson equation. + + $$ + \nabla^2 \varphi = C + $$ + + Args: + dim (int): Dimension of equation. + detach_keys (Optional[Tuple[str, ...]]): Keys used for detach during computing. + Defaults to None. + + Examples: + >>> import ppsci + >>> pde = ppsci.equation.Poisson(2) + """ + + def __init__(self, dim: int, detach_keys: Optional[Tuple[str, ...]] = None): + super().__init__() + self.detach_keys = detach_keys + invars = self.create_symbols("x y z")[:dim] + p = self.create_function("p", invars) + self.dim = dim + + poisson = 0 + for invar in invars: + poisson += p.diff(invar, 2) + + self.add_equation("poisson", poisson) + + self._apply_detach() diff --git a/examples/smc_reac/ppsci/equation/pde/viv.py b/examples/smc_reac/ppsci/equation/pde/viv.py new file mode 100644 index 0000000000..c3d85895f1 --- /dev/null +++ b/examples/smc_reac/ppsci/equation/pde/viv.py @@ -0,0 +1,64 @@ +# 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. + +from __future__ import annotations + +import paddle +import sympy as sp +from paddle.nn import initializer + +from ppsci.equation.pde import base + + +class Vibration(base.PDE): + r"""Vortex induced vibration equation. + + $$ + \rho \dfrac{\partial^2 \eta}{\partial t^2} + e^{k1} \dfrac{\partial \eta}{\partial t} + e^{k2} \eta = f + $$ + + Args: + rho (float): Generalized mass. + k1 (float): Learnable parameter for modal damping. + k2 (float): Learnable parameter for generalized stiffness. + + Examples: + >>> import ppsci + >>> pde = ppsci.equation.Vibration(1.0, 4.0, -1.0) + """ + + def __init__(self, rho: float, k1: float, k2: float): + super().__init__() + self.rho = rho + self.k1 = paddle.create_parameter( + shape=[], + dtype=paddle.get_default_dtype(), + default_initializer=initializer.Constant(k1), + ) + self.k2 = paddle.create_parameter( + shape=[], + dtype=paddle.get_default_dtype(), + default_initializer=initializer.Constant(k2), + ) + self.learnable_parameters.append(self.k1) + self.learnable_parameters.append(self.k2) + + t_f = self.create_symbols("t_f") + eta = self.create_function("eta", (t_f,)) + k1 = self.create_symbols(self.k1.name) + k2 = self.create_symbols(self.k2.name) + f = self.rho * eta.diff(t_f, 2) + sp.exp(k1) * eta.diff(t_f) + sp.exp(k2) * eta + self.add_equation("f", f) + + self._apply_detach() diff --git a/examples/smc_reac/ppsci/experimental/__init__.py b/examples/smc_reac/ppsci/experimental/__init__.py new file mode 100644 index 0000000000..842f19428a --- /dev/null +++ b/examples/smc_reac/ppsci/experimental/__init__.py @@ -0,0 +1,37 @@ +# 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. + +""" +This module is for experimental API +""" + +from ppsci.experimental.math_module import bessel_i0 +from ppsci.experimental.math_module import bessel_i0e +from ppsci.experimental.math_module import bessel_i1 +from ppsci.experimental.math_module import bessel_i1e +from ppsci.experimental.math_module import fractional_diff +from ppsci.experimental.math_module import gaussian_integrate +from ppsci.experimental.math_module import montecarlo_integrate +from ppsci.experimental.math_module import trapezoid_integrate + +__all__ = [ + "bessel_i0", + "bessel_i0e", + "bessel_i1", + "bessel_i1e", + "fractional_diff", + "gaussian_integrate", + "trapezoid_integrate", + "montecarlo_integrate", +] diff --git a/examples/smc_reac/ppsci/experimental/math_module.py b/examples/smc_reac/ppsci/experimental/math_module.py new file mode 100644 index 0000000000..fc255c5671 --- /dev/null +++ b/examples/smc_reac/ppsci/experimental/math_module.py @@ -0,0 +1,646 @@ +# 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. + +from __future__ import annotations + +import functools +from typing import Any +from typing import Callable +from typing import List +from typing import Optional +from typing import Tuple +from typing import Union + +import numpy as np +import paddle +from typing_extensions import Literal + + +def bessel_i0(x: paddle.Tensor) -> paddle.Tensor: + """Zero-order modified Bézier curve functions of the first kind. + + Args: + x (paddle.Tensor): Input data of the formula. + + Examples: + >>> import paddle + >>> import ppsci + >>> res = ppsci.experimental.bessel_i0(paddle.to_tensor([0, 1, 2, 3, 4], dtype="float32")) + """ + return paddle.i0(x) + + +def bessel_i0e(x: paddle.Tensor) -> paddle.Tensor: + """Exponentially scaled zero-order modified Bézier curve functions of the first kind. + + Args: + x (paddle.Tensor): Input data of the formula. + + Examples: + >>> import paddle + >>> import ppsci + >>> res = ppsci.experimental.bessel_i0e(paddle.to_tensor([0, 1, 2, 3, 4], dtype="float32")) + """ + return paddle.i0e(x) + + +def bessel_i1(x: paddle.Tensor) -> paddle.Tensor: + """First-order modified Bézier curve functions of the first kind. + + Args: + x (paddle.Tensor): Input data of the formula. + + Examples: + >>> import paddle + >>> import ppsci + >>> res = ppsci.experimental.bessel_i1(paddle.to_tensor([0, 1, 2, 3, 4], dtype="float32")) + """ + return paddle.i1(x) + + +def bessel_i1e(x: paddle.Tensor) -> paddle.Tensor: + """Exponentially scaled first-order modified Bézier curve functions of the first kind. + + Args: + x (paddle.Tensor): Input data of the formula. + + Examples: + >>> import paddle + >>> import ppsci + >>> res = ppsci.experimental.bessel_i1e(paddle.to_tensor([0, 1, 2, 3, 4], dtype="float32")) + """ + return paddle.i1e(x) + + +def expand_func_values_and_squeeze_integral(f: Callable): + """This decorator ensures that the trailing dimension of integrands is indeed the integrand dimension. + This is pertinent in the 1d case when the sampled values are often of shape `(N,)`. Then, to maintain backward + consistency, we squeeze the result in the 1d case so it does not have any trailing dimensions. + + Args: + f (Callable): The wrapped function. + """ + + @functools.wraps(f) + def wrap(*args, **kwargs): + # i.e we only have one dimension, or the second dimension (that of the integrand) is 1 + is_1d = len(args[0].shape) == 1 or ( + len(args[0].shape) == 2 and args[0].shape[1] == 1 + ) + if is_1d: + return paddle.squeeze( + f(paddle.unsqueeze(args[0], axis=1), *args[1:], **kwargs) + ) + return f(*args, **kwargs) + + return wrap + + +def gaussian_integrate( + fn: Callable[[Any], paddle.Tensor], + dim: int, + N: int, + integration_domains: List[List[float]], + dtype: Literal["float32", "float64"] = "float64", +) -> paddle.Tensor: + """Integrate given function using gaussian quadrature. + + Args: + fn (Callable[[Any], paddle.Tensor]): Function to be integrated. + dim (int): Dimensionality of the integrand. + N (int): Number of dicretization points. + integration_domains (List[List[float]]): Integration domains. + dtype (Literal["float32", "float64"], optional): Dtype used during computation. Defaults to "float64". + + Returns: + paddle.Tensor: Integral result. + + Examples: + >>> import numpy as np + >>> import paddle + >>> import ppsci.experimental + >>> func = lambda x: paddle.sin(x) + >>> dim = 1 + >>> N = 500 + >>> integration_domains = [[0, np.pi]] + >>> result = ppsci.experimental.gaussian_integrate(func, dim, N, integration_domains) + >>> np.testing.assert_allclose(float(result), 2.0, 1e-6) + >>> print(float(result)) + 1.9999999999999576 + """ + + def _compatible_meshgrid(*args: paddle.Tensor, **kwargs: paddle.Tensor): + # TODO(HydrogenSulfate): paddle.meshgrid do not support single Tensor, + # which will be fixed in paddle framework. + if len(args) == 1: + return args + else: + return paddle.meshgrid(*args, **kwargs) + + def _roots(N: int) -> np.ndarray: + return np.polynomial.legendre.leggauss(N)[0] + + def _calculate_grid( + N: int, + integration_domains: paddle.Tensor, + ) -> Tuple[paddle.Tensor, paddle.Tensor, int]: + """Calculate grid points, widths and N per dim + + Args: + N (int): Number of points. + integration_domain (paddle.Tensor): Integration domain. + + Returns: + Tuple[paddle.Tensor, paddle.Tensor, int]: Grid points, grid widths and + Number of grid slices per dimension. + """ + # Create grid and assemble evaluation points + grid_1d = [] + _dim = integration_domains.shape[0] + n_per_dim = int(N ** (1.0 / _dim) + 1e-8) + + # Determine for each dimension grid points and mesh width + def _resize_roots( + integration_domain: Tuple[float, float], roots: np.ndarray + ): # scale from [-1,1] to [a,b] + a = integration_domain[0] + b = integration_domain[1] + return ((b - a) / 2) * roots + ((a + b) / 2) + + for dim in range(_dim): + grid_1d.append(_resize_roots(integration_domains[dim], _roots(n_per_dim))) + h = paddle.stack([grid_1d[dim][1] - grid_1d[dim][0] for dim in range(_dim)]) + + # Get grid points + points = _compatible_meshgrid(*grid_1d) + points = paddle.stack([mg.reshape([-1]) for mg in points], axis=1) + + return points, h, n_per_dim + + def _evaluate_integrand(fn, points, weights=None, fn_args=None) -> paddle.Tensor: + """Evaluate the integrand function at the passed points. + + Args: + fn (function): Integrand function. + points (paddle.Tensor): Integration points. + weights (paddle.Tensor, optional): Integration weights. Defaults to None. + fn_args (list or tuple, optional): Any arguments required by the function. Defaults to None. + + Returns: + paddle.Tensor: Integral result. + """ + if fn_args is None: + fn_args = () + + result = fn(points, *fn_args) + if not str(result.dtype).endswith(dtype): + result = result.astype(dtype) + + if result.shape[0] != points.shape[0]: + raise ValueError( + f"The passed function was given {points.shape[0]} points but only returned {result.shape[0]} value(s)." + f"Please ensure that your function is vectorized, i.e. can be called with multiple evaluation points at once. It should return a tensor " + f"where first dimension matches length of passed elements. " + ) + + if weights is not None: + if ( + len(result.shape) > 1 + ): # if the the integrand is multi-dimensional, we need to reshape/repeat weights so they can be broadcast in the *= + integrand_shape = result.shape[1:] + weights = paddle.repeat_interleave( + paddle.unsqueeze(weights, axis=1), np.prod(integrand_shape) + ).reshape((weights.shape[0], *(integrand_shape))) + result *= weights + + return result + + def _weights(N, dim): + """Return the weights, broadcast across the dimensions, generated from the polynomial of choice. + + Args: + N (int): Number of nodes. + dim (int): Number of dimensions. + + Returns: + paddle.Tensor: Integration weights. + """ + weights = paddle.to_tensor(np.polynomial.legendre.leggauss(N)[1], dtype=dtype) + return paddle.prod( + paddle.stack(_compatible_meshgrid(*([weights] * dim)), axis=0), + axis=0, + ).reshape([-1]) + + def _apply_composite_rule(cur_dim_areas, dim, hs, domain): + """Apply "composite" rule for gaussian integrals + + cur_dim_areas will contain the areas per dimension + """ + # We collapse dimension by dimension + for cur_dim in range(dim): + cur_dim_areas = ( + 0.5 + * (domain[cur_dim][1] - domain[cur_dim][0]) + * paddle.sum( + cur_dim_areas, axis=len(cur_dim_areas.shape) - 1, dtype=dtype + ) + ) + return cur_dim_areas + + @expand_func_values_and_squeeze_integral + def _calculate_result( + function_values: paddle.Tensor, + dim: int, + n_per_dim: int, + hs: paddle.Tensor, + integration_domains: paddle.Tensor, + ) -> paddle.Tensor: + """Apply the "composite rule" to calculate a result from the evaluated integrand. + + Args: + function_values (paddle.Tensor): Output of the integrand. + dim (int): Dimensionality. + n_per_dim (int): Number of grid slices per dimension. + hs (paddle.Tensor): Distances between grid slices for each dimension. + + Returns: + paddle.Tensor: Quadrature result. + """ + # Reshape the output to be [integrand_dim,N,N,...] points instead of [integrand_dim,dim*N] points + integrand_shape = function_values.shape[1:] + dim_shape = [n_per_dim] * dim + new_shape = [*integrand_shape, *dim_shape] + + perm = list(range(len(function_values.shape))) + if len(perm) >= 2: + perm.append(perm.pop(0)) + reshaped_function_values = paddle.transpose(function_values, perm) + reshaped_function_values = reshaped_function_values.reshape(new_shape) + + assert new_shape == list( + reshaped_function_values.shape + ), f"reshaping produced shape {reshaped_function_values.shape}, expected shape was {new_shape}" + + result = _apply_composite_rule( + reshaped_function_values, dim, hs, integration_domains + ) + return result + + assert dtype in [ + "float32", + "float64", + ], f"dtype must be either 'float32' or 'float64', but got {dtype}" + + neg = False + for i, (a, b) in enumerate(integration_domains): + if a > b: + neg = not neg + integration_domains[i] = [b, a] + + integration_domains = paddle.to_tensor( + integration_domains, + dtype=dtype, + ) + + if integration_domains.shape[0] != dim: + raise ValueError( + f"The number of integration domain({integration_domains.shape[0]}) " + f"must be equal to the given 'dim'({dim})." + ) + if integration_domains.shape[1] != 2: + raise ValueError( + f"integration_domain should be in format of [[a_1, b_1], [a_2, b_2], ..., " + f"[a_dim, b_dim]], but got each range of integration is {integration_domains[0]}" + ) + grid_points, hs, n_per_dim = _calculate_grid(N, integration_domains) + + function_values = _evaluate_integrand( + fn, grid_points, weights=_weights(n_per_dim, dim) + ) + + result = _calculate_result(function_values, dim, n_per_dim, hs, integration_domains) + return result if (not neg) else -result + + +def fractional_diff( + func: Callable, alpha: float, a: float, t: float, h: float, dtype="float64" +) -> paddle.Tensor: + r"""Compute fractional derivative of given function at point t with fractional order + alpha using [Caputo derivative of fractional](https://en.wikipedia.org/wiki/Fractional_calculus#Caputo_fractional_derivative). + + $$ + D_t^\alpha f(t)=\frac{1}{\Gamma(n-\alpha)} \int_0^t \frac{f^{(n)}(s)}{(t-s)^{\alpha+1-n}} d s . + $$ + + $$ + s.t. 0 \lt \alpha \lt 1 . + $$ + + Args: + func (Callable): Function to compute the fractional derivative of. + alpha (float): Fractional order. + t (float): Point to compute the fractional derivative at. + a (float): Start point of the fractional integral. + h (float): Step size for finite difference. + dtype (str, optional): Data dtype during computation. Defaults to "float64". + + Returns: + paddle.Tensor: Fractional derivative result of the function at t. + + Examples: + >>> from ppsci.experimental import fractional_diff + >>> import numpy as np + >>> # define f(x) = x^2 + >>> def f(x): + ... return x * x + >>> # compute 0.5-order fractional derivative of f(x) at t=1.0 with step size h=1e-6 + >>> res = fractional_diff(f, alpha=0.5, a=0, t=1.0, h=1e-6, dtype="float64") + >>> np.testing.assert_allclose(float(res), 1.503547, 1e-6) + """ + + if not (0 < alpha < 1): + raise NotImplementedError( + f"Given alpha should be in range (0, 1), but got {alpha}" + ) + + def _finite_derivative( + func: Callable, x: paddle.Tensor, dx: float + ) -> paddle.Tensor: + """Compute the finite difference of a function at x using centered difference. + + Args: + func (Callable): Function to compute the finite difference of. + x (paddle.Tensor): Point to compute the finite difference at. + dx (float): Delta to use for the finite difference. + + Returns: + paddle.Tensor: First-order Finite difference of the function at x. + """ + return (func(x + dx) - func(x - dx)) / (2 * dx) + + def int_func(s): + return _finite_derivative(func, s, dx=h) / (t - s) ** (alpha) + + result = ( + 1.0 / paddle.exp(paddle.lgamma(paddle.to_tensor(1.0 - alpha, dtype=dtype))) + ) * gaussian_integrate( + int_func, dim=1, N=2**10 + 1, integration_domains=[[a, t]], dtype=dtype + ) + return result + + +def trapezoid_integrate( + y: paddle.Tensor, + x: paddle.Tensor = None, + dx: float = None, + axis: int = -1, + mode: Literal["sum", "cumsum"] = "sum", +) -> paddle.Tensor: + """ + Integrate along the given axis using the composite trapezoidal rule. Use the sum method. + + Args: + y (paddle.Tensor): Input to be integrated. + x (paddle.Tensor, optional): The sample points corresponding to the input samples. its shape should be + (1) input.shape; (2) the input.shape[axis] if axis is not default. Defaults to None. + dx (float, optional): The sample points are assumed to be evenly spaced and it is the spacing between sample points. + If 'x' and 'dx' are both default, 'dx' is set to 1 by default. Defaults to None. + axis (int, optional): The axis along which to integrate. Defaults to -1. + mode (Literal["sum", "cumsum"], optional): Which type cumulative sum function used. Defaults to "sum". + + Returns: + paddle.Tensor: Integral result. If dim of input is N, return is N-1 dim. + + Examples: + >>> import paddle + >>> import ppsci + >>> y = paddle.to_tensor([[0, 1, 2], [3, 4, 5]], dtype="float32") + >>> res = ppsci.experimental.trapezoid_integrate(y) + >>> print(res) + Tensor(shape=[2], dtype=float32, place=Place(gpu:0), stop_gradient=True, + [2., 8.]) + >>> res = ppsci.experimental.trapezoid_integrate(y, mode="cumsum") + >>> print(res) + Tensor(shape=[2, 2], dtype=float32, place=Place(gpu:0), stop_gradient=True, + [[0.50000000, 2. ], + [3.50000000, 8. ]]) + >>> res = ppsci.experimental.trapezoid_integrate( + ... y, x=paddle.to_tensor([[0, 1, 2], [3, 4, 5]], dtype="float32") + ... ) + >>> print(res) + Tensor(shape=[2], dtype=float32, place=Place(gpu:0), stop_gradient=True, + [2., 8.]) + >>> res = ppsci.experimental.trapezoid_integrate( + ... y, x=paddle.to_tensor([0, 1], dtype="float32"), axis=0 + ... ) + >>> print(res) + Tensor(shape=[3], dtype=float32, place=Place(gpu:0), stop_gradient=True, + [1.50000000, 2.50000000, 3.50000000]) + >>> res = ppsci.experimental.trapezoid_integrate( + ... y, x=paddle.to_tensor([0, 1, 2], dtype="float32"), axis=1 + ... ) + >>> print(res) + Tensor(shape=[2], dtype=float32, place=Place(gpu:0), stop_gradient=True, + [2., 8.]) + >>> res = ppsci.experimental.trapezoid_integrate(y, dx=2) + >>> print(res) + Tensor(shape=[2], dtype=float32, place=Place(gpu:0), stop_gradient=True, + [4. , 16.]) + """ + if mode == "sum": + return paddle.trapezoid(y, x, dx, axis) + elif mode == "cumsum": + return paddle.cumulative_trapezoid(y, x, dx, axis) + else: + raise ValueError(f'mode should be "sum" or "cumsum", but got {mode}') + + +def montecarlo_integrate( + fn: Callable, + dim: int, + N: int = 1000, + integration_domain: Union[List[List[float]], paddle.Tensor] = None, + seed: int = None, +) -> paddle.Tensor: + """Integrates the passed function on the passed domain using vanilla Monte + Carlo Integration. + + Args: + fn (Callable): The function to integrate over. + dim (int): Dimensionality of the function's domain over which to + integrate. + N (Optional[int]): Number of sample points to use for the integration. + Defaults to 1000. + integration_domain (Union[List[List[float]], paddle.Tensor]): Integration + domain, e.g. [[-1,1],[0,1]]. Defaults to [-1,1]^dim. + seed (Optional[int]): Random number generation seed to the sampling + point creation, only set if provided. Defaults to None. + + Raises: + ValueError: If len(integration_domain) != dim + + Returns: + paddle.Tensor: Integral result. + + Examples: + >>> import paddle + >>> import ppsci + + >>> _ = paddle.seed(1024) + >>> # The function we want to integrate, in this example + >>> # f(x0,x1) = sin(x0) + e^x1 for x0=[0,1] and x1=[-1,1] + >>> # Note that the function needs to support multiple evaluations at once (first + >>> # dimension of x here) + >>> # Expected result here is ~3.2698 + >>> def some_function(x): + ... return paddle.sin(x[:, 0]) + paddle.exp(x[:, 1]) + + >>> # Compute the function integral by sampling 10000 points over domain + >>> integral_value = ppsci.experimental.montecarlo_integrate( + ... some_function, + ... dim=2, + ... N=10000, + ... integration_domain=[[0, 1], [-1, 1]], + ... ) + + >>> print(integral_value) + Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, + 3.25152588) + """ + + @expand_func_values_and_squeeze_integral + def calculate_result(function_values, integration_domain): + """Calculate an integral result from the function evaluations + + Args: + function_values (paddle.Tensor): Output of the integrand + integration_domain (paddle.Tensor): Integration domain + + Returns: + Quadrature result + """ + scales = integration_domain[:, 1] - integration_domain[:, 0] + volume = paddle.prod(scales) + + # Integral = V / N * sum(func values) + N = function_values.shape[0] + integral = volume * paddle.sum(function_values, axis=0) / N + return integral + + def calculate_sample_points( + N: int, integration_domain: paddle.Tensor, seed: Optional[int] = None + ): + """Calculate random points for the integrand evaluation. + + Args: + N (int): Number of points + integration_domain (paddle.Tensor): Integration domain. + seed (int, optional): Random number generation seed for the sampling point creation, only set if provided. Defaults to None. + Returns: + Sample points. + """ + dim = integration_domain.shape[0] + domain_starts = integration_domain[:, 0] + domain_sizes = integration_domain[:, 1] - domain_starts + # Scale and translate random numbers via broadcasting + return ( + paddle.uniform( + shape=[N, dim], + dtype=domain_sizes.dtype, + min=0.0, + max=1.0, + seed=seed or 0, + ) + * domain_sizes + + domain_starts + ) + + if dim is not None: + if dim < 1: + raise ValueError("Dimension needs to be 1 or larger.") + if N is not None: + if N < 1 or type(N) is not int: + raise ValueError("N has to be a positive integer.") + + integration_domain = _setup_integration_domain(dim, integration_domain) + sample_points = calculate_sample_points(N, integration_domain, seed) + function_values, _ = _evaluate_integrand(fn, sample_points) + return calculate_result(function_values, integration_domain) + + +def _setup_integration_domain( + dim: int, integration_domain: Union[List[List[float]], paddle.Tensor] +) -> paddle.Tensor: + """Sets up the integration domain if unspecified by the user. + Args: + dim (int): Dimensionality of the integration domain. + integration_domain (List or Tensor): Integration domain, e.g. [[-1,1],[0,1]]. Defaults to [-1,1]^dim. + + Returns: + Integration domain. + """ + # If no integration_domain is specified, create [-1,1]^d bounds + if integration_domain is None: + integration_domain = [[-1.0, 1.0]] * dim + + integration_domain = [[float(b) for b in bounds] for bounds in integration_domain] + + integration_domain = paddle.to_tensor(integration_domain) + + if tuple(integration_domain.shape) != (dim, 2): + raise ValueError( + "The integration domain has an unexpected shape. " + f"Expected {(dim, 2)}, got {integration_domain.shape}" + ) + return integration_domain + + +def _evaluate_integrand(fn, points, weights=None, args=None): + """Evaluate the integrand function at the passed points. + + Args: + fn (Callable): Integrand function. + points (paddle.Tensor): Integration points. + weights (Optional[paddle.Tensor]): Integration weights. Defaults to None. + args (Optional[List, Tuple]): Any arguments required by the function. Defaults to None. + + Returns: + padlde.Tensor: Integrand function output. + int: Number of evaluated points. + """ + num_points = points.shape[0] + + if args is None: + args = () + + result = fn(points, *args) + num_results = result.shape[0] + if num_results != num_points: + raise ValueError( + f"The passed function was given {num_points} points but only returned {num_results} value(s)." + f"Please ensure that your function is vectorized, i.e. can be called with multiple evaluation points at once. It should return a tensor " + f"where first dimension matches length of passed elements. " + ) + + if weights is not None: + if ( + len(result.shape) > 1 + ): # if the the integrand is multi-dimensional, we need to reshape/repeat weights so they can be broadcast in the *= + integrand_shape = paddle.to_tensor(result.shape[1:]) + weights = paddle.tile( + paddle.unsqueeze(weights, axis=1), paddle.prod(integrand_shape) + ).reshape((weights.shape[0], *(integrand_shape))) + result *= weights + + return result, num_points diff --git a/examples/smc_reac/ppsci/externals/__init__.py b/examples/smc_reac/ppsci/externals/__init__.py new file mode 100644 index 0000000000..67e62f29f3 --- /dev/null +++ b/examples/smc_reac/ppsci/externals/__init__.py @@ -0,0 +1,20 @@ +"""External Development Packages""" + +import importlib.util + +EXTERNAL_PACKAGES_LIST = [ + "deepali", + "neuraloperator", + "open3d", + "paddle_harmonics", + "paddle_scatter", + "paddle_sparse", + "tensorly", + "warp", +] + +__all__ = [] +for package_name in EXTERNAL_PACKAGES_LIST: + if importlib.util.find_spec(package_name): + globals()[package_name] = __import__(package_name) + __all__.append(package_name) diff --git a/examples/smc_reac/ppsci/geometry/__init__.py b/examples/smc_reac/ppsci/geometry/__init__.py new file mode 100644 index 0000000000..30b4ad0859 --- /dev/null +++ b/examples/smc_reac/ppsci/geometry/__init__.py @@ -0,0 +1,83 @@ +# 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. + +import copy + +from ppsci.geometry.geometry import Geometry +from ppsci.geometry.geometry_1d import Interval +from ppsci.geometry.geometry_2d import Disk +from ppsci.geometry.geometry_2d import Polygon +from ppsci.geometry.geometry_2d import Rectangle +from ppsci.geometry.geometry_2d import Triangle +from ppsci.geometry.geometry_3d import Cuboid +from ppsci.geometry.geometry_3d import Sphere +from ppsci.geometry.geometry_nd import Hypercube +from ppsci.geometry.geometry_nd import Hypersphere +from ppsci.geometry.mesh import Mesh +from ppsci.geometry.mesh import SDFMesh +from ppsci.geometry.pointcloud import PointCloud +from ppsci.geometry.timedomain import TimeDomain +from ppsci.geometry.timedomain import TimeXGeometry +from ppsci.utils import logger +from ppsci.utils import misc + +__all__ = [ + "build_geometry", + "Cuboid", + "Disk", + "Geometry", + "Hypercube", + "Hypersphere", + "Interval", + "Mesh", + "SDFMesh", + "Polygon", + "Rectangle", + "Sphere", + "TimeDomain", + "TimeXGeometry", + "Triangle", + "PointCloud", +] + + +def build_geometry(cfg): + """Build geometry(ies) + + Args: + cfg (List[DictConfig]): Geometry config list. + + Returns: + Dict[str, Geometry]: Geometry(ies) in dict. + """ + if cfg is None: + return None + cfg = copy.deepcopy(cfg) + + geom_dict = misc.PrettyOrderedDict() + for _item in cfg: + geom_cls = next(iter(_item.keys())) + geom_cfg = _item[geom_cls] + geom_name = geom_cfg.pop("name", geom_cls) + if geom_cls == "TimeXGeometry": + time_cfg = geom_cfg.pop("TimeDomain") + geom_cls = next(iter(geom_cfg.keys())) + geom_dict[geom_name] = TimeXGeometry( + TimeDomain(**time_cfg), eval(geom_cls)(**geom_cfg[geom_cls]) + ) + else: + geom_dict[geom_name] = eval(geom_cls)(**geom_cfg) + + logger.debug(str(geom_dict[geom_name])) + return geom_dict diff --git a/examples/smc_reac/ppsci/geometry/csg.py b/examples/smc_reac/ppsci/geometry/csg.py new file mode 100644 index 0000000000..87534bedd6 --- /dev/null +++ b/examples/smc_reac/ppsci/geometry/csg.py @@ -0,0 +1,337 @@ +# 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. + +""" +Code below is heavily based on [https://github.com/lululxvi/deepxde](https://github.com/lululxvi/deepxde) +""" + +from __future__ import annotations + +import numpy as np +import paddle + +from ppsci.geometry import geometry + + +class CSGUnion(geometry.Geometry): + """Construct an object by CSG Union(except for Mesh).""" + + def __init__(self, geom1, geom2): + if geom1.ndim != geom2.ndim: + raise ValueError( + f"{geom1}.ndim({geom1.ndim}) should be equal to " + f"{geom2}.ndim({geom1.ndim})" + ) + super().__init__( + geom1.ndim, + ( + np.minimum(geom1.bbox[0], geom2.bbox[0]), + np.maximum(geom1.bbox[1], geom2.bbox[1]), + ), + geom1.diam + geom2.diam, + ) + self.geom1 = geom1 + self.geom2 = geom2 + + def is_inside(self, x): + return np.logical_or(self.geom1.is_inside(x), self.geom2.is_inside(x)) + + def on_boundary(self, x): + return np.logical_or( + np.logical_and(self.geom1.on_boundary(x), ~self.geom2.is_inside(x)), + np.logical_and(self.geom2.on_boundary(x), ~self.geom1.is_inside(x)), + ) + + def boundary_normal(self, x): + return np.logical_and(self.geom1.on_boundary(x), ~self.geom2.is_inside(x))[ + :, np.newaxis + ] * self.geom1.boundary_normal(x) + np.logical_and( + self.geom2.on_boundary(x), ~self.geom1.is_inside(x) + )[ + :, np.newaxis + ] * self.geom2.boundary_normal( + x + ) + + def random_points(self, n, random="pseudo"): + x = np.empty(shape=(n, self.ndim), dtype=paddle.get_default_dtype()) + _size = 0 + while _size < n: + points = ( + np.random.rand(n, self.ndim) * (self.bbox[1] - self.bbox[0]) + + self.bbox[0] + ) + points = points[self.is_inside(points)] + + if len(points) > n - _size: + points = points[: n - _size] + x[_size : _size + len(points)] = points + _size += len(points) + return x + + def random_boundary_points(self, n, random="pseudo"): + x = np.empty(shape=(n, self.ndim), dtype=paddle.get_default_dtype()) + _size = 0 + while _size < n: + geom1_boundary_points = self.geom1.random_boundary_points(n, random=random) + geom1_boundary_points = geom1_boundary_points[ + ~self.geom2.is_inside(geom1_boundary_points) + ] + + geom2_boundary_points = self.geom2.random_boundary_points(n, random=random) + geom2_boundary_points = geom2_boundary_points[ + ~self.geom1.is_inside(geom2_boundary_points) + ] + + points = np.concatenate((geom1_boundary_points, geom2_boundary_points)) + points = np.random.permutation(points) + + if len(points) > n - _size: + points = points[: n - _size] + x[_size : _size + len(points)] = points + _size += len(points) + return x + + def periodic_point(self, x, component): + x = np.copy(x) + on_boundary_geom1 = np.logical_and( + self.geom1.on_boundary(x), ~self.geom2.is_inside(x) + ) + x[on_boundary_geom1] = self.geom1.periodic_point(x, component)[ + on_boundary_geom1 + ] + on_boundary_geom2 = np.logical_and( + self.geom2.on_boundary(x), ~self.geom1.is_inside(x) + ) + x[on_boundary_geom2] = self.geom2.periodic_point(x, component)[ + on_boundary_geom2 + ] + return x + + def sdf_func(self, points: np.ndarray) -> np.ndarray: + """Compute signed distance field of CSG union of two geometries. + ref: https://iquilezles.org/articles/distfunctions/ + + Args: + points (np.ndarray): The coordinate points used to calculate the SDF + value, the shape is [N, D]. + + Returns: + np.ndarray: SDF values of input points without squared, the shape is [N, 1]. + """ + sdf1 = self.geom1.sdf_func(points) + sdf2 = self.geom2.sdf_func(points) + return np.minimum(sdf1, sdf2) + + +class CSGDifference(geometry.Geometry): + """Construct an object by CSG Difference.""" + + def __init__(self, geom1, geom2): + if geom1.ndim != geom2.ndim: + raise ValueError( + f"{geom1}.ndim({geom1.ndim}) should be equal to " + f"{geom2}.ndim({geom1.ndim})." + ) + super().__init__(geom1.ndim, geom1.bbox, geom1.diam) + self.geom1 = geom1 + self.geom2 = geom2 + + def is_inside(self, x): + return np.logical_and(self.geom1.is_inside(x), ~self.geom2.is_inside(x)) + + def on_boundary(self, x): + return np.logical_or( + np.logical_and(self.geom1.on_boundary(x), ~self.geom2.is_inside(x)), + np.logical_and(self.geom1.is_inside(x), self.geom2.on_boundary(x)), + ) + + def boundary_normal(self, x): + return np.logical_and(self.geom1.on_boundary(x), ~self.geom2.is_inside(x))[ + :, np.newaxis + ] * self.geom1.boundary_normal(x) + np.logical_and( + self.geom1.is_inside(x), self.geom2.on_boundary(x) + )[ + :, np.newaxis + ] * -self.geom2.boundary_normal( + x + ) + + def random_points(self, n, random="pseudo"): + x = np.empty(shape=(n, self.ndim), dtype=paddle.get_default_dtype()) + _size = 0 + while _size < n: + tmp = self.geom1.random_points(n, random=random) + tmp = tmp[~self.geom2.is_inside(tmp)] + + if len(tmp) > n - _size: + tmp = tmp[: n - _size] + x[_size : _size + len(tmp)] = tmp + _size += len(tmp) + return x + + def random_boundary_points(self, n, random="pseudo"): + x = np.empty(shape=(n, self.ndim), dtype=paddle.get_default_dtype()) + _size = 0 + while _size < n: + geom1_boundary_points = self.geom1.random_boundary_points(n, random=random) + geom1_boundary_points = geom1_boundary_points[ + ~self.geom2.is_inside(geom1_boundary_points) + ] + + geom2_boundary_points = self.geom2.random_boundary_points(n, random=random) + geom2_boundary_points = geom2_boundary_points[ + self.geom1.is_inside(geom2_boundary_points) + ] + + points = np.concatenate((geom1_boundary_points, geom2_boundary_points)) + points = np.random.permutation(points) + + if len(points) > n - _size: + points = points[: n - _size] + x[_size : _size + len(points)] = points + _size += len(points) + return x + + def periodic_point(self, x, component): + x = np.copy(x) + on_boundary_geom1 = np.logical_and( + self.geom1.on_boundary(x), ~self.geom2.is_inside(x) + ) + x[on_boundary_geom1] = self.geom1.periodic_point(x, component)[ + on_boundary_geom1 + ] + return x + + def sdf_func(self, points: np.ndarray) -> np.ndarray: + """Compute signed distance field of CSG difference of two geometries. + + Args: + points (np.ndarray): The coordinate points used to calculate the SDF + value, the shape is [N, D]. + + Returns: + np.ndarray: SDF values of input points without squared, the shape is [N, 1]. + """ + sdf1 = self.geom1.sdf_func(points) + sdf2 = self.geom2.sdf_func(points) + return np.maximum(sdf1, -sdf2) + + +class CSGIntersection(geometry.Geometry): + """Construct an object by CSG Intersection.""" + + def __init__(self, geom1, geom2): + if geom1.ndim != geom2.ndim: + raise ValueError( + f"{geom1}.ndim({geom1.ndim}) should be equal to " + f"{geom2}.ndim({geom1.ndim})" + ) + super().__init__( + geom1.ndim, + ( + np.maximum(geom1.bbox[0], geom2.bbox[0]), + np.minimum(geom1.bbox[1], geom2.bbox[1]), + ), + min(geom1.diam, geom2.diam), + ) + self.geom1 = geom1 + self.geom2 = geom2 + + def is_inside(self, x): + return np.logical_and(self.geom1.is_inside(x), self.geom2.is_inside(x)) + + def on_boundary(self, x): + return np.logical_or( + np.logical_and(self.geom1.on_boundary(x), self.geom2.is_inside(x)), + np.logical_and(self.geom1.is_inside(x), self.geom2.on_boundary(x)), + ) + + def boundary_normal(self, x): + return np.logical_and(self.geom1.on_boundary(x), self.geom2.is_inside(x))[ + :, np.newaxis + ] * self.geom1.boundary_normal(x) + np.logical_and( + self.geom1.is_inside(x), self.geom2.on_boundary(x) + )[ + :, np.newaxis + ] * self.geom2.boundary_normal( + x + ) + + def random_points(self, n, random="pseudo"): + x = np.empty(shape=(n, self.ndim), dtype=paddle.get_default_dtype()) + _size = 0 + while _size < n: + points = self.geom1.random_points(n, random=random) + points = points[self.geom2.is_inside(points)] + + if len(points) > n - _size: + points = points[: n - _size] + x[_size : _size + len(points)] = points + _size += len(points) + return x + + def random_boundary_points(self, n, random="pseudo"): + x = np.empty(shape=(n, self.ndim), dtype=paddle.get_default_dtype()) + _size = 0 + while _size < n: + geom1_boundary_points = self.geom1.random_boundary_points(n, random=random) + geom1_boundary_points = geom1_boundary_points[ + self.geom2.is_inside(geom1_boundary_points) + ] + + geom2_boundary_points = self.geom2.random_boundary_points(n, random=random) + geom2_boundary_points = geom2_boundary_points[ + self.geom1.is_inside(geom2_boundary_points) + ] + + points = np.concatenate((geom1_boundary_points, geom2_boundary_points)) + points = np.random.permutation(points) + + if len(points) > n - _size: + points = points[: n - _size] + x[_size : _size + len(points)] = points + _size += len(points) + return x + + def periodic_point(self, x, component): + x = np.copy(x) + on_boundary_geom1 = np.logical_and( + self.geom1.on_boundary(x), self.geom2.is_inside(x) + ) + x[on_boundary_geom1] = self.geom1.periodic_point(x, component)[ + on_boundary_geom1 + ] + on_boundary_geom2 = np.logical_and( + self.geom2.on_boundary(x), self.geom1.is_inside(x) + ) + x[on_boundary_geom2] = self.geom2.periodic_point(x, component)[ + on_boundary_geom2 + ] + return x + + def sdf_func(self, points: np.ndarray) -> np.ndarray: + """Compute signed distance field of CSG intersection of two geometries. + ref: https://iquilezles.org/articles/distfunctions/ + + Args: + points (np.ndarray): The coordinate points used to calculate the SDF + value the shape is [N, D]. + + Returns: + np.ndarray: SDF values of input points without squared, the shape is [N, 1]. + """ + sdf1 = self.geom1.sdf_func(points) + sdf2 = self.geom2.sdf_func(points) + return np.maximum(sdf1, sdf2) diff --git a/examples/smc_reac/ppsci/geometry/geometry.py b/examples/smc_reac/ppsci/geometry/geometry.py new file mode 100644 index 0000000000..5bda675414 --- /dev/null +++ b/examples/smc_reac/ppsci/geometry/geometry.py @@ -0,0 +1,696 @@ +# 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. + +""" +Code below is heavily based on [https://github.com/lululxvi/deepxde](https://github.com/lululxvi/deepxde) +""" +from __future__ import annotations + +import abc +from typing import Callable +from typing import Dict +from typing import Optional +from typing import Tuple + +import numpy as np +import paddle +from typing_extensions import Literal + +from ppsci.utils import logger +from ppsci.utils import misc + + +class Geometry: + """Base class for geometry. + + Args: + ndim (int): Number of geometry dimension. + bbox (Tuple[np.ndarray, np.ndarray]): Bounding box of upper and lower. + diam (float): Diameter of geometry. + """ + + def __init__(self, ndim: int, bbox: Tuple[np.ndarray, np.ndarray], diam: float): + self.ndim = ndim + self.bbox = bbox + self.diam = min(diam, np.linalg.norm(bbox[1] - bbox[0])) + + @property + def dim_keys(self): + return ("x", "y", "z")[: self.ndim] + + @abc.abstractmethod + def is_inside(self, x: np.ndarray) -> np.ndarray: + """Returns a boolean array where x is inside the geometry. + + Args: + x (np.ndarray): Points to check if inside the geometry. The shape is [N, D], + where D is the number of dimension of geometry. + + Returns: + np.ndarray: Boolean array where x is inside the geometry. The shape is [N]. + + Examples: + >>> import numpy as np + >>> import ppsci + >>> interval = ppsci.geometry.Interval(0, 1) + >>> x = np.array([[0], [0.5], [1.5]]) + >>> interval.is_inside(x) + array([ True, True, False]) + >>> rectangle = ppsci.geometry.Rectangle((0, 0), (1, 1)) + >>> x = np.array([[0.0, 0.0], [0.5, 0.5], [1.5, 1.5]]) + >>> rectangle.is_inside(x) + array([ True, True, False]) + >>> cuboid = ppsci.geometry.Cuboid((0, 0, 0), (1, 1, 1)) + >>> x = np.array([[0, 0, 0], [0.5, 0.5, 0.5], [1.5, 1.5, 1.5]]) + >>> cuboid.is_inside(x) + array([ True, True, False]) + """ + + @abc.abstractmethod + def on_boundary(self, x: np.ndarray) -> np.ndarray: + """Returns a boolean array where x is on geometry boundary. + + Args: + x (np.ndarray): Points to check if on the geometry boundary. The shape is [N, D], + where D is the number of dimension of geometry. + + Returns: + np.ndarray: Boolean array where x is on the geometry boundary. The shape is [N]. + + Examples: + >>> import numpy as np + >>> import ppsci + >>> interval = ppsci.geometry.Interval(0, 1) + >>> x = np.array([[0], [0.5], [1.5]]) + >>> interval.on_boundary(x) + array([ True, False, False]) + >>> rectangle = ppsci.geometry.Rectangle((0, 0), (1, 1)) + >>> x = np.array([[0, 0], [0.5, 0.5], [1, 1.5]]) + >>> rectangle.on_boundary(x) + array([ True, False, False]) + >>> cuboid = ppsci.geometry.Cuboid((0, 0, 0), (1, 1, 1)) + >>> x = np.array([[0, 0, 0], [0.5, 0.5, 0.5], [1, 1, 1.5]]) + >>> cuboid.on_boundary(x) + array([ True, False, False]) + """ + + def boundary_normal(self, x): + """Compute the unit normal at x.""" + raise NotImplementedError(f"{self}.boundary_normal is not implemented") + + def uniform_points(self, n: int, boundary: bool = True) -> np.ndarray: + """Compute the equi-spaced points in the geometry. + + Args: + n (int): Number of points. + boundary (bool): Include boundary points. Defaults to True. + + Returns: + np.ndarray: Random points in the geometry. The shape is [N, D]. + """ + logger.warning( + f"{self}.uniform_points not implemented. " f"Use random_points instead." + ) + return self.random_points(n) + + def sample_interior( + self, + n: int, + random: Literal["pseudo", "Halton", "LHS"] = "pseudo", + criteria: Optional[Callable[..., np.ndarray]] = None, + evenly: bool = False, + compute_sdf_derivatives: bool = False, + ) -> Dict[str, np.ndarray]: + """Sample random points in the geometry and return those meet criteria. + + Args: + n (int): Number of points. + random (Literal["pseudo", "Halton", "LHS"]): Random method. Defaults to "pseudo". + pseudo: Pseudo random. + Halton: Halton sequence. + LHS: Latin Hypercube Sampling. + criteria (Optional[Callable[..., np.ndarray]]): Criteria function. Given + coords from different dimension and return a boolean array with shape [n,]. + Defaults to None. + evenly (bool): Evenly sample points. Defaults to False. + compute_sdf_derivatives (bool): Compute SDF derivatives. Defaults to False. + + Returns: + Dict[str, np.ndarray]: Random points in the geometry. The shape is [N, D]. + their signed distance function. The shape is [N, 1]. + their derivatives of SDF(optional). The shape is [N, D]. + + Examples: + >>> import numpy as np + >>> import ppsci + >>> np.random.seed(42) + >>> interval = ppsci.geometry.Interval(0, 1) + >>> interval.sample_interior(2) + {'x': array([[0.37454012], + [0.9507143 ]], dtype=float32), 'sdf': array([[0.37454012], + [0.04928571]], dtype=float32)} + >>> rectangle = ppsci.geometry.Rectangle((0, 0), (1, 1)) + >>> rectangle.sample_interior(2, "pseudo", None, False, True) + {'x': array([[0.7319939 ], + [0.15601864]], dtype=float32), 'y': array([[0.5986585 ], + [0.15599452]], dtype=float32), 'sdf': array([[0.2680061 ], + [0.15599453]], dtype=float32), 'sdf__x': array([[-1.0001659 ], + [ 0.25868416]], dtype=float32), 'sdf__y': array([[-0. ], + [ 0.74118376]], dtype=float32)} + >>> cuboid = ppsci.geometry.Cuboid((0, 0, 0), (1, 1, 1)) + >>> cuboid.sample_interior(2, "pseudo", None, True, True) + {'x': array([[0.], + [0.]], dtype=float32), 'y': array([[0.], + [0.]], dtype=float32), 'z': array([[0.], + [1.]], dtype=float32), 'sdf': array([[0.], + [0.]], dtype=float32), 'sdf__x': array([[0.50008297], + [0.50008297]], dtype=float32), 'sdf__y': array([[0.50008297], + [0.50008297]], dtype=float32), 'sdf__z': array([[ 0.50008297], + [-0.49948692]], dtype=float32)} + """ + x = np.empty(shape=(n, self.ndim), dtype=paddle.get_default_dtype()) + _size, _ntry, _nsuc = 0, 0, 0 + while _size < n: + if evenly: + points = self.uniform_points(n) + else: + if misc.typename(self) == "TimeXGeometry": + points = self.random_points(n, random, criteria) + else: + points = self.random_points(n, random) + + if criteria is not None: + criteria_mask = criteria(*np.split(points, self.ndim, axis=1)).flatten() + points = points[criteria_mask] + + if len(points) > n - _size: + points = points[: n - _size] + x[_size : _size + len(points)] = points + + _size += len(points) + _ntry += 1 + if len(points) > 0: + _nsuc += 1 + + if _ntry >= 1000 and _nsuc == 0: + raise ValueError( + "Sample interior points failed, " + "please check correctness of geometry and given criteria." + ) + + # if sdf_func added, return x_dict and sdf_dict, else, only return the x_dict + if hasattr(self, "sdf_func"): + sdf = -self.sdf_func(x) + sdf_dict = misc.convert_to_dict(sdf, ("sdf",)) + sdf_derives_dict = {} + if compute_sdf_derivatives: + sdf_derives = -self.sdf_derivatives(x) + sdf_derives_dict = misc.convert_to_dict( + sdf_derives, tuple(f"sdf__{key}" for key in self.dim_keys) + ) + else: + sdf_dict = {} + sdf_derives_dict = {} + x_dict = misc.convert_to_dict(x, self.dim_keys) + + return {**x_dict, **sdf_dict, **sdf_derives_dict} + + def sample_boundary( + self, + n: int, + random: Literal["pseudo", "Halton", "LHS"] = "pseudo", + criteria: Optional[Callable[..., np.ndarray]] = None, + evenly: bool = False, + ) -> Dict[str, np.ndarray]: + """Compute the random points in the geometry and return those meet criteria. + + Args: + n (int): Number of points. + random (Literal["pseudo", "Halton", "LHS"]): Random method. Defaults to "pseudo". + pseudo: Pseudo random. + Halton: Halton sequence. + LHS: Latin Hypercube Sampling. + criteria (Optional[Callable[..., np.ndarray]]): Criteria function. Given + coords from different dimension and return a boolean array with shape [n,]. + Defaults to None. + evenly (bool): Evenly sample points. Defaults to False. + + Returns: + Dict[str, np.ndarray]: Random points in the geometry. The shape is [N, D]. + their normal vectors. The shape is [N, D]. + their area. The shape is [N, 1].(only if the geometry is a mesh) + + Examples: + >>> import numpy as np + >>> import ppsci + >>> np.random.seed(42) + >>> interval = ppsci.geometry.Interval(0, 1) + >>> interval.sample_boundary(2) + {'x': array([[0.], + [1.]], dtype=float32), 'normal_x': array([[-1.], + [ 1.]], dtype=float32)} + >>> rectangle = ppsci.geometry.Rectangle((0, 0), (1, 1)) + >>> rectangle.sample_boundary(2) + {'x': array([[1.], + [0.]], dtype=float32), 'y': array([[0.49816048], + [0.19714284]], dtype=float32), 'normal_x': array([[ 1.], + [-1.]], dtype=float32), 'normal_y': array([[0.], + [0.]], dtype=float32)} + >>> cuboid = ppsci.geometry.Cuboid((0, 0, 0), (1, 1, 1)) + >>> cuboid.sample_boundary(2) + {'x': array([[0.83244264], + [0.18182497]], dtype=float32), 'y': array([[0.21233912], + [0.1834045 ]], dtype=float32), 'z': array([[0.], + [1.]], dtype=float32), 'normal_x': array([[0.], + [0.]], dtype=float32), 'normal_y': array([[0.], + [0.]], dtype=float32), 'normal_z': array([[-1.], + [ 1.]], dtype=float32)} + """ + x = np.empty(shape=(n, self.ndim), dtype=paddle.get_default_dtype()) + _size, _ntry, _nsuc = 0, 0, 0 + while _size < n: + if evenly: + if ( + misc.typename(self) == "TimeXGeometry" + and misc.typename(self.geometry) == "Mesh" + ): + points, normal, area = self.uniform_boundary_points(n) + else: + points = self.uniform_boundary_points(n) + else: + if ( + misc.typename(self) == "TimeXGeometry" + and misc.typename(self.geometry) == "Mesh" + ): + points, normal, area = self.random_boundary_points(n, random) + else: + if misc.typename(self) == "TimeXGeometry": + points = self.random_boundary_points(n, random, criteria) + else: + points = self.random_boundary_points(n, random) + + if criteria is not None: + criteria_mask = criteria(*np.split(points, self.ndim, axis=1)).flatten() + points = points[criteria_mask] + + if len(points) > n - _size: + points = points[: n - _size] + x[_size : _size + len(points)] = points + + _size += len(points) + _ntry += 1 + if len(points) > 0: + _nsuc += 1 + + if _ntry >= 10000 and _nsuc == 0: + raise ValueError( + "Sample boundary points failed, " + "please check correctness of geometry and given criteria." + ) + + if not ( + misc.typename(self) == "TimeXGeometry" + and misc.typename(self.geometry) == "Mesh" + ): + normal = self.boundary_normal(x) + + normal_dict = misc.convert_to_dict( + normal[:, 1:] if "t" in self.dim_keys else normal, + [f"normal_{key}" for key in self.dim_keys if key != "t"], + ) + x_dict = misc.convert_to_dict(x, self.dim_keys) + if ( + misc.typename(self) == "TimeXGeometry" + and misc.typename(self.geometry) == "Mesh" + ): + area_dict = misc.convert_to_dict(area[:, 1:], ["area"]) + return {**x_dict, **normal_dict, **area_dict} + + return {**x_dict, **normal_dict} + + @abc.abstractmethod + def random_points( + self, n: int, random: Literal["pseudo", "Halton", "LHS"] = "pseudo" + ) -> np.ndarray: + """Compute the random points in the geometry. + + Args: + n (int): Number of points. + random (Literal["pseudo", "Halton", "LHS"]): Random method. Defaults to "pseudo". + pseudo: Pseudo random. + Halton: Halton sequence. + LHS: Latin Hypercube Sampling. + + Returns: + np.ndarray: Random points in the geometry. The shape is [N, D]. + + Examples: + >>> import numpy as np + >>> import ppsci + >>> np.random.seed(42) + >>> interval = ppsci.geometry.Interval(0, 1) + >>> interval.random_points(2) + array([[0.37454012], + [0.9507143 ]], dtype=float32) + >>> rectangle = ppsci.geometry.Rectangle((0, 0), (1, 1)) + >>> rectangle.random_points(2) + array([[0.7319939 , 0.5986585 ], + [0.15601864, 0.15599452]], dtype=float32) + >>> cuboid = ppsci.geometry.Cuboid((0, 0, 0), (1, 1, 1)) + >>> cuboid.random_points(2) + array([[0.05808361, 0.8661761 , 0.601115 ], + [0.7080726 , 0.02058449, 0.96990985]], dtype=float32) + """ + + def uniform_boundary_points(self, n: int) -> np.ndarray: + """Compute the equi-spaced points on the boundary(not implemented). + + Args: + n (int): Number of points. + + Returns: + np.ndarray: Random points on the boundary. The shape is [N, D]. + """ + logger.warning( + f"{self}.uniform_boundary_points not implemented. " + f"Use random_boundary_points instead." + ) + return self.random_boundary_points(n) + + @abc.abstractmethod + def random_boundary_points( + self, n: int, random: Literal["pseudo", "Halton", "LHS"] = "pseudo" + ) -> np.ndarray: + """Compute the random points on the boundary. + + Args: + n (int): Number of points. + random (Literal["pseudo", "Halton", "LHS"]): Random method. Defaults to "pseudo". + pseudo: Pseudo random. + Halton: Halton sequence. + LHS: Latin Hypercube Sampling. + + Returns: + np.ndarray: Random points on the boundary. The shape is [N, D]. + + Examples: + >>> import numpy as np + >>> import ppsci + >>> np.random.seed(42) + >>> interval = ppsci.geometry.Interval(0, 1) + >>> interval.random_boundary_points(2) + array([[0.], + [1.]], dtype=float32) + >>> rectangle = ppsci.geometry.Rectangle((0, 0), (1, 1)) + >>> rectangle.random_boundary_points(2) + array([[1. , 0.49816048], + [0. , 0.19714284]], dtype=float32) + >>> cuboid = ppsci.geometry.Cuboid((0, 0, 0), (1, 1, 1)) + >>> cuboid.random_boundary_points(2) + array([[0.83244264, 0.21233912, 0. ], + [0.18182497, 0.1834045 , 1. ]], dtype=float32) + """ + + def periodic_point(self, x: np.ndarray, component: int): + """Compute the periodic image of x(not implemented).""" + raise NotImplementedError(f"{self}.periodic_point to be implemented") + + def sdf_derivatives(self, x: np.ndarray, epsilon: float = 1e-4) -> np.ndarray: + """Compute derivatives of SDF function. + + Args: + x (np.ndarray): Points for computing SDF derivatives using central + difference. The shape is [N, D], D is the number of dimension of + geometry. + epsilon (float): Derivative step. Defaults to 1e-4. + + Returns: + np.ndarray: Derivatives of corresponding SDF function. + The shape is [N, D]. D is the number of dimension of geometry. + + Examples: + >>> import numpy as np + >>> import ppsci + >>> interval = ppsci.geometry.Interval(0, 1) + >>> x = np.array([[0], [0.5], [1.5]]) + >>> interval.sdf_derivatives(x) + array([[-1.], + [ 0.], + [ 1.]]) + >>> rectangle = ppsci.geometry.Rectangle((0, 0), (1, 1)) + >>> x = np.array([[0.0, 0.0], [0.5, 0.5], [1.5, 1.5]]) + >>> rectangle.sdf_derivatives(x) + array([[-0.5 , -0.5 ], + [ 0. , 0. ], + [ 0.70710678, 0.70710678]]) + >>> cuboid = ppsci.geometry.Cuboid((0, 0, 0), (1, 1, 1)) + >>> x = np.array([[0, 0, 0], [0.5, 0.5, 0.5], [1, 1, 1]]) + >>> cuboid.sdf_derivatives(x) + array([[-0.5, -0.5, -0.5], + [ 0. , 0. , 0. ], + [ 0.5, 0.5, 0.5]]) + """ + if not hasattr(self, "sdf_func"): + raise NotImplementedError( + f"{misc.typename(self)}.sdf_func should be implemented " + "when using 'sdf_derivatives'." + ) + # Only compute sdf derivatives for those already implement `sdf_func` method. + sdf_derives = np.empty_like(x) + for i in range(self.ndim): + h = np.zeros_like(x) + h[:, i] += epsilon / 2 + derives_at_i = (self.sdf_func(x + h) - self.sdf_func(x - h)) / epsilon + sdf_derives[:, i : i + 1] = derives_at_i + return sdf_derives + + def union(self, other: "Geometry") -> "Geometry": + """CSG Union. + + Args: + other (Geometry): The other geometry. + + Returns: + Geometry: The union of two geometries. + + Examples: + >>> import numpy as np + >>> import ppsci + >>> interval1 = ppsci.geometry.Interval(0, 1) + >>> interval2 = ppsci.geometry.Interval(0.5, 1.5) + >>> union = interval1.union(interval2) + >>> union.bbox + (array([[0.]]), array([[1.5]])) + >>> rectangle1 = ppsci.geometry.Rectangle((0, 0), (2, 3)) + >>> rectangle2 = ppsci.geometry.Rectangle((0, 0), (3, 2)) + >>> union = rectangle1.union(rectangle2) + >>> union.bbox + (array([0., 0.], dtype=float32), array([3., 3.], dtype=float32)) + >>> cuboid1 = ppsci.geometry.Cuboid((0, 0, 0), (1, 2, 2)) + >>> cuboid2 = ppsci.geometry.Cuboid((0, 0, 0), (2, 1, 1)) + >>> union = cuboid1 | cuboid2 + >>> union.bbox + (array([0., 0., 0.], dtype=float32), array([2., 2., 2.], dtype=float32)) + """ + from ppsci.geometry import csg + + return csg.CSGUnion(self, other) + + def __or__(self, other: "Geometry") -> "Geometry": + """CSG Union. + + Args: + other (Geometry): The other geometry. + + Returns: + Geometry: The union of two geometries. + + Examples: + >>> import numpy as np + >>> import ppsci + >>> interval1 = ppsci.geometry.Interval(0, 1) + >>> interval2 = ppsci.geometry.Interval(0.5, 1.5) + >>> union = interval1.__or__(interval2) + >>> union.bbox + (array([[0.]]), array([[1.5]])) + >>> rectangle1 = ppsci.geometry.Rectangle((0, 0), (2, 3)) + >>> rectangle2 = ppsci.geometry.Rectangle((0, 0), (3, 2)) + >>> union = rectangle1.__or__(rectangle2) + >>> union.bbox + (array([0., 0.], dtype=float32), array([3., 3.], dtype=float32)) + >>> cuboid1 = ppsci.geometry.Cuboid((0, 0, 0), (1, 2, 2)) + >>> cuboid2 = ppsci.geometry.Cuboid((0, 0, 0), (2, 1, 1)) + >>> union = cuboid1 | cuboid2 + >>> union.bbox + (array([0., 0., 0.], dtype=float32), array([2., 2., 2.], dtype=float32)) + """ + from ppsci.geometry import csg + + return csg.CSGUnion(self, other) + + def difference(self, other: "Geometry") -> "Geometry": + """CSG Difference. + + Args: + other (Geometry): The other geometry. + + Returns: + Geometry: The difference of two geometries. + + Examples: + >>> import numpy as np + >>> import ppsci + >>> interval1 = ppsci.geometry.Interval(0.0, 2.0) + >>> interval2 = ppsci.geometry.Interval(1.0, 3.0) + >>> difference = interval1.difference(interval2) + >>> difference.bbox + (array([[0.]]), array([[2.]])) + >>> rectangle1 = ppsci.geometry.Rectangle((0.0, 0.0), (2.0, 3.0)) + >>> rectangle2 = ppsci.geometry.Rectangle((1.0, 1.0), (2.0, 2.0)) + >>> difference = rectangle1.difference(rectangle2) + >>> difference.bbox + (array([0., 0.], dtype=float32), array([2., 3.], dtype=float32)) + >>> cuboid1 = ppsci.geometry.Cuboid((0, 0, 0), (1, 2, 2)) + >>> cuboid2 = ppsci.geometry.Cuboid((0, 0, 0), (2, 1, 1)) + >>> difference = cuboid1 - cuboid2 + >>> difference.bbox + (array([0., 0., 0.], dtype=float32), array([1., 2., 2.], dtype=float32)) + """ + from ppsci.geometry import csg + + return csg.CSGDifference(self, other) + + def __sub__(self, other: "Geometry") -> "Geometry": + """CSG Difference. + + Args: + other (Geometry): The other geometry. + + Returns: + Geometry: The difference of two geometries. + + Examples: + >>> import numpy as np + >>> import ppsci + >>> interval1 = ppsci.geometry.Interval(0.0, 2.0) + >>> interval2 = ppsci.geometry.Interval(1.0, 3.0) + >>> difference = interval1.__sub__(interval2) + >>> difference.bbox + (array([[0.]]), array([[2.]])) + >>> rectangle1 = ppsci.geometry.Rectangle((0.0, 0.0), (2.0, 3.0)) + >>> rectangle2 = ppsci.geometry.Rectangle((1.0, 1.0), (2.0, 2.0)) + >>> difference = rectangle1.__sub__(rectangle2) + >>> difference.bbox + (array([0., 0.], dtype=float32), array([2., 3.], dtype=float32)) + >>> cuboid1 = ppsci.geometry.Cuboid((0, 0, 0), (1, 2, 2)) + >>> cuboid2 = ppsci.geometry.Cuboid((0, 0, 0), (2, 1, 1)) + >>> difference = cuboid1 - cuboid2 + >>> difference.bbox + (array([0., 0., 0.], dtype=float32), array([1., 2., 2.], dtype=float32)) + """ + from ppsci.geometry import csg + + return csg.CSGDifference(self, other) + + def intersection(self, other: "Geometry") -> "Geometry": + """CSG Intersection. + + Args: + other (Geometry): The other geometry. + + Returns: + Geometry: The intersection of two geometries. + + Examples: + >>> import numpy as np + >>> import ppsci + >>> interval1 = ppsci.geometry.Interval(0.0, 1.0) + >>> interval2 = ppsci.geometry.Interval(0.5, 1.5) + >>> intersection = interval1.intersection(interval2) + >>> intersection.bbox + (array([[0.5]]), array([[1.]])) + >>> rectangle1 = ppsci.geometry.Rectangle((0.0, 0.0), (2.0, 3.0)) + >>> rectangle2 = ppsci.geometry.Rectangle((0.0, 0.0), (3.0, 2.0)) + >>> intersection = rectangle1.intersection(rectangle2) + >>> intersection.bbox + (array([0., 0.], dtype=float32), array([2., 2.], dtype=float32)) + >>> cuboid1 = ppsci.geometry.Cuboid((0, 0, 0), (1, 2, 2)) + >>> cuboid2 = ppsci.geometry.Cuboid((0, 0, 0), (2, 1, 1)) + >>> intersection = cuboid1 & cuboid2 + >>> intersection.bbox + (array([0., 0., 0.], dtype=float32), array([1., 1., 1.], dtype=float32)) + """ + from ppsci.geometry import csg + + return csg.CSGIntersection(self, other) + + def __and__(self, other: "Geometry") -> "Geometry": + """CSG Intersection. + + Args: + other (Geometry): The other geometry. + + Returns: + Geometry: The intersection of two geometries. + + Examples: + >>> import numpy as np + >>> import ppsci + >>> interval1 = ppsci.geometry.Interval(0.0, 1.0) + >>> interval2 = ppsci.geometry.Interval(0.5, 1.5) + >>> intersection = interval1.__and__(interval2) + >>> intersection.bbox + (array([[0.5]]), array([[1.]])) + >>> rectangle1 = ppsci.geometry.Rectangle((0.0, 0.0), (2.0, 3.0)) + >>> rectangle2 = ppsci.geometry.Rectangle((0.0, 0.0), (3.0, 2.0)) + >>> intersection = rectangle1.__and__(rectangle2) + >>> intersection.bbox + (array([0., 0.], dtype=float32), array([2., 2.], dtype=float32)) + >>> cuboid1 = ppsci.geometry.Cuboid((0, 0, 0), (1, 2, 2)) + >>> cuboid2 = ppsci.geometry.Cuboid((0, 0, 0), (2, 1, 1)) + >>> intersection = cuboid1 & cuboid2 + >>> intersection.bbox + (array([0., 0., 0.], dtype=float32), array([1., 1., 1.], dtype=float32)) + """ + from ppsci.geometry import csg + + return csg.CSGIntersection(self, other) + + def __str__(self) -> str: + """Return the name of class. + + Returns: + str: Meta information of geometry. + + Examples: + >>> import ppsci + >>> interval = ppsci.geometry.Interval(0, 1) + >>> interval.__str__() + "Interval, ndim = 1, bbox = (array([[0]]), array([[1]])), diam = 1, dim_keys = ('x',)" + >>> rectangle = ppsci.geometry.Rectangle((0, 0), (1, 1)) + >>> rectangle.__str__() + "Rectangle, ndim = 2, bbox = (array([0., 0.], dtype=float32), array([1., 1.], dtype=float32)), diam = 1.4142135381698608, dim_keys = ('x', 'y')" + >>> cuboid = ppsci.geometry.Cuboid((0, 0, 0), (1, 1, 1)) + >>> cuboid.__str__() + "Cuboid, ndim = 3, bbox = (array([0., 0., 0.], dtype=float32), array([1., 1., 1.], dtype=float32)), diam = 1.7320507764816284, dim_keys = ('x', 'y', 'z')" + """ + return ", ".join( + [ + self.__class__.__name__, + f"ndim = {self.ndim}", + f"bbox = {self.bbox}", + f"diam = {self.diam}", + f"dim_keys = {self.dim_keys}", + ] + ) diff --git a/examples/smc_reac/ppsci/geometry/geometry_1d.py b/examples/smc_reac/ppsci/geometry/geometry_1d.py new file mode 100644 index 0000000000..d5de01fe56 --- /dev/null +++ b/examples/smc_reac/ppsci/geometry/geometry_1d.py @@ -0,0 +1,119 @@ +# 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. + +""" +Code below is heavily based on [https://github.com/lululxvi/deepxde](https://github.com/lululxvi/deepxde) +""" + +from __future__ import annotations + +import numpy as np +import paddle + +from ppsci.geometry import geometry +from ppsci.geometry.sampler import sample +from ppsci.utils import misc + + +class Interval(geometry.Geometry): + """Class for interval. + + Args: + l (float): Left position of interval. + r (float): Right position of interval. + + Examples: + >>> import ppsci + >>> geom = ppsci.geometry.Interval(-1, 1) + """ + + def __init__(self, l: float, r: float): + super().__init__(1, (np.array([[l]]), np.array([[r]])), r - l) + self.l = l + self.r = r + + def is_inside(self, x: np.ndarray): + return ((self.l <= x) & (x <= self.r)).flatten() + + def on_boundary(self, x: np.ndarray): + return (np.isclose(x, self.l) | np.isclose(x, self.r)).flatten() + + def boundary_normal(self, x: np.ndarray): + return -np.isclose(x, self.l).astype(paddle.get_default_dtype()) + np.isclose( + x, self.r + ).astype(paddle.get_default_dtype()) + + def uniform_points(self, n: int, boundary: bool = True): + if boundary: + return np.linspace( + self.l, self.r, n, dtype=paddle.get_default_dtype() + ).reshape([-1, 1]) + return np.linspace( + self.l, self.r, n + 1, endpoint=False, dtype=paddle.get_default_dtype() + )[1:].reshape([-1, 1]) + + def random_points(self, n: int, random: str = "pseudo"): + x = sample(n, 1, random) + return (self.l + x * self.diam).astype(paddle.get_default_dtype()) + + def uniform_boundary_points(self, n: int): + if n == 1: + return np.array([[self.l]], dtype=paddle.get_default_dtype()) + xl = np.full([n // 2, 1], self.l, dtype=paddle.get_default_dtype()) + xr = np.full([n - n // 2, 1], self.r, dtype=paddle.get_default_dtype()) + return np.concatenate((xl, xr), axis=0) + + def random_boundary_points(self, n: int, random: str = "pseudo"): + if n == 2: + return np.array([[self.l], [self.r]], dtype=paddle.get_default_dtype()) + return ( + np.random.choice([self.l, self.r], n) + .reshape([-1, 1]) + .astype(paddle.get_default_dtype()) + ) + + def periodic_point(self, x: np.ndarray, component: int = 0): + x_array = misc.convert_to_array(x, self.dim_keys) + periodic_x = x_array + periodic_x[np.isclose(x_array, self.l)] = self.r + periodic_x[np.isclose(x_array, self.r)] = self.l + periodic_x_normal = self.boundary_normal(periodic_x) + + periodic_x = misc.convert_to_dict(periodic_x, self.dim_keys) + periodic_x_normal = misc.convert_to_dict( + periodic_x_normal, [f"normal_{k}" for k in self.dim_keys] + ) + return {**periodic_x, **periodic_x_normal} + + def sdf_func(self, points: np.ndarray) -> np.ndarray: + """Compute signed distance field + + Args: + points (np.ndarray): The coordinate points used to calculate the SDF value, + the shape is [N, 1] + + Returns: + np.ndarray: SDF values of input points without squared, the shape is [N, 1]. + + NOTE: This function usually returns ndarray with negative values, because + according to the definition of SDF, the SDF value of the coordinate point inside + the object(interior points) is negative, the outside is positive, and the edge + is 0. Therefore, when used for weighting, a negative sign is often added before + the result of this function. + """ + if points.shape[1] != self.ndim: + raise ValueError( + f"Shape of given points should be [*, {self.ndim}], but got {points.shape}" + ) + return -((self.r - self.l) / 2 - np.abs(points - (self.l + self.r) / 2)) diff --git a/examples/smc_reac/ppsci/geometry/geometry_2d.py b/examples/smc_reac/ppsci/geometry/geometry_2d.py new file mode 100644 index 0000000000..2df6293b27 --- /dev/null +++ b/examples/smc_reac/ppsci/geometry/geometry_2d.py @@ -0,0 +1,706 @@ +# 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. + +""" +Code below is heavily based on [https://github.com/lululxvi/deepxde](https://github.com/lululxvi/deepxde) +""" + +from __future__ import annotations + +from typing import Tuple + +import numpy as np +import paddle +from scipy import spatial + +from ppsci.geometry import geometry +from ppsci.geometry import geometry_nd +from ppsci.geometry import sampler + + +class Disk(geometry.Geometry): + """Class for disk geometry + + Args: + center (Tuple[float, float]): Center point of disk [x0, y0]. + radius (float): Radius of disk. + + Examples: + >>> import ppsci + >>> geom = ppsci.geometry.Disk((0.0, 0.0), 1.0) + """ + + def __init__(self, center: Tuple[float, float], radius: float): + self.center = np.array(center, dtype=paddle.get_default_dtype()) + self.radius = radius + super().__init__(2, (self.center - radius, self.center + radius), 2 * radius) + + def is_inside(self, x): + return np.linalg.norm(x - self.center, axis=1) <= self.radius + + def on_boundary(self, x): + return np.isclose(np.linalg.norm(x - self.center, axis=1), self.radius) + + def boundary_normal(self, x): + ox = x - self.center + ox_len = np.linalg.norm(ox, axis=1, keepdims=True) + ox = (ox / ox_len) * np.isclose(ox_len, self.radius).astype( + paddle.get_default_dtype() + ) + return ox + + def random_points(self, n, random="pseudo"): + # http://mathworld.wolfram.com/DiskPointPicking.html + rng = sampler.sample(n, 2, random) + r, theta = rng[:, 0], 2 * np.pi * rng[:, 1] + x = np.sqrt(r) * np.cos(theta) + y = np.sqrt(r) * np.sin(theta) + return self.radius * np.stack((x, y), axis=1) + self.center + + def uniform_boundary_points(self, n): + theta = np.linspace( + 0, 2 * np.pi, num=n, endpoint=False, dtype=paddle.get_default_dtype() + ) + X = np.stack((np.cos(theta), np.sin(theta)), axis=1) + return self.radius * X + self.center + + def random_boundary_points(self, n, random="pseudo"): + theta = 2 * np.pi * sampler.sample(n, 1, random) + X = np.concatenate((np.cos(theta), np.sin(theta)), axis=1) + return self.radius * X + self.center + + def sdf_func(self, points: np.ndarray) -> np.ndarray: + """Compute signed distance field. + + Args: + points (np.ndarray): The coordinate points used to calculate the SDF value, + the shape is [N, 2] + + Returns: + np.ndarray: SDF values of input points without squared, the shape is [N, 1]. + + NOTE: This function usually returns ndarray with negative values, because + according to the definition of SDF, the SDF value of the coordinate point inside + the object(interior points) is negative, the outside is positive, and the edge + is 0. Therefore, when used for weighting, a negative sign is often added before + the result of this function. + """ + if points.shape[1] != self.ndim: + raise ValueError( + f"Shape of given points should be [*, {self.ndim}], but got {points.shape}" + ) + sdf = self.radius - np.linalg.norm(points - self.center, axis=1) + sdf = -sdf[..., np.newaxis] + return sdf + + +class Rectangle(geometry_nd.Hypercube): + """Class for rectangle geometry + + Args: + xmin (Tuple[float, float]): Bottom left corner point, [x0, y0]. + xmax (Tuple[float, float]): Top right corner point, [x1, y1]. + + Examples: + >>> import ppsci + >>> geom = ppsci.geometry.Rectangle((0.0, 0.0), (1.0, 1.0)) + """ + + def __init__(self, xmin, xmax): + super().__init__(xmin, xmax) + self.perimeter = 2 * np.sum(self.xmax - self.xmin) + self.area = np.prod(self.xmax - self.xmin) + + def uniform_boundary_points(self, n): + nx, ny = np.ceil(n / self.perimeter * (self.xmax - self.xmin)).astype(int) + bottom = np.hstack( + ( + np.linspace( + self.xmin[0], + self.xmax[0], + nx, + endpoint=False, + dtype=paddle.get_default_dtype(), + ).reshape([nx, 1]), + np.full([nx, 1], self.xmin[1], dtype=paddle.get_default_dtype()), + ) + ) + right = np.hstack( + ( + np.full([ny, 1], self.xmax[0], dtype=paddle.get_default_dtype()), + np.linspace( + self.xmin[1], + self.xmax[1], + ny, + endpoint=False, + dtype=paddle.get_default_dtype(), + ).reshape([ny, 1]), + ) + ) + top = np.hstack( + ( + np.linspace( + self.xmin[0], self.xmax[0], nx + 1, dtype=paddle.get_default_dtype() + )[1:].reshape([nx, 1]), + np.full([nx, 1], self.xmax[1], dtype=paddle.get_default_dtype()), + ) + ) + left = np.hstack( + ( + np.full([ny, 1], self.xmin[0], dtype=paddle.get_default_dtype()), + np.linspace( + self.xmin[1], self.xmax[1], ny + 1, dtype=paddle.get_default_dtype() + )[1:].reshape([ny, 1]), + ) + ) + x = np.vstack((bottom, right, top, left)) + if len(x) > n: + x = x[0:n] + return x + + def random_boundary_points(self, n, random="pseudo"): + l1 = self.xmax[0] - self.xmin[0] + l2 = l1 + self.xmax[1] - self.xmin[1] + l3 = l2 + l1 + u = np.ravel(sampler.sample(n + 10, 1, random)) + # Remove the possible points very close to the corners + u = u[~np.isclose(u, l1 / self.perimeter)] + u = u[~np.isclose(u, l3 / self.perimeter)] + u = u[0:n] + + u *= self.perimeter + x = [] + for l in u: + if l < l1: + x.append([self.xmin[0] + l, self.xmin[1]]) + elif l < l2: + x.append([self.xmax[0], self.xmin[1] + (l - l1)]) + elif l < l3: + x.append([self.xmax[0] - (l - l2), self.xmax[1]]) + else: + x.append([self.xmin[0], self.xmax[1] - (l - l3)]) + return np.vstack(x) + + @staticmethod + def is_valid(vertices): + """Check if the geometry is a Rectangle.""" + return ( + len(vertices) == 4 + and np.isclose(np.prod(vertices[1] - vertices[0]), 0) + and np.isclose(np.prod(vertices[2] - vertices[1]), 0) + and np.isclose(np.prod(vertices[3] - vertices[2]), 0) + and np.isclose(np.prod(vertices[0] - vertices[3]), 0) + ) + + def sdf_func(self, points: np.ndarray) -> np.ndarray: + """Compute signed distance field. + + Args: + points (np.ndarray): The coordinate points used to calculate the SDF value, + the shape of the array is [N, 2]. + + Returns: + np.ndarray: SDF values of input points without squared, the shape is [N, 1]. + + NOTE: This function usually returns ndarray with negative values, because + according to the definition of SDF, the SDF value of the coordinate point inside + the object(interior points) is negative, the outside is positive, and the edge + is 0. Therefore, when used for weighting, a negative sign is often added before + the result of this function. + """ + if points.shape[1] != self.ndim: + raise ValueError( + f"Shape of given points should be [*, {self.ndim}], but got {points.shape}" + ) + center = (self.xmin + self.xmax) / 2 + dist_to_boundary = ( + np.abs(points - center) - np.array([self.xmax - self.xmin]) / 2 + ) + return ( + np.linalg.norm(np.maximum(dist_to_boundary, 0), axis=1) + + np.minimum(np.max(dist_to_boundary, axis=1), 0) + ).reshape(-1, 1) + + +class Triangle(geometry.Geometry): + """Class for Triangle + + The order of vertices can be in a clockwise or counterclockwise direction. The + vertices will be re-ordered in counterclockwise (right hand rule). + + Args: + x1 (Tuple[float, float]): First point of Triangle [x0, y0]. + x2 (Tuple[float, float]): Second point of Triangle [x1, y1]. + x3 (Tuple[float, float]): Third point of Triangle [x2, y2]. + + Examples: + >>> import ppsci + >>> geom = ppsci.geometry.Triangle((0, 0), (1, 0), (0, 1)) + """ + + def __init__(self, x1, x2, x3): + self.area = polygon_signed_area([x1, x2, x3]) + # Clockwise + if self.area < 0: + self.area = -self.area + x2, x3 = x3, x2 + + self.x1 = np.array(x1, dtype=paddle.get_default_dtype()) + self.x2 = np.array(x2, dtype=paddle.get_default_dtype()) + self.x3 = np.array(x3, dtype=paddle.get_default_dtype()) + + self.v12 = self.x2 - self.x1 + self.v23 = self.x3 - self.x2 + self.v31 = self.x1 - self.x3 + self.l12 = np.linalg.norm(self.v12) + self.l23 = np.linalg.norm(self.v23) + self.l31 = np.linalg.norm(self.v31) + self.n12 = self.v12 / self.l12 + self.n23 = self.v23 / self.l23 + self.n31 = self.v31 / self.l31 + self.n12_normal = clockwise_rotation_90(self.n12) + self.n23_normal = clockwise_rotation_90(self.n23) + self.n31_normal = clockwise_rotation_90(self.n31) + self.perimeter = self.l12 + self.l23 + self.l31 + + super().__init__( + 2, + (np.minimum(x1, np.minimum(x2, x3)), np.maximum(x1, np.maximum(x2, x3))), + self.l12 + * self.l23 + * self.l31 + / ( + self.perimeter + * (self.l12 + self.l23 - self.l31) + * (self.l23 + self.l31 - self.l12) + * (self.l31 + self.l12 - self.l23) + ) + ** 0.5, + ) + + def is_inside(self, x): + # https://stackoverflow.com/a/2049593/12679294 + _sign = np.stack( + [ + np.cross(self.v12, x - self.x1), + np.cross(self.v23, x - self.x2), + np.cross(self.v31, x - self.x3), + ], + axis=1, + ) + return ~(np.any(_sign > 0, axis=-1) & np.any(_sign < 0, axis=-1)) + + def on_boundary(self, x): + l1 = np.linalg.norm(x - self.x1, axis=-1) + l2 = np.linalg.norm(x - self.x2, axis=-1) + l3 = np.linalg.norm(x - self.x3, axis=-1) + return np.any( + np.isclose( + [l1 + l2 - self.l12, l2 + l3 - self.l23, l3 + l1 - self.l31], + 0, + atol=1e-6, + ), + axis=0, + ) + + def boundary_normal(self, x): + l1 = np.linalg.norm(x - self.x1, axis=-1, keepdims=True) + l2 = np.linalg.norm(x - self.x2, axis=-1, keepdims=True) + l3 = np.linalg.norm(x - self.x3, axis=-1, keepdims=True) + on12 = np.isclose(l1 + l2, self.l12) + on23 = np.isclose(l2 + l3, self.l23) + on31 = np.isclose(l3 + l1, self.l31) + # Check points on the vertexes + if np.any(np.count_nonzero(np.hstack([on12, on23, on31]), axis=-1) > 1): + raise ValueError( + "{}.boundary_normal do not accept points on the vertexes.".format( + self.__class__.__name__ + ) + ) + return self.n12_normal * on12 + self.n23_normal * on23 + self.n31_normal * on31 + + def random_points(self, n, random="pseudo"): + # There are two methods for triangle point picking. + # Method 1 (used here): + # - https://math.stackexchange.com/questions/18686/uniform-random-point-in-triangle + # Method 2: + # - http://mathworld.wolfram.com/TrianglePointPicking.html + # - https://hbfs.wordpress.com/2010/10/05/random-points-in-a-triangle-generating-random-sequences-ii/ + # - https://stackoverflow.com/questions/19654251/random-point-inside-triangle-inside-java + sqrt_r1 = np.sqrt(np.random.rand(n, 1)) + r2 = np.random.rand(n, 1) + return ( + (1 - sqrt_r1) * self.x1 + + sqrt_r1 * (1 - r2) * self.x2 + + r2 * sqrt_r1 * self.x3 + ) + + def uniform_boundary_points(self, n): + density = n / self.perimeter + x12 = ( + np.linspace( + 0, + 1, + num=int(np.ceil(density * self.l12)), + endpoint=False, + dtype=paddle.get_default_dtype(), + )[:, None] + * self.v12 + + self.x1 + ) + x23 = ( + np.linspace( + 0, + 1, + num=int(np.ceil(density * self.l23)), + endpoint=False, + dtype=paddle.get_default_dtype(), + )[:, None] + * self.v23 + + self.x2 + ) + x31 = ( + np.linspace( + 0, + 1, + num=int(np.ceil(density * self.l31)), + endpoint=False, + dtype=paddle.get_default_dtype(), + )[:, None] + * self.v31 + + self.x3 + ) + x = np.vstack((x12, x23, x31)) + if len(x) > n: + x = x[0:n] + return x + + def random_boundary_points(self, n, random="pseudo"): + u = np.ravel(sampler.sample(n + 2, 1, random)) + # Remove the possible points very close to the corners + u = u[np.logical_not(np.isclose(u, self.l12 / self.perimeter))] + u = u[np.logical_not(np.isclose(u, (self.l12 + self.l23) / self.perimeter))] + u = u[:n] + + u *= self.perimeter + x = [] + for l in u: + if l < self.l12: + x.append(l * self.n12 + self.x1) + elif l < self.l12 + self.l23: + x.append((l - self.l12) * self.n23 + self.x2) + else: + x.append((l - self.l12 - self.l23) * self.n31 + self.x3) + return np.vstack(x) + + def sdf_func(self, points: np.ndarray) -> np.ndarray: + """Compute signed distance field. + + Args: + points (np.ndarray): The coordinate points used to calculate the SDF value, + the shape of the array is [N, 2]. + + Returns: + np.ndarray: SDF values of input points without squared, the shape is [N, 1]. + + NOTE: This function usually returns ndarray with negative values, because + according to the definition of SDF, the SDF value of the coordinate point inside + the object(interior points) is negative, the outside is positive, and the edge + is 0. Therefore, when used for weighting, a negative sign is often added before + the result of this function. + """ + if points.shape[1] != self.ndim: + raise ValueError( + f"Shape of given points should be [*, {self.ndim}], but got {points.shape}" + ) + v1p = points - self.x1 # v1p: vector from x1 to points + v2p = points - self.x2 + v3p = points - self.x3 + # vv12_p: vertical vector of points to v12(If the vertical point is in the extension of v12, + # the vector will be the vector from x1 to points) + vv12_p = ( + self.v12 + * np.clip(np.dot(v1p, self.v12.reshape(2, -1)) / self.l12**2, 0, 1) + - v1p + ) + vv23_p = ( + self.v23 + * np.clip(np.dot(v2p, self.v23.reshape(2, -1)) / self.l23**2, 0, 1) + - v2p + ) + vv31_p = ( + self.v31 + * np.clip(np.dot(v3p, self.v31.reshape(2, -1)) / self.l31**2, 0, 1) + - v3p + ) + is_inside = self.is_inside(points).reshape(-1, 1) * 2 - 1 + len_vv12_p = np.linalg.norm(vv12_p, axis=1, keepdims=True) + len_vv23_p = np.linalg.norm(vv23_p, axis=1, keepdims=True) + len_vv31_p = np.linalg.norm(vv31_p, axis=1, keepdims=True) + mini_dist = np.minimum(np.minimum(len_vv12_p, len_vv23_p), len_vv31_p) + return is_inside * mini_dist + + +class Polygon(geometry.Geometry): + """Class for simple polygon. + + Args: + vertices (Tuple[Tuple[float, float], ...]): The order of vertices can be in a + clockwise or counter-clockwise direction. The vertices will be re-ordered in + counterclockwise (right hand rule). + + Examples: + >>> import ppsci + >>> geom = ppsci.geometry.Polygon(((0, 0), (1, 0), (2, 1), (2, 2), (0, 2))) + """ + + def __init__(self, vertices): + self.vertices = np.array(vertices, dtype=paddle.get_default_dtype()) + if len(vertices) == 3: + raise ValueError("The polygon is a triangle. Use Triangle instead.") + if Rectangle.is_valid(self.vertices): + raise ValueError("The polygon is a rectangle. Use Rectangle instead.") + + self.area = polygon_signed_area(self.vertices) + # Clockwise + if self.area < 0: + self.area = -self.area + self.vertices = np.flipud(self.vertices) + + self.diagonals = spatial.distance.squareform( + spatial.distance.pdist(self.vertices) + ) + super().__init__( + 2, + (np.amin(self.vertices, axis=0), np.amax(self.vertices, axis=0)), + np.max(self.diagonals), + ) + self.nvertices = len(self.vertices) + self.perimeter = np.sum( + [self.diagonals[i, i + 1] for i in range(-1, self.nvertices - 1)] + ) + self.bbox = np.array( + [np.min(self.vertices, axis=0), np.max(self.vertices, axis=0)], + dtype=paddle.get_default_dtype(), + ) + + self.segments = self.vertices[1:] - self.vertices[:-1] + self.segments = np.vstack((self.vertices[0] - self.vertices[-1], self.segments)) + self.normal = clockwise_rotation_90(self.segments.T).T + self.normal = self.normal / np.linalg.norm(self.normal, axis=1).reshape(-1, 1) + + def is_inside(self, x): + def wn_PnPoly(P, V): + """Winding number algorithm. + + https://en.wikipedia.org/wiki/Point_in_polygon + http://geomalgorithms.com/a03-_inclusion.html + + Args: + P: A point. + V: Vertex points of a polygon. + + Returns: + wn: Winding number (=0 only if P is outside polygon). + """ + wn = np.zeros(len(P)) # Winding number counter + + # Repeat the first vertex at end + # Loop through all edges of the polygon + for i in range(-1, self.nvertices - 1): # Edge from V[i] to V[i+1] + tmp = np.all( + np.hstack( + [ + V[i, 1] <= P[:, 1:2], # Start y <= P[1] + V[i + 1, 1] > P[:, 1:2], # An upward crossing + is_left(V[i], V[i + 1], P) > 0, # P left of edge + ] + ), + axis=-1, + ) + wn[tmp] += 1 # Have a valid up intersect + tmp = np.all( + np.hstack( + [ + V[i, 1] > P[:, 1:2], # Start y > P[1] + V[i + 1, 1] <= P[:, 1:2], # A downward crossing + is_left(V[i], V[i + 1], P) < 0, # P right of edge + ] + ), + axis=-1, + ) + wn[tmp] -= 1 # Have a valid down intersect + return wn + + return wn_PnPoly(x, self.vertices) != 0 + + def on_boundary(self, x): + _on = np.zeros(shape=len(x), dtype=np.int) + for i in range(-1, self.nvertices - 1): + l1 = np.linalg.norm(self.vertices[i] - x, axis=-1) + l2 = np.linalg.norm(self.vertices[i + 1] - x, axis=-1) + _on[np.isclose(l1 + l2, self.diagonals[i, i + 1])] += 1 + return _on > 0 + + def random_points(self, n, random="pseudo"): + x = np.empty((0, 2), dtype=paddle.get_default_dtype()) + vbbox = self.bbox[1] - self.bbox[0] + while len(x) < n: + x_new = sampler.sample(n, 2, "pseudo") * vbbox + self.bbox[0] + x = np.vstack((x, x_new[self.is_inside(x_new)])) + return x[:n] + + def uniform_boundary_points(self, n): + density = n / self.perimeter + x = [] + for i in range(-1, self.nvertices - 1): + x.append( + np.linspace( + 0, + 1, + num=int(np.ceil(density * self.diagonals[i, i + 1])), + endpoint=False, + dtype=paddle.get_default_dtype(), + )[:, None] + * (self.vertices[i + 1] - self.vertices[i]) + + self.vertices[i] + ) + x = np.vstack(x) + if len(x) > n: + x = x[0:n] + return x + + def random_boundary_points(self, n, random="pseudo"): + u = np.ravel(sampler.sample(n + self.nvertices, 1, random)) + # Remove the possible points very close to the corners + l = 0 + for i in range(0, self.nvertices - 1): + l += self.diagonals[i, i + 1] + u = u[np.logical_not(np.isclose(u, l / self.perimeter))] + u = u[:n] + u *= self.perimeter + u.sort() + + x = [] + i = -1 + l0 = 0 + l1 = l0 + self.diagonals[i, i + 1] + v = (self.vertices[i + 1] - self.vertices[i]) / self.diagonals[i, i + 1] + for l in u: + if l > l1: + i += 1 + l0, l1 = l1, l1 + self.diagonals[i, i + 1] + v = (self.vertices[i + 1] - self.vertices[i]) / self.diagonals[i, i + 1] + x.append((l - l0) * v + self.vertices[i]) + return np.vstack(x) + + def sdf_func(self, points: np.ndarray) -> np.ndarray: + """Compute signed distance field. + + Args: + points (np.ndarray): The coordinate points used to calculate the SDF value, + the shape is [N, 2] + Returns: + np.ndarray: SDF values of input points without squared, the shape is [N, 1]. + + NOTE: This function usually returns ndarray with negative values, because + according to the definition of SDF, the SDF value of the coordinate point inside + the object(interior points) is negative, the outside is positive, and the edge + is 0. Therefore, when used for weighting, a negative sign is often added before + the result of this function. + """ + if points.shape[1] != self.ndim: + raise ValueError( + f"Shape of given points should be [*, {self.ndim}], but got {points.shape}" + ) + sdf_value = np.empty((points.shape[0], 1), dtype=paddle.get_default_dtype()) + for n in range(points.shape[0]): + distance = np.dot( + points[n] - self.vertices[0], points[n] - self.vertices[0] + ) + inside_tag = 1.0 + for i in range(self.vertices.shape[0]): + j = (self.vertices.shape[0] - 1) if i == 0 else (i - 1) + # Calculate the shortest distance from point P to each edge. + vector_ij = self.vertices[j] - self.vertices[i] + vector_in = points[n] - self.vertices[i] + distance_vector = vector_in - vector_ij * np.clip( + np.dot(vector_in, vector_ij) / np.dot(vector_ij, vector_ij), + 0.0, + 1.0, + ) + distance = np.minimum( + distance, np.dot(distance_vector, distance_vector) + ) + # Calculate the inside and outside using the Odd-even rule + odd_even_rule_number = np.array( + [ + points[n][1] >= self.vertices[i][1], + points[n][1] < self.vertices[j][1], + vector_ij[0] * vector_in[1] > vector_ij[1] * vector_in[0], + ] + ) + if odd_even_rule_number.all() or np.all(~odd_even_rule_number): + inside_tag *= -1.0 + sdf_value[n] = inside_tag * np.sqrt(distance) + return -sdf_value + + +def polygon_signed_area(vertices): + """The (signed) area of a simple polygon. + + If the vertices are in the counterclockwise direction, then the area is positive; if + they are in the clockwise direction, the area is negative. + + Shoelace formula: https://en.wikipedia.org/wiki/Shoelace_formula + + Args: + vertices (np.ndarray): Polygon vertices with shape of [N, 2]. + + Returns: + float: The (signed) area of a simple polygon. + """ + x, y = zip(*vertices) + x = np.array(list(x) + [x[0]], dtype=paddle.get_default_dtype()) + y = np.array(list(y) + [y[0]], dtype=paddle.get_default_dtype()) + return 0.5 * (np.sum(x[:-1] * y[1:]) - np.sum(x[1:] * y[:-1])) + + +def clockwise_rotation_90(v): + """Rotate a vector of 90 degrees clockwise about the origin. + + Args: + v (np.ndarray): Vector with shape of [2, N]. + + Returns: + np.ndarray: Rotated vector with shape of [2, N]. + """ + return np.array([v[1], -v[0]], dtype=paddle.get_default_dtype()) + + +def is_left(P0, P1, P2): + """Test if a point is Left|On|Right of an infinite line. + + See: the January 2001 Algorithm "Area of 2D and 3D Triangles and Polygons". + + Args: + P0 (np.ndarray): One point in the line. + P1 (np.ndarray): One point in the line. + P2 (np.ndarray): A array of point to be tested with shape of [N, 2]. + + Returns: + np.ndarray: >0 if P2 left of the line through P0 and P1, =0 if P2 on the line, <0 if P2 + right of the line. + """ + return np.cross(P1 - P0, P2 - P0, axis=-1).reshape((-1, 1)) diff --git a/examples/smc_reac/ppsci/geometry/geometry_3d.py b/examples/smc_reac/ppsci/geometry/geometry_3d.py new file mode 100644 index 0000000000..8af958b1b9 --- /dev/null +++ b/examples/smc_reac/ppsci/geometry/geometry_3d.py @@ -0,0 +1,203 @@ +# 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. + +""" +Code below is heavily based on [https://github.com/lululxvi/deepxde](https://github.com/lululxvi/deepxde) +""" + +from __future__ import annotations + +import itertools +from typing import Tuple + +import numpy as np +import paddle + +from ppsci.geometry import geometry_2d +from ppsci.geometry import geometry_nd + + +class Cuboid(geometry_nd.Hypercube): + """Class for Cuboid + + Args: + xmin (Tuple[float, float, float]): Bottom left corner point [x0, y0, z0]. + xmax (Tuple[float, float, float]): Top right corner point [x1, y1, z1]. + + Examples: + >>> import ppsci + >>> geom = ppsci.geometry.Cuboid((0, 0, 0), (1, 1, 1)) + """ + + def __init__( + self, xmin: Tuple[float, float, float], xmax: Tuple[float, float, float] + ): + super().__init__(xmin, xmax) + dx = self.xmax - self.xmin + self.area = 2 * np.sum(dx * np.roll(dx, 2)) + + def random_boundary_points(self, n, random="pseudo"): + pts = [] + density = n / self.area + rect = geometry_2d.Rectangle(self.xmin[:-1], self.xmax[:-1]) + for z in [self.xmin[-1], self.xmax[-1]]: + u = rect.random_points(int(np.ceil(density * rect.area)), random=random) + pts.append( + np.hstack( + (u, np.full((len(u), 1), z, dtype=paddle.get_default_dtype())) + ) + ) + rect = geometry_2d.Rectangle(self.xmin[::2], self.xmax[::2]) + for y in [self.xmin[1], self.xmax[1]]: + u = rect.random_points(int(np.ceil(density * rect.area)), random=random) + pts.append( + np.hstack( + ( + u[:, 0:1], + np.full((len(u), 1), y, dtype=paddle.get_default_dtype()), + u[:, 1:], + ) + ) + ) + rect = geometry_2d.Rectangle(self.xmin[1:], self.xmax[1:]) + for x in [self.xmin[0], self.xmax[0]]: + u = rect.random_points(int(np.ceil(density * rect.area)), random=random) + pts.append( + np.hstack( + (np.full((len(u), 1), x, dtype=paddle.get_default_dtype()), u) + ) + ) + pts = np.vstack(pts) + if len(pts) > n: + return pts[np.random.choice(len(pts), size=n, replace=False)] + return pts + + def uniform_boundary_points(self, n): + h = (self.area / n) ** 0.5 + nx, ny, nz = np.ceil((self.xmax - self.xmin) / h).astype(int) + 1 + x = np.linspace( + self.xmin[0], self.xmax[0], num=nx, dtype=paddle.get_default_dtype() + ) + y = np.linspace( + self.xmin[1], self.xmax[1], num=ny, dtype=paddle.get_default_dtype() + ) + z = np.linspace( + self.xmin[2], self.xmax[2], num=nz, dtype=paddle.get_default_dtype() + ) + + pts = [] + for v in [self.xmin[-1], self.xmax[-1]]: + u = list(itertools.product(x, y)) + pts.append( + np.hstack( + (u, np.full((len(u), 1), v, dtype=paddle.get_default_dtype())) + ) + ) + if nz > 2: + for v in [self.xmin[1], self.xmax[1]]: + u = np.array( + list(itertools.product(x, z[1:-1])), + dtype=paddle.get_default_dtype(), + ) + pts.append( + np.hstack( + ( + u[:, 0:1], + np.full((len(u), 1), v, dtype=paddle.get_default_dtype()), + u[:, 1:], + ) + ) + ) + if ny > 2 and nz > 2: + for v in [self.xmin[0], self.xmax[0]]: + u = list(itertools.product(y[1:-1], z[1:-1])) + pts.append( + np.hstack( + (np.full((len(u), 1), v, dtype=paddle.get_default_dtype()), u) + ) + ) + pts = np.vstack(pts) + if len(pts) > n: + return pts[np.random.choice(len(pts), size=n, replace=False)] + return pts + + def sdf_func(self, points: np.ndarray) -> np.ndarray: + """Compute signed distance field. + + Args: + points (np.ndarray): The coordinate points used to calculate the SDF value, + the shape is [N, 3] + + Returns: + np.ndarray: SDF values of input points without squared, the shape is [N, 1]. + + NOTE: This function usually returns ndarray with negative values, because + according to the definition of SDF, the SDF value of the coordinate point inside + the object(interior points) is negative, the outside is positive, and the edge + is 0. Therefore, when used for weighting, a negative sign is often added before + the result of this function. + """ + if points.shape[1] != self.ndim: + raise ValueError( + f"Shape of given points should be [*, {self.ndim}], but got {points.shape}" + ) + sdf = ( + ((self.xmax - self.xmin) / 2 - abs(points - (self.xmin + self.xmax) / 2)) + ).min(axis=1) + sdf = -sdf[..., np.newaxis] + return sdf + + +class Sphere(geometry_nd.Hypersphere): + """Class for Sphere + + Args: + center (Tuple[float, float, float]): Center of the sphere [x0, y0, z0]. + radius (float): Radius of the sphere. + """ + + def __init__(self, center, radius): + super().__init__(center, radius) + + def uniform_boundary_points(self, n: int): + nl = np.arange(1, n + 1).astype(paddle.get_default_dtype()) + g = (np.sqrt(5) - 1) / 2 + z = (2 * nl - 1) / n - 1 + x = np.sqrt(1 - z**2) * np.cos(2 * np.pi * nl * g) + y = np.sqrt(1 - z**2) * np.sin(2 * np.pi * nl * g) + return np.stack((x, y, z), axis=-1) + + def sdf_func(self, points: np.ndarray) -> np.ndarray: + """Compute signed distance field. + + Args: + points (np.ndarray): The coordinate points used to calculate the SDF value, + the shape is [N, 3] + + Returns: + np.ndarray: SDF values of input points without squared, the shape is [N, 1]. + + NOTE: This function usually returns ndarray with negative values, because + according to the definition of SDF, the SDF value of the coordinate point inside + the object(interior points) is negative, the outside is positive, and the edge + is 0. Therefore, when used for weighting, a negative sign is often added before + the result of this function. + """ + if points.shape[1] != self.ndim: + raise ValueError( + f"Shape of given points should be [*, {self.ndim}], but got {points.shape}" + ) + sdf = self.radius - (((points - self.center) ** 2).sum(axis=1)) ** 0.5 + sdf = -sdf[..., np.newaxis] + return sdf diff --git a/examples/smc_reac/ppsci/geometry/geometry_nd.py b/examples/smc_reac/ppsci/geometry/geometry_nd.py new file mode 100644 index 0000000000..84a8a3edbc --- /dev/null +++ b/examples/smc_reac/ppsci/geometry/geometry_nd.py @@ -0,0 +1,196 @@ +# 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. + +""" +Code below is heavily based on [https://github.com/lululxvi/deepxde](https://github.com/lululxvi/deepxde) +""" +from __future__ import annotations + +import itertools +from typing import Tuple + +import numpy as np +import paddle +from scipy import stats +from sklearn import preprocessing + +from ppsci.geometry import geometry +from ppsci.geometry import sampler +from ppsci.utils import misc + + +class Hypercube(geometry.Geometry): + """Multi-dimensional hyper cube. + + Args: + xmin (Tuple[float, ...]): Lower corner point. + xmax (Tuple[float, ...]): Upper corner point. + + Examples: + >>> import ppsci + >>> geom = ppsci.geometry.Hypercube((0, 0, 0, 0), (1, 1, 1, 1)) + """ + + def __init__(self, xmin: Tuple[float, ...], xmax: Tuple[float, ...]): + if len(xmin) != len(xmax): + raise ValueError("Dimensions of xmin and xmax do not match.") + + self.xmin = np.array(xmin, dtype=paddle.get_default_dtype()) + self.xmax = np.array(xmax, dtype=paddle.get_default_dtype()) + if np.any(self.xmin >= self.xmax): + raise ValueError("xmin >= xmax") + + self.side_length = self.xmax - self.xmin + super().__init__( + len(xmin), (self.xmin, self.xmax), np.linalg.norm(self.side_length) + ) + self.volume = np.prod(self.side_length, dtype=paddle.get_default_dtype()) + + def is_inside(self, x): + return np.logical_and( + np.all(x >= self.xmin, axis=-1), np.all(x <= self.xmax, axis=-1) + ) + + def on_boundary(self, x): + _on_boundary = np.logical_or( + np.any(np.isclose(x, self.xmin), axis=-1), + np.any(np.isclose(x, self.xmax), axis=-1), + ) + return np.logical_and(self.is_inside(x), _on_boundary) + + def boundary_normal(self, x): + _n = -np.isclose(x, self.xmin).astype(paddle.get_default_dtype()) + np.isclose( + x, self.xmax + ) + # For vertices, the normal is averaged for all directions + idx = np.count_nonzero(_n, axis=-1) > 1 + if np.any(idx): + l = np.linalg.norm(_n[idx], axis=-1, keepdims=True) + _n[idx] /= l + return _n + + def uniform_points(self, n, boundary=True): + dx = (self.volume / n) ** (1 / self.ndim) + xi = [] + for i in range(self.ndim): + ni = int(np.ceil(self.side_length[i] / dx)) + if boundary: + xi.append( + np.linspace( + self.xmin[i], + self.xmax[i], + num=ni, + dtype=paddle.get_default_dtype(), + ) + ) + else: + xi.append( + np.linspace( + self.xmin[i], + self.xmax[i], + num=ni + 1, + endpoint=False, + dtype=paddle.get_default_dtype(), + )[1:] + ) + x = np.array(list(itertools.product(*xi)), dtype=paddle.get_default_dtype()) + if len(x) > n: + x = x[0:n] + return x + + def random_points(self, n, random="pseudo"): + x = sampler.sample(n, self.ndim, random) + # print(f"Hypercube's range: {self.__class__.__name__}", self.xmin, self.xmax) + return (self.xmax - self.xmin) * x + self.xmin + + def random_boundary_points(self, n, random="pseudo"): + x = sampler.sample(n, self.ndim, random) + # Randomly pick a dimension + rand_dim = np.random.randint(self.ndim, size=n) + # Replace value of the randomly picked dimension with the nearest boundary value (0 or 1) + x[np.arange(n), rand_dim] = np.round(x[np.arange(n), rand_dim]) + return (self.xmax - self.xmin) * x + self.xmin + + def periodic_point(self, x, component): + y = misc.convert_to_array(x, self.dim_keys) + _on_xmin = np.isclose(y[:, component], self.xmin[component]) + _on_xmax = np.isclose(y[:, component], self.xmax[component]) + y[:, component][_on_xmin] = self.xmax[component] + y[:, component][_on_xmax] = self.xmin[component] + y_normal = self.boundary_normal(y) + + y = misc.convert_to_dict(y, self.dim_keys) + y_normal = misc.convert_to_dict( + y_normal, [f"normal_{k}" for k in self.dim_keys] + ) + return {**y, **y_normal} + + +class Hypersphere(geometry.Geometry): + """Multi-dimensional hyper sphere. + + Args: + center (Tuple[float, ...]): Center point coordinate. + radius (Tuple[float, ...]): Radius along each dimension. + + Examples: + >>> import ppsci + >>> geom = ppsci.geometry.Hypersphere((0, 0, 0, 0), 1.0) + """ + + def __init__(self, center, radius): + self.center = np.array(center, dtype=paddle.get_default_dtype()) + self.radius = radius + super().__init__( + len(center), (self.center - radius, self.center + radius), 2 * radius + ) + + self._r2 = radius**2 + + def is_inside(self, x): + return np.linalg.norm(x - self.center, axis=-1) <= self.radius + + def on_boundary(self, x): + return np.isclose(np.linalg.norm(x - self.center, axis=-1), self.radius) + + def boundary_normal(self, x): + _n = x - self.center + l = np.linalg.norm(_n, axis=-1, keepdims=True) + _n = _n / l * np.isclose(l, self.radius) + return _n + + def random_points(self, n, random="pseudo"): + # https://math.stackexchange.com/questions/87230/picking-random-points-in-the-volume-of-sphere-with-uniform-probability + if random == "pseudo": + U = np.random.rand(n, 1).astype(paddle.get_default_dtype()) + X = np.random.normal(size=(n, self.ndim)).astype(paddle.get_default_dtype()) + else: + rng = sampler.sample(n, self.ndim + 1, random) + U, X = rng[:, 0:1], rng[:, 1:] # Error if X = [0, 0, ...] + X = stats.norm.ppf(X).astype(paddle.get_default_dtype()) + X = preprocessing.normalize(X) + X = U ** (1 / self.ndim) * X + return self.radius * X + self.center + + def random_boundary_points(self, n, random="pseudo"): + # http://mathworld.wolfram.com/HyperspherePointPicking.html + if random == "pseudo": + X = np.random.normal(size=(n, self.ndim)).astype(paddle.get_default_dtype()) + else: + U = sampler.sample( + n, self.ndim, random + ) # Error for [0, 0, ...] or [0.5, 0.5, ...] + X = stats.norm.ppf(U).astype(paddle.get_default_dtype()) + X = preprocessing.normalize(X) + return self.radius * X + self.center diff --git a/examples/smc_reac/ppsci/geometry/inflation.py b/examples/smc_reac/ppsci/geometry/inflation.py new file mode 100644 index 0000000000..198a7cd223 --- /dev/null +++ b/examples/smc_reac/ppsci/geometry/inflation.py @@ -0,0 +1,192 @@ +# 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. + +from __future__ import annotations + +import numpy as np +import paddle + +from ppsci.utils import checker + +if not checker.dynamic_import_to_globals(["pymesh", "open3d"]): + raise ModuleNotFoundError + +__all__ = [ + "pymesh_inflation", +] + + +def open3d_inflation( + mesh: open3d.geometry.TriangleMesh, distance: float, direction: int = 1 +) -> open3d.geometry.TriangleMesh: + """Inflate mesh geometry. + + Args: + mesh (open3d.geometry.TriangleMesh): Open3D mesh object. + distance (float): Distance along exterior normal to inflate. + direction (int): 1 for exterior normal, -1 for interior normal. Defaults to 1. + + Returns: + open3d.geometry.TriangleMesh: Inflated mesh. + """ + mesh.remove_duplicated_vertices() + mesh.remove_degenerate_triangles() + mesh.remove_duplicated_triangles() + mesh.remove_unreferenced_vertices() + triangles = np.asarray(mesh.triangles) + points = np.asarray(mesh.vertices) + + remove_ids = [] + for i, point in enumerate(points): + boolean_index = np.argwhere(triangles == i)[:, 0] + if len(boolean_index) < 3: + remove_ids.append(i) + mesh.remove_vertices_by_index(remove_ids) + + points = np.asarray(mesh.vertices, dtype=paddle.get_default_dtype()) + mesh.compute_triangle_normals() + normals = np.asarray(mesh.triangle_normals, dtype=paddle.get_default_dtype()) + mesh.orient_triangles() + triangles = np.asarray(mesh.triangles, dtype=paddle.get_default_dtype()) + new_points = [] + for i, point in enumerate(points): + boolean_index = np.argwhere(triangles == i)[:, 0] + normal = normals[boolean_index] * direction + d = np.ones(len(normal), dtype=paddle.get_default_dtype()) * distance + + new_point = np.linalg.lstsq(normal, d, rcond=None)[0].squeeze() + new_point = point + new_point + if np.linalg.norm(new_point - point) > distance * 2: + # TODO : Find a better way to solve the bad inflation + new_point = point + distance * normal.mean(axis=0) + + new_points.append(new_point) + + new_points = np.array(new_points, dtype=paddle.get_default_dtype()) + new_mesh = open3d.geometry.TriangleMesh( + open3d.utility.Vector3dVector(new_points), + open3d.utility.Vector3iVector(triangles), + ) + + new_mesh.remove_duplicated_vertices() + new_mesh.remove_degenerate_triangles() + new_mesh.remove_duplicated_triangles() + new_mesh.remove_unreferenced_vertices() + new_mesh.compute_triangle_normals() + return new_mesh + + +def pymesh_inflation(mesh: pymesh.Mesh, distance: float) -> pymesh.Mesh: + """Inflate mesh by distance. + + Args: + mesh (pymesh.Mesh): PyMesh object. + distance (float): Inflation distance. + + Returns: + pymesh.Mesh: Inflated mesh. + """ + vertices = np.array(mesh.vertices, dtype=paddle.get_default_dtype()) + faces = np.array(mesh.faces) + open3d_mesh = open3d.geometry.TriangleMesh( + open3d.utility.Vector3dVector(vertices), open3d.utility.Vector3iVector(faces) + ) + inflated_open3d_mesh = open3d_inflation( + open3d_mesh, abs(distance), 1.0 if distance >= 0.0 else -1.0 + ) + vertices = np.array(inflated_open3d_mesh.vertices, dtype=paddle.get_default_dtype()) + faces = np.array(inflated_open3d_mesh.triangles) + inflated_pymesh = pymesh.form_mesh(vertices, faces) + return inflated_pymesh + + +def offset(mesh, distance) -> open3d.geometry.TriangleMesh: + """Offset the 2D mesh + + Args: + mesh (open3d.geometry.TriangleMesh): The mesh to be offset. + distance (float): The distance to offset. + + Returns: + open3d.geometry.TriangleMesh: Result mesh. + """ + # check if the mesh is 2D + mesh.compute_triangle_normals() + normals = np.asarray(mesh.triangle_normals, dtype=paddle.get_default_dtype()) + if not np.allclose(normals[:, :-1], 0): + raise ValueError("The mesh is not 2D") + + mesh.remove_duplicated_vertices() + mesh.remove_degenerate_triangles() + mesh.remove_duplicated_triangles() + mesh.remove_unreferenced_vertices() + triangles = np.asarray(mesh.triangles, dtype=paddle.get_default_dtype()) + + edges = np.vstack( + [triangles[:, [0, 1]], triangles[:, [1, 2]], triangles[:, [2, 0]]] + ) + edges = set(map(tuple, edges)) + edges = np.array(list(edges)) + + vertices = np.asarray(mesh.vertices, dtype=paddle.get_default_dtype())[:, :-1] + edges_in_triangle = np.array( + [ + np.intersect1d( + np.argwhere(triangles == edge[0])[:, 0], + np.argwhere(triangles == edge[1])[:, 0], + ) + for edge in edges + ], + dtype=object, + ) + surface_edges = edges[[len(i) == 1 for i in edges_in_triangle]] + edges_in_triangle = [i for i in edges_in_triangle if len(i) == 1] + + edges_normals = [] + for edge, triangle in zip(surface_edges, edges_in_triangle): + triangle = triangles[triangle].squeeze() + other_point = vertices[np.setdiff1d(triangle, edge)].squeeze() + edge = vertices[edge] + u = (other_point[0] - edge[0][0]) * (edge[0][0] - edge[1][0]) + ( + other_point[1] - edge[0][1] + ) * (edge[0][1] - edge[1][1]) + u = u / np.sum((edge[0] - edge[1]) ** 2) + edge_normal = edge[0] + u * (edge[0] - edge[1]) + edge_normal = edge_normal - other_point + edges_normals.append(edge_normal) + + edges_normals = np.array(edges_normals, dtype=paddle.get_default_dtype()) + edges_normals = edges_normals / np.linalg.norm(edges_normals, axis=1)[:, None] + + new_mesh = open3d.geometry.TriangleMesh() + new_vertices = [] + for point in set(surface_edges.reshape(-1)): + index = np.argwhere(surface_edges == point)[:, 0] + normal = edges_normals[index] + d = np.ones(len(index), dtype=paddle.get_default_dtype()) * distance + new_point = np.linalg.lstsq(normal, d, rcond=None)[0] + new_point = vertices[point] + new_point + new_vertices.append(new_point) + + new_vertices = np.hstack( + ( + np.array(new_vertices, dtype=paddle.get_default_dtype()), + np.zeros((len(new_vertices), 1), dtype=paddle.get_default_dtype()), + ) + ) + new_mesh.vertices = open3d.utility.Vector3dVector(new_vertices) + new_mesh.triangles = open3d.utility.Vector3iVector(triangles) + new_mesh.compute_triangle_normals() + new_mesh.compute_vertex_normals() + return new_mesh diff --git a/examples/smc_reac/ppsci/geometry/mesh.py b/examples/smc_reac/ppsci/geometry/mesh.py new file mode 100644 index 0000000000..8a363ca36b --- /dev/null +++ b/examples/smc_reac/ppsci/geometry/mesh.py @@ -0,0 +1,1392 @@ +# 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. + +from __future__ import annotations + +from typing import TYPE_CHECKING +from typing import Callable +from typing import Dict +from typing import Optional +from typing import Tuple +from typing import Union + +import numpy as np +import paddle + +try: + from stl import mesh as np_mesh_module +except ModuleNotFoundError: + pass +except ImportError: + pass + +from typing_extensions import Literal + +from ppsci.geometry import geometry +from ppsci.geometry import geometry_3d +from ppsci.geometry import sampler +from ppsci.geometry import sdf as sdf_module +from ppsci.utils import checker +from ppsci.utils import misc + +if TYPE_CHECKING: + import pymesh + + +class Mesh(geometry.Geometry): + """Class for mesh geometry. + + Args: + mesh (Union[str, Mesh]): Mesh file path or mesh object, such as "/path/to/mesh.stl". + + Examples: + >>> import ppsci + >>> geom = ppsci.geometry.Mesh("/path/to/mesh.stl") # doctest: +SKIP + """ + + def __init__(self, mesh: Union["pymesh.Mesh", str]): + # check if pymesh is installed when using Mesh Class + if not checker.dynamic_import_to_globals(["pymesh"]): + raise ImportError( + "Could not import pymesh python package." + "Please install it as https://pymesh.readthedocs.io/en/latest/installation.html." + ) + import pymesh + + if isinstance(mesh, str): + self.py_mesh = pymesh.meshio.load_mesh(mesh) + elif isinstance(mesh, pymesh.Mesh): + self.py_mesh = mesh + else: + raise ValueError("arg `mesh` should be path string or `pymesh.Mesh`") + + self.init_mesh() + + @classmethod + def from_pymesh(cls, mesh: "pymesh.Mesh") -> "Mesh": + """Instantiate Mesh object with given PyMesh object. + + Args: + mesh (pymesh.Mesh): PyMesh object. + + Returns: + Mesh: Instantiated ppsci.geometry.Mesh object. + + Examples: + >>> import ppsci + >>> import pymesh # doctest: +SKIP + >>> import numpy as np # doctest: +SKIP + >>> box = pymesh.generate_box_mesh(np.array([0, 0, 0]), np.array([1, 1, 1])) # doctest: +SKIP + >>> mesh = ppsci.geometry.Mesh.from_pymesh(box) # doctest: +SKIP + >>> print(mesh.vertices) # doctest: +SKIP + [[0. 0. 0.] + [1. 0. 0.] + [1. 1. 0.] + [0. 1. 0.] + [0. 0. 1.] + [1. 0. 1.] + [1. 1. 1.] + [0. 1. 1.]] + """ + # check if pymesh is installed when using Mesh Class + if not checker.dynamic_import_to_globals(["pymesh"]): + raise ImportError( + "Could not import pymesh python package." + "Please install it as https://pymesh.readthedocs.io/en/latest/installation.html." + ) + import pymesh + + if isinstance(mesh, pymesh.Mesh): + return cls(mesh) + else: + raise ValueError( + f"arg `mesh` should be type of `pymesh.Mesh`, but got {type(mesh)}" + ) + + def init_mesh(self): + """Initialize necessary variables for mesh""" + if "face_normal" not in self.py_mesh.get_attribute_names(): + self.py_mesh.add_attribute("face_normal") + self.face_normal = self.py_mesh.get_attribute("face_normal").reshape([-1, 3]) + + if not checker.dynamic_import_to_globals(["open3d"]): + raise ImportError( + "Could not import open3d python package. " + "Please install it with `pip install open3d`." + ) + import open3d + + self.open3d_mesh = open3d.geometry.TriangleMesh( + open3d.utility.Vector3dVector(np.array(self.py_mesh.vertices)), + open3d.utility.Vector3iVector(np.array(self.py_mesh.faces)), + ) + self.open3d_mesh.compute_vertex_normals() + + self.vertices = self.py_mesh.vertices + self.faces = self.py_mesh.faces + self.vectors = self.vertices[self.faces] + super().__init__( + self.vertices.shape[-1], + (np.amin(self.vertices, axis=0), np.amax(self.vertices, axis=0)), + np.inf, + ) + self.v0 = self.vectors[:, 0] + self.v1 = self.vectors[:, 1] + self.v2 = self.vectors[:, 2] + self.num_vertices = self.py_mesh.num_vertices + self.num_faces = self.py_mesh.num_faces + + if not checker.dynamic_import_to_globals(["pysdf"]): + raise ImportError( + "Could not import pysdf python package. " + "Please install open3d with `pip install pysdf`." + ) + import pysdf + + self.pysdf = pysdf.SDF(self.vertices, self.faces) + self.bounds = ( + ((np.min(self.vectors[:, :, 0])), np.max(self.vectors[:, :, 0])), + ((np.min(self.vectors[:, :, 1])), np.max(self.vectors[:, :, 1])), + ((np.min(self.vectors[:, :, 2])), np.max(self.vectors[:, :, 2])), + ) + + def sdf_func(self, points: np.ndarray) -> np.ndarray: + """Compute signed distance field. + + Args: + points (np.ndarray): The coordinate points used to calculate the SDF value, + the shape is [N, 3] + + Returns: + np.ndarray: SDF values of input points without squared, the shape is [N, 1]. + + NOTE: This function usually returns ndarray with negative values, because + according to the definition of SDF, the SDF value of the coordinate point inside + the object(interior points) is negative, the outside is positive, and the edge + is 0. Therefore, when used for weighting, a negative sign is often added before + the result of this function. + """ + if not checker.dynamic_import_to_globals(["pymesh"]): + raise ImportError( + "Could not import pymesh python package." + "Please install it as https://pymesh.readthedocs.io/en/latest/installation.html." + ) + import pymesh + + sdf, _, _, _ = pymesh.signed_distance_to_mesh(self.py_mesh, points) + sdf = sdf[..., np.newaxis].astype(paddle.get_default_dtype()) + return sdf + + def is_inside(self, x): + # NOTE: point on boundary is included + return self.pysdf.contains(x) + + def on_boundary(self, x): + return np.isclose(self.sdf_func(x), 0.0).ravel() + + def translate(self, translation: np.ndarray, relative: bool = True) -> "Mesh": + """Translate by given offsets. + + NOTE: This API generate a completely new Mesh object with translated geometry, + without modifying original Mesh object inplace. + + Args: + translation (np.ndarray): Translation offsets, numpy array of shape (3,): + [offset_x, offset_y, offset_z]. + relative (bool, optional): Whether translate relatively. Defaults to True. + + Returns: + Mesh: Translated Mesh object. + + Examples: + >>> import ppsci + >>> import pymesh # doctest: +SKIP + >>> import numpy as np + >>> box = pymesh.generate_box_mesh(np.array([0, 0, 0]), np.array([1, 1, 1])) # doctest: +SKIP + >>> mesh = ppsci.geometry.Mesh(box) # doctest: +SKIP + >>> print(mesh.vertices) # doctest: +SKIP + [[0. 0. 0.] + [1. 0. 0.] + [1. 1. 0.] + [0. 1. 0.] + [0. 0. 1.] + [1. 0. 1.] + [1. 1. 1.] + [0. 1. 1.]] + >>> print(mesh.translate((-0.5, 0, 0.5), False).vertices) # the center is moved to the translation vector. # doctest: +SKIP + [[-1. -0.5 0. ] + [ 0. -0.5 0. ] + [ 0. 0.5 0. ] + [-1. 0.5 0. ] + [-1. -0.5 1. ] + [ 0. -0.5 1. ] + [ 0. 0.5 1. ] + [-1. 0.5 1. ]] + >>> print(mesh.translate((-0.5, 0, 0.5), True).vertices) # the translation vector is directly added to the geometry coordinates # doctest: +SKIP + [[-0.5 0. 0.5] + [ 0.5 0. 0.5] + [ 0.5 1. 0.5] + [-0.5 1. 0.5] + [-0.5 0. 1.5] + [ 0.5 0. 1.5] + [ 0.5 1. 1.5] + [-0.5 1. 1.5]] + """ + vertices = np.array(self.vertices, dtype=paddle.get_default_dtype()) + faces = np.array(self.faces) + + if not checker.dynamic_import_to_globals(("open3d", "pymesh")): + raise ImportError( + "Could not import open3d and pymesh python package. " + "Please install open3d with `pip install open3d` and " + "pymesh as https://paddlescience-docs.readthedocs.io/zh/latest/zh/install_setup/#__tabbed_4_1" + ) + import open3d # isort:skip + import pymesh # isort:skip + + open3d_mesh = open3d.geometry.TriangleMesh( + open3d.utility.Vector3dVector(vertices), + open3d.utility.Vector3iVector(faces), + ) + open3d_mesh = open3d_mesh.translate(translation, relative) + translated_mesh = pymesh.form_mesh( + np.asarray(open3d_mesh.vertices, dtype=paddle.get_default_dtype()), faces + ) + # Generate a new Mesh object using class method + return Mesh.from_pymesh(translated_mesh) + + def scale( + self, scale: float, center: Tuple[float, float, float] = (0, 0, 0) + ) -> "Mesh": + """Scale by given scale coefficient and center coordinate. + + NOTE: This API generate a completely new Mesh object with scaled geometry, + without modifying original Mesh object inplace. + + Args: + scale (float): Scale coefficient. + center (Tuple[float,float,float], optional): Center coordinate, [x, y, z]. + Defaults to (0, 0, 0). + + Returns: + Mesh: Scaled Mesh object. + + Examples: + >>> import ppsci + >>> import pymesh # doctest: +SKIP + >>> import numpy as np + >>> box = pymesh.generate_box_mesh(np.array([0, 0, 0]), np.array([1, 1, 1])) # doctest: +SKIP + >>> mesh = ppsci.geometry.Mesh(box) # doctest: +SKIP + >>> print(mesh.vertices) # doctest: +SKIP + [[0. 0. 0.] + [1. 0. 0.] + [1. 1. 0.] + [0. 1. 0.] + [0. 0. 1.] + [1. 0. 1.] + [1. 1. 1.] + [0. 1. 1.]] + >>> mesh = mesh.scale(2, (0.25, 0.5, 0.75)) # doctest: +SKIP + >>> print(mesh.vertices) # doctest: +SKIP + [[-0.25 -0.5 -0.75] + [ 1.75 -0.5 -0.75] + [ 1.75 1.5 -0.75] + [-0.25 1.5 -0.75] + [-0.25 -0.5 1.25] + [ 1.75 -0.5 1.25] + [ 1.75 1.5 1.25] + [-0.25 1.5 1.25]] + """ + vertices = np.array(self.vertices, dtype=paddle.get_default_dtype()) + faces = np.array(self.faces, dtype=paddle.get_default_dtype()) + + if not checker.dynamic_import_to_globals(("open3d", "pymesh")): + raise ImportError( + "Could not import open3d and pymesh python package. " + "Please install open3d with `pip install open3d` and " + "pymesh as https://pymesh.readthedocs.io/en/latest/installation.html." + ) + import open3d # isort:skip + import pymesh # isort:skip + + open3d_mesh = open3d.geometry.TriangleMesh( + open3d.utility.Vector3dVector(vertices), + open3d.utility.Vector3iVector(faces), + ) + open3d_mesh = open3d_mesh.scale(scale, center) + scaled_pymesh = pymesh.form_mesh( + np.asarray(open3d_mesh.vertices, dtype=paddle.get_default_dtype()), faces + ) + # Generate a new Mesh object using class method + return Mesh.from_pymesh(scaled_pymesh) + + def uniform_boundary_points(self, n: int): + """Compute the equi-spaced points on the boundary.""" + return self.pysdf.sample_surface(n) + + def inflated_random_points(self, n, distance, random="pseudo", criteria=None): + if not isinstance(n, (tuple, list)): + n = [n] + if not isinstance(distance, (tuple, list)): + distance = [distance] + if len(n) != len(distance): + raise ValueError( + f"len(n)({len(n)}) should be equal to len(distance)({len(distance)})" + ) + + from ppsci.geometry import inflation + + all_points = [] + all_areas = [] + for _n, _dist in zip(n, distance): + inflated_mesh = Mesh(inflation.pymesh_inflation(self.py_mesh, _dist)) + points, areas = inflated_mesh.random_points(_n, random, criteria) + all_points.append(points) + all_areas.append(areas) + + all_points = np.concatenate(all_points, axis=0) + all_areas = np.concatenate(all_areas, axis=0) + return all_points, all_areas + + def _approximate_area( + self, + random: Literal["pseudo"] = "pseudo", + criteria: Optional[Callable] = None, + n_appr: int = 10000, + ) -> float: + """Approximate area with given `criteria` and `n_appr` points by Monte Carlo + algorithm. + + Args: + random (str, optional): Random method. Defaults to "pseudo". + criteria (Optional[Callable]): Criteria function. Defaults to None. + n_appr (int): Number of points for approximating area. Defaults to 10000. + + Returns: + float: Approximation area with given criteria. + """ + triangle_areas = area_of_triangles(self.v0, self.v1, self.v2) + triangle_probabilities = triangle_areas / np.linalg.norm(triangle_areas, ord=1) + triangle_index = np.arange(triangle_probabilities.shape[0]) + npoint_per_triangle = np.random.choice( + triangle_index, n_appr, p=triangle_probabilities + ) + npoint_per_triangle, _ = np.histogram( + npoint_per_triangle, + np.arange(triangle_probabilities.shape[0] + 1) - 0.5, + ) + + appr_areas = [] + if criteria is not None: + aux_points = [] + + for i, npoint in enumerate(npoint_per_triangle): + if npoint == 0: + continue + # sample points for computing criteria mask if criteria is given + if criteria is not None: + points_at_triangle_i = sample_in_triangle( + self.v0[i], self.v1[i], self.v2[i], npoint, random + ) + aux_points.append(points_at_triangle_i) + + appr_areas.append( + np.full( + (npoint, 1), triangle_areas[i] / npoint, paddle.get_default_dtype() + ) + ) + appr_areas = np.concatenate(appr_areas, axis=0) # [n_appr, 1] + + # set invalid area to 0 by computing criteria mask with auxiliary points + if criteria is not None: + aux_points = np.concatenate(aux_points, axis=0) # [n_appr, 3] + criteria_mask = criteria(*np.split(aux_points, self.ndim, 1)) + appr_areas *= criteria_mask + return appr_areas.sum() + + def random_boundary_points(self, n, random="pseudo"): + triangle_area = area_of_triangles(self.v0, self.v1, self.v2) + triangle_prob = triangle_area / np.linalg.norm(triangle_area, ord=1) + npoint_per_triangle = np.random.choice( + np.arange(len(triangle_prob)), n, p=triangle_prob + ) + npoint_per_triangle, _ = np.histogram( + npoint_per_triangle, np.arange(len(triangle_prob) + 1) - 0.5 + ) + + points = [] + normal = [] + areas = [] + for i, npoint in enumerate(npoint_per_triangle): + if npoint == 0: + continue + points_at_triangle_i = sample_in_triangle( + self.v0[i], self.v1[i], self.v2[i], npoint, random + ) + normal_at_triangle_i = np.tile(self.face_normal[i], (npoint, 1)).astype( + paddle.get_default_dtype() + ) + areas_at_triangle_i = np.full( + (npoint, 1), + triangle_area[i] / npoint, + dtype=paddle.get_default_dtype(), + ) + + points.append(points_at_triangle_i) + normal.append(normal_at_triangle_i) + areas.append(areas_at_triangle_i) + + points = np.concatenate(points, axis=0) + normal = np.concatenate(normal, axis=0) + areas = np.concatenate(areas, axis=0) + + return points, normal, areas + + def sample_boundary( + self, + n: int, + random: Literal["pseudo"] = "pseudo", + criteria: Optional[Callable[..., np.ndarray]] = None, + evenly: bool = False, + inflation_dist: Union[float, Tuple[float, ...]] = None, + ) -> Dict[str, np.ndarray]: + # TODO(sensen): Support for time-dependent points(repeat data in time) + if inflation_dist is not None: + if not isinstance(n, (tuple, list)): + n = [n] + if not isinstance(inflation_dist, (tuple, list)): + inflation_dist = [inflation_dist] + if len(n) != len(inflation_dist): + raise ValueError( + f"len(n)({len(n)}) should be equal to len(inflation_dist)({len(inflation_dist)})" + ) + + from ppsci.geometry import inflation + + inflated_data_dict = {} + for _n, _dist in zip(n, inflation_dist): + # 1. manually inflate mesh at first + inflated_mesh = Mesh(inflation.pymesh_inflation(self.py_mesh, _dist)) + # 2. compute all data by sample_boundary with `inflation_dist=None` + data_dict = inflated_mesh.sample_boundary( + _n, + random, + criteria, + evenly, + inflation_dist=None, + ) + for key, value in data_dict.items(): + if key not in inflated_data_dict: + inflated_data_dict[key] = value + else: + inflated_data_dict[key] = np.concatenate( + (inflated_data_dict[key], value), axis=0 + ) + return inflated_data_dict + else: + if evenly: + raise ValueError( + "Can't sample evenly on mesh now, please set evenly=False." + ) + _size, _ntry, _nsuc = 0, 0, 0 + all_points = [] + all_normal = [] + while _size < n: + points, normal, _ = self.random_boundary_points(n, random) + if criteria is not None: + criteria_mask = criteria( + *np.split(points, self.ndim, axis=1) + ).ravel() + points = points[criteria_mask] + normal = normal[criteria_mask] + + if len(points) > n - _size: + points = points[: n - _size] + normal = normal[: n - _size] + + all_points.append(points) + all_normal.append(normal) + + _size += len(points) + _ntry += 1 + if len(points) > 0: + _nsuc += 1 + + if _ntry >= 1000 and _nsuc == 0: + raise ValueError( + "Sample boundary points failed, " + "please check correctness of geometry and given criteria." + ) + + all_points = np.concatenate(all_points, axis=0) + all_normal = np.concatenate(all_normal, axis=0) + appr_area = self._approximate_area(random, criteria) + all_areas = np.full((n, 1), appr_area / n, paddle.get_default_dtype()) + + x_dict = misc.convert_to_dict(all_points, self.dim_keys) + normal_dict = misc.convert_to_dict( + all_normal, [f"normal_{key}" for key in self.dim_keys if key != "t"] + ) + area_dict = misc.convert_to_dict(all_areas, ["area"]) + return {**x_dict, **normal_dict, **area_dict} + + def random_points(self, n, random="pseudo", criteria=None): + _size = 0 + all_points = [] + cuboid = geometry_3d.Cuboid( + [bound[0] for bound in self.bounds], + [bound[1] for bound in self.bounds], + ) + _nsample, _nvalid = 0, 0 + while _size < n: + random_points = cuboid.random_points(n, random) + valid_mask = self.is_inside(random_points) + + if criteria: + valid_mask &= criteria( + *np.split(random_points, self.ndim, axis=1) + ).ravel() + valid_points = random_points[valid_mask] + _nvalid += len(valid_points) + + if len(valid_points) > n - _size: + valid_points = valid_points[: n - _size] + + all_points.append(valid_points) + _size += len(valid_points) + _nsample += n + + all_points = np.concatenate(all_points, axis=0) + cuboid_volume = np.prod([b[1] - b[0] for b in self.bounds]) + all_areas = np.full( + (n, 1), cuboid_volume * (_nvalid / _nsample) / n, paddle.get_default_dtype() + ) + return all_points, all_areas + + def sample_interior( + self, + n: int, + random: Literal["pseudo"] = "pseudo", + criteria: Optional[Callable[..., np.ndarray]] = None, + evenly: bool = False, + compute_sdf_derivatives: bool = False, + ): + """Sample random points in the geometry and return those meet criteria.""" + if evenly: + # TODO(sensen): Implement uniform sample for mesh interior. + raise NotImplementedError( + "uniformly sample for interior in mesh is not support yet, " + "you may need to set evenly=False in config dict of constraint" + ) + points, areas = self.random_points(n, random, criteria) + + x_dict = misc.convert_to_dict(points, self.dim_keys) + area_dict = misc.convert_to_dict(areas, ("area",)) + + # NOTE: add negative to the sdf values because weight should be positive. + sdf = -self.sdf_func(points) + sdf_dict = misc.convert_to_dict(sdf, ("sdf",)) + + sdf_derives_dict = {} + if compute_sdf_derivatives: + sdf_derives = -self.sdf_derivatives(points) + sdf_derives_dict = misc.convert_to_dict( + sdf_derives, tuple(f"sdf__{key}" for key in self.dim_keys) + ) + + return {**x_dict, **area_dict, **sdf_dict, **sdf_derives_dict} + + def union(self, other: "Mesh"): + if not checker.dynamic_import_to_globals(["pymesh"]): + raise ImportError( + "Could not import pymesh python package. " + "Please install it as https://pymesh.readthedocs.io/en/latest/installation.html." + ) + import pymesh + + csg = pymesh.CSGTree( + {"union": [{"mesh": self.py_mesh}, {"mesh": other.py_mesh}]} + ) + return Mesh(csg.mesh) + + def __or__(self, other: "Mesh"): + return self.union(other) + + def __add__(self, other: "Mesh"): + return self.union(other) + + def difference(self, other: "Mesh"): + if not checker.dynamic_import_to_globals(["pymesh"]): + raise ImportError( + "Could not import pymesh python package. " + "Please install it as https://pymesh.readthedocs.io/en/latest/installation.html." + ) + import pymesh + + csg = pymesh.CSGTree( + {"difference": [{"mesh": self.py_mesh}, {"mesh": other.py_mesh}]} + ) + return Mesh(csg.mesh) + + def __sub__(self, other: "Mesh"): + return self.difference(other) + + def intersection(self, other: "Mesh"): + if not checker.dynamic_import_to_globals(["pymesh"]): + raise ImportError( + "Could not import pymesh python package. " + "Please install it as https://pymesh.readthedocs.io/en/latest/installation.html." + ) + import pymesh + + csg = pymesh.CSGTree( + {"intersection": [{"mesh": self.py_mesh}, {"mesh": other.py_mesh}]} + ) + return Mesh(csg.mesh) + + def __and__(self, other: "Mesh"): + return self.intersection(other) + + def __str__(self) -> str: + """Return the name of class""" + return ", ".join( + [ + self.__class__.__name__, + f"num_vertices = {self.num_vertices}", + f"num_faces = {self.num_faces}", + f"bounds = {self.bounds}", + f"dim_keys = {self.dim_keys}", + ] + ) + + +class SDFMesh(geometry.Geometry): + """Class for SDF geometry, a kind of implicit surface mesh. + + Args: + vectors (np.ndarray): Vectors of triangles of mesh with shape [M, 3, 3]. + normals (np.ndarray): Unit normals of each triangle face with shape [M, 3]. + sdf_func (Callable[[np.ndarray, bool], np.ndarray]): Signed distance function + of the triangle mesh. + + Examples: + >>> import ppsci + >>> geom = ppsci.geometry.SDFMesh.from_stl("/path/to/mesh.stl") # doctest: +SKIP + """ + + eps = 1e-6 + + def __init__( + self, + vectors: np.ndarray, + normals: np.ndarray, + sdf_func: Callable[[np.ndarray, bool], np.ndarray], + ): + if vectors.shape[1:] != (3, 3): + raise ValueError( + f"The shape of `vectors` must be [M, 3, 3], but got {vectors.shape}" + ) + if normals.shape[1] != 3: + raise ValueError( + f"The shape of `normals` must be [M, 3], but got {normals.shape}" + ) + self.vectors = vectors + self.face_normal = normals + self.sdf_func = sdf_func # overwrite sdf_func + self.bounds = ( + ((np.min(self.vectors[:, :, 0])), np.max(self.vectors[:, :, 0])), + ((np.min(self.vectors[:, :, 1])), np.max(self.vectors[:, :, 1])), + ((np.min(self.vectors[:, :, 2])), np.max(self.vectors[:, :, 2])), + ) + self.ndim = 3 + super().__init__( + self.vectors.shape[-1], + (np.amin(self.vectors, axis=(0, 1)), np.amax(self.vectors, axis=(0, 1))), + np.inf, + ) + + @property + def v0(self) -> np.ndarray: + return self.vectors[:, 0] + + @property + def v1(self) -> np.ndarray: + return self.vectors[:, 1] + + @property + def v2(self) -> np.ndarray: + return self.vectors[:, 2] + + @classmethod + def from_stl(cls, mesh_file: str) -> "SDFMesh": + """Instantiate SDFMesh from given mesh file. + + Args: + mesh_file (str): Path to triangle mesh file. + + Returns: + SDFMesh: Instantiated ppsci.geometry.SDFMesh object. + + Examples: + >>> import ppsci + >>> import pymesh # doctest: +SKIP + >>> import numpy as np # doctest: +SKIP + >>> box = pymesh.generate_box_mesh(np.array([0, 0, 0]), np.array([1, 1, 1])) # doctest: +SKIP + >>> pymesh.save_mesh("box.stl", box) # doctest: +SKIP + >>> mesh = ppsci.geometry.SDFMesh.from_stl("box.stl") # doctest: +SKIP + >>> print(sdfmesh.vectors.shape) # doctest: +SKIP + (12, 3, 3) + """ + # check if pymesh is installed when using Mesh Class + if not checker.dynamic_import_to_globals(["stl"]): + raise ImportError( + "Could not import stl python package. " + "Please install numpy-stl with: pip install 'numpy-stl>=2.16,<2.17'" + ) + + np_mesh_obj = np_mesh_module.Mesh.from_file(mesh_file) + return cls( + np_mesh_obj.vectors, + np_mesh_obj.get_unit_normals(), + make_sdf(np_mesh_obj.vectors), + ) + + def sdf_func( + self, points: np.ndarray, compute_sdf_derivatives: bool = False + ) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]: + """Compute signed distance field. + + Args: + points (np.ndarray): The coordinate points used to calculate the SDF value, + the shape is [N, 3] + compute_sdf_derivatives (bool): Whether to compute SDF derivatives. + Defaults to False. + + Returns: + Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]: + If compute_sdf_derivatives is True, then return both SDF values([N, 1]) + and their derivatives([N, 3]); otherwise only return SDF values([N, 1]). + + NOTE: This function usually returns ndarray with negative values, because + according to the definition of SDF, the SDF value of the coordinate point inside + the object(interior points) is negative, the outside is positive, and the edge + is 0. Therefore, when used for weighting, a negative sign is often added before + the result of this function. + """ + # normalize triangles + x_min, y_min, z_min = np.min(points, axis=0) + x_max, y_max, z_max = np.max(points, axis=0) + max_dis = max(max((x_max - x_min), (y_max - y_min)), (z_max - z_min)) + store_triangles = np.array(self.vectors, dtype=np.float64) + store_triangles[:, :, 0] -= x_min + store_triangles[:, :, 1] -= y_min + store_triangles[:, :, 2] -= z_min + store_triangles *= 1 / max_dis + store_triangles = store_triangles.reshape([-1, 3]) + + # normalize query points + points = points.copy() + points[:, 0] -= x_min + points[:, 1] -= y_min + points[:, 2] -= z_min + points *= 1 / max_dis + points = points.astype(np.float64).ravel() + + # compute sdf values for query points + sdf = sdf_module.signed_distance_field( + store_triangles, + np.arange((store_triangles.shape[0])), + points, + include_hit_points=compute_sdf_derivatives, + ) + if compute_sdf_derivatives: + sdf, hit_points = sdf + + sdf = sdf.numpy() # [N] + sdf = np.expand_dims(max_dis * sdf, axis=1) # [N, 1] + + if compute_sdf_derivatives: + hit_points = hit_points.numpy() # [N, 3] + # Gradient of SDF is the unit vector from the query point to the hit point. + sdf_derives = hit_points - points + sdf_derives /= np.linalg.norm(sdf_derives, axis=1, keepdims=True) + return sdf, sdf_derives + + return sdf + + def is_inside(self, x): + # NOTE: point on boundary is included + return np.less(self.sdf_func(x), 0.0).ravel() + + def on_boundary(self, x: np.ndarray, normal: np.ndarray) -> np.ndarray: + x_plus = x + self.eps * normal + x_minus = x - self.eps * normal + + sdf_x_plus = self.sdf_func(x_plus) + sdf_x_minus = self.sdf_func(x_minus) + mask_on_boundary = np.less_equal(sdf_x_plus * sdf_x_minus, 0) + return mask_on_boundary.ravel() + + def translate(self, translation: np.ndarray) -> "SDFMesh": + """Translate by given offsets. + + NOTE: This API generate a completely new Mesh object with translated geometry, + without modifying original Mesh object inplace. + + Args: + translation (np.ndarray): Translation offsets, numpy array of shape (3,): + [offset_x, offset_y, offset_z]. + + Returns: + Mesh: Translated Mesh object. + + Examples: + >>> import ppsci + >>> import pymesh # doctest: +SKIP + >>> mesh = ppsci.geometry.SDFMesh.from_stl('/path/to/mesh.stl') # doctest: +SKIP + >>> mesh = mesh.translate(np.array([1, -1, 2])) # doctest: +SKIP + """ + new_vectors = self.vectors + translation.reshape([1, 1, 3]) + + return SDFMesh( + new_vectors, + self.face_normal, + make_sdf(new_vectors), + ) + + def scale(self, scale: float) -> "SDFMesh": + """Scale by given scale coefficient and center coordinate. + + NOTE: This API generate a completely new Mesh object with scaled geometry, + without modifying original Mesh object inplace. + + Args: + scale (float): Scale coefficient. + + Returns: + Mesh: Scaled Mesh object. + + Examples: + >>> import ppsci + >>> import pymesh # doctest: +SKIP + >>> mesh = ppsci.geometry.SDFMesh.from_stl('/path/to/mesh.stl') # doctest: +SKIP + >>> mesh = mesh.scale(np.array([1.3, 1.5, 2.0])) # doctest: +SKIP + """ + new_vectors = self.vectors * scale + return SDFMesh( + new_vectors, + self.face_normal, + make_sdf(new_vectors), + ) + + def uniform_boundary_points(self, n: int): + """Compute the equi-spaced points on the boundary.""" + raise NotImplementedError( + "'uniform_boundary_points' is not available in SDFMesh." + ) + + def inflated_random_points(self, n, distance, random="pseudo", criteria=None): + raise NotImplementedError( + "'inflated_random_points' is not available in SDFMesh." + ) + + def _approximate_area( + self, + random: Literal["pseudo"] = "pseudo", + criteria: Optional[Callable] = None, + n_appr: int = 10000, + ) -> float: + """Approximate area with given `criteria` and `n_appr` points by Monte Carlo + algorithm. + + Args: + random (str, optional): Random method. Defaults to "pseudo". + criteria (Optional[Callable]): Criteria function. Defaults to None. + n_appr (int): Number of points for approximating area. Defaults to 10000. + + Returns: + float: Approximation area with given criteria. + """ + triangle_areas = area_of_triangles(self.v0, self.v1, self.v2) + triangle_probabilities = triangle_areas / np.linalg.norm(triangle_areas, ord=1) + triangle_index = np.arange(triangle_probabilities.shape[0]) + npoint_per_triangle = np.random.choice( + triangle_index, n_appr, p=triangle_probabilities + ) + npoint_per_triangle, _ = np.histogram( + npoint_per_triangle, + np.arange(triangle_probabilities.shape[0] + 1) - 0.5, + ) + + aux_points = [] + aux_normals = [] + appr_areas = [] + + for i, npoint in enumerate(npoint_per_triangle): + if npoint == 0: + continue + # sample points for computing criteria mask if criteria is given + points_at_triangle_i = sample_in_triangle( + self.v0[i], self.v1[i], self.v2[i], npoint, random + ) + normal_at_triangle_i = np.tile( + self.face_normal[i].reshape(1, 3), (npoint, 1) + ) + aux_points.append(points_at_triangle_i) + aux_normals.append(normal_at_triangle_i) + appr_areas.append( + np.full( + (npoint, 1), triangle_areas[i] / npoint, paddle.get_default_dtype() + ) + ) + + aux_points = np.concatenate(aux_points, axis=0) # [n_appr, 3] + aux_normals = np.concatenate(aux_normals, axis=0) # [n_appr, 3] + appr_areas = np.concatenate(appr_areas, axis=0) # [n_appr, 1] + valid_mask = self.on_boundary(aux_points, aux_normals)[:, None] + # set invalid area to 0 by computing criteria mask with auxiliary points + if criteria is not None: + criteria_mask = criteria(*np.split(aux_points, self.ndim, 1)) + assert valid_mask.shape == criteria_mask.shape + valid_mask = np.logical_and(valid_mask, criteria_mask) + + appr_areas *= valid_mask + + return appr_areas.sum() + + def random_boundary_points( + self, n, random="pseudo" + ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: + triangle_area = area_of_triangles(self.v0, self.v1, self.v2) + triangle_prob = triangle_area / np.linalg.norm(triangle_area, ord=1) + npoint_per_triangle = np.random.choice( + np.arange(len(triangle_prob)), n, p=triangle_prob + ) + npoint_per_triangle, _ = np.histogram( + npoint_per_triangle, np.arange(len(triangle_prob) + 1) - 0.5 + ) + + points = [] + normal = [] + areas = [] + for i, npoint in enumerate(npoint_per_triangle): + if npoint == 0: + continue + points_at_triangle_i = sample_in_triangle( + self.v0[i], self.v1[i], self.v2[i], npoint, random + ) + normal_at_triangle_i = np.tile(self.face_normal[i], (npoint, 1)).astype( + paddle.get_default_dtype() + ) + areas_at_triangle_i = np.full( + (npoint, 1), + triangle_area[i] / npoint, + dtype=paddle.get_default_dtype(), + ) + + points.append(points_at_triangle_i) + normal.append(normal_at_triangle_i) + areas.append(areas_at_triangle_i) + + points = np.concatenate(points, axis=0) + normal = np.concatenate(normal, axis=0) + areas = np.concatenate(areas, axis=0) + + return points, normal, areas + + def sample_boundary( + self, + n: int, + random: Literal["pseudo"] = "pseudo", + criteria: Optional[Callable[..., np.ndarray]] = None, + evenly: bool = False, + inflation_dist: Union[float, Tuple[float, ...]] = None, + ) -> Dict[str, np.ndarray]: + # TODO(sensen): Support for time-dependent points(repeat data in time) + if inflation_dist is not None: + raise NotImplementedError("Not implemented yet") + else: + if evenly: + raise ValueError( + "Can't sample evenly on mesh now, please set evenly=False." + ) + _size, _ntry, _nsuc = 0, 0, 0 + all_points = [] + all_normal = [] + while _size < n: + points, normal, _ = self.random_boundary_points(n, random) + valid_mask = self.on_boundary(points, normal) + + if criteria is not None: + criteria_mask = criteria( + *np.split(points, self.ndim, axis=1) + ).ravel() + assert valid_mask.shape == criteria_mask.shape + valid_mask = np.logical_and(valid_mask, criteria_mask) + + points = points[valid_mask] + normal = normal[valid_mask] + + if len(points) > n - _size: + points = points[: n - _size] + normal = normal[: n - _size] + + all_points.append(points) + all_normal.append(normal) + + _size += len(points) + _ntry += 1 + if len(points) > 0: + _nsuc += 1 + + if _ntry >= 1000 and _nsuc == 0: + raise ValueError( + "Sample boundary points failed, " + "please check correctness of geometry and given criteria." + ) + + all_points = np.concatenate(all_points, axis=0) + all_normal = np.concatenate(all_normal, axis=0) + _appr_area = self._approximate_area(random, criteria) + all_areas = np.full((n, 1), _appr_area / n, paddle.get_default_dtype()) + + x_dict = misc.convert_to_dict(all_points, self.dim_keys) + normal_dict = misc.convert_to_dict( + all_normal, [f"normal_{key}" for key in self.dim_keys if key != "t"] + ) + area_dict = misc.convert_to_dict(all_areas, ["area"]) + return {**x_dict, **normal_dict, **area_dict} + + def random_points(self, n, random="pseudo", criteria=None): + _size = 0 + all_points = [] + cuboid = geometry_3d.Cuboid( + [bound[0] for bound in self.bounds], + [bound[1] for bound in self.bounds], + ) + _nsample, _nvalid = 0, 0 + while _size < n: + random_points = cuboid.random_points(n, random) + valid_mask = self.is_inside(random_points) + + if criteria: + criteria_mask = criteria( + *np.split(random_points, self.ndim, axis=1) + ).ravel() + assert valid_mask.shape == criteria_mask.shape + valid_mask = np.logical_and(valid_mask, criteria_mask) + + valid_points = random_points[valid_mask] + _nvalid += len(valid_points) + + if len(valid_points) > n - _size: + valid_points = valid_points[: n - _size] + + all_points.append(valid_points) + _size += len(valid_points) + _nsample += n + + all_points = np.concatenate(all_points, axis=0) + cuboid_volume = np.prod([b[1] - b[0] for b in self.bounds]) + all_areas = np.full( + (n, 1), cuboid_volume * (_nvalid / _nsample) / n, paddle.get_default_dtype() + ) + return all_points, all_areas + + def sample_interior( + self, + n: int, + random: Literal["pseudo"] = "pseudo", + criteria: Optional[Callable[..., np.ndarray]] = None, + evenly: bool = False, + compute_sdf_derivatives: bool = False, + ): + """Sample random points in the geometry and return those meet criteria.""" + if evenly: + # TODO(sensen): Implement uniform sample for mesh interior. + raise NotImplementedError( + "uniformly sample for interior in mesh is not support yet, " + "you may need to set evenly=False in config dict of constraint" + ) + points, areas = self.random_points(n, random, criteria) + + x_dict = misc.convert_to_dict(points, self.dim_keys) + area_dict = misc.convert_to_dict(areas, ("area",)) + + sdf = self.sdf_func(points, compute_sdf_derivatives) + if compute_sdf_derivatives: + sdf, sdf_derives = sdf + + # NOTE: Negate sdf because weight should be positive. + sdf_dict = misc.convert_to_dict(-sdf, ("sdf",)) + + sdf_derives_dict = {} + if compute_sdf_derivatives: + # NOTE: Negate sdf derivatives + sdf_derives_dict = misc.convert_to_dict( + -sdf_derives, tuple(f"sdf__{key}" for key in self.dim_keys) + ) + + return {**x_dict, **area_dict, **sdf_dict, **sdf_derives_dict} + + def union(self, other: "SDFMesh"): + new_vectors = np.concatenate([self.vectors, other.vectors], axis=0) + new_normals = np.concatenate([self.face_normal, other.face_normal], axis=0) + + def make_union_new_sdf(sdf_func1, sdf_func2): + def new_sdf_func(points: np.ndarray, compute_sdf_derivatives: bool = False): + # Invert definition of sdf to make boolean operation accurate + # see: https://iquilezles.org/articles/interiordistance/ + sdf_self = sdf_func1(points, compute_sdf_derivatives) + sdf_other = sdf_func2(points, compute_sdf_derivatives) + if compute_sdf_derivatives: + sdf_self, sdf_derives_self = sdf_self + sdf_other, sdf_derives_other = sdf_other + + computed_sdf = -np.maximum(-sdf_self, -sdf_other) + + if compute_sdf_derivatives: + computed_sdf_derives = -np.where( + sdf_self < sdf_other, + sdf_derives_self, + sdf_derives_other, + ) + return computed_sdf, computed_sdf_derives + + return computed_sdf + + return new_sdf_func + + return SDFMesh( + new_vectors, + new_normals, + make_union_new_sdf(self.sdf_func, other.sdf_func), + ) + + def __or__(self, other: "SDFMesh"): + return self.union(other) + + def __add__(self, other: "SDFMesh"): + return self.union(other) + + def difference(self, other: "SDFMesh"): + new_vectors = np.concatenate([self.vectors, other.vectors], axis=0) + new_normals = np.concatenate([self.face_normal, -other.face_normal], axis=0) + + def make_difference_new_sdf(sdf_func1, sdf_func2): + def new_sdf_func(points: np.ndarray, compute_sdf_derivatives: bool = False): + # Invert definition of sdf to make boolean operation accurate + # see: https://iquilezles.org/articles/interiordistance/ + sdf_self = sdf_func1(points, compute_sdf_derivatives) + sdf_other = sdf_func2(points, compute_sdf_derivatives) + if compute_sdf_derivatives: + sdf_self, sdf_derives_self = sdf_self + sdf_other, sdf_derives_other = sdf_other + + computed_sdf = -np.minimum(-sdf_self, sdf_other) + + if compute_sdf_derivatives: + computed_sdf_derives = np.where( + -sdf_self < sdf_other, + -sdf_derives_self, + sdf_derives_other, + ) + return computed_sdf, computed_sdf_derives + + return computed_sdf + + return new_sdf_func + + return SDFMesh( + new_vectors, + new_normals, + make_difference_new_sdf(self.sdf_func, other.sdf_func), + ) + + def __sub__(self, other: "SDFMesh"): + return self.difference(other) + + def intersection(self, other: "SDFMesh"): + new_vectors = np.concatenate([self.vectors, other.vectors], axis=0) + new_normals = np.concatenate([self.face_normal, other.face_normal], axis=0) + + def make_intersection_new_sdf(sdf_func1, sdf_func2): + def new_sdf_func(points: np.ndarray, compute_sdf_derivatives: bool = False): + # Invert definition of sdf to make boolean operation accurate + # see: https://iquilezles.org/articles/interiordistance/ + sdf_self = sdf_func1(points, compute_sdf_derivatives) + sdf_other = sdf_func2(points, compute_sdf_derivatives) + if compute_sdf_derivatives: + sdf_self, sdf_derives_self = sdf_self + sdf_other, sdf_derives_other = sdf_other + + computed_sdf = -np.minimum(-sdf_self, -sdf_other) + + if compute_sdf_derivatives: + computed_sdf_derives = np.where( + sdf_self > sdf_other, + -sdf_derives_self, + -sdf_derives_other, + ) + return computed_sdf, computed_sdf_derives + + return computed_sdf + + return new_sdf_func + + return SDFMesh( + new_vectors, + new_normals, + make_intersection_new_sdf(self.sdf_func, other.sdf_func), + ) + + def __and__(self, other: "SDFMesh"): + return self.intersection(other) + + def __str__(self) -> str: + """Return the name of class""" + return ", ".join( + [ + self.__class__.__name__, + f"num_faces = {self.vectors.shape[0]}", + f"bounds = {self.bounds}", + f"dim_keys = {self.dim_keys}", + ] + ) + + +def area_of_triangles(v0, v1, v2): + """Ref https://math.stackexchange.com/questions/128991/how-to-calculate-the-area-of-a-3d-triangle + + Args: + v0 (np.ndarray): Coordinates of the first vertex of the triangle surface with shape of [N, 3]. + v1 (np.ndarray): Coordinates of the second vertex of the triangle surface with shape of [N, 3]. + v2 (np.ndarray): Coordinates of the third vertex of the triangle surface with shape of [N, 3]. + + Returns: + np.ndarray: Area of each triangle with shape of [N, ]. + """ + a = np.sqrt( + (v0[:, 0] - v1[:, 0]) ** 2 + + (v0[:, 1] - v1[:, 1]) ** 2 + + (v0[:, 2] - v1[:, 2]) ** 2 + + 1e-10 + ) + b = np.sqrt( + (v1[:, 0] - v2[:, 0]) ** 2 + + (v1[:, 1] - v2[:, 1]) ** 2 + + (v1[:, 2] - v2[:, 2]) ** 2 + + 1e-10 + ) + c = np.sqrt( + (v0[:, 0] - v2[:, 0]) ** 2 + + (v0[:, 1] - v2[:, 1]) ** 2 + + (v0[:, 2] - v2[:, 2]) ** 2 + + 1e-10 + ) + p = (a + b + c) / 2 + area = np.sqrt(p * (p - a) * (p - b) * (p - c) + 1e-10) + return area + + +def sample_in_triangle(v0, v1, v2, n, random="pseudo", criteria=None): + """ + Uniformly sample n points in an 3D triangle defined by 3 vertices v0, v1, v2 + https://math.stackexchange.com/questions/18686/uniform-random-point-in-triangle + + Args: + v0 (np.ndarray): Coordinates of the first vertex of an triangle with shape of [3, ]. + v1 (np.ndarray): Coordinates of the second vertex of an triangle with shape of [3, ]. + v2 (np.ndarray): Coordinates of the third vertex of an triangle with shape of [3, ]. + n (int): Number of points to be sampled. + + Returns: + np.ndarray: Coordinates of sampled n points with shape of [n, 3]. + """ + xs, ys, zs = [], [], [] + _size = 0 + while _size < n: + r1 = sampler.sample(n, 1, random).ravel() + r2 = sampler.sample(n, 1, random).ravel() + s1 = np.sqrt(r1) + x = v0[0] * (1.0 - s1) + v1[0] * (1.0 - r2) * s1 + v2[0] * r2 * s1 + y = v0[1] * (1.0 - s1) + v1[1] * (1.0 - r2) * s1 + v2[1] * r2 * s1 + z = v0[2] * (1.0 - s1) + v1[2] * (1.0 - r2) * s1 + v2[2] * r2 * s1 + + if criteria is not None: + criteria_mask = criteria(x, y, z).ravel() + x = x[criteria_mask] + y = y[criteria_mask] + z = z[criteria_mask] + + if len(x) > n - _size: + x = x[: n - _size] + y = y[: n - _size] + z = z[: n - _size] + + xs.append(x) + ys.append(y) + zs.append(z) + _size += len(x) + + xs = np.concatenate(xs, axis=0) + ys = np.concatenate(ys, axis=0) + zs = np.concatenate(zs, axis=0) + + return np.stack([xs, ys, zs], axis=1) + + +def make_sdf(vectors: np.ndarray): + def sdf_func(points: np.ndarray, compute_sdf_derivatives=False): + points = points.copy() + x_min, y_min, z_min = np.min(points, axis=0) + x_max, y_max, z_max = np.max(points, axis=0) + max_dis = max(max((x_max - x_min), (y_max - y_min)), (z_max - z_min)) + store_triangles = vectors.copy() + store_triangles[:, :, 0] -= x_min + store_triangles[:, :, 1] -= y_min + store_triangles[:, :, 2] -= z_min + store_triangles *= 1 / max_dis + store_triangles = store_triangles.reshape([-1, 3]) + points[:, 0] -= x_min + points[:, 1] -= y_min + points[:, 2] -= z_min + points *= 1 / max_dis + points = points.astype(np.float64).ravel() + + # compute sdf values + sdf = sdf_module.signed_distance_field( + store_triangles, + np.arange((store_triangles.shape[0])), + points, + include_hit_points=compute_sdf_derivatives, + ) + if compute_sdf_derivatives: + sdf, sdf_derives = sdf + + sdf = sdf.numpy() + sdf = np.expand_dims(max_dis * sdf, axis=1) + + if compute_sdf_derivatives: + sdf_derives = sdf_derives.numpy().reshape(-1) + sdf_derives = -(sdf_derives - points) + sdf_derives = np.reshape(sdf_derives, (sdf_derives.shape[0] // 3, 3)) + sdf_derives = sdf_derives / np.linalg.norm( + sdf_derives, axis=1, keepdims=True + ) + return sdf, sdf_derives + + return sdf + + return sdf_func diff --git a/examples/smc_reac/ppsci/geometry/pointcloud.py b/examples/smc_reac/ppsci/geometry/pointcloud.py new file mode 100644 index 0000000000..ae4d4fb4cc --- /dev/null +++ b/examples/smc_reac/ppsci/geometry/pointcloud.py @@ -0,0 +1,312 @@ +# 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. + +from __future__ import annotations + +from typing import Dict +from typing import Optional +from typing import Tuple + +import numpy as np + +from ppsci.geometry import geometry +from ppsci.utils import misc + + +class PointCloud(geometry.Geometry): + """Class for point cloud geometry, i.e. a set of points from given file or array. + + Args: + interior (Dict[str, np.ndarray]): Filepath or dict data, which store interior points of a point cloud, such as {"x": np.ndarray, "y": np.ndarray}. + coord_keys (Tuple[str, ...]): Tuple of coordinate keys, such as ("x", "y"). + boundary (Dict[str, np.ndarray]): Boundary points of a point cloud. Defaults to None. + boundary_normal (Dict[str, np.ndarray]): Boundary normal points of a point cloud. Defaults to None. + + Examples: + >>> import ppsci + >>> import numpy as np + >>> interior_points = {"x": np.linspace(-1, 1, dtype="float32").reshape((-1, 1))} + >>> geom = ppsci.geometry.PointCloud(interior_points, ("x",)) + """ + + def __init__( + self, + interior: Dict[str, np.ndarray], + coord_keys: Tuple[str, ...], + boundary: Optional[Dict[str, np.ndarray]] = None, + boundary_normal: Optional[Dict[str, np.ndarray]] = None, + ): + # Interior points + self.interior = misc.convert_to_array(interior, coord_keys) + self.len = self.interior.shape[0] + + # Boundary points + self.boundary = boundary + if self.boundary is not None: + self.boundary = misc.convert_to_array(self.boundary, coord_keys) + + # Boundary normal points + self.normal = boundary_normal + if self.normal is not None: + self.normal = misc.convert_to_array( + self.normal, tuple(f"{key}_normal" for key in coord_keys) + ) + if list(self.normal.shape) != list(self.boundary.shape): + raise ValueError( + f"boundary's shape({self.boundary.shape}) must equal " + f"to normal's shape({self.normal.shape})" + ) + + self.input_keys = coord_keys + super().__init__( + len(coord_keys), + (np.amin(self.interior, axis=0), np.amax(self.interior, axis=0)), + np.inf, + ) + + @property + def dim_keys(self): + return self.input_keys + + def is_inside(self, x): + # NOTE: point on boundary is included + return ( + np.isclose((x[:, None, :] - self.interior[None, :, :]), 0, atol=1e-6) + .all(axis=2) + .any(axis=1) + ) + + def on_boundary(self, x): + if not self.boundary: + raise ValueError( + "self.boundary must be initialized" " when call 'on_boundary' function" + ) + return ( + np.isclose( + (x[:, None, :] - self.boundary[None, :, :]), + 0, + atol=1e-6, + ) + .all(axis=2) + .any(axis=1) + ) + + def translate(self, translation: np.ndarray) -> "PointCloud": + """ + Translate the geometry by the given offset. + + Args: + translation (np.ndarray): Translation offset.The shape of translation must be the same as the shape of the interior points. + + Returns: + PointCloud: Translated point cloud. + + Examples: + >>> import ppsci + >>> import numpy as np + >>> interior_points = {"x": np.linspace(0, 2, 5, dtype="float32").reshape((-1, 1))} + >>> geom = ppsci.geometry.PointCloud(interior_points, ("x",)) + >>> translation = np.array([1.0]) + >>> print(geom.translate(translation).interior) + [[1. ] + [1.5] + [2. ] + [2.5] + [3. ]] + >>> interior_points_2d = {"x": np.linspace(0, 2, 5, dtype="float32").reshape((-1, 1)), + ... "y": np.linspace(0, 2, 5, dtype="float32").reshape((-1, 1))} + >>> geom_2d = ppsci.geometry.PointCloud(interior_points_2d, ("x", "y")) + >>> translation_2d = np.array([1.0, 3.0]) + >>> print(geom_2d.translate(translation_2d).interior) + [[1. 3. ] + [1.5 3.5] + [2. 4. ] + [2.5 4.5] + [3. 5. ]] + """ + for i, offset in enumerate(translation): + self.interior[:, i] += offset + if self.boundary: + self.boundary += offset + return self + + def scale(self, scale: np.ndarray) -> "PointCloud": + """ + Scale the geometry by the given factor. + + Args: + scale (np.ndarray): Scale factor.The shape of scale must be the same as the shape of the interior points. + + Returns: + PointCloud: Scaled point cloud. + + Examples: + >>> import ppsci + >>> import numpy as np + >>> interior_points = {"x": np.linspace(0, 2, 5, dtype="float32").reshape((-1, 1))} + >>> geom = ppsci.geometry.PointCloud(interior_points, ("x",)) + >>> scale = np.array([2.0]) + >>> print(geom.scale(scale).interior) + [[0.] + [1.] + [2.] + [3.] + [4.]] + >>> interior_points_2d = {"x": np.linspace(0, 2, 5, dtype="float32").reshape((-1, 1)), + ... "y": np.linspace(0, 2, 5, dtype="float32").reshape((-1, 1))} + >>> geom_2d = ppsci.geometry.PointCloud(interior_points_2d, ("x", "y")) + >>> scale_2d = np.array([2.0, 0.5]) + >>> print(geom_2d.scale(scale_2d).interior) + [[0. 0. ] + [1. 0.25] + [2. 0.5 ] + [3. 0.75] + [4. 1. ]] + """ + for i, _scale in enumerate(scale): + self.interior[:, i] *= _scale + if self.boundary: + self.boundary[:, i] *= _scale + if self.normal: + self.normal[:, i] *= _scale + return self + + def uniform_boundary_points(self, n: int): + """Compute the equi-spaced points on the boundary.""" + raise NotImplementedError( + "PointCloud do not have 'uniform_boundary_points' method" + ) + + def random_boundary_points(self, n: int, random: str = "pseudo") -> np.ndarray: + """Randomly sample points on the boundary. + + Args: + n (int): Number of sample points. + random (str): Random method. Defaults to "pseudo". + + Returns: + np.ndarray: Randomly sampled points on the boundary.The shape of the returned array is (n, ndim). + + Examples: + >>> import ppsci + >>> import numpy as np + >>> np.random.seed(0) + >>> interior_points = {"x": np.linspace(0, 2, 5, dtype="float32").reshape((-1, 1))} + >>> boundary_points = {"x": np.array([0.0, 2.0], dtype="float32").reshape((-1, 1))} + >>> geom = ppsci.geometry.PointCloud(interior_points, ("x",), boundary_points) + >>> print(geom.random_boundary_points(1)) + [[2.]] + """ + assert self.boundary is not None, ( + "boundary points can't be empty when call " + "'random_boundary_points' method" + ) + assert n <= len(self.boundary), ( + f"number of sample points({n}) " + f"can't be more than that in boundary({len(self.boundary)})" + ) + return self.boundary[ + np.random.choice(len(self.boundary), size=n, replace=False) + ] + + def random_points(self, n: int, random: str = "pseudo") -> np.ndarray: + """Randomly sample points in the geometry. + + Args: + n (int): Number of sample points. + random (str): Random method. Defaults to "pseudo". + + Returns: + np.ndarray: Randomly sampled points in the geometry.The shape of the returned array is (n, ndim). + + Examples: + >>> import ppsci + >>> import numpy as np + >>> np.random.seed(0) + >>> interior_points = {"x": np.linspace(0, 2, 5, dtype="float32").reshape((-1, 1))} + >>> geom = ppsci.geometry.PointCloud(interior_points, ("x",)) + >>> print(geom.random_points(2)) + [[1.] + [0.]] + """ + assert n <= len(self.interior), ( + f"number of sample points({n}) " + f"can't be more than that in points({len(self.interior)})" + ) + return self.interior[ + np.random.choice(len(self.interior), size=n, replace=False) + ] + + def uniform_points(self, n: int, boundary: bool = True) -> np.ndarray: + """Compute the equi-spaced points in the geometry. + + Args: + n (int): Number of sample points. + boundary (bool): Whether to include boundary points. Defaults to True. + + Returns: + np.ndarray: Equi-spaced points in the geometry.The shape of the returned array is (n, ndim). + + Examples: + >>> import ppsci + >>> import numpy as np + >>> interior_points = {"x": np.linspace(0, 2, 5, dtype="float32").reshape((-1, 1))} + >>> geom = ppsci.geometry.PointCloud(interior_points, ("x",)) + >>> print(geom.uniform_points(2)) + [[0. ] + [0.5]] + """ + return self.interior[:n] + + def union(self, other): + raise NotImplementedError( + "Union operation for PointCloud is not supported yet." + ) + + def __or__(self, other): + raise NotImplementedError( + "Union operation for PointCloud is not supported yet." + ) + + def difference(self, other): + raise NotImplementedError( + "Subtraction operation for PointCloud is not supported yet." + ) + + def __sub__(self, other): + raise NotImplementedError( + "Subtraction operation for PointCloud is not supported yet." + ) + + def intersection(self, other): + raise NotImplementedError( + "Intersection operation for PointCloud is not supported yet." + ) + + def __and__(self, other): + raise NotImplementedError( + "Intersection operation for PointCloud is not supported yet." + ) + + def __str__(self) -> str: + """Return the name of class.""" + return ", ".join( + [ + self.__class__.__name__, + f"num_points = {len(self.interior)}", + f"ndim = {self.ndim}", + f"bbox = {self.bbox}", + f"dim_keys = {self.dim_keys}", + ] + ) diff --git a/examples/smc_reac/ppsci/geometry/sampler.py b/examples/smc_reac/ppsci/geometry/sampler.py new file mode 100644 index 0000000000..a6de5015ff --- /dev/null +++ b/examples/smc_reac/ppsci/geometry/sampler.py @@ -0,0 +1,92 @@ +# 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. + +""" +Code below is heavily based on [https://github.com/lululxvi/deepxde](https://github.com/lululxvi/deepxde) +""" + +from __future__ import annotations + +import numpy as np +import paddle +import skopt +from typing_extensions import Literal + + +def sample( + n_samples: int, ndim: int, method: Literal["pseudo", "Halton", "LHS"] = "pseudo" +) -> np.ndarray: + """Generate pseudorandom or quasi-random samples in [0, 1]^ndim. + + Args: + n_samples (int): The number of samples. + ndim (int): Number of dimension. + method (str): One of the following: "pseudo" (pseudorandom), "LHS" (Latin + hypercube sampling), "Halton" (Halton sequence), "Hammersley" (Hammersley + sequence), or "Sobol" (Sobol sequence). + + Returns: + np.ndarray: Generated random samples with shape of [n_samples, ndim]. + """ + if method == "pseudo": + return pseudorandom(n_samples, ndim) + if method in ["LHS", "Halton", "Hammersley", "Sobol"]: + return quasirandom(n_samples, ndim, method) + raise ValueError(f"Sampling method({method}) is not available.") + + +def pseudorandom(n_samples: int, ndim: int) -> np.ndarray: + """Pseudo random.""" + # If random seed is set, then the rng based code always returns the same random + # number, which may not be what we expect. + # rng = np.random.default_rng(config.random_seed) + # return rng.random(size=(n_samples, ndim), dtype=dtype=paddle.get_default_dtype()) + return np.random.random(size=(n_samples, ndim)).astype( + dtype=paddle.get_default_dtype() + ) + + +def quasirandom( + n_samples: int, ndim: int, method: Literal["pseudo", "LHS"] +) -> np.ndarray: + """Quasi random""" + # Certain points should be removed: + # - Boundary points such as [..., 0, ...] + # - Special points [0, 0, 0, ...] and [0.5, 0.5, 0.5, ...], which cause error in + # Hypersphere.random_points() and Hypersphere.random_boundary_points() + skip = 0 + if method == "LHS": + sampler = skopt.sampler.Lhs() + elif method == "Halton": + # 1st point: [0, 0, ...] + sampler = skopt.sampler.Halton(min_skip=1, max_skip=1) + elif method == "Hammersley": + # 1st point: [0, 0, ...] + if ndim == 1: + sampler = skopt.sampler.Hammersly(min_skip=1, max_skip=1) + else: + sampler = skopt.sampler.Hammersly() + skip = 1 + elif method == "Sobol": + # 1st point: [0, 0, ...], 2nd point: [0.5, 0.5, ...] + sampler = skopt.sampler.Sobol(randomize=False) + if ndim < 3: + skip = 1 + else: + skip = 2 + space = [(0.0, 1.0)] * ndim + return np.asarray( + sampler.generate(space, n_samples + skip)[skip:], + dtype=paddle.get_default_dtype(), + ) diff --git a/examples/smc_reac/ppsci/geometry/sdf.py b/examples/smc_reac/ppsci/geometry/sdf.py new file mode 100644 index 0000000000..bb3e260540 --- /dev/null +++ b/examples/smc_reac/ppsci/geometry/sdf.py @@ -0,0 +1,198 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +# ruff: noqa: F401 + +# modified from: https://github.com/NVIDIA/modulus/blob/main/modulus/utils/sdf.py + +from __future__ import annotations + +import importlib.util +from typing import Tuple +from typing import overload + +from numpy import ndarray + +try: + import warp as wp + + @wp.kernel + def _bvh_query_distance( + mesh: wp.uint64, + points: wp.array(dtype=wp.vec3f), + max_dist: wp.float32, + sdf: wp.array(dtype=wp.float32), + sdf_hit_point: wp.array(dtype=wp.vec3f), + sdf_hit_point_id: wp.array(dtype=wp.int32), + ): + + """ + Computes the signed distance from each point in the given array `points` + to the mesh represented by `mesh`,within the maximum distance `max_dist`, + and stores the result in the array `sdf`. + + Parameters: + mesh (wp.uint64): The identifier of the mesh. + points (wp.array): An array of 3D points for which to compute the + signed distance. + max_dist (wp.float32): The maximum distance within which to search + for the closest point on the mesh. + sdf (wp.array): An array to store the computed signed distances. + sdf_hit_point (wp.array): An array to store the computed hit points. + sdf_hit_point_id (wp.array): An array to store the computed hit point ids. + + Returns: + None + """ + tid = wp.tid() + + res = wp.mesh_query_point_sign_normal(mesh, points[tid], max_dist) + + mesh_ = wp.mesh_get(mesh) + + p0 = mesh_.points[mesh_.indices[3 * res.face + 0]] + p1 = mesh_.points[mesh_.indices[3 * res.face + 1]] + p2 = mesh_.points[mesh_.indices[3 * res.face + 2]] + + p_closest = res.u * p0 + res.v * p1 + (1.0 - res.u - res.v) * p2 + + sdf[tid] = res.sign * wp.abs(wp.length(points[tid] - p_closest)) + sdf_hit_point[tid] = p_closest + sdf_hit_point_id[tid] = res.face + +except ModuleNotFoundError: + pass +except Exception: + raise + + +@overload +def signed_distance_field( + mesh_vertices: list[tuple[float, float, float]], + mesh_indices: ndarray, + input_points: list[tuple[float, float, float]], + max_dist: float = 1e8, + include_hit_points: bool = False, + include_hit_points_id: bool = False, +) -> wp.array: + ... + + +@overload +def signed_distance_field( + mesh_vertices: list[tuple[float, float, float]], + mesh_indices: ndarray, + input_points: list[tuple[float, float, float]], + max_dist: float = 1e8, + include_hit_points: bool = True, + include_hit_points_id: bool = False, +) -> Tuple[wp.array, wp.array]: + ... + + +@overload +def signed_distance_field( + mesh_vertices: list[tuple[float, float, float]], + mesh_indices: ndarray, + input_points: list[tuple[float, float, float]], + max_dist: float = 1e8, + include_hit_points: bool = False, + include_hit_points_id: bool = True, +) -> Tuple[wp.array, wp.array]: + ... + + +@overload +def signed_distance_field( + mesh_vertices: list[tuple[float, float, float]], + mesh_indices: ndarray, + input_points: list[tuple[float, float, float]], + max_dist: float = 1e8, + include_hit_points: bool = True, + include_hit_points_id: bool = True, +) -> Tuple[wp.array, wp.array, wp.array]: + ... + + +def signed_distance_field( + mesh_vertices: list[tuple[float, float, float]], + mesh_indices: ndarray, + input_points: list[tuple[float, float, float]], + max_dist: float = 1e8, + include_hit_points: bool = False, + include_hit_points_id: bool = False, +) -> wp.array: + """ + Computes the signed distance field (SDF) for a given mesh and input points. + + Args: + mesh_vertices (list[tuple[float, float, float]]): List of vertices defining the mesh. + mesh_indices (list[tuple[int, int, int]]): List of indices defining the triangles of the mesh. + input_points (list[tuple[float, float, float]]): List of input points for which to compute the SDF. + max_dist (float, optional): Maximum distance within which to search for + the closest point on the mesh. Default is 1e8. + include_hit_points (bool, optional): Whether to include hit points in + the output. Default is False. + include_hit_points_id (bool, optional): Whether to include hit point + IDs in the output. Default is False. + + Returns: + wp.array: An array containing the computed signed distance field. + + Example: + >>> mesh_vertices = [(0, 0, 0), (1, 0, 0), (0, 1, 0)] + >>> mesh_indices = np.array((0, 1, 2)) + >>> input_points = [(0.5, 0.5, 0.5)] + >>> signed_distance_field(mesh_vertices, mesh_indices, input_points).numpy() + Module modulus.utils.sdf load on device 'cuda:0' took ... + array([0.5], dtype=float32) + """ + if not importlib.util.find_spec("warp"): + raise ModuleNotFoundError("Please install warp with: pip install warp-lang") + + wp.init() + mesh = wp.Mesh( + wp.array(mesh_vertices, dtype=wp.vec3), + wp.array(mesh_indices, dtype=wp.int32), + ) + + sdf_points = wp.array(input_points, dtype=wp.vec3) + + sdf = wp.zeros(shape=sdf_points.shape, dtype=wp.float32) + sdf_hit_point = wp.zeros(shape=sdf_points.shape, dtype=wp.vec3f) + sdf_hit_point_id = wp.zeros(shape=sdf_points.shape, dtype=wp.int32) + + wp.launch( + kernel=_bvh_query_distance, + dim=len(sdf_points), + inputs=[ + mesh.id, + sdf_points, + max_dist, + sdf, + sdf_hit_point, + sdf_hit_point_id, + ], + ) + + if include_hit_points and include_hit_points_id: + return (sdf, sdf_hit_point, sdf_hit_point_id) + elif include_hit_points: + return (sdf, sdf_hit_point) + elif include_hit_points_id: + return (sdf, sdf_hit_point_id) + else: + return sdf diff --git a/examples/smc_reac/ppsci/geometry/timedomain.py b/examples/smc_reac/ppsci/geometry/timedomain.py new file mode 100644 index 0000000000..909944a9b4 --- /dev/null +++ b/examples/smc_reac/ppsci/geometry/timedomain.py @@ -0,0 +1,793 @@ +# 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. + +""" +Code below is heavily based on [https://github.com/lululxvi/deepxde](https://github.com/lululxvi/deepxde) +""" + +from __future__ import annotations + +import itertools +from typing import Callable +from typing import Dict +from typing import Optional +from typing import Tuple + +import numpy as np +import paddle + +from ppsci.geometry import geometry +from ppsci.geometry import geometry_1d +from ppsci.geometry import geometry_2d +from ppsci.geometry import geometry_3d +from ppsci.geometry import geometry_nd +from ppsci.geometry import mesh +from ppsci.utils import misc + + +class TimeDomain(geometry_1d.Interval): + """Class for timedomain, an special interval geometry. + + Args: + t0 (float): Start of time. + t1 (float): End of time. + time_step (Optional[float]): Step interval of time. Defaults to None. + timestamps (Optional[Tuple[float, ...]]): List of timestamps. + Defaults to None. + + Examples: + >>> import ppsci + >>> geom = ppsci.geometry.TimeDomain(0, 1) + """ + + def __init__( + self, + t0: float, + t1: float, + time_step: Optional[float] = None, + timestamps: Optional[Tuple[float, ...]] = None, + ): + super().__init__(t0, t1) + self.t0 = t0 + self.t1 = t1 + self.time_step = time_step + if timestamps is None: + self.timestamps = None + else: + self.timestamps = np.array( + timestamps, dtype=paddle.get_default_dtype() + ).reshape([-1]) + if time_step is not None: + if time_step <= 0: + raise ValueError(f"time_step({time_step}) must be larger than 0.") + self.num_timestamps = int(np.ceil((t1 - t0) / time_step)) + 1 + elif timestamps is not None: + self.num_timestamps = len(timestamps) + + def on_initial(self, t: np.ndarray) -> np.ndarray: + """Check if a specific time is on the initial time point. + + Args: + t (np.ndarray): The time to be checked. + + Returns: + np.ndarray: Bool numpy array of whether the specific time is on the initial time point. + + Examples: + >>> import paddle + >>> import ppsci + >>> geom = ppsci.geometry.TimeDomain(0, 1) + >>> T = [0, 0.01, 0.126, 0.2, 0.3] + >>> check = geom.on_initial(T) + >>> print(check) + [ True False False False False] + """ + return np.isclose(t, self.t0).flatten() + + +class TimeXGeometry(geometry.Geometry): + """Class for combination of time and geometry. + + Args: + timedomain (TimeDomain): TimeDomain object. + geometry (geometry.Geometry): Geometry object. + + Examples: + >>> import ppsci + >>> timedomain = ppsci.geometry.TimeDomain(0, 1) + >>> geom = ppsci.geometry.Rectangle((0, 0), (1, 1)) + >>> time_geom = ppsci.geometry.TimeXGeometry(timedomain, geom) + """ + + def __init__(self, timedomain: TimeDomain, geometry: geometry.Geometry): + self.timedomain = timedomain + self.geometry = geometry + self.ndim = geometry.ndim + timedomain.ndim + + @property + def dim_keys(self): + return ("t",) + self.geometry.dim_keys + + def on_boundary(self, x): + # [N, ndim(txyz)] + return self.geometry.on_boundary(x[:, 1:]) + + def on_initial(self, x): + # [N, 1(t)] + return self.timedomain.on_initial(x[:, :1]) + + def boundary_normal(self, x): + # x: [N, ndim(txyz)] + normal = self.geometry.boundary_normal(x[:, 1:]) + return np.hstack((x[:, :1], normal)) + + def uniform_points(self, n: int, boundary: bool = True) -> np.ndarray: + """Uniform points on the spatial-temporal domain. + Geometry volume ~ bbox. + Time volume ~ diam. + + Args: + n (int): The total number of sample points to be generated. + boundary (bool): Indicates whether boundary points are included, default is True. + + Returns: + np.ndarray: a set of spatial-temporal coordinate points 'tx' that represent sample points evenly distributed within the spatial-temporal domain. + + Examples: + >>> import ppsci + >>> timedomain = ppsci.geometry.TimeDomain(0, 1, 0.001) + >>> geom = ppsci.geometry.Rectangle((0, 0), (1, 1)) + >>> time_geom = ppsci.geometry.TimeXGeometry(timedomain, geom) + >>> ts = time_geom.uniform_points(1000) + >>> print(ts.shape) + (1000, 3) + """ + if self.timedomain.time_step is not None: + # exclude start time t0 + nt = int(np.ceil(self.timedomain.diam / self.timedomain.time_step)) + nx = int(np.ceil(n / nt)) + elif self.timedomain.timestamps is not None: + # exclude start time t0 + nt = self.timedomain.num_timestamps - 1 + nx = int(np.ceil(n / nt)) + else: + nx = int( + np.ceil( + ( + n + * np.prod(self.geometry.bbox[1] - self.geometry.bbox[0]) + / self.timedomain.diam + ) + ** 0.5 + ) + ) + nt = int(np.ceil(n / nx)) + x = self.geometry.uniform_points(nx, boundary=boundary) + nx = len(x) + if boundary and ( + self.timedomain.time_step is None and self.timedomain.timestamps is None + ): + t = self.timedomain.uniform_points(nt, boundary=True) + else: + if self.timedomain.time_step is not None: + t = np.linspace( + self.timedomain.t1, + self.timedomain.t0, + num=nt, + endpoint=boundary, + dtype=paddle.get_default_dtype(), + )[:, None][::-1] + else: + t = self.timedomain.timestamps[1:] + tx = [] + for ti in t: + tx.append( + np.hstack((np.full([nx, 1], ti, dtype=paddle.get_default_dtype()), x)) + ) + tx = np.vstack(tx) + if len(tx) > n: + tx = tx[:n] + return tx + + def random_points( + self, n: int, random: str = "pseudo", criteria: Optional[Callable] = None + ) -> np.ndarray: + """Generate random points on the spatial-temporal domain. + + Args: + n (int): The total number of random points to generate. + random (str): Specifies the way to generate random points, default is "pseudo" , which means that a pseudo-random number generator is used. + criteria (Optional[Callable]): A method that filters on the generated random points. Defaults to None. + + Returns: + np.ndarray: A set of random spatial-temporal points. + + Examples: + >>> import ppsci + >>> timedomain = ppsci.geometry.TimeDomain(0, 1, 0.001) + >>> geom = ppsci.geometry.Rectangle((0, 0), (1, 1)) + >>> time_geom = ppsci.geometry.TimeXGeometry(timedomain, geom) + >>> ts = time_geom.random_points(1000) + >>> print(ts.shape) + (1000, 3) + """ + if self.timedomain.time_step is None and self.timedomain.timestamps is None: + raise ValueError("Either time_step or timestamps must be provided.") + # time evenly and geometry random, if time_step if specified + if self.timedomain.time_step is not None: + nt = int(np.ceil(self.timedomain.diam / self.timedomain.time_step)) + t = np.linspace( + self.timedomain.t1, + self.timedomain.t0, + num=nt, + endpoint=False, + dtype=paddle.get_default_dtype(), + )[:, None][ + ::-1 + ] # [nt, 1] + # 1. sample nx points in static geometry with criteria + nx = int(np.ceil(n / nt)) + _size, _ntry, _nsuc = 0, 0, 0 + x = np.empty( + shape=(nx, self.geometry.ndim), dtype=paddle.get_default_dtype() + ) + while _size < nx: + _x = self.geometry.random_points(nx, random) + if criteria is not None: + # fix arg 't' to None in criteria there + criteria_mask = criteria( + None, *np.split(_x, self.geometry.ndim, axis=1) + ).flatten() + _x = _x[criteria_mask] + if len(_x) > nx - _size: + _x = _x[: nx - _size] + x[_size : _size + len(_x)] = _x + + _size += len(_x) + _ntry += 1 + if len(_x) > 0: + _nsuc += 1 + + if _ntry >= 1000 and _nsuc == 0: + raise ValueError( + "Sample points failed, " + "please check correctness of geometry and given criteria." + ) + + # 2. repeat spatial points along time + tx = [] + for ti in t: + tx.append( + np.hstack( + (np.full([nx, 1], ti, dtype=paddle.get_default_dtype()), x) + ) + ) + tx = np.vstack(tx) + if len(tx) > n: + tx = tx[:n] + return tx + elif self.timedomain.timestamps is not None: + nt = self.timedomain.num_timestamps - 1 + t = self.timedomain.timestamps[1:] + nx = int(np.ceil(n / nt)) + + _size, _ntry, _nsuc = 0, 0, 0 + x = np.empty( + shape=(nx, self.geometry.ndim), dtype=paddle.get_default_dtype() + ) + while _size < nx: + _x = self.geometry.random_points(nx, random) + if criteria is not None: + # fix arg 't' to None in criteria there + criteria_mask = criteria( + None, *np.split(_x, self.geometry.ndim, axis=1) + ).flatten() + _x = _x[criteria_mask] + if len(_x) > nx - _size: + _x = _x[: nx - _size] + x[_size : _size + len(_x)] = _x + + _size += len(_x) + _ntry += 1 + if len(_x) > 0: + _nsuc += 1 + + if _ntry >= 1000 and _nsuc == 0: + raise ValueError( + "Sample interior points failed, " + "please check correctness of geometry and given criteria." + ) + + tx = [] + for ti in t: + tx.append( + np.hstack( + (np.full([nx, 1], ti, dtype=paddle.get_default_dtype()), x) + ) + ) + tx = np.vstack(tx) + if len(tx) > n: + tx = tx[:n] + return tx + + if isinstance(self.geometry, geometry_1d.Interval): + geom = geometry_2d.Rectangle( + [self.timedomain.t0, self.geometry.l], + [self.timedomain.t1, self.geometry.r], + ) + return geom.random_points(n, random=random) + + if isinstance(self.geometry, geometry_2d.Rectangle): + geom = geometry_3d.Cuboid( + [self.timedomain.t0, self.geometry.xmin[0], self.geometry.xmin[1]], + [self.timedomain.t1, self.geometry.xmax[0], self.geometry.xmax[1]], + ) + return geom.random_points(n, random=random) + + if isinstance(self.geometry, (geometry_3d.Cuboid, geometry_nd.Hypercube)): + geom = geometry_nd.Hypercube( + np.append(self.timedomain.t0, self.geometry.xmin), + np.append(self.timedomain.t1, self.geometry.xmax), + ) + return geom.random_points(n, random=random) + + x = self.geometry.random_points(n, random=random) + t = self.timedomain.random_points(n, random=random) + t = np.random.permutation(t) + return np.hstack((t, x)) + + def uniform_boundary_points( + self, n: int, criteria: Optional[Callable] = None + ) -> np.ndarray: + """Uniform boundary points on the spatial-temporal domain. + Geometry surface area ~ bbox. + Time surface area ~ diam. + + Args: + n (int): The total number of boundary points on the spatial-temporal domain to be generated that are evenly distributed across geometry boundaries. + criteria (Optional[Callable]): Used to filter the generated boundary points, only points that meet certain conditions are retained. Default is None. + + Returns: + np.ndarray: A set of point coordinates evenly distributed across geometry boundaries on the spatial-temporal domain. + + Examples: + >>> import ppsci + >>> timedomain = ppsci.geometry.TimeDomain(0, 1) + >>> geom = ppsci.geometry.Rectangle((0, 0), (1, 1)) + >>> time_geom = ppsci.geometry.TimeXGeometry(timedomain, geom) + >>> ts = time_geom.uniform_boundary_points(1000) + >>> print(ts.shape) + (1000, 3) + """ + if self.geometry.ndim == 1: + nx = 2 + else: + s = 2 * sum( + map( + lambda l: l[0] * l[1], + itertools.combinations( + self.geometry.bbox[1] - self.geometry.bbox[0], 2 + ), + ) + ) + nx = int((n * s / self.timedomain.diam) ** 0.5) + nt = int(np.ceil(n / nx)) + + _size, _ntry, _nsuc = 0, 0, 0 + x = np.empty(shape=(nx, self.geometry.ndim), dtype=paddle.get_default_dtype()) + while _size < nx: + _x = self.geometry.uniform_boundary_points(nx) + if criteria is not None: + # fix arg 't' to None in criteria there + criteria_mask = criteria( + None, *np.split(_x, self.geometry.ndim, axis=1) + ).flatten() + _x = _x[criteria_mask] + if len(_x) > nx - _size: + _x = _x[: nx - _size] + x[_size : _size + len(_x)] = _x + + _size += len(_x) + _ntry += 1 + if len(_x) > 0: + _nsuc += 1 + + if _ntry >= 1000 and _nsuc == 0: + raise ValueError( + "Sample boundary points failed, " + "please check correctness of geometry and given criteria." + ) + + nx = len(x) + t = np.linspace( + self.timedomain.t1, + self.timedomain.t0, + num=nt, + endpoint=False, + dtype=paddle.get_default_dtype(), + )[:, None][::-1] + tx = [] + for ti in t: + tx.append( + np.hstack((np.full([nx, 1], ti, dtype=paddle.get_default_dtype()), x)) + ) + tx = np.vstack(tx) + if len(tx) > n: + tx = tx[:n] + return tx + + def random_boundary_points( + self, n: int, random: str = "pseudo", criteria: Optional[Callable] = None + ) -> np.ndarray: + """Random boundary points on the spatial-temporal domain. + + Args: + n (int): The total number of spatial-temporal points generated on a given geometry boundary. + random (str): Controls the way to generate random points. Default is "pseudo". + criteria (Optional[Callable]): Used to filter the generated boundary points, only points that meet certain conditions are retained. Default is None. + + Returns: + np.ndarray: A set of point coordinates randomly distributed across geometry boundaries on the spatial-temporal domain. + + Examples: + >>> import ppsci + >>> timedomain = ppsci.geometry.TimeDomain(0, 1, 0.001) + >>> geom = ppsci.geometry.Rectangle((0, 0), (1, 1)) + >>> time_geom = ppsci.geometry.TimeXGeometry(timedomain, geom) + >>> ts = time_geom.random_boundary_points(1000) + >>> print(ts.shape) + (1000, 3) + """ + if self.timedomain.time_step is None and self.timedomain.timestamps is None: + raise ValueError("Either time_step or timestamps must be provided.") + if self.timedomain.time_step is not None: + # exclude start time t0 + nt = int(np.ceil(self.timedomain.diam / self.timedomain.time_step)) + t = np.linspace( + self.timedomain.t1, + self.timedomain.t0, + num=nt, + endpoint=False, + dtype=paddle.get_default_dtype(), + )[:, None][::-1] + nx = int(np.ceil(n / nt)) + + if isinstance(self.geometry, mesh.Mesh): + x, _n, a = self.geometry.random_boundary_points(nx, random=random) + else: + _size, _ntry, _nsuc = 0, 0, 0 + x = np.empty( + shape=(nx, self.geometry.ndim), dtype=paddle.get_default_dtype() + ) + while _size < nx: + _x = self.geometry.random_boundary_points(nx, random) + if criteria is not None: + # fix arg 't' to None in criteria there + criteria_mask = criteria( + None, *np.split(_x, self.geometry.ndim, axis=1) + ).flatten() + _x = _x[criteria_mask] + if len(_x) > nx - _size: + _x = _x[: nx - _size] + x[_size : _size + len(_x)] = _x + + _size += len(_x) + _ntry += 1 + if len(_x) > 0: + _nsuc += 1 + + if _ntry >= 1000 and _nsuc == 0: + raise ValueError( + "Sample boundary points failed, " + "please check correctness of geometry and given criteria." + ) + + t_x = [] + if isinstance(self.geometry, mesh.Mesh): + t_normal = [] + t_area = [] + + for ti in t: + t_x.append( + np.hstack( + (np.full([nx, 1], ti, dtype=paddle.get_default_dtype()), x) + ) + ) + if isinstance(self.geometry, mesh.Mesh): + t_normal.append( + np.hstack( + (np.full([nx, 1], ti, dtype=paddle.get_default_dtype()), _n) + ) + ) + t_area.append( + np.hstack( + (np.full([nx, 1], ti, dtype=paddle.get_default_dtype()), a) + ) + ) + + t_x = np.vstack(t_x) + if isinstance(self.geometry, mesh.Mesh): + t_normal = np.vstack(t_normal) + t_area = np.vstack(t_area) + + if len(t_x) > n: + t_x = t_x[:n] + if isinstance(self.geometry, mesh.Mesh): + t_normal = t_normal[:n] + t_area = t_area[:n] + + if isinstance(self.geometry, mesh.Mesh): + return t_x, t_normal, t_area + else: + return t_x + elif self.timedomain.timestamps is not None: + # exclude start time t0 + nt = self.timedomain.num_timestamps - 1 + t = self.timedomain.timestamps[1:] + nx = int(np.ceil(n / nt)) + + if isinstance(self.geometry, mesh.Mesh): + x, _n, a = self.geometry.random_boundary_points(nx, random=random) + else: + _size, _ntry, _nsuc = 0, 0, 0 + x = np.empty( + shape=(nx, self.geometry.ndim), dtype=paddle.get_default_dtype() + ) + while _size < nx: + _x = self.geometry.random_boundary_points(nx, random) + if criteria is not None: + # fix arg 't' to None in criteria there + criteria_mask = criteria( + None, *np.split(_x, self.geometry.ndim, axis=1) + ).flatten() + _x = _x[criteria_mask] + if len(_x) > nx - _size: + _x = _x[: nx - _size] + x[_size : _size + len(_x)] = _x + + _size += len(_x) + _ntry += 1 + if len(_x) > 0: + _nsuc += 1 + + if _ntry >= 1000 and _nsuc == 0: + raise ValueError( + "Sample boundary points failed, " + "please check correctness of geometry and given criteria." + ) + + t_x = [] + if isinstance(self.geometry, mesh.Mesh): + t_normal = [] + t_area = [] + + for ti in t: + t_x.append( + np.hstack( + (np.full([nx, 1], ti, dtype=paddle.get_default_dtype()), x) + ) + ) + if isinstance(self.geometry, mesh.Mesh): + t_normal.append( + np.hstack( + (np.full([nx, 1], ti, dtype=paddle.get_default_dtype()), _n) + ) + ) + t_area.append( + np.hstack( + (np.full([nx, 1], ti, dtype=paddle.get_default_dtype()), a) + ) + ) + + t_x = np.vstack(t_x) + if isinstance(self.geometry, mesh.Mesh): + t_normal = np.vstack(t_normal) + t_area = np.vstack(t_area) + + if len(t_x) > n: + t_x = t_x[:n] + if isinstance(self.geometry, mesh.Mesh): + t_normal = t_normal[:n] + t_area = t_area[:n] + + if isinstance(self.geometry, mesh.Mesh): + return t_x, t_normal, t_area + else: + return t_x + else: + if isinstance(self.geometry, mesh.Mesh): + x, _n, a = self.geometry.random_boundary_points(n, random=random) + else: + x = self.geometry.random_boundary_points(n, random=random) + + t = self.timedomain.random_points(n, random=random) + t = np.random.permutation(t) + + t_x = np.hstack((t, x)) + + if isinstance(self.geometry, mesh.Mesh): + t_normal = np.hstack((_n, t)) + t_area = np.hstack((_n, t)) + return t_x, t_normal, t_area + else: + return t_x + + def uniform_initial_points(self, n: int) -> np.ndarray: + """Generate evenly distributed point coordinates on the spatial-temporal domain at the initial moment. + + Args: + n (int): The total number of generated points. + + Returns: + np.ndarray: A set of point coordinates evenly distributed on the spatial-temporal domain at the initial moment. + + Examples: + >>> import ppsci + >>> timedomain = ppsci.geometry.TimeDomain(0, 1) + >>> geom = ppsci.geometry.Rectangle((0, 0), (1, 1)) + >>> time_geom = ppsci.geometry.TimeXGeometry(timedomain, geom) + >>> ts = time_geom.uniform_initial_points(1000) + >>> print(ts.shape) + (1000, 3) + """ + x = self.geometry.uniform_points(n, True) + t = self.timedomain.t0 + if len(x) > n: + x = x[:n] + return np.hstack((np.full([n, 1], t, dtype=paddle.get_default_dtype()), x)) + + def random_initial_points(self, n: int, random: str = "pseudo") -> np.ndarray: + """Generate randomly distributed point coordinates on the spatial-temporal domain at the initial moment. + + Args: + n (int): The total number of generated points. + random (str): Controls the way to generate random points. Default is "pseudo". + + Returns: + np.ndarray: A set of point coordinates randomly distributed on the spatial-temporal domain at the initial moment. + + Examples: + >>> import ppsci + >>> timedomain = ppsci.geometry.TimeDomain(0, 1) + >>> geom = ppsci.geometry.Rectangle((0, 0), (1, 1)) + >>> time_geom = ppsci.geometry.TimeXGeometry(timedomain, geom) + >>> ts = time_geom.random_initial_points(1000) + >>> print(ts.shape) + (1000, 3) + """ + x = self.geometry.random_points(n, random=random) + t = self.timedomain.t0 + return np.hstack((np.full([n, 1], t, dtype=paddle.get_default_dtype()), x)) + + def periodic_point( + self, x: Dict[str, np.ndarray], component: int + ) -> Dict[str, np.ndarray]: + """Process given point coordinates to satisfy the periodic boundary conditions of the geometry. + + Args: + x (Dict[str, np.ndarray]): Contains the coordinates and timestamps of the points. It represents the coordinates of the point to be processed. + component (int): Specifies the components or dimensions of specific spatial coordinates that are periodically processed. + + Returns: + Dict[str, np.ndarray] : contains the original timestamps and the coordinates of the spatial point after periodic processing. + + Examples: + >>> import ppsci + >>> timedomain = ppsci.geometry.TimeDomain(0, 1, 0.1) + >>> geom = ppsci.geometry.Rectangle((0, 0), (1, 1)) + >>> time_geom = ppsci.geometry.TimeXGeometry(timedomain, geom) + >>> ts = time_geom.sample_boundary(1000) + >>> result = time_geom.periodic_point(ts, 0) + >>> for k,v in result.items(): + ... print(k, v.shape) + t (1000, 1) + x (1000, 1) + y (1000, 1) + normal_x (1000, 1) + normal_y (1000, 1) + """ + xp = self.geometry.periodic_point(x, component) + txp = {"t": x["t"], **xp} + return txp + + def sample_initial_interior( + self, + n: int, + random: str = "pseudo", + criteria: Optional[Callable] = None, + evenly: bool = False, + compute_sdf_derivatives: bool = False, + ) -> Dict[str, np.ndarray]: + """Sample random points in the time-geometry and return those meet criteria. + + Args: + n (int): The total number of interior points generated. + random (str): The method used to specify the initial point of generation. Default is "pseudo". + criteria (Optional[Callable]): Used to filter the generated interior points, only points that meet certain conditions are retained. Default is None. + evenly (bool): Indicates whether the initial points are generated evenly. Default is False. + compute_sdf_derivatives (bool): Indicates whether to calculate the derivative of signed distance function or not. Default is False. + + Returns: + np.ndarray: Contains the coordinates of the initial internal point generated, as well as the potentially computed signed distance function and its derivative. + + Examples: + >>> import ppsci + >>> timedomain = ppsci.geometry.TimeDomain(0, 1) + >>> geom = ppsci.geometry.Rectangle((0, 0), (1, 1)) + >>> time_geom = ppsci.geometry.TimeXGeometry(timedomain, geom) + >>> ts = time_geom.sample_initial_interior(1000) + >>> for k,v in ts.items(): + ... print(k, v.shape) + t (1000, 1) + x (1000, 1) + y (1000, 1) + sdf (1000, 1) + """ + x = np.empty(shape=(n, self.ndim), dtype=paddle.get_default_dtype()) + _size, _ntry, _nsuc = 0, 0, 0 + while _size < n: + if evenly: + points = self.uniform_initial_points(n) + else: + points = self.random_initial_points(n, random) + + if criteria is not None: + criteria_mask = criteria(*np.split(points, self.ndim, axis=1)).flatten() + points = points[criteria_mask] + + if len(points) > n - _size: + points = points[: n - _size] + x[_size : _size + len(points)] = points + + _size += len(points) + _ntry += 1 + if len(points) > 0: + _nsuc += 1 + + if _ntry >= 1000 and _nsuc == 0: + raise ValueError( + "Sample initial interior points failed, " + "please check correctness of geometry and given criteria." + ) + + # if sdf_func added, return x_dict and sdf_dict, else, only return the x_dict + if hasattr(self.geometry, "sdf_func"): + # compute sdf excluding time t + sdf = -self.geometry.sdf_func(x[..., 1:]) + sdf_dict = misc.convert_to_dict(sdf, ("sdf",)) + sdf_derives_dict = {} + if compute_sdf_derivatives: + # compute sdf derivatives excluding time t + sdf_derives = -self.geometry.sdf_derivatives(x[..., 1:]) + sdf_derives_dict = misc.convert_to_dict( + sdf_derives, tuple(f"sdf__{key}" for key in self.geometry.dim_keys) + ) + else: + sdf_dict = {} + sdf_derives_dict = {} + x_dict = misc.convert_to_dict(x, self.dim_keys) + + return {**x_dict, **sdf_dict, **sdf_derives_dict} + + def __str__(self) -> str: + """Return the name of class""" + return ", ".join( + [ + self.__class__.__name__, + f"ndim = {self.ndim}", + f"bbox = (time){self.timedomain.bbox} x (space){self.geometry.bbox}", + f"diam = (time){self.timedomain.diam} x (space){self.geometry.diam}", + f"dim_keys = {self.dim_keys}", + ] + ) diff --git a/examples/smc_reac/ppsci/loss/__init__.py b/examples/smc_reac/ppsci/loss/__init__.py new file mode 100644 index 0000000000..d9502e8248 --- /dev/null +++ b/examples/smc_reac/ppsci/loss/__init__.py @@ -0,0 +1,67 @@ +# 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. + +import copy + +from ppsci.loss import mtl +from ppsci.loss.base import Loss +from ppsci.loss.chamfer import ChamferLoss +from ppsci.loss.func import FunctionalLoss +from ppsci.loss.integral import IntegralLoss +from ppsci.loss.kl import KLLoss +from ppsci.loss.l1 import L1Loss +from ppsci.loss.l1 import PeriodicL1Loss +from ppsci.loss.l2 import L2Loss +from ppsci.loss.l2 import L2RelLoss +from ppsci.loss.l2 import PeriodicL2Loss +from ppsci.loss.mae import MAELoss +from ppsci.loss.mse import CausalMSELoss +from ppsci.loss.mse import MSELoss +from ppsci.loss.mse import MSELossWithL2Decay +from ppsci.loss.mse import PeriodicMSELoss + +__all__ = [ + "Loss", + "FunctionalLoss", + "IntegralLoss", + "L1Loss", + "PeriodicL1Loss", + "L2Loss", + "L2RelLoss", + "PeriodicL2Loss", + "MAELoss", + "CausalMSELoss", + "ChamferLoss", + "MSELoss", + "MSELossWithL2Decay", + "PeriodicMSELoss", + "KLLoss", + "mtl", +] + + +def build_loss(cfg): + """Build loss. + + Args: + cfg (DictConfig): Loss config. + + Returns: + Loss: Callable loss object. + """ + cfg = copy.deepcopy(cfg) + + loss_cls = cfg.pop("name") + loss = eval(loss_cls)(**cfg) + return loss diff --git a/examples/smc_reac/ppsci/loss/base.py b/examples/smc_reac/ppsci/loss/base.py new file mode 100644 index 0000000000..378013bb9e --- /dev/null +++ b/examples/smc_reac/ppsci/loss/base.py @@ -0,0 +1,38 @@ +# 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. + +from __future__ import annotations + +from typing import Dict +from typing import Optional +from typing import Union + +from paddle import nn +from typing_extensions import Literal + + +class Loss(nn.Layer): + """Base class for loss.""" + + def __init__( + self, + reduction: Literal["mean", "sum"], + weight: Optional[Union[float, Dict[str, float]]] = None, + ): + super().__init__() + self.reduction = reduction + self.weight = weight + + def __str__(self): + return f"{self.__class__.__name__}(reduction={self.reduction}, weight={self.weight})" diff --git a/examples/smc_reac/ppsci/loss/chamfer.py b/examples/smc_reac/ppsci/loss/chamfer.py new file mode 100644 index 0000000000..8740f2a18b --- /dev/null +++ b/examples/smc_reac/ppsci/loss/chamfer.py @@ -0,0 +1,92 @@ +# 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. + +from __future__ import annotations + +from typing import Dict +from typing import Optional +from typing import Union + +import paddle + +from ppsci.loss import base + + +class ChamferLoss(base.Loss): + r"""Class for Chamfe distance loss. + + $$ + L = \dfrac{1}{S_1} \sum_{x \in S_1} \min_{y \in S_2} \Vert x - y \Vert_2^2 + \dfrac{1}{S_2} \sum_{y \in S_2} \min_{x \in S_1} \Vert y - x \Vert_2^2 + $$ + + $$ + \text{where } S_1 \text{ and } S_2 \text{ is the coordinate matrix of two point clouds}. + $$ + + Args: + weight (Optional[Union[float, Dict[str, float]]]): Weight for loss. Defaults to None. + + Examples: + >>> import paddle + >>> from ppsci.loss import ChamferLoss + >>> _ = paddle.seed(42) + >>> batch_point_cloud1 = paddle.rand([2, 100, 3]) + >>> batch_point_cloud2 = paddle.rand([2, 50, 3]) + >>> output_dict = {"s1": batch_point_cloud1} + >>> label_dict = {"s1": batch_point_cloud2} + >>> weight = {"s1": 0.8} + >>> loss = ChamferLoss(weight=weight) + >>> result = loss(output_dict, label_dict) + >>> print(result) + {'s1': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, + 0.04415882)} + """ + + def __init__( + self, + weight: Optional[Union[float, Dict[str, float]]] = None, + ): + super().__init__("mean", weight) + + def forward( + self, output_dict, label_dict, weight_dict=None + ) -> Dict[str, "paddle.Tensor"]: + losses = {} + + for key in label_dict: + s1 = output_dict[key] + s2 = label_dict[key] + N1, N2 = s1.shape[1], s2.shape[1] + + # [B, N1, N2, 3] + s1_expand = paddle.expand(s1.reshape([-1, N1, 1, 3]), shape=[-1, N1, N2, 3]) + # [B, N1, N2, 3] + s2_expand = paddle.expand(s2.reshape([-1, 1, N2, 3]), shape=[-1, N1, N2, 3]) + + dis = ((s1_expand - s2_expand) ** 2).sum(axis=3) # [B, N1, N2] + loss_s12 = dis.min(axis=2) # [B, N1] + loss_s21 = dis.min(axis=1) # [B, N2] + loss = loss_s12.mean() + loss_s21.mean() + + if weight_dict and key in weight_dict: + loss *= weight_dict[key] + + if isinstance(self.weight, (float, int)): + loss *= self.weight + elif isinstance(self.weight, dict) and key in self.weight: + loss *= self.weight[key] + + losses[key] = loss + + return losses diff --git a/examples/smc_reac/ppsci/loss/func.py b/examples/smc_reac/ppsci/loss/func.py new file mode 100644 index 0000000000..49992c5f7f --- /dev/null +++ b/examples/smc_reac/ppsci/loss/func.py @@ -0,0 +1,94 @@ +# 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. + +from __future__ import annotations + +from typing import Callable +from typing import Dict +from typing import Optional +from typing import Union + +import paddle + +from ppsci.loss import base + + +class FunctionalLoss(base.Loss): + r"""Functional loss class, which allows to use custom loss computing function from given loss_expr for complex computation cases. + + $$ + L = f(\mathbf{x}, \mathbf{y}) + $$ + + $$ + \mathbf{x}, \mathbf{y} \in \mathcal{R}^{N} + $$ + + Args: + loss_expr (Callable[..., paddle.Tensor]): Function for custom loss computation. + weight (Optional[Union[float, Dict[str, float]]]): Weight for loss. Defaults to None. + + Examples: + >>> import paddle + >>> from ppsci.loss import FunctionalLoss + >>> import paddle.nn.functional as F + >>> def mse_sum_loss(output_dict, label_dict, weight_dict=None): + ... losses = 0 + ... for key in output_dict.keys(): + ... loss = F.mse_loss(output_dict[key], label_dict[key], "sum") + ... if weight_dict: + ... loss *= weight_dict[key] + ... losses += loss + ... return {"mse_loss": losses} + >>> loss = FunctionalLoss(mse_sum_loss) + >>> output_dict = {'u': paddle.to_tensor([[0.5, 0.9], [1.1, -1.3]]), + ... 'v': paddle.to_tensor([[0.5, 0.9], [1.1, -1.3]])} + >>> label_dict = {'u': paddle.to_tensor([[-1.8, 1.0], [-0.2, 2.5]]), + ... 'v': paddle.to_tensor([[0.1, 0.1], [0.1, 0.1]])} + >>> weight_dict = {'u': 0.8, 'v': 0.2} + >>> result = loss(output_dict, label_dict, weight_dict) + >>> print(result) + {'mse_loss': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, + 17.89600182)} + """ + + def __init__( + self, + loss_expr: Callable[..., paddle.Tensor], + weight: Optional[Union[float, Dict[str, float]]] = None, + ): + super().__init__(None, weight) + self.loss_expr = loss_expr + + def forward( + self, output_dict, label_dict=None, weight_dict=None + ) -> Dict[str, "paddle.Tensor"]: + losses = self.loss_expr(output_dict, label_dict, weight_dict) + + assert isinstance(losses, dict), ( + "Loss computed by custom function should be type of 'dict', " + f"but got {type(losses)}." + " Please check the return type of custom loss function." + ) + + for key in losses: + assert isinstance( + losses[key], (paddle.Tensor, paddle.static.Variable, paddle.pir.Value) + ), ( + "Loss computed by custom function should be type of 'paddle.Tensor', " + f"'paddle.static.Variable' or 'paddle.pir.Value', but got {type(losses[key])}." + " Please check the return type of custom loss function." + ) + + return losses diff --git a/examples/smc_reac/ppsci/loss/integral.py b/examples/smc_reac/ppsci/loss/integral.py new file mode 100644 index 0000000000..74223c73fc --- /dev/null +++ b/examples/smc_reac/ppsci/loss/integral.py @@ -0,0 +1,112 @@ +# 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. + +from __future__ import annotations + +from typing import TYPE_CHECKING +from typing import Dict +from typing import Optional +from typing import Union + +import paddle.nn.functional as F +from typing_extensions import Literal + +from ppsci.loss import base + +if TYPE_CHECKING: + import paddle + + +class IntegralLoss(base.Loss): + r"""Class for integral loss with Monte-Carlo integration algorithm. + + $$ + L = + \begin{cases} + \dfrac{1}{N} \Vert \displaystyle\sum_{i=1}^{M}{\mathbf{s}_i \cdot \mathbf{x}_i} - \mathbf{y} \Vert_2^2, & \text{if reduction='mean'} \\ + \Vert \displaystyle\sum_{i=0}^{M}{\mathbf{s}_i \cdot \mathbf{x}_i} - \mathbf{y} \Vert_2^2, & \text{if reduction='sum'} + \end{cases} + $$ + + $$ + \mathbf{x}, \mathbf{s} \in \mathcal{R}^{M \times N}, \mathbf{y} \in \mathcal{R}^{N} + $$ + + Args: + reduction (Literal["mean", "sum"], optional): Reduction method. Defaults to "mean". + weight (Optional[Union[float, Dict[str, float]]]): Weight for loss. Defaults to None. + + Examples: + >>> import paddle + >>> from ppsci.loss import IntegralLoss + + >>> output_dict = {'u': paddle.to_tensor([[0.5, 2.2, 0.9], [1.1, 0.8, -1.3]]), + ... 'v': paddle.to_tensor([[0.5, 2.2, 0.9], [1.1, 0.8, -1.3]]), + ... 'area': paddle.to_tensor([[0.01, 0.02, 0.03], [0.01, 0.02, 0.03]])} + >>> label_dict = {'u': paddle.to_tensor([-1.8, 0.0]), + ... 'v': paddle.to_tensor([0.1, 0.1])} + >>> weight = {'u': 0.8, 'v': 0.2} + >>> loss = IntegralLoss(weight=weight) + >>> result = loss(output_dict, label_dict) + >>> print(result) + {'u': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, + 1.40780795), 'v': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, + 0.00131200)} + + >>> loss = IntegralLoss(reduction="sum", weight=weight) + >>> result = loss(output_dict, label_dict) + >>> print(result) + {'u': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, + 2.81561589), 'v': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, + 0.00262400)} + """ + + def __init__( + self, + reduction: Literal["mean", "sum"] = "mean", + weight: Optional[Union[float, Dict[str, float]]] = None, + ): + if reduction not in ["mean", "sum"]: + raise ValueError( + f"reduction should be 'mean' or 'sum', but got {reduction}" + ) + super().__init__(reduction, weight) + + def forward( + self, output_dict, label_dict, weight_dict=None + ) -> Dict[str, "paddle.Tensor"]: + losses = {} + + for key in label_dict: + loss = F.mse_loss( + (output_dict[key] * output_dict["area"]).sum(axis=1), + label_dict[key], + "none", + ) + if weight_dict and key in weight_dict: + loss *= weight_dict[key] + + if self.reduction == "sum": + loss = loss.sum() + elif self.reduction == "mean": + loss = loss.mean() + + if isinstance(self.weight, (float, int)): + loss *= self.weight + elif isinstance(self.weight, dict) and key in self.weight: + loss *= self.weight[key] + + losses[key] = loss + + return losses diff --git a/examples/smc_reac/ppsci/loss/kl.py b/examples/smc_reac/ppsci/loss/kl.py new file mode 100644 index 0000000000..c07c3ed2c6 --- /dev/null +++ b/examples/smc_reac/ppsci/loss/kl.py @@ -0,0 +1,51 @@ +# 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. + +from __future__ import annotations + +from typing import Dict +from typing import Optional +from typing import Union + +import paddle +from typing_extensions import Literal + +from ppsci.loss import base + + +class KLLoss(base.Loss): + def __init__( + self, + reduction: Literal["mean", "sum"] = "mean", + weight: Optional[Union[float, Dict[str, float]]] = None, + ): + if reduction not in ["mean", "sum"]: + raise ValueError( + f"reduction should be 'mean' or 'sum', but got {reduction}" + ) + super().__init__(reduction, weight) + + def forward( + self, output_dict, label_dict=None, weight_dict=None + ) -> Dict[str, "paddle.Tensor"]: + losses = {} + + mu, log_sigma = output_dict["mu"], output_dict["log_sigma"] + + base = paddle.exp(2.0 * log_sigma) + paddle.pow(mu, 2) - 1.0 - 2.0 * log_sigma + loss = 0.5 * paddle.sum(base) / mu.shape[0] + + losses["kl_loss"] = loss + + return loss diff --git a/examples/smc_reac/ppsci/loss/l1.py b/examples/smc_reac/ppsci/loss/l1.py new file mode 100644 index 0000000000..3edbc2e102 --- /dev/null +++ b/examples/smc_reac/ppsci/loss/l1.py @@ -0,0 +1,219 @@ +# 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. + +from __future__ import annotations + +from typing import TYPE_CHECKING +from typing import Dict +from typing import Optional +from typing import Union + +import paddle.nn.functional as F +from typing_extensions import Literal + +from ppsci.loss import base + +if TYPE_CHECKING: + + import paddle + + +class L1Loss(base.Loss): + r"""Class for l1 loss. + + $$ + L = \Vert \mathbf{x} - \mathbf{y} \Vert_1 + $$ + + $$ + \mathbf{x}, \mathbf{y} \in \mathcal{R}^{N} + $$ + + when `reduction` is set to "mean" + + $$ + L = MEAN \left( \Vert \mathbf{x} - \mathbf{y} \Vert_1 \right) + $$ + + when `reduction` is set to "sum" + + $$ + L = SUM \left( \Vert \mathbf{x} - \mathbf{y} \Vert_1 \right) + $$ + + Args: + reduction (Literal["mean", "sum"], optional): Reduction method. Defaults to "mean". + weight (Optional[Union[float, Dict[str, float]]]): Weight for loss. Defaults to None. + + Examples: + >>> import paddle + >>> from ppsci.loss import L1Loss + >>> output_dict = {"u": paddle.to_tensor([[0.5, 0.9], [1.1, -1.3]]), + ... "v": paddle.to_tensor([[0.5, 0.9], [1.1, -1.3]])} + >>> label_dict = {"u": paddle.to_tensor([[-1.8, 1.0], [-0.2, 2.5]]), + ... "v": paddle.to_tensor([[0.1, 0.1], [0.1, 0.1]])} + >>> weight = {"u": 0.8, "v": 0.2} + >>> loss = L1Loss(weight=weight) + >>> result = loss(output_dict, label_dict) + >>> print(result) + {'u': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, + 3.), 'v': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, + 0.35999998)} + + >>> loss = L1Loss(reduction="sum", weight=weight) + >>> result = loss(output_dict, label_dict) + >>> print(result) + {'u': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, + 6.), 'v': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, + 0.71999997)} + """ + + def __init__( + self, + reduction: Literal["mean", "sum"] = "mean", + weight: Optional[Union[float, Dict[str, float]]] = None, + ): + if reduction not in ["mean", "sum"]: + raise ValueError( + f"reduction should be 'mean' or 'sum', but got {reduction}" + ) + super().__init__(reduction, weight) + + def forward( + self, output_dict, label_dict, weight_dict=None + ) -> Dict[str, "paddle.Tensor"]: + losses = {} + + for key in label_dict: + loss = F.l1_loss(output_dict[key], label_dict[key], "none") + if weight_dict and key in weight_dict: + loss *= weight_dict[key] + + if "area" in output_dict: + loss *= output_dict["area"] + + loss = loss.sum(axis=1) + + if self.reduction == "sum": + loss = loss.sum() + elif self.reduction == "mean": + loss = loss.mean() + + if isinstance(self.weight, (float, int)): + loss *= self.weight + elif isinstance(self.weight, dict) and key in self.weight: + loss *= self.weight[key] + + losses[key] = loss + + return losses + + +class PeriodicL1Loss(base.Loss): + r"""Class for periodic l1 loss. + + $$ + L = \Vert \mathbf{x_l}-\mathbf{x_r} \Vert_1 + $$ + + $\mathbf{x_l} \in \mathcal{R}^{N}$ is the first half of batch output, + $\mathbf{x_r} \in \mathcal{R}^{N}$ is the second half of batch output. + + when `reduction` is set to "mean" + + $$ + L = MEAN \left( \Vert \mathbf{x_l}-\mathbf{x_r} \Vert_1 \right) + $$ + + when `reduction` is set to "sum" + + $$ + L = SUM \left( \Vert \mathbf{x_l}-\mathbf{x_r} \Vert_1 \right) + $$ + + Args: + reduction (Literal["mean", "sum"], optional): Reduction method. Defaults to "mean". + weight (Optional[Union[float, Dict[str, float]]]): Weight for loss. Defaults to None. + + Examples: + >>> import paddle + >>> from ppsci.loss import PeriodicL1Loss + + >>> output_dict = {'u': paddle.to_tensor([[0.5, 2.2, 0.9], [1.1, 0.8, -1.3]]), + ... 'v': paddle.to_tensor([[0.5, 2.2, 0.9], [1.1, 0.8, -1.3]])} + >>> label_dict = {'u': paddle.to_tensor([[-1.8, 0.0, 1.0], [-0.2, 0.2, 2.5]]), + ... 'v': paddle.to_tensor([[0.1, 0.1, 0.1], [0.1, 0.1, 0.1]])} + >>> weight = {'u': 0.8, 'v': 0.2} + >>> loss = PeriodicL1Loss(weight=weight) + >>> result = loss(output_dict, label_dict) + >>> print(result) + {'u': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, + 3.35999990), 'v': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, + 0.83999997)} + + >>> loss = PeriodicL1Loss(reduction="sum", weight=weight) + >>> result = loss(output_dict, label_dict) + >>> print(result) + {'u': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, + 3.35999990), 'v': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, + 0.83999997)} + """ + + def __init__( + self, + reduction: Literal["mean", "sum"] = "mean", + weight: Optional[Union[float, Dict[str, float]]] = None, + ): + if reduction not in ["mean", "sum"]: + raise ValueError( + f"reduction should be 'mean' or 'sum', but got {reduction}" + ) + super().__init__(reduction, weight) + + def forward( + self, output_dict, label_dict, weight_dict=None + ) -> Dict[str, "paddle.Tensor"]: + losses = {} + + for key in label_dict: + n_output = len(output_dict[key]) + if n_output % 2 > 0: + raise ValueError( + f"Length of output({n_output}) of key({key}) should be even." + ) + + n_output //= 2 + loss = F.l1_loss( + output_dict[key][:n_output], output_dict[key][n_output:], "none" + ) + if weight_dict and key in weight_dict: + loss *= weight_dict[key] + if "area" in output_dict: + loss *= output_dict["area"] + + loss = loss.sum(axis=1) + + if self.reduction == "sum": + loss = loss.sum() + elif self.reduction == "mean": + loss = loss.mean() + + if isinstance(self.weight, (float, int)): + loss *= self.weight + elif isinstance(self.weight, dict) and key in self.weight: + loss *= self.weight[key] + + losses[key] = loss + + return losses diff --git a/examples/smc_reac/ppsci/loss/l2.py b/examples/smc_reac/ppsci/loss/l2.py new file mode 100644 index 0000000000..7b65a937c6 --- /dev/null +++ b/examples/smc_reac/ppsci/loss/l2.py @@ -0,0 +1,310 @@ +# 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. + +from __future__ import annotations + +from typing import Dict +from typing import Optional +from typing import Union + +import paddle +import paddle.nn.functional as F +from typing_extensions import Literal + +from ppsci.loss import base + + +class L2Loss(base.Loss): + r"""Class for l2 loss. + + $$ + L =\Vert \mathbf{x} - \mathbf{y} \Vert_2 + $$ + + $$ + \mathbf{x}, \mathbf{y} \in \mathcal{R}^{N} + $$ + + when `reduction` is set to "mean" + + $$ + L = MEAN \left( \Vert \mathbf{x} - \mathbf{y} \Vert_2 \right) + $$ + + when `reduction` is set to "sum" + + $$ + L = SUM \left( \Vert \mathbf{x} - \mathbf{y} \Vert_2 \right) + $$ + + Args: + reduction (Literal["mean", "sum"], optional): Reduction method. Defaults to "mean". + weight (Optional[Union[float, Dict[str, float]]]): Weight for loss. Defaults to None. + + Examples: + >>> import paddle + >>> from ppsci.loss import L2Loss + >>> output_dict = {'u': paddle.to_tensor([[0.5, 0.9], [1.1, -1.3]]), + ... 'v': paddle.to_tensor([[0.5, 0.9], [1.1, -1.3]])} + >>> label_dict = {'u': paddle.to_tensor([[-1.8, 1.0], [-0.2, 2.5]]), + ... 'v': paddle.to_tensor([[0.1, 0.1], [0.1, 0.1]])} + >>> weight = {'u': 0.8, 'v': 0.2} + >>> loss = L2Loss(weight=weight) + >>> result = loss(output_dict, label_dict) + >>> print(result) + {'u': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, + 2.52735591), 'v': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, + 0.26148924)} + >>> loss = L2Loss(reduction="sum", weight=weight) + >>> result = loss(output_dict, label_dict) + >>> print(result) + {'u': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, + 5.05471182), 'v': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, + 0.52297848)} + """ + + def __init__( + self, + reduction: Literal["mean", "sum"] = "mean", + weight: Optional[Union[float, Dict[str, float]]] = None, + ): + if reduction not in ["mean", "sum"]: + raise ValueError( + f"reduction should be 'mean' or 'sum', but got {reduction}" + ) + super().__init__(reduction, weight) + + def forward( + self, output_dict, label_dict, weight_dict=None + ) -> Dict[str, "paddle.Tensor"]: + losses = {} + + for key in label_dict: + loss = F.mse_loss(output_dict[key], label_dict[key], "none") + if weight_dict and key in weight_dict: + loss *= weight_dict[key] + + if "area" in output_dict: + loss *= output_dict["area"] + + loss = loss.sum(axis=1).sqrt() + + if self.reduction == "sum": + loss = loss.sum() + elif self.reduction == "mean": + loss = loss.mean() + + if isinstance(self.weight, (float, int)): + loss *= self.weight + elif isinstance(self.weight, dict) and key in self.weight: + loss *= self.weight[key] + + losses[key] = loss + + return losses + + +class PeriodicL2Loss(base.Loss): + r"""Class for Periodic l2 loss. + + $$ + L = \Vert \mathbf{x_l}-\mathbf{x_r} \Vert_2 + $$ + + $\mathbf{x_l} \in \mathcal{R}^{N}$ is the first half of batch output, + $\mathbf{x_r} \in \mathcal{R}^{N}$ is the second half of batch output. + + when `reduction` is set to "mean" + + $$ + L = MEAN \left( \Vert \mathbf{x_l}-\mathbf{x_r} \Vert_2 \right) + $$ + + when `reduction` is set to "sum" + + $$ + L = SUM \left( \Vert \mathbf{x_l}-\mathbf{x_r} \Vert_2 \right) + $$ + + Args: + reduction (Literal["mean", "sum"], optional): Reduction method. Defaults to "mean". + weight (Optional[Union[float, Dict[str, float]]]): Weight for loss. Defaults to None. + + Examples: + >>> import paddle + >>> from ppsci.loss import PeriodicL2Loss + + >>> output_dict = {'u': paddle.to_tensor([[0.5, 2.2, 0.9], [1.1, 0.8, -1.3]]), + ... 'v': paddle.to_tensor([[0.5, 2.2, 0.9], [1.1, 0.8, -1.3]])} + >>> label_dict = {'u': paddle.to_tensor([[-1.8, 0.0, 1.0], [-0.2, 0.2, 2.5]]), + ... 'v': paddle.to_tensor([[0.1, 0.1, 0.1], [0.1, 0.1, 0.1]])} + >>> weight = {'u': 0.8, 'v': 0.2} + >>> loss = PeriodicL2Loss(weight=weight) + >>> result = loss(output_dict, label_dict) + >>> print(result) + {'u': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, + 2.14065409), 'v': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, + 0.53516352)} + + >>> loss = PeriodicL2Loss(reduction="sum", weight=weight) + >>> result = loss(output_dict, label_dict) + >>> print(result) + {'u': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, + 2.14065409), 'v': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, + 0.53516352)} + """ + + def __init__( + self, + reduction: Literal["mean", "sum"] = "mean", + weight: Optional[Union[float, Dict[str, float]]] = None, + ): + if reduction not in ["mean", "sum"]: + raise ValueError( + f"reduction should be 'mean' or 'sum', but got {reduction}" + ) + super().__init__(reduction, weight) + + def forward( + self, output_dict, label_dict, weight_dict=None + ) -> Dict[str, "paddle.Tensor"]: + losses = {} + + for key in label_dict: + n_output = len(output_dict[key]) + if n_output % 2 > 0: + raise ValueError( + f"Length of output({n_output}) of key({key}) should be even." + ) + n_output //= 2 + + loss = F.mse_loss( + output_dict[key][:n_output], output_dict[key][n_output:], "none" + ) + if weight_dict and key in weight_dict: + loss *= weight_dict[key] + + if "area" in output_dict: + loss *= output_dict["area"] + + loss = loss.sum(axis=1).sqrt() + + if self.reduction == "sum": + loss = loss.sum() + elif self.reduction == "mean": + loss = loss.mean() + + if isinstance(self.weight, (float, int)): + loss *= self.weight + elif isinstance(self.weight, dict) and key in self.weight: + loss *= self.weight[key] + + losses[key] = loss + + return losses + + +class L2RelLoss(base.Loss): + r"""Class for l2 relative loss. + + $$ + L = \dfrac{\Vert \mathbf{x} - \mathbf{y} \Vert_2}{\Vert \mathbf{y} \Vert_2} + $$ + + $$ + \mathbf{x}, \mathbf{y} \in \mathcal{R}^{N} + $$ + + when `reduction` is set to "mean" + + $$ + L = MEAN \left( \dfrac{\Vert \mathbf{x} - \mathbf{y} \Vert_2}{\Vert \mathbf{y} \Vert_2} \right) + $$ + + when `reduction` is set to "sum" + + $$ + L = SUM \left( \dfrac{\Vert \mathbf{x} - \mathbf{y} \Vert_2}{\Vert \mathbf{y} \Vert_2} \right) + $$ + + Args: + reduction (Literal["mean", "sum"], optional): Specifies the reduction to apply to the output: 'mean' | 'sum'. Defaults to "mean". + weight (Optional[Union[float, Dict[str, float]]]): Weight for loss. Defaults to None. + + Examples: + >>> import paddle + >>> from ppsci.loss import L2RelLoss + + >>> output_dict = {'u': paddle.to_tensor([[0.5, 0.9], [1.1, -1.3]]), + ... 'v': paddle.to_tensor([[0.5, 0.9], [1.1, -1.3]])} + >>> label_dict = {'u': paddle.to_tensor([[-1.8, 1.0], [-0.2, 2.5]]), + ... 'v': paddle.to_tensor([[0.1, 0.1], [0.1, 0.1]])} + >>> weight = {'u': 0.8, 'v': 0.2} + >>> loss = L2RelLoss(weight=weight) + >>> result = loss(output_dict, label_dict) + >>> print(result) + {'u': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, + 1.08776188), 'v': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, + 1.84900820)} + + >>> loss = L2RelLoss(reduction="sum", weight=weight) + >>> result = loss(output_dict, label_dict) + >>> print(result) + {'u': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, + 2.17552376), 'v': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, + 3.69801641)} + """ + + def __init__( + self, + reduction: Literal["mean", "sum"] = "mean", + weight: Optional[Union[float, Dict[str, float]]] = None, + ): + if reduction not in ["mean", "sum"]: + raise ValueError( + f"reduction should be 'mean' or 'sum', but got {reduction}" + ) + super().__init__(reduction, weight) + + def rel_loss(self, x, y): + batch_size = x.shape[0] + x_ = x.reshape((batch_size, -1)) + y_ = y.reshape((batch_size, -1)) + diff_norms = paddle.norm(x_ - y_, p=2, axis=1) + y_norms = paddle.norm(y_, p=2, axis=1) + return diff_norms / y_norms + + def forward( + self, output_dict, label_dict, weight_dict=None + ) -> Dict[str, "paddle.Tensor"]: + losses = {} + + for key in label_dict: + loss = self.rel_loss(output_dict[key], label_dict[key]) + if weight_dict: + loss *= weight_dict[key] + + if self.reduction == "sum": + loss = loss.sum() + elif self.reduction == "mean": + loss = loss.mean() + + if isinstance(self.weight, float): + loss *= self.weight + elif isinstance(self.weight, dict) and key in self.weight: + loss *= self.weight[key] + + losses[key] = loss + + return losses diff --git a/examples/smc_reac/ppsci/loss/mae.py b/examples/smc_reac/ppsci/loss/mae.py new file mode 100644 index 0000000000..ff7869535c --- /dev/null +++ b/examples/smc_reac/ppsci/loss/mae.py @@ -0,0 +1,109 @@ +# 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. + +from __future__ import annotations + +from typing import TYPE_CHECKING +from typing import Dict +from typing import Optional +from typing import Union + +import paddle.nn.functional as F +from typing_extensions import Literal + +from ppsci.loss import base + +if TYPE_CHECKING: + import paddle + + +class MAELoss(base.Loss): + r"""Class for mean absolute error loss. + + $$ + L = + \begin{cases} + \dfrac{1}{N} \Vert {\mathbf{x}-\mathbf{y}} \Vert_1, & \text{if reduction='mean'} \\ + \Vert {\mathbf{x}-\mathbf{y}} \Vert_1, & \text{if reduction='sum'} + \end{cases} + $$ + + $$ + \mathbf{x}, \mathbf{y} \in \mathcal{R}^{N} + $$ + + Args: + reduction (Literal["mean", "sum"], optional): Reduction method. Defaults to "mean". + weight (Optional[Union[float, Dict[str, float]]]): Weight for loss. Defaults to None. + + Examples: + >>> import paddle + >>> from ppsci.loss import MAELoss + + >>> output_dict = {'u': paddle.to_tensor([[0.5, 0.9], [1.1, -1.3]]), + ... 'v': paddle.to_tensor([[0.5, 0.9], [1.1, -1.3]])} + >>> label_dict = {'u': paddle.to_tensor([[-1.8, 1.0], [-0.2, 2.5]]), + ... 'v': paddle.to_tensor([[0.1, 0.1], [0.1, 0.1]])} + >>> weight = {'u': 0.8, 'v': 0.2} + >>> loss = MAELoss(weight=weight) + >>> result = loss(output_dict, label_dict) + >>> print(result) + {'u': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, + 1.50000000), 'v': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, + 0.17999999)} + + >>> loss = MAELoss(reduction="sum", weight=weight) + >>> result = loss(output_dict, label_dict) + >>> print(result) + {'u': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, + 6.), 'v': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, + 0.71999997)} + """ + + def __init__( + self, + reduction: Literal["mean", "sum"] = "mean", + weight: Optional[Union[float, Dict[str, float]]] = None, + ): + if reduction not in ["mean", "sum"]: + raise ValueError( + f"reduction should be 'mean' or 'sum', but got {reduction}" + ) + super().__init__(reduction, weight) + + def forward( + self, output_dict, label_dict, weight_dict=None + ) -> Dict[str, "paddle.Tensor"]: + losses = {} + + for key in label_dict: + loss = F.l1_loss(output_dict[key], label_dict[key], "none") + if weight_dict and key in weight_dict: + loss *= weight_dict[key] + + if "area" in output_dict: + loss *= output_dict["area"] + + if self.reduction == "sum": + loss = loss.sum() + elif self.reduction == "mean": + loss = loss.mean() + if isinstance(self.weight, (float, int)): + loss *= self.weight + elif isinstance(self.weight, dict) and key in self.weight: + loss *= self.weight[key] + + losses[key] = loss + + return losses diff --git a/examples/smc_reac/ppsci/loss/mse.py b/examples/smc_reac/ppsci/loss/mse.py new file mode 100644 index 0000000000..184f11cc40 --- /dev/null +++ b/examples/smc_reac/ppsci/loss/mse.py @@ -0,0 +1,355 @@ +# 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. + +from __future__ import annotations + +from typing import Dict +from typing import Optional +from typing import Union + +import paddle +import paddle.nn.functional as F +from typing_extensions import Literal + +from ppsci.loss import base + + +class MSELoss(base.Loss): + r"""Class for mean squared error loss. + + $$ + L = + \begin{cases} + \dfrac{1}{N} \Vert {\mathbf{x}-\mathbf{y}} \Vert_2^2, & \text{if reduction='mean'} \\ + \Vert {\mathbf{x}-\mathbf{y}} \Vert_2^2, & \text{if reduction='sum'} + \end{cases} + $$ + + $$ + \mathbf{x}, \mathbf{y} \in \mathcal{R}^{N} + $$ + + Args: + reduction (Literal["mean", "sum"], optional): Reduction method. Defaults to "mean". + weight (Optional[Union[float, Dict[str, float]]]): Weight for loss. Defaults to None. + + Examples: + >>> import paddle + >>> from ppsci.loss import MSELoss + + >>> output_dict = {'u': paddle.to_tensor([[0.5, 0.9], [1.1, -1.3]]), + ... 'v': paddle.to_tensor([[0.5, 0.9], [1.1, -1.3]])} + >>> label_dict = {'u': paddle.to_tensor([[-1.8, 1.0], [-0.2, 2.5]]), + ... 'v': paddle.to_tensor([[0.1, 0.1], [0.1, 0.1]])} + >>> weight = {'u': 0.8, 'v': 0.2} + >>> loss = MSELoss(weight=weight) + >>> result = loss(output_dict, label_dict) + >>> print(result) + {'u': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, + 4.28600025), 'v': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, + 0.18800001)} + + >>> loss = MSELoss(reduction="sum", weight=weight) + >>> result = loss(output_dict, label_dict) + >>> print(result) + {'u': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, + 17.14400101), 'v': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, + 0.75200003)} + """ + + def __init__( + self, + reduction: Literal["mean", "sum"] = "mean", + weight: Optional[Union[float, Dict[str, float]]] = None, + ): + if reduction not in ["mean", "sum"]: + raise ValueError( + f"reduction should be 'mean' or 'sum', but got {reduction}" + ) + super().__init__(reduction, weight) + + def forward( + self, output_dict, label_dict, weight_dict=None + ) -> Dict[str, "paddle.Tensor"]: + losses = {} + + for key in label_dict: + loss = F.mse_loss(output_dict[key], label_dict[key], "none") + if weight_dict and key in weight_dict: + loss *= weight_dict[key] + + if "area" in output_dict: + loss *= output_dict["area"] + + if self.reduction == "sum": + loss = loss.sum() + elif self.reduction == "mean": + loss = loss.mean() + if isinstance(self.weight, (float, int)): + loss *= self.weight + elif isinstance(self.weight, dict) and key in self.weight: + loss *= self.weight[key] + + losses[key] = loss + + return losses + + +class CausalMSELoss(base.Loss): + r"""Class for mean squared error loss. + + $$ + L = \frac{1}{M} \displaystyle\sum_{i=1}^M{w_i} \mathcal{L}_r^i, + $$ + + where $w_i=\exp (-\epsilon \displaystyle\sum_{k=1}^{i-1} \mathcal{L}_r^k), i=2,3, \ldots, M.$ + + Args: + n_chunks (int): $M$, Number of split time windows. + reduction (Literal["mean", "sum"], optional): Reduction method. Defaults to "mean". + weight (Optional[Union[float, Dict[str, float]]]): Weight for loss. Defaults to None. + tol (float, optional): Causal tolerance, i.e. $\epsilon$ in paper. Defaults to 1.0. + + Examples: + >>> import paddle + >>> from ppsci.loss import CausalMSELoss + + >>> output_dict = {'u': paddle.to_tensor([[0.5, 0.9, 1.0], [1.1, -1.3, 0.0]])} + >>> label_dict = {'u': paddle.to_tensor([[-1.8, 1.0, -0.1], [-0.2, 2.5, 2.0]])} + >>> loss = CausalMSELoss(n_chunks=3) + >>> result = loss(output_dict, label_dict) + >>> print(result) + {'u': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, + 0.96841478)} + """ + + def __init__( + self, + n_chunks: int, + reduction: Literal["mean", "sum"] = "mean", + weight: Optional[Union[float, Dict[str, float]]] = None, + tol: float = 1.0, + ): + if n_chunks <= 0: + raise ValueError(f"n_chunks should be positive, but got {n_chunks}") + if reduction not in ["mean", "sum"]: + raise ValueError( + f"reduction should be 'mean' or 'sum', but got {reduction}" + ) + super().__init__(reduction, weight) + self.n_chunks = n_chunks + self.tol = tol + self.register_buffer( + "acc_mat", paddle.tril(paddle.ones([n_chunks, n_chunks]), -1) + ) + + def forward( + self, output_dict, label_dict, weight_dict=None + ) -> Dict[str, "paddle.Tensor"]: + losses = {} + + for key in label_dict: + loss = F.mse_loss(output_dict[key], label_dict[key], "none") + if weight_dict and key in weight_dict: + loss *= weight_dict[key] + + if "area" in output_dict: + loss *= output_dict["area"] + + # causal weighting + loss_t = loss.reshape([self.n_chunks, -1]) # [nt, nx] + weight_t = paddle.exp( + -self.tol * (self.acc_mat @ loss_t.mean(-1, keepdim=True)) + ) # [nt, nt] x [nt, 1] ==> [nt, 1] + assert weight_t.shape[0] == self.n_chunks + loss = loss_t * weight_t.detach() + + if self.reduction == "sum": + loss = loss.sum() + elif self.reduction == "mean": + loss = loss.mean() + if isinstance(self.weight, (float, int)): + loss *= self.weight + elif isinstance(self.weight, dict) and key in self.weight: + loss *= self.weight[key] + + losses[key] = loss + + return losses + + +class MSELossWithL2Decay(MSELoss): + r"""MSELoss with L2 decay. + + $$ + L = + \begin{cases} + \dfrac{1}{N} \Vert {\mathbf{x}-\mathbf{y}} \Vert_2^2 + \displaystyle\sum_{i=1}^{M}{\Vert \mathbf{K_i} \Vert_F^2}, & \text{if reduction='mean'} \\ + \Vert {\mathbf{x}-\mathbf{y}} \Vert_2^2 + \displaystyle\sum_{i=1}^{M}{\Vert \mathbf{K_i} \Vert_F^2}, & \text{if reduction='sum'} + \end{cases} + $$ + + $$ + \mathbf{x}, \mathbf{y} \in \mathcal{R}^{N}, \mathbf{K_i} \in \mathcal{R}^{O_i \times P_i} + $$ + + $M$ is the number of which apply regularization on. + + Args: + reduction (Literal["mean", "sum"], optional): Specifies the reduction to apply to the output: 'mean' | 'sum'. Defaults to "mean". + regularization_dict (Optional[Dict[str, float]]): Regularization dictionary. Defaults to None. + weight (Optional[Union[float, Dict[str, float]]]): Weight for loss. Defaults to None. + + Raises: + ValueError: reduction should be 'mean' or 'sum'. + + Examples: + >>> import paddle + >>> from ppsci.loss import MSELossWithL2Decay + + >>> output_dict = {'u': paddle.to_tensor([[0.5, 0.9], [1.1, -1.3]]), + ... 'v': paddle.to_tensor([[0.5, 0.9], [1.1, -1.3]])} + >>> label_dict = {'u': paddle.to_tensor([[-1.8, 1.0], [-0.2, 2.5]]), + ... 'v': paddle.to_tensor([[0.1, 0.1], [0.1, 0.1]])} + >>> weight = {'u': 0.8, 'v': 0.2} + >>> regularization_dict = {'u': 2.0} + >>> loss = MSELossWithL2Decay(regularization_dict=regularization_dict, weight=weight) + >>> result = loss(output_dict, label_dict) + >>> print(result) + {'u': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, + 7.91999960), 'v': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, + 0.18800001)} + + >>> regularization_dict = {'v': 1.0} + >>> loss = MSELossWithL2Decay(reduction="sum", regularization_dict=regularization_dict, weight=weight) + >>> result = loss(output_dict, label_dict) + >>> print(result) + {'u': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, + 17.14400101), 'v': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, + 3.95999980)} + """ + + def __init__( + self, + reduction: Literal["mean", "sum"] = "mean", + regularization_dict: Optional[Dict[str, float]] = None, + weight: Optional[Union[float, Dict[str, float]]] = None, + ): + if reduction not in ["mean", "sum"]: + raise ValueError( + f"reduction should be 'mean' or 'sum', but got {reduction}" + ) + super().__init__(reduction, weight) + self.regularization_dict = regularization_dict + + def forward( + self, output_dict, label_dict, weight_dict=None + ) -> Dict[str, "paddle.Tensor"]: + losses = super().forward(output_dict, label_dict, weight_dict) + + if self.regularization_dict is not None: + for reg_key, reg_weight in self.regularization_dict.items(): + loss = output_dict[reg_key].pow(2).sum() + losses[reg_key] = loss * reg_weight + + return losses + + +class PeriodicMSELoss(base.Loss): + r"""Class for periodic mean squared error loss. + + $$ + L = + \begin{cases} + \dfrac{1}{N} \Vert \mathbf{x_l}-\mathbf{x_r} \Vert_2^2, & \text{if reduction='mean'} \\ + \Vert \mathbf{x_l}-\mathbf{x_r} \Vert_2^2, & \text{if reduction='sum'} + \end{cases} + $$ + + $\mathbf{x_l} \in \mathcal{R}^{N}$ is the first half of batch output, + $\mathbf{x_r} \in \mathcal{R}^{N}$ is the second half of batch output. + + Args: + reduction (Literal["mean", "sum"], optional): Reduction method. Defaults to "mean". + weight (Optional[Union[float, Dict[str, float]]]): Weight for loss. Defaults to None. + + Examples: + >>> import paddle + >>> from ppsci.loss import PeriodicMSELoss + + >>> output_dict = {'u': paddle.to_tensor([[0.5, 0.9], [1.1, -1.3]]), + ... 'v': paddle.to_tensor([[0.5, 0.9], [1.1, -1.3]])} + >>> label_dict = {'u': paddle.to_tensor([[-1.8, 1.0], [-0.2, 2.5]]), + ... 'v': paddle.to_tensor([[0.1, 0.1], [0.1, 0.1]])} + >>> weight = {'u': 0.8, 'v': 0.2} + >>> loss = PeriodicMSELoss(weight=weight) + >>> result = loss(output_dict, label_dict) + >>> print(result) + {'u': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, + 2.07999969), 'v': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, + 0.51999992)} + + >>> loss = PeriodicMSELoss(reduction="sum", weight=weight) + >>> result = loss(output_dict, label_dict) + >>> print(result) + {'u': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, + 4.15999937), 'v': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, + 1.03999984)} + """ + + def __init__( + self, + reduction: Literal["mean", "sum"] = "mean", + weight: Optional[Union[float, Dict[str, float]]] = None, + ): + if reduction not in ["mean", "sum"]: + raise ValueError( + f"reduction should be 'mean' or 'sum', but got {reduction}" + ) + super().__init__(reduction, weight) + + def forward( + self, output_dict, label_dict, weight_dict=None + ) -> Dict[str, "paddle.Tensor"]: + losses = {} + + for key in label_dict: + n_output = len(output_dict[key]) + if n_output % 2 > 0: + raise ValueError( + f"Length of output({n_output}) of key({key}) should be even." + ) + + n_output //= 2 + loss = F.mse_loss( + output_dict[key][:n_output], output_dict[key][n_output:], "none" + ) + if weight_dict: + loss *= weight_dict[key] + if "area" in output_dict: + loss *= output_dict["area"] + + if self.reduction == "sum": + loss = loss.sum() + elif self.reduction == "mean": + loss = loss.mean() + + if isinstance(self.weight, (float, int)): + loss *= self.weight + elif isinstance(self.weight, dict) and key in self.weight: + loss *= self.weight[key] + + losses[key] = loss + + return losses diff --git a/examples/smc_reac/ppsci/loss/mtl/__init__.py b/examples/smc_reac/ppsci/loss/mtl/__init__.py new file mode 100644 index 0000000000..3bff2aaa7f --- /dev/null +++ b/examples/smc_reac/ppsci/loss/mtl/__init__.py @@ -0,0 +1,49 @@ +# 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. + +import copy + +from ppsci.loss.mtl.agda import AGDA +from ppsci.loss.mtl.base import LossAggregator +from ppsci.loss.mtl.grad_norm import GradNorm +from ppsci.loss.mtl.ntk import NTK +from ppsci.loss.mtl.pcgrad import PCGrad +from ppsci.loss.mtl.relobralo import Relobralo +from ppsci.loss.mtl.sum import Sum + +__all__ = [ + "AGDA", + "GradNorm", + "LossAggregator", + "PCGrad", + "Relobralo", + "Sum", + "NTK", +] + + +def build_mtl_aggregator(cfg): + """Build loss aggregator with multi-task learning method. + + Args: + cfg (DictConfig): Aggregator config. + + Returns: + Loss: Callable loss aggregator object. + """ + cfg = copy.deepcopy(cfg) + + aggregator_cls = cfg.pop("name") + aggregator = eval(aggregator_cls)(**cfg) + return aggregator diff --git a/examples/smc_reac/ppsci/loss/mtl/agda.py b/examples/smc_reac/ppsci/loss/mtl/agda.py new file mode 100644 index 0000000000..31193f5b76 --- /dev/null +++ b/examples/smc_reac/ppsci/loss/mtl/agda.py @@ -0,0 +1,161 @@ +# 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. + +from __future__ import annotations + +from typing import ClassVar +from typing import List + +import paddle +from paddle import nn + +from ppsci.loss.mtl import base + + +class AGDA(base.LossAggregator): + r""" + **A**daptive **G**radient **D**escent **A**lgorithm + + [Physics-informed neural network based on a new adaptive gradient descent algorithm for solving partial differential equations of flow problems](https://pubs.aip.org/aip/pof/article-abstract/35/6/063608/2899773/Physics-informed-neural-network-based-on-a-new) + + NOTE: This loss aggregator is only suitable for two-task learning and the first task loss must be PDE loss. + + Attributes: + should_persist(bool): Whether to persist the loss aggregator when saving. + Those loss aggregators with parameters and/or buffers should be persisted. + + Args: + model (nn.Layer): Training model. + M (int, optional): Smoothing period. Defaults to 100. + gamma (float, optional): Smooth factor. Defaults to 0.999. + + Examples: + >>> import paddle + >>> from ppsci.loss import mtl + >>> model = paddle.nn.Linear(3, 4) + >>> loss_aggregator = mtl.AGDA(model) + >>> for i in range(5): + ... x1 = paddle.randn([8, 3]) + ... x2 = paddle.randn([8, 3]) + ... y1 = model(x1) + ... y2 = model(x2) + ... pde_loss = paddle.sum(y1) + ... bc_loss = paddle.sum((y2 - 2) ** 2) + ... loss_aggregator({'pde_loss': pde_loss, 'bc_loss': bc_loss}).backward() + """ + should_persist: ClassVar[bool] = False + + def __init__(self, model: nn.Layer, M: int = 100, gamma: float = 0.999) -> None: + super().__init__(model) + self.M = M + self.gamma = gamma + self.Lf_smooth = 0 + self.Lu_smooth = 0 + self.Lf_tilde_acc = 0.0 + self.Lu_tilde_acc = 0.0 + + def __call__(self, losses, step: int = 0) -> "AGDA": + if len(losses) != 2: + raise ValueError( + f"Number of losses(tasks) for AGDA should be 2, but got {len(losses)}" + ) + return super().__call__(losses, step) + + def backward(self) -> None: + grads_list = self._compute_grads() + with paddle.no_grad(): + refined_grads = self._refine_grads(grads_list) + self._set_grads(refined_grads) + + def _compute_grads(self) -> List[paddle.Tensor]: + # compute all gradients derived by each loss + grads_list = [] # num_params x num_losses + for key in self.losses: + # backward with current loss + self.losses[key].backward() + grads_list.append( + paddle.concat( + [ + param.grad.clone().reshape([-1]) + for param in self.model.parameters() + if param.grad is not None + ], + axis=0, + ) + ) + # clear gradients for current loss for not affecting other loss + self.model.clear_gradients() + + return grads_list + + def _refine_grads(self, grads_list: List[paddle.Tensor]) -> List[paddle.Tensor]: + # compute moving average of L^smooth_i(n) - eq.(16) + losses_seq = list(self.losses.values()) + self.Lf_smooth = ( + self.gamma * self.Lf_smooth + (1 - self.gamma) * losses_seq[0].item() + ) + self.Lu_smooth = ( + self.gamma * self.Lu_smooth + (1 - self.gamma) * losses_seq[1].item() + ) + + # compute L^smooth_i(kM) - eq.(17) + if self.step % self.M == 0: + Lf_smooth_kM = self.Lf_smooth + Lu_smooth_kM = self.Lu_smooth + Lf_tilde = self.Lf_smooth / Lf_smooth_kM + Lu_tilde = self.Lu_smooth / Lu_smooth_kM + + # compute r_i(n) - eq.(18) + self.Lf_tilde_acc += Lf_tilde + self.Lu_tilde_acc += Lu_tilde + rf = Lf_tilde / self.Lf_tilde_acc + ru = Lu_tilde / self.Lu_tilde_acc + + # compute E(g(n)) - step1(1) + gf_magn = (grads_list[0] * grads_list[0]).sum().sqrt() + gu_magn = (grads_list[1] * grads_list[1]).sum().sqrt() + Eg = (gf_magn + gu_magn) / 2 + + # compute \omega_f(n) - step1(2) + omega_f = (rf * (Eg - gf_magn) + gf_magn) / gf_magn + omega_u = (ru * (Eg - gu_magn) + gu_magn) / gu_magn + + # compute g_bar(n) - step1(3) + gf_bar = omega_f * grads_list[0] + gu_bar = omega_u * grads_list[1] + + # compute gradient projection - step2(1) + dot_product = (gf_bar * gu_bar).sum() + if dot_product < 0: + gu_bar = gu_bar - (dot_product / (gf_bar * gf_bar).sum()) * gf_bar + grads_list = [gf_bar, gu_bar] + + proj_grads: List[paddle.Tensor] = [] + for j in range(len(self.losses)): + start_idx = 0 + for idx, var in enumerate(self.model.parameters()): + grad_shape = var.shape + flatten_dim = var.numel() + refined_grad = grads_list[j][start_idx : start_idx + flatten_dim] + refined_grad = paddle.reshape(refined_grad, grad_shape) + if len(proj_grads) < self.param_num: + proj_grads.append(refined_grad) + else: + proj_grads[idx] += refined_grad + start_idx += flatten_dim + return proj_grads + + def _set_grads(self, grads_list: List[paddle.Tensor]) -> None: + for i, param in enumerate(self.model.parameters()): + param.grad = grads_list[i] diff --git a/examples/smc_reac/ppsci/loss/mtl/base.py b/examples/smc_reac/ppsci/loss/mtl/base.py new file mode 100644 index 0000000000..eec88c9c00 --- /dev/null +++ b/examples/smc_reac/ppsci/loss/mtl/base.py @@ -0,0 +1,68 @@ +# 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. + +from __future__ import annotations + +from typing import TYPE_CHECKING +from typing import ClassVar +from typing import Dict +from typing import Union + +from paddle import nn + +if TYPE_CHECKING: + import paddle + + +class LossAggregator(nn.Layer): + """Base class of loss aggregator mainly for multitask learning. + + Attributes: + should_persist(bool): Whether to persist the loss aggregator when saving. + Those loss aggregators with parameters and/or buffers should be persisted. + + Args: + model (nn.Layer): Training model. + """ + + should_persist: ClassVar[bool] = False + + def __init__(self, model: nn.Layer) -> None: + super().__init__() + self.model = model + self.step = 0 + self.param_num = 0 + for param in self.model.parameters(): + if not param.stop_gradient: + self.param_num += 1 + + def forward( + self, losses: Dict[str, "paddle.Tensor"], step: int = 0 + ) -> Union["paddle.Tensor", "LossAggregator"]: + self.losses = losses + self.loss_num = len(losses) + self.step = step + return self + + def backward(self) -> None: + raise NotImplementedError( + f"'backward' should be implemented in subclass {self.__class__.__name__}" + ) + + def state_dict(self): + agg_state = super().state_dict() + model_state = self.model.state_dict() + # remove model parameters from state dict for already in pdparams + agg_state = {k: v for k, v in agg_state.items() if k not in model_state} + return agg_state diff --git a/examples/smc_reac/ppsci/loss/mtl/grad_norm.py b/examples/smc_reac/ppsci/loss/mtl/grad_norm.py new file mode 100644 index 0000000000..309d3c257c --- /dev/null +++ b/examples/smc_reac/ppsci/loss/mtl/grad_norm.py @@ -0,0 +1,145 @@ +# 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. + +from __future__ import annotations + +from typing import ClassVar +from typing import Dict +from typing import List + +import paddle +from paddle import nn + +from ppsci.loss.mtl import base + +# from ppsci.utils import logger + + +class GradNorm(base.LossAggregator): + r"""GradNorm loss weighting algorithm. + + reference: [https://github.com/PredictiveIntelligenceLab/jaxpi/blob/main/jaxpi/models.py#L132-L146](https://github.com/PredictiveIntelligenceLab/jaxpi/blob/main/jaxpi/models.py#L132-L146) + + $$ + \begin{align*} + L^t &= \sum_{i=1}^{N}{\tilde{w}_i^t\cdot L_i^t}, \\ + \text{where } \\ + \tilde{w}_i^0&=1, \\ + \tilde{w}_i^t&=\tilde{w}_i^{t-1}\cdot m+w_i^t\cdot (1-m), t\ge1\\ + w_i^t&=\dfrac{\overline{\Vert \nabla_{\theta}{L_i^t} \Vert_2}}{\Vert \nabla_{\theta}{L_i^t} \Vert_2}, \\ + \overline{\Vert \nabla_{\theta}{L_i^t} \Vert_2}&=\dfrac{1}{N}\sum_{i=1}^N{\Vert \nabla_{\theta}{L_i^t} \Vert_2}, \\ + &t \text{ is the training step started from 0}. + \end{align*} + $$ + + Attributes: + should_persist(bool): Whether to persist the loss aggregator when saving. + Those loss aggregators with parameters and/or buffers should be persisted. + + Args: + model (nn.Layer): Training model. + num_losses (int, optional): Number of losses. Defaults to 1. + update_freq (int, optional): Weight updating frequency. Defaults to 1000. + momentum (float, optional): Momentum $m$ for moving weight. Defaults to 0.9. + init_weights (List[float]): Initial weights list. Defaults to None. + + Examples: + >>> import paddle + >>> from ppsci.loss import mtl + >>> model = paddle.nn.Linear(3, 4) + >>> loss_aggregator = mtl.GradNorm(model, num_losses=2) + >>> for i in range(5): + ... x1 = paddle.randn([8, 3]) + ... x2 = paddle.randn([8, 3]) + ... y1 = model(x1) + ... y2 = model(x2) + ... loss1 = paddle.sum(y1) + ... loss2 = paddle.sum((y2 - 2) ** 2) + ... loss_aggregator({'loss1': loss1, 'loss2': loss2}).backward() + """ + should_persist: ClassVar[bool] = True + weight: paddle.Tensor + + def __init__( + self, + model: nn.Layer, + num_losses: int = 1, + update_freq: int = 1000, + momentum: float = 0.9, + init_weights: List[float] = None, + ) -> None: + super().__init__(model) + self.step = 0 + self.num_losses = num_losses + self.update_freq = update_freq + self.momentum = momentum + if init_weights is not None and num_losses != len(init_weights): + raise ValueError( + f"Length of init_weights({len(init_weights)}) should be equal to " + f"num_losses({num_losses})." + ) + self.register_buffer( + "weight", + paddle.to_tensor(init_weights, dtype="float32") + if init_weights is not None + else paddle.ones([num_losses]), + ) + + def _compute_weight(self, losses: List["paddle.Tensor"]) -> List["paddle.Tensor"]: + grad_norms = [] + for loss in losses: + loss.backward(retain_graph=True) # NOTE: Keep graph for loss backward + with paddle.no_grad(): + grad_vector = paddle.concat( + [ + p.grad.reshape([-1]) + for p in self.model.parameters() + if p.grad is not None + ] + ) + grad_norms.append(paddle.linalg.norm(grad_vector, p=2)) + self.model.clear_gradients() + + mean_grad_norm = paddle.mean(paddle.stack(grad_norms)) + weight = [(mean_grad_norm / x) for x in grad_norms] + + return weight + + def __call__( + self, losses: Dict[str, "paddle.Tensor"], step: int = 0 + ) -> "paddle.Tensor": + assert len(losses) == self.num_losses, ( + f"Length of given losses({len(losses)}) should be equal to " + f"num_losses({self.num_losses})." + ) + self.step = step + + # compute current loss with moving weights + loss = 0.0 + for i, key in enumerate(losses): + if i == 0: + loss = self.weight[i] * losses[key] + else: + loss += self.weight[i] * losses[key] + + # update moving weights every 'update_freq' steps + if self.step % self.update_freq == 0: + weight = self._compute_weight(list(losses.values())) + for i in range(self.num_losses): + self.weight[i].set_value( + self.momentum * self.weight[i] + (1 - self.momentum) * weight[i] + ) + # logger.message(f"weight at step {self.step}: {self.weight.numpy()}") + + return loss diff --git a/examples/smc_reac/ppsci/loss/mtl/ntk.py b/examples/smc_reac/ppsci/loss/mtl/ntk.py new file mode 100644 index 0000000000..b2dab91fc7 --- /dev/null +++ b/examples/smc_reac/ppsci/loss/mtl/ntk.py @@ -0,0 +1,118 @@ +# 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. + +from __future__ import annotations + +from typing import TYPE_CHECKING +from typing import ClassVar +from typing import Dict + +import paddle + +from ppsci.loss.mtl import base + +if TYPE_CHECKING: + from paddle import nn + + +class NTK(base.LossAggregator): + r"""Weighted Neural Tangent Kernel. + + reference: [https://github.com/PredictiveIntelligenceLab/jaxpi/blob/main/jaxpi/models.py#L148-L158](https://github.com/PredictiveIntelligenceLab/jaxpi/blob/main/jaxpi/models.py#L148-L158) + + Attributes: + should_persist(bool): Whether to persist the loss aggregator when saving. + Those loss aggregators with parameters and/or buffers should be persisted. + + Args: + model (nn.Layer): Training model. + num_losses (int, optional): Number of losses. Defaults to 1. + update_freq (int, optional): Weight updating frequency. Defaults to 1000. + + Examples: + >>> import paddle + >>> from ppsci.loss import mtl + >>> model = paddle.nn.Linear(3, 4) + >>> loss_aggregator = mtl.NTK(model, num_losses=2) + >>> for i in range(5): + ... x1 = paddle.randn([8, 3]) + ... x2 = paddle.randn([8, 3]) + ... y1 = model(x1) + ... y2 = model(x2) + ... loss1 = paddle.sum(y1) + ... loss2 = paddle.sum((y2 - 2) ** 2) + ... loss_aggregator({'loss1': loss1, 'loss2': loss2}).backward() + """ + should_persist: ClassVar[bool] = True + weight: paddle.Tensor + + def __init__( + self, + model: nn.Layer, + num_losses: int = 1, + update_freq: int = 1000, + ) -> None: + super().__init__(model) + self.step = 0 + self.num_losses = num_losses + self.update_freq = update_freq + self.register_buffer("weight", paddle.ones([num_losses])) + + def _compute_weight(self, losses: Dict[str, paddle.Tensor]): + ntk_sum = 0 + ntk_value = [] + for loss in losses.values(): + grads = paddle.grad( + loss, + self.model.parameters(), + create_graph=False, + retain_graph=True, + allow_unused=True, + ) + with paddle.no_grad(): + grad = paddle.concat( + [grad.reshape([-1]) for grad in grads if grad is not None] + ) + ntk_value.append( + paddle.sqrt( + paddle.sum(grad.detach() ** 2), + ) + ) + + ntk_sum += paddle.sum(paddle.stack(ntk_value, axis=0)) + ntk_weight = [(ntk_sum / x) for x in ntk_value] + + return ntk_weight + + def __call__( + self, losses: Dict[str, "paddle.Tensor"], step: int = 0 + ) -> "paddle.Tensor": + assert len(losses) == self.num_losses, ( + f"Length of given losses({len(losses)}) should be equal to " + f"num_losses({self.num_losses})." + ) + self.step = step + + # compute current loss with moving weights + loss = 0 + for i, (k, v) in enumerate(losses.items()): + loss = loss + self.weight[i] * v + + # update moving weights every 'update_freq' steps + if self.step % self.update_freq == 0: + computed_weight = self._compute_weight(losses) + for i in range(self.num_losses): + self.weight[i].set_value(computed_weight[i]) + + return loss diff --git a/examples/smc_reac/ppsci/loss/mtl/pcgrad.py b/examples/smc_reac/ppsci/loss/mtl/pcgrad.py new file mode 100644 index 0000000000..45b5923110 --- /dev/null +++ b/examples/smc_reac/ppsci/loss/mtl/pcgrad.py @@ -0,0 +1,124 @@ +# 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. + +from __future__ import annotations + +from typing import ClassVar +from typing import List + +import numpy as np +import paddle +from paddle import nn + +from ppsci.loss.mtl import base + + +class PCGrad(base.LossAggregator): + r""" + **P**rojecting **C**onflicting Gradients + + [Gradient Surgery for Multi-Task Learning](https://papers.nips.cc/paper/2020/hash/3fe78a8acf5fda99de95303940a2420c-Abstract.html) + + Code reference: [https://github.com/tianheyu927/PCGrad/blob/master/PCGrad_tf.py](https://github.com/tianheyu927/PCGrad/blob/master/PCGrad_tf.py) + + Attributes: + should_persist(bool): Whether to persist the loss aggregator when saving. + Those loss aggregators with parameters and/or buffers should be persisted. + + Args: + model (nn.Layer): Training model. + + Examples: + >>> import paddle + >>> from ppsci.loss import mtl + >>> model = paddle.nn.Linear(3, 4) + >>> loss_aggregator = mtl.PCGrad(model) + >>> for i in range(5): + ... x1 = paddle.randn([8, 3]) + ... x2 = paddle.randn([8, 3]) + ... y1 = model(x1) + ... y2 = model(x2) + ... loss1 = paddle.sum(y1) + ... loss2 = paddle.sum((y2 - 2) ** 2) + ... loss_aggregator({'loss1': loss1, 'loss2': loss2}).backward() + """ + should_persist: ClassVar[bool] = False + + def __init__(self, model: nn.Layer) -> None: + super().__init__(model) + self._zero = paddle.zeros([]) + + def backward(self) -> None: + # shuffle order of losses + keys = list(self.losses.keys()) + np.random.shuffle(keys) + self.losses = {key: self.losses[key] for key in keys} + + grads_list = self._compute_grads() + with paddle.no_grad(): + refined_grads = self._refine_grads(grads_list) + self._set_grads(refined_grads) + + def _compute_grads(self) -> List[paddle.Tensor]: + # compute all gradients derived by each loss + grads_list = [] # num_params x num_losses + for key in self.losses: + # backward with current loss + self.losses[key].backward() + grads_list.append( + paddle.concat( + [ + param.grad.clone().reshape([-1]) + for param in self.model.parameters() + if param.grad is not None + ], + axis=0, + ) + ) + # clear gradients for current loss for not affecting other loss + self.model.clear_gradients() + + return grads_list + + def _refine_grads(self, grads_list: List[paddle.Tensor]) -> List[paddle.Tensor]: + def proj_grad(grad: paddle.Tensor): + for k in range(self.loss_num): + inner_product = paddle.sum(grad * grads_list[k]) + proj_direction = inner_product / paddle.sum( + grads_list[k] * grads_list[k] + ) + grad = grad - paddle.minimum(proj_direction, self._zero) * grads_list[k] + return grad + + grads_list = [proj_grad(grad) for grad in grads_list] + + # Unpack flattened projected gradients back to their original shapes. + proj_grads: List[paddle.Tensor] = [] + for j in range(self.loss_num): + start_idx = 0 + for idx, var in enumerate(self.model.parameters()): + grad_shape = var.shape + flatten_dim = var.numel() + refined_grad = grads_list[j][start_idx : start_idx + flatten_dim] + refined_grad = paddle.reshape(refined_grad, grad_shape) + if len(proj_grads) < self.param_num: + proj_grads.append(refined_grad) + else: + proj_grads[idx] += refined_grad + start_idx += flatten_dim + return proj_grads + + def _set_grads(self, grads_list: List[paddle.Tensor]) -> None: + for i, param in enumerate(self.model.parameters()): + param.grad = grads_list[i] diff --git a/examples/smc_reac/ppsci/loss/mtl/relobralo.py b/examples/smc_reac/ppsci/loss/mtl/relobralo.py new file mode 100644 index 0000000000..02ec8f1339 --- /dev/null +++ b/examples/smc_reac/ppsci/loss/mtl/relobralo.py @@ -0,0 +1,127 @@ +# 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. + +from __future__ import annotations + +from typing import ClassVar +from typing import Dict + +import paddle +from paddle import nn + + +class Relobralo(nn.Layer): + r""" + **Re**lative **Lo**ss **B**alancing with **Ra**ndom **Lo**okback + + [Multi-Objective Loss Balancing for Physics-Informed Deep Learning](https://arxiv.org/abs/2110.09813) + + Attributes: + should_persist(bool): Whether to persist the loss aggregator when saving. + Those loss aggregators with parameters and/or buffers should be persisted. + + Args: + num_losses (int): Number of losses. + alpha (float, optional): Ability for remembering past in paper. Defaults to 0.95. + beta (float, optional): Parameter for generating $\rho$ from bernoulli distribution, + and $E[\rho](=\beta)$ should be close to 1. Defaults to 0.99. + tau (float, optional): Temperature factor. Equivalent to softmax when $\tau$=1.0, + equivalent to argmax when $\tau$=0. Defaults to 1.0. + eps (float, optional): $\epsilon$ to avoid divided by 0 in losses. Defaults to 1e-8. + + Examples: + >>> import paddle + >>> from ppsci.loss import mtl + >>> model = paddle.nn.Linear(3, 4) + >>> loss_aggregator = mtl.Relobralo(num_losses=2) + >>> for i in range(5): + ... x1 = paddle.randn([8, 3]) + ... x2 = paddle.randn([8, 3]) + ... y1 = model(x1) + ... y2 = model(x2) + ... loss1 = paddle.sum(y1) + ... loss2 = paddle.sum((y2 - 2) ** 2) + ... loss_aggregator({'loss1': loss1, 'loss2': loss2}).backward() + """ + should_persist: ClassVar[bool] = True + + def __init__( + self, + num_losses: int, + alpha: float = 0.95, + beta: float = 0.99, + tau: float = 1.0, + eps: float = 1e-8, + ) -> None: + super().__init__() + self.step = 0 + self.num_losses: int = num_losses + self.alpha: float = alpha + self.beta: float = beta + self.tau: float = tau + self.eps: float = eps + self.register_buffer("losses_init", paddle.zeros([self.num_losses])) + self.register_buffer("losses_prev", paddle.zeros([self.num_losses])) + self.register_buffer("lmbda", paddle.ones([self.num_losses])) + + def _softmax(self, vec: "paddle.Tensor") -> "paddle.Tensor": + max_item = vec.max() + result = paddle.exp(vec - max_item) / paddle.exp(vec - max_item).sum() + return result + + def _compute_bal( + self, losses_vec1: "paddle.Tensor", losses_vec2: "paddle.Tensor" + ) -> "paddle.Tensor": + return self.num_losses * ( + self._softmax(losses_vec1 / (self.tau * losses_vec2 + self.eps)) + ) + + def __call__( + self, losses: Dict[str, "paddle.Tensor"], step: int = 0 + ) -> "paddle.Tensor": + assert len(losses) == self.num_losses, ( + f"Length of given losses({len(losses)}) should be equal to " + f"num_losses({self.num_losses})." + ) + self.step = step + losses_stacked = paddle.stack(list(losses.values())) # [num_losses, ] + + if self.step == 0: + loss = losses_stacked.sum() + with paddle.no_grad(): + paddle.assign(losses_stacked.detach(), self.losses_init) + else: + with paddle.no_grad(): + # 1. update lambda_hist + rho = paddle.bernoulli(paddle.to_tensor(self.beta)) + lmbda_hist = rho * self.lmbda + (1 - rho) * self._compute_bal( + losses_stacked, self.losses_init + ) + + # 2. update lambda + paddle.assign( + self.alpha * lmbda_hist + + (1 - self.alpha) + * self._compute_bal(losses_stacked, self.losses_prev), + self.lmbda, + ) + + # 3. compute reweighted total loss with lambda + loss = (losses_stacked * self.lmbda).sum() + + # update losses_prev at the end of each step + with paddle.no_grad(): + paddle.assign(losses_stacked.detach(), self.losses_prev) + + return loss diff --git a/examples/smc_reac/ppsci/loss/mtl/sum.py b/examples/smc_reac/ppsci/loss/mtl/sum.py new file mode 100644 index 0000000000..d2c9a7bd50 --- /dev/null +++ b/examples/smc_reac/ppsci/loss/mtl/sum.py @@ -0,0 +1,60 @@ +# 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. + +from __future__ import annotations + +from typing import TYPE_CHECKING +from typing import Dict + +if TYPE_CHECKING: + import paddle + +from typing import ClassVar + +from ppsci.loss.mtl.base import LossAggregator + + +class Sum(LossAggregator): + r""" + **Default loss aggregator** which do simple summation for given losses as below. + + $$ + loss = \sum_i^N losses_i + $$ + + Attributes: + should_persist(bool): Whether to persist the loss aggregator when saving. + Those loss aggregators with parameters and/or buffers should be persisted. + """ + should_persist: ClassVar[bool] = False + + def __init__(self) -> None: + self.step = 0 + + def __call__( + self, losses: Dict[str, "paddle.Tensor"], step: int = 0 + ) -> "paddle.Tensor": + assert ( + len(losses) > 0 + ), f"Number of given losses({len(losses)}) can not be empty." + self.step = step + + total_loss = 0.0 + for i, key in enumerate(losses): + if i == 0: + total_loss = losses[key] + else: + total_loss += losses[key] + + return total_loss diff --git a/examples/smc_reac/ppsci/metric/__init__.py b/examples/smc_reac/ppsci/metric/__init__.py new file mode 100644 index 0000000000..0a1d069aa9 --- /dev/null +++ b/examples/smc_reac/ppsci/metric/__init__.py @@ -0,0 +1,63 @@ +# 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. + +import copy + +from ppsci.metric.anomaly_coef import LatitudeWeightedACC +from ppsci.metric.base import Metric +from ppsci.metric.func import FunctionalMetric +from ppsci.metric.l2_rel import L2Rel +from ppsci.metric.l2_rel import MeanL2Rel +from ppsci.metric.mae import MAE +from ppsci.metric.max_ae import MaxAE +from ppsci.metric.mse import MSE +from ppsci.metric.r2_score import R2Score +from ppsci.metric.rmse import RMSE +from ppsci.metric.rmse import LatitudeWeightedRMSE +from ppsci.utils import misc + +__all__ = [ + "LatitudeWeightedACC", + "Metric", + "FunctionalMetric", + "L2Rel", + "MeanL2Rel", + "MAE", + "MaxAE", + "MSE", + "RMSE", + "LatitudeWeightedRMSE", + "R2Score", + "build_metric", +] + + +def build_metric(cfg): + """Build metric. + + Args: + cfg (List[DictConfig]): List of metric config. + + Returns: + Dict[str, Metric]: Dict of callable metric object. + """ + cfg = copy.deepcopy(cfg) + + metric_dict = misc.PrettyOrderedDict() + for _item in cfg: + metric_cls = next(iter(_item.keys())) + metric_cfg = _item.pop(metric_cls) + metric = eval(metric_cls)(**metric_cfg) + metric_dict[metric_cls] = metric + return metric_dict diff --git a/examples/smc_reac/ppsci/metric/anomaly_coef.py b/examples/smc_reac/ppsci/metric/anomaly_coef.py new file mode 100644 index 0000000000..33633228ca --- /dev/null +++ b/examples/smc_reac/ppsci/metric/anomaly_coef.py @@ -0,0 +1,122 @@ +# 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. + +from __future__ import annotations + +from typing import Dict +from typing import Optional +from typing import Tuple +from typing import Union + +import numpy as np +import paddle + +from ppsci.metric import base + + +class LatitudeWeightedACC(base.Metric): + r"""Latitude weighted anomaly correlation coefficient. + + $$ + metric = + \dfrac{\sum\limits_{m,n}{L_mX_{mn}Y_{mn}}}{\sqrt{\sum\limits_{m,n}{L_mX_{mn}^{2}}\sum\limits_{m,n}{L_mY_{mn}^{2}}}} + $$ + + $$ + L_m = N_{lat}\dfrac{\cos(lat_m)}{\sum\limits_{j=1}^{N_{lat}}\cos(lat_j)} + $$ + + $lat_m$ is the latitude at m. + $N_{lat}$ is the number of latitude set by `num_lat`. + + Args: + num_lat (int): Number of latitude. + mean (Optional[Union[np.array, Tuple[float, ...]]]): Mean of training data. Defaults to None. + keep_batch (bool, optional): Whether keep batch axis. Defaults to False. + variable_dict (Optional[Dict[str, int]]): Variable dictionary, the key is the name of a variable and + the value is its index. Defaults to None. + unlog (bool, optional): Whether calculate expm1 for all elements in the array. Defaults to False. + scale (float, optional): The scale value used after expm1. Defaults to 1e-5. + + Examples: + >>> import numpy as np + >>> import ppsci + >>> mean = np.random.randn(20, 720, 1440) + >>> metric = ppsci.metric.LatitudeWeightedACC(720, mean=mean) + """ + + def __init__( + self, + num_lat: int, + mean: Optional[Union[np.array, Tuple[float, ...]]], + keep_batch: bool = False, + variable_dict: Optional[Dict[str, int]] = None, + unlog: bool = False, + scale: float = 1e-5, + ): + super().__init__(keep_batch) + self.num_lat = num_lat + self.mean = ( + None if mean is None else paddle.to_tensor(mean, paddle.get_default_dtype()) + ) + self.variable_dict = variable_dict + self.unlog = unlog + self.scale = scale + + self.weight = self.get_latitude_weight(num_lat) + + def get_latitude_weight(self, num_lat: int = 720): + lat_t = paddle.linspace(start=0, stop=1, num=num_lat) + lat_t = paddle.cos(3.1416 * (0.5 - lat_t)) + weight = num_lat * lat_t / paddle.sum(lat_t) + weight = weight.reshape((1, 1, -1, 1)) + return weight + + def scale_expm1(self, x: paddle.Tensor): + return self.scale * paddle.expm1(x) + + @paddle.no_grad() + def forward(self, output_dict, label_dict) -> Dict[str, "paddle.Tensor"]: + metric_dict = {} + + for key in label_dict: + output = ( + self.scale_expm1(output_dict[key]) if self.unlog else output_dict[key] + ) + label = self.scale_expm1(label_dict[key]) if self.unlog else label_dict[key] + + if self.mean is not None: + output = output - self.mean + label = label - self.mean + + rmse = paddle.sum( + self.weight * output * label, axis=(-1, -2) + ) / paddle.sqrt( + paddle.sum(self.weight * output**2, axis=(-1, -2)) + * paddle.sum(self.weight * label**2, axis=(-1, -2)) + ) + + if self.variable_dict is not None: + for variable_name, idx in self.variable_dict.items(): + if self.keep_batch: + metric_dict[f"{key}.{variable_name}"] = rmse[:, idx] + else: + metric_dict[f"{key}.{variable_name}"] = rmse[:, idx].mean() + else: + if self.keep_batch: + metric_dict[key] = rmse.mean(axis=1) + else: + metric_dict[key] = rmse.mean() + + return metric_dict diff --git a/examples/smc_reac/ppsci/metric/base.py b/examples/smc_reac/ppsci/metric/base.py new file mode 100644 index 0000000000..750e629882 --- /dev/null +++ b/examples/smc_reac/ppsci/metric/base.py @@ -0,0 +1,25 @@ +# 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. + +from __future__ import annotations + +from paddle import nn + + +class Metric(nn.Layer): + """Base class for metric.""" + + def __init__(self, keep_batch: bool = False): + super().__init__() + self.keep_batch = keep_batch diff --git a/examples/smc_reac/ppsci/metric/func.py b/examples/smc_reac/ppsci/metric/func.py new file mode 100644 index 0000000000..bee646b656 --- /dev/null +++ b/examples/smc_reac/ppsci/metric/func.py @@ -0,0 +1,66 @@ +# 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. + +from __future__ import annotations + +from typing import TYPE_CHECKING +from typing import Callable +from typing import Dict + +from ppsci.metric import base + +if TYPE_CHECKING: + import paddle + + +class FunctionalMetric(base.Metric): + r"""Functional metric class, which allows to use custom metric computing function from given metric_expr for complex computation cases. + + Args: + metric_expr (Callable): Expression of metric calculation. + keep_batch (bool, optional): Whether keep batch axis. Defaults to False. + + Examples: + >>> import paddle + >>> from ppsci.metric import FunctionalMetric + >>> def metric_expr(output_dict, *args): + ... rel_l2 = 0 + ... for key in output_dict: + ... length = int(len(output_dict[key])/2) + ... out_dict = output_dict[key][:length] + ... label_dict = output_dict[key][length:] + ... rel_l2 += paddle.norm(out_dict - label_dict) / paddle.norm(label_dict) + ... return {"rel_l2": rel_l2} + >>> metric_dict = FunctionalMetric(metric_expr) + >>> output_dict = {'u': paddle.to_tensor([[0.5, 0.9], [1.1, -1.3], [-0.2, 1.5], [-0.1, -0.3]]), + ... 'v': paddle.to_tensor([[0.5, 0.9], [1.1, -1.3], [-1.8, 1.0], [-0.2, 2.5]])} + >>> result = metric_dict(output_dict) + >>> print(result) + {'rel_l2': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, + 2.59985542)} + """ + + def __init__( + self, + metric_expr: Callable[ + [Dict["str", "paddle.Tensor"], Dict["str", "paddle.Tensor"]], + Dict["str", "paddle.Tensor"], + ], + keep_batch: bool = False, + ): + super().__init__(keep_batch) + self.metric_expr = metric_expr + + def forward(self, output_dict, label_dict=None) -> Dict[str, "paddle.Tensor"]: + return self.metric_expr(output_dict, label_dict) diff --git a/examples/smc_reac/ppsci/metric/l2_rel.py b/examples/smc_reac/ppsci/metric/l2_rel.py new file mode 100644 index 0000000000..2a64e9befc --- /dev/null +++ b/examples/smc_reac/ppsci/metric/l2_rel.py @@ -0,0 +1,139 @@ +# 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. + +from __future__ import annotations + +from typing import Dict + +import numpy as np +import paddle + +from ppsci.metric import base + + +class L2Rel(base.Metric): + r"""Class for l2 relative error. + + NOTE: This metric API is slightly different from `MeanL2Rel`, difference is as below: + + - `L2Rel` regards the input sample as a whole and calculates the l2 relative error of the whole; + - `MeanL2Rel` will calculate L2Rel separately for each input sample and return the average of l2 relative error for all samples. + + $$ + metric = \dfrac{\Vert \mathbf{x} - \mathbf{y} \Vert_2}{\max(\Vert \mathbf{y} \Vert_2, \epsilon)} + $$ + + $$ + \mathbf{x}, \mathbf{y} \in \mathcal{R}^{N} + $$ + + Args: + keep_batch (bool, optional): Whether keep batch axis. Defaults to False. + + Examples: + >>> import paddle + >>> from ppsci.metric import L2Rel + >>> output_dict = {'u': paddle.to_tensor([[0.5, 0.9], [1.1, -1.3]]), + ... 'v': paddle.to_tensor([[0.5, 0.9], [1.1, -1.3]])} + >>> label_dict = {'u': paddle.to_tensor([[-1.8, 1.0], [-0.2, 2.5]]), + ... 'v': paddle.to_tensor([[0.1, 0.1], [0.1, 0.1]])} + >>> loss = L2Rel() + >>> result = loss(output_dict, label_dict) + >>> print(result) + {'u': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, + 1.42658269), 'v': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, + 9.69535923)} + """ + + # NOTE: Avoid divide by zero in result + # see https://github.com/scikit-learn/scikit-learn/pull/15007 + EPS: float = np.finfo(np.float32).eps + + def __init__(self, keep_batch: bool = False): + if keep_batch: + raise ValueError(f"keep_batch should be False, but got {keep_batch}.") + super().__init__(keep_batch) + + @paddle.no_grad() + def forward(self, output_dict, label_dict) -> Dict[str, "paddle.Tensor"]: + metric_dict = {} + for key in label_dict: + rel_l2 = paddle.norm(label_dict[key] - output_dict[key], p=2) / paddle.norm( + label_dict[key], p=2 + ).clip(min=self.EPS) + metric_dict[key] = rel_l2 + + return metric_dict + + +class MeanL2Rel(base.Metric): + r"""Class for mean l2 relative error. + + NOTE: This metric API is slightly different from `L2Rel`, difference is as below: + + - `MeanL2Rel` will calculate L2Rel separately for each input sample and return the average of l2 relative error for all samples. + - `L2Rel` regards the input sample as a whole and calculates the l2 relative error of the whole; + + $$ + metric = \dfrac{1}{M} \sum_{i=1}^{M}\dfrac{\Vert \mathbf{x_i} - \mathbf{y_i} \Vert_2}{\max(\Vert \mathbf{y_i} \Vert_2, \epsilon) } + $$ + + $$ + \mathbf{x_i}, \mathbf{y_i} \in \mathcal{R}^{N} + $$ + + Args: + keep_batch (bool, optional): Whether keep batch axis. Defaults to False. + + Examples: + >>> import paddle + >>> from ppsci.metric import MeanL2Rel + >>> output_dict = {'u': paddle.to_tensor([[0.5, 0.9], [1.1, -1.3]]), + ... 'v': paddle.to_tensor([[0.5, 0.9], [1.1, -1.3]])} + >>> label_dict = {'u': paddle.to_tensor([[-1.8, 1.0], [-0.2, 2.5]]), + ... 'v': paddle.to_tensor([[0.1, 0.1], [0.1, 0.1]])} + >>> loss = MeanL2Rel() + >>> result = loss(output_dict, label_dict) + >>> print(result) + {'u': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, + 1.35970235), 'v': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, + 9.24504089)} + >>> loss = MeanL2Rel(keep_batch=True) + >>> result = loss(output_dict, label_dict) + >>> print(result) + {'u': Tensor(shape=[2], dtype=float32, place=Place(gpu:0), stop_gradient=True, + [1.11803389, 1.60137081]), 'v': Tensor(shape=[2], dtype=float32, place=Place(gpu:0), stop_gradient=True, + [6.32455540 , 12.16552544])} + """ + + # NOTE: Avoid divide by zero in result + # see https://github.com/scikit-learn/scikit-learn/pull/15007 + EPS: float = np.finfo(np.float32).eps + + def __init__(self, keep_batch: bool = False): + super().__init__(keep_batch) + + @paddle.no_grad() + def forward(self, output_dict, label_dict) -> Dict[str, "paddle.Tensor"]: + metric_dict = {} + for key in label_dict: + rel_l2 = paddle.norm( + label_dict[key] - output_dict[key], p=2, axis=1 + ) / paddle.norm(label_dict[key], p=2, axis=1).clip(min=self.EPS) + if self.keep_batch: + metric_dict[key] = rel_l2 + else: + metric_dict[key] = rel_l2.mean() + + return metric_dict diff --git a/examples/smc_reac/ppsci/metric/mae.py b/examples/smc_reac/ppsci/metric/mae.py new file mode 100644 index 0000000000..3b6ebdedbb --- /dev/null +++ b/examples/smc_reac/ppsci/metric/mae.py @@ -0,0 +1,73 @@ +# 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. + +from __future__ import annotations + +from typing import Dict + +import paddle +import paddle.nn.functional as F + +from ppsci.metric import base + + +class MAE(base.Metric): + r"""Mean absolute error. + + $$ + metric = \dfrac{1}{N} \Vert \mathbf{x} - \mathbf{y} \Vert_1 + $$ + + $$ + \mathbf{x}, \mathbf{y} \in \mathcal{R}^{N} + $$ + + Args: + keep_batch (bool, optional): Whether keep batch axis. Defaults to False. + + Examples: + >>> import paddle + >>> from ppsci.metric import MAE + >>> output_dict = {'u': paddle.to_tensor([[0.5, 0.9], [1.1, -1.3]]), + ... 'v': paddle.to_tensor([[0.5, 0.9], [1.1, -1.3]])} + >>> label_dict = {'u': paddle.to_tensor([[-1.8, 1.0], [-0.2, 2.5]]), + ... 'v': paddle.to_tensor([[0.1, 0.1], [0.1, 0.1]])} + >>> loss = MAE() + >>> result = loss(output_dict, label_dict) + >>> print(result) + {'u': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, + 1.87500000), 'v': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, + 0.89999998)} + >>> loss = MAE(keep_batch=True) + >>> result = loss(output_dict, label_dict) + >>> print(result) + {'u': Tensor(shape=[2], dtype=float32, place=Place(gpu:0), stop_gradient=True, + [1.20000005, 2.54999995]), 'v': Tensor(shape=[2], dtype=float32, place=Place(gpu:0), stop_gradient=True, + [0.59999996, 1.20000005])} + """ + + def __init__(self, keep_batch: bool = False): + super().__init__(keep_batch) + + @paddle.no_grad() + def forward(self, output_dict, label_dict) -> Dict[str, "paddle.Tensor"]: + metric_dict = {} + for key in label_dict: + mae = F.l1_loss(output_dict[key], label_dict[key], "none") + if self.keep_batch: + metric_dict[key] = mae.mean(axis=tuple(range(1, mae.ndim))) + else: + metric_dict[key] = mae.mean() + + return metric_dict diff --git a/examples/smc_reac/ppsci/metric/max_ae.py b/examples/smc_reac/ppsci/metric/max_ae.py new file mode 100644 index 0000000000..7d94544b97 --- /dev/null +++ b/examples/smc_reac/ppsci/metric/max_ae.py @@ -0,0 +1,77 @@ +# Copyright (c) 2024 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. + +from __future__ import annotations + +from typing import Dict + +import paddle + +from ppsci.metric import base + + +class MaxAE(base.Metric): + r"""Maximum Absolute Error (MaxAE). + + $$ + \text{MaxAE} = \max_i \left( |x_i - y_i| \right) + $$ + + $$ + \mathbf{x}, \mathbf{y} \in \mathcal{R}^{N} + $$ + + Args: + keep_batch (bool, optional): Whether keep batch axis. Defaults to False. + + Examples: + >>> import paddle + >>> from ppsci.metric import MaxAE + >>> output_dict = {'u': paddle.to_tensor([[0.5, 0.9], [1.1, -1.3]]), + ... 'v': paddle.to_tensor([[0.5, 0.9], [1.1, -1.3]])} + >>> label_dict = {'u': paddle.to_tensor([[-1.8, 1.0], [-0.2, 2.5]]), + ... 'v': paddle.to_tensor([[0.1, 0.1], [0.1, 0.1]])} + >>> metric = MaxAE() + >>> result = metric(output_dict, label_dict) + >>> print(result) + {'u': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, + 3.80000007), 'v': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, + 0.80000001)} + >>> metric = MaxAE(keep_batch=True) + >>> result = metric(output_dict, label_dict) + >>> print(result) + {'u': Tensor(shape=[2], dtype=float32, place=Place(gpu:0), stop_gradient=True, + [1.30000002, 3.80000007]), 'v': Tensor(shape=[2], dtype=float32, place=Place(gpu:0), stop_gradient=True, + [0.40000001, 0.80000001])} + """ + + def __init__(self, keep_batch: bool = False): + super().__init__(keep_batch) + + @paddle.no_grad() + def forward(self, output_dict, label_dict) -> Dict[str, "paddle.Tensor"]: + maxae_dict = {} + + for key in label_dict: + # Calculate absolute error + ae = paddle.abs(output_dict[key] - label_dict[key]) + + if self.keep_batch: + # Take the maximum AE within each batch + maxae_dict[key] = paddle.amax(ae, axis=tuple(range(1, ae.ndim))) + else: + # Take the global maximum AE across all elements + maxae_dict[key] = paddle.amax(ae) + + return maxae_dict diff --git a/examples/smc_reac/ppsci/metric/mse.py b/examples/smc_reac/ppsci/metric/mse.py new file mode 100644 index 0000000000..9e47a7cf5a --- /dev/null +++ b/examples/smc_reac/ppsci/metric/mse.py @@ -0,0 +1,73 @@ +# 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. + +from __future__ import annotations + +from typing import Dict + +import paddle +import paddle.nn.functional as F + +from ppsci.metric import base + + +class MSE(base.Metric): + r"""Mean square error + + $$ + metric = \dfrac{1}{N} \Vert \mathbf{x} - \mathbf{y} \Vert_2^2 + $$ + + $$ + \mathbf{x}, \mathbf{y} \in \mathcal{R}^{N} + $$ + + Args: + keep_batch (bool, optional): Whether keep batch axis. Defaults to False. + + Examples: + >>> import paddle + >>> from ppsci.metric import MSE + >>> output_dict = {'u': paddle.to_tensor([[0.5, 0.9], [1.1, -1.3]]), + ... 'v': paddle.to_tensor([[0.5, 0.9], [1.1, -1.3]])} + >>> label_dict = {'u': paddle.to_tensor([[-1.8, 1.0], [-0.2, 2.5]]), + ... 'v': paddle.to_tensor([[0.1, 0.1], [0.1, 0.1]])} + >>> loss = MSE() + >>> result = loss(output_dict, label_dict) + >>> print(result) + {'u': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, + 5.35750008), 'v': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, + 0.94000000)} + >>> loss = MSE(keep_batch=True) + >>> result = loss(output_dict, label_dict) + >>> print(result) + {'u': Tensor(shape=[2], dtype=float32, place=Place(gpu:0), stop_gradient=True, + [2.65000010, 8.06499958]), 'v': Tensor(shape=[2], dtype=float32, place=Place(gpu:0), stop_gradient=True, + [0.39999998, 1.48000002])} + """ + + def __init__(self, keep_batch: bool = False): + super().__init__(keep_batch) + + @paddle.no_grad() + def forward(self, output_dict, label_dict) -> Dict[str, "paddle.Tensor"]: + metric_dict = {} + for key in label_dict: + mse = F.mse_loss(output_dict[key], label_dict[key], "none") + if self.keep_batch: + metric_dict[key] = mse.mean(axis=tuple(range(1, mse.ndim))) + else: + metric_dict[key] = mse.mean() + + return metric_dict diff --git a/examples/smc_reac/ppsci/metric/r2_score.py b/examples/smc_reac/ppsci/metric/r2_score.py new file mode 100644 index 0000000000..4ad3eb608d --- /dev/null +++ b/examples/smc_reac/ppsci/metric/r2_score.py @@ -0,0 +1,97 @@ +# Copyright (c) 2024 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. + +from __future__ import annotations + +from typing import Dict + +import paddle + +from ppsci.metric import base + + +class R2Score(base.Metric): + r"""Coefficient of Determination (R^2 Score). + + $$ + R^2 = 1 - \frac{\sum_{i=1}^{N} (y_i - \hat{y}_i)^2}{\sum_{i=1}^{N} (y_i - \bar{y})^2} + $$ + + $$ + \mathbf{y}, \mathbf{\hat{y}} \in \mathcal{R}^{N} + $$ + + Args: + keep_batch (bool, optional): Whether keep batch axis. Defaults to False. + + Examples: + >>> import paddle + >>> from ppsci.metric import R2Score + >>> output_dict = {'u': paddle.to_tensor([[0.5, 0.9], [1.1, -1.3]]), + ... 'v': paddle.to_tensor([[0.5, 0.9], [1.1, -1.3]])} + >>> label_dict = {'u': paddle.to_tensor([[-1.8, 1.0], [-0.2, 2.5]]), + ... 'v': paddle.to_tensor([[0.1, 0.1], [0.1, 0.1]])} + >>> metric = R2Score() + >>> result = metric(output_dict, label_dict) + >>> print(result) + {'u': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, + -3.75000000), 'v': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, + 0.00000000)} + >>> metric = R2Score(keep_batch=True) + >>> result = metric(output_dict, label_dict) + >>> print(result) + {'u': Tensor(shape=[2], dtype=float32, place=Place(gpu:0), stop_gradient=True, + [-0.64000008, -6.88000011]), 'v': Tensor(shape=[2], dtype=float32, place=Place(gpu:0), stop_gradient=True, + [0.00000000, 0.00000000])} + """ + + def __init__(self, keep_batch: bool = False): + super().__init__(keep_batch) + + @paddle.no_grad() + def forward(self, output_dict, label_dict) -> Dict[str, "paddle.Tensor"]: + r2score_dict = {} + + for key in label_dict: + + output = output_dict[key] + target = label_dict[key] + + # Ensure the shapes of output and target match + if output.shape != target.shape: + raise ValueError( + f"Output and target shapes do not match for key '{key}'. Output shape: {output.shape}, Target shape: {target.shape}" + ) + + output = output.flatten() + target = target.flatten() + # Calculate mean of target along the last dimension (features) + target_mean = target.mean(axis=-1, keepdim=True) + + # Calculate total sum of squares (TSS) + ss_tot = paddle.sum(x=(target - target_mean) ** 2) + + # Calculate residual sum of squares (RSS) + ss_res = paddle.sum(x=(target - output) ** 2) + + # Calculate R^2 score with numerical stability + r2 = 1 - (ss_res / (ss_tot + 1e-8)) + + # Handle batch-wise or mean R^2 based on self.keep_batch + if self.keep_batch: + r2score_dict[key] = r2.unsqueeze(axis=0) + else: + r2score_dict[key] = r2 + + return r2score_dict diff --git a/examples/smc_reac/ppsci/metric/rmse.py b/examples/smc_reac/ppsci/metric/rmse.py new file mode 100644 index 0000000000..55be8e9102 --- /dev/null +++ b/examples/smc_reac/ppsci/metric/rmse.py @@ -0,0 +1,155 @@ +# 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. + +from __future__ import annotations + +from typing import Dict +from typing import Optional +from typing import Tuple +from typing import Union + +import numpy as np +import paddle +import paddle.nn.functional as F + +from ppsci.metric import base + + +class RMSE(base.Metric): + r"""Root mean square error + + $$ + metric = \sqrt{\dfrac{1}{N} \Vert \mathbf{x} - \mathbf{y} \Vert_2^2} + $$ + + $$ + \mathbf{x}, \mathbf{y} \in \mathcal{R}^{N} + $$ + + Args: + keep_batch (bool, optional): Whether keep batch axis. Defaults to False. + + Examples: + >>> import paddle + >>> from ppsci.metric import RMSE + >>> output_dict = {'u': paddle.to_tensor([[0.5, 0.9], [1.1, -1.3]]), + ... 'v': paddle.to_tensor([[0.5, 0.9], [1.1, -1.3]])} + >>> label_dict = {'u': paddle.to_tensor([[-1.8, 1.0], [-0.2, 2.5]]), + ... 'v': paddle.to_tensor([[0.1, 0.1], [0.1, 0.1]])} + >>> loss = RMSE() + >>> result = loss(output_dict, label_dict) + >>> print(result) + {'u': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, + 2.31462741), 'v': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, + 0.96953595)} + """ + + def __init__(self, keep_batch: bool = False): + if keep_batch: + raise ValueError(f"keep_batch should be False, but got {keep_batch}.") + super().__init__(keep_batch) + + @paddle.no_grad() + def forward(self, output_dict, label_dict) -> Dict[str, "paddle.Tensor"]: + metric_dict = {} + for key in label_dict: + rmse = F.mse_loss(output_dict[key], label_dict[key], "mean") ** 0.5 + metric_dict[key] = rmse + + return metric_dict + + +class LatitudeWeightedRMSE(base.Metric): + r"""Latitude weighted root mean square error. + + $$ + metric =\sqrt{\dfrac{1}{MN}\sum\limits_{m=1}^{M}\sum\limits_{n=1}^{N}L_m(X_{mn}-Y_{mn})^{2}} + $$ + + $$ + L_m = N_{lat}\dfrac{\cos(lat_m)}{\sum\limits_{j=1}^{N_{lat}}\cos(lat_j)} + $$ + + $lat_m$ is the latitude at m. + $N_{lat}$ is the number of latitude set by `num_lat`. + + Args: + num_lat (int): Number of latitude. + std (Optional[Union[np.array, Tuple[float, ...]]]): Standard Deviation of training dataset. Defaults to None. + keep_batch (bool, optional): Whether keep batch axis. Defaults to False. + variable_dict (Optional[Dict[str, int]]): Variable dictionary, the key is the name of a variable and + the value is its index. Defaults to None. + unlog (bool, optional): Whether calculate expm1 for all elements in the array. Defaults to False. + scale (float, optional): The scale value used after expm1. Defaults to 1e-5. + + Examples: + >>> import numpy as np + >>> import ppsci + >>> std = np.random.randn(20, 1, 1) + >>> metric = ppsci.metric.LatitudeWeightedRMSE(720, std=std) + """ + + def __init__( + self, + num_lat: int, + std: Optional[Union[np.array, Tuple[float, ...]]] = None, + keep_batch: bool = False, + variable_dict: Dict[str, int] = None, + unlog: bool = False, + scale: float = 1e-5, + ): + super().__init__(keep_batch) + self.num_lat = num_lat + self.std = ( + None + if std is None + else paddle.to_tensor(std, paddle.get_default_dtype()).reshape((1, -1)) + ) + self.variable_dict = variable_dict + self.unlog = unlog + self.scale = scale + self.weight = self.get_latitude_weight(num_lat) + + def get_latitude_weight(self, num_lat: int = 720): + lat_t = paddle.linspace(start=0, stop=1, num=num_lat) + lat_t = paddle.cos(3.1416 * (0.5 - lat_t)) + weight = num_lat * lat_t / paddle.sum(lat_t) + weight = weight.reshape((1, 1, -1, 1)) + return weight + + def scale_expm1(self, x: paddle.Tensor): + return self.scale * paddle.expm1(x) + + @paddle.no_grad() + def forward(self, output_dict, label_dict) -> Dict[str, "paddle.Tensor"]: + metric_dict = {} + for key in label_dict: + output = ( + self.scale_expm1(output_dict[key]) if self.unlog else output_dict[key] + ) + label = self.scale_expm1(label_dict[key]) if self.unlog else label_dict[key] + + mse = F.mse_loss(output, label, "none") + rmse = (mse * self.weight).mean(axis=(-1, -2)) ** 0.5 + if self.std is not None: + rmse = rmse * self.std + if self.variable_dict is not None: + for variable_name, idx in self.variable_dict.items(): + metric_dict[f"{key}.{variable_name}"] = ( + rmse[:, idx] if self.keep_batch else rmse[:, idx].mean() + ) + else: + metric_dict[key] = rmse.mean(axis=1) if self.keep_batch else rmse.mean() + + return metric_dict diff --git a/examples/smc_reac/ppsci/optimizer/__init__.py b/examples/smc_reac/ppsci/optimizer/__init__.py new file mode 100644 index 0000000000..c03b0717ee --- /dev/null +++ b/examples/smc_reac/ppsci/optimizer/__init__.py @@ -0,0 +1,84 @@ +# 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. + +import copy + +from ppsci.optimizer import lr_scheduler +from ppsci.optimizer.optimizer import LBFGS +from ppsci.optimizer.optimizer import SGD +from ppsci.optimizer.optimizer import SOAP +from ppsci.optimizer.optimizer import Adam +from ppsci.optimizer.optimizer import AdamW +from ppsci.optimizer.optimizer import Momentum +from ppsci.optimizer.optimizer import OptimizerList +from ppsci.optimizer.optimizer import RMSProp + +__all__ = [ + "LBFGS", + "SGD", + "Adam", + "AdamW", + "Momentum", + "RMSProp", + "OptimizerList", + "lr_scheduler", + "SOAP", +] + + +def build_lr_scheduler(cfg, epochs, iters_per_epoch): + """Build learning rate scheduler. + + Args: + cfg (DictConfig): Learning rate scheduler config. + epochs (int): Total epochs. + iters_per_epoch (int): Number of iterations of one epoch. + + Returns: + LRScheduler: Learning rate scheduler. + """ + cfg = copy.deepcopy(cfg) + cfg.update({"epochs": epochs, "iters_per_epoch": iters_per_epoch}) + lr_scheduler_cls = cfg.pop("name") + lr_scheduler_ = eval(lr_scheduler_cls)(**cfg) + return lr_scheduler_() + + +def build_optimizer(cfg, model_list, epochs, iters_per_epoch): + """Build optimizer and learning rate scheduler + + Args: + cfg (DictConfig): Learning rate scheduler config. + model_list (Tuple[nn.Layer, ...]): Tuple of model(s). + epochs (int): Total epochs. + iters_per_epoch (int): Number of iterations of one epoch. + + Returns: + Optimizer, LRScheduler: Optimizer and learning rate scheduler. + """ + # build lr_scheduler + cfg = copy.deepcopy(cfg) + lr_cfg = cfg.pop("lr") + if isinstance(lr_cfg, float): + lr_scheduler = lr_cfg + else: + lr_scheduler = build_lr_scheduler(lr_cfg, epochs, iters_per_epoch) + + # build optimizer + opt_cls = cfg.pop("name") + optimizer = eval(opt_cls)(learning_rate=lr_scheduler, **cfg)(model_list) + + if isinstance(lr_scheduler, float): + return optimizer, None + return optimizer, lr_scheduler diff --git a/examples/smc_reac/ppsci/optimizer/lr_scheduler.py b/examples/smc_reac/ppsci/optimizer/lr_scheduler.py new file mode 100644 index 0000000000..cad7da3503 --- /dev/null +++ b/examples/smc_reac/ppsci/optimizer/lr_scheduler.py @@ -0,0 +1,911 @@ +# 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. + +from __future__ import annotations + +import abc +import math +from typing import Callable +from typing import List +from typing import Tuple +from typing import Union + +from paddle.optimizer import lr + +from ppsci.utils import logger + +__all__ = [ + "Linear", + "Cosine", + "Step", + "Piecewise", + "MultiStepDecay", + "ExponentialDecay", + "CosineWarmRestarts", + "OneCycleLR", + "LambdaDecay", + "ReduceOnPlateau", +] + + +class LRBase: + """Base class for custom learning rates. + + Args: + epochs (int): Total epoch(s). + iters_per_epoch (int): Number of iterations within an epoch. + learning_rate (float): Learning rate. + warmup_epoch (int): Number of warmup epochs. + warmup_start_lr (float): Start learning rate within warmup. + last_epoch (int): Last epoch. + by_epoch (bool): Learning rate decays by epoch when by_epoch is True, else by iter. + verbose (bool): If True, prints a message to stdout for each update. Defaults to False. + """ + + def __init__( + self, + epochs: int, + iters_per_epoch: int, + learning_rate: float, + warmup_epoch: int, + warmup_start_lr: float, + last_epoch: int, + by_epoch: bool, + verbose: bool = False, + ) -> None: + """Initialize and record the necessary parameters.""" + super().__init__() + if warmup_epoch >= epochs: + msg = ( + "When using warm up, the value of 'Global.epochs' should be greater " + "than value of 'Optimizer.lr.warmup_epoch'. The value of " + f"'Optimizer.lr.warmup_epoch' has been set to {epochs}." + ) + logger.warning(msg) + warmup_epoch = epochs + self.epochs = epochs + self.iters_per_epoch = iters_per_epoch + self.learning_rate = learning_rate + self.warmup_epoch = warmup_epoch + self.warmup_steps = ( + self.warmup_epoch + if by_epoch + else round(self.warmup_epoch * self.iters_per_epoch) + ) + self.warmup_start_lr = warmup_start_lr + self.last_epoch = last_epoch + self.by_epoch = by_epoch + self.verbose = verbose + + @abc.abstractmethod + def __call__(self, *args, **kwargs) -> lr.LRScheduler: + """Generate an learning rate scheduler. + + Returns: + lr.LinearWarmup: learning rate scheduler. + """ + pass + + def linear_warmup( + self, learning_rate: Union[float, lr.LRScheduler] + ) -> lr.LinearWarmup: + """Add an Linear Warmup before learning_rate. + + Args: + learning_rate (Union[float, lr.LRScheduler]): Original learning rate without + warmup. + + Returns: + lr.LinearWarmup: learning rate scheduler with warmup. + """ + warmup_lr = lr.LinearWarmup( + learning_rate=learning_rate, + warmup_steps=self.warmup_steps, + start_lr=self.warmup_start_lr, + end_lr=self.learning_rate, + last_epoch=self.last_epoch, + verbose=self.verbose, + ) + return warmup_lr + + +class Constant(lr.LRScheduler): + """Constant learning rate Class implementation. + + Args: + learning_rate (float): The initial learning rate. + last_epoch (int, optional): The index of last epoch. Default: -1. + """ + + def __init__(self, learning_rate: float, last_epoch: int = -1): + self.learning_rate = learning_rate + self.last_epoch = last_epoch + super().__init__() + + def get_lr(self) -> float: + """Always return the same learning rate""" + return self.learning_rate + + +class Linear(LRBase): + """Linear learning rate decay. + + Args: + epochs (int): Total epoch(s). + iters_per_epoch (int): Number of iterations within an epoch. + learning_rate (float): Learning rate. + end_lr (float, optional): The minimum final learning rate. Defaults to 0.0. + power (float, optional): Power of polynomial. Defaults to 1.0. + cycle (bool, optional): Whether the learning rate rises again. If True, then the learning rate will rise when it decrease + to ``end_lr`` . If False, the learning rate is monotone decreasing. Defaults to False. + warmup_epoch (int): Number of warmup epochs. + warmup_start_lr (float): Start learning rate within warmup. + last_epoch (int): Last epoch. + by_epoch (bool): Learning rate decays by epoch when by_epoch is True, else by iter. + + Examples: + >>> import ppsci + >>> lr = ppsci.optimizer.lr_scheduler.Linear(10, 2, 0.001)() + """ + + def __init__( + self, + epochs: int, + iters_per_epoch: int, + learning_rate: float, + end_lr: float = 0.0, + power: float = 1.0, + cycle: bool = False, + warmup_epoch: int = 0, + warmup_start_lr: float = 0.0, + last_epoch: int = -1, + by_epoch: bool = False, + ): + super().__init__( + epochs, + iters_per_epoch, + learning_rate, + warmup_epoch, + warmup_start_lr, + last_epoch, + by_epoch, + ) + self.decay_steps = (epochs - self.warmup_epoch) * iters_per_epoch + self.end_lr = end_lr + self.power = power + self.cycle = cycle + self.warmup_steps = round(self.warmup_epoch * iters_per_epoch) + if self.by_epoch: + self.decay_steps = self.epochs - self.warmup_epoch + + def __call__(self): + learning_rate = ( + lr.PolynomialDecay( + learning_rate=self.learning_rate, + decay_steps=self.decay_steps, + end_lr=self.end_lr, + power=self.power, + cycle=self.cycle, + last_epoch=self.last_epoch, + ) + if self.decay_steps > 0 + else Constant(self.learning_rate) + ) + + if self.warmup_steps > 0: + learning_rate = self.linear_warmup(learning_rate) + + setattr(learning_rate, "by_epoch", self.by_epoch) + return learning_rate + + +class ExponentialDecay(LRBase): + """ExponentialDecay learning rate decay. + + Args: + epochs (int): Total epoch(s). + iters_per_epoch (int): Number of iterations within an epoch. + learning_rate (float): Learning rate. + gamma (float): The decay rate. + decay_steps (int): The number of steps to decay. + warmup_epoch (int): Number of warmup epochs. + warmup_start_lr (float): Start learning rate within warmup. + last_epoch (int): Last epoch. + by_epoch (bool): Learning rate decays by epoch when by_epoch is True, else by iter. + + Examples: + >>> import ppsci + >>> lr = ppsci.optimizer.lr_scheduler.ExponentialDecay(10, 2, 1e-3, 0.95, 3)() + """ + + def __init__( + self, + epochs: int, + iters_per_epoch: int, + learning_rate: float, + gamma: float, + decay_steps: int, + warmup_epoch: int = 0, + warmup_start_lr: float = 0.0, + last_epoch: int = -1, + by_epoch: bool = False, + ): + super().__init__( + epochs, + iters_per_epoch, + learning_rate, + warmup_epoch, + warmup_start_lr, + last_epoch, + by_epoch, + ) + self.decay_steps = decay_steps + self.gamma = gamma + self.warmup_steps = round(self.warmup_epoch * iters_per_epoch) + if self.by_epoch: + self.decay_steps /= iters_per_epoch + + def __call__(self): + learning_rate = lr.ExponentialDecay( + learning_rate=self.learning_rate, + gamma=self.gamma ** (1 / self.decay_steps), + last_epoch=self.last_epoch, + ) + + if self.warmup_steps > 0: + learning_rate = self.linear_warmup(learning_rate) + + setattr(learning_rate, "by_epoch", self.by_epoch) + return learning_rate + + +class Cosine(LRBase): + """Cosine learning rate decay. + + lr = 0.05 * (math.cos(epoch * (math.pi / epochs)) + 1) + + Args: + epochs (int): Total epoch(s). + iters_per_epoch (int): Number of iterations within an epoch. + learning_rate (float): Learning rate. + eta_min (float, optional): Minimum learning rate. Defaults to 0.0. + warmup_epoch (int, optional): The epoch numbers for LinearWarmup. Defaults to 0. + warmup_start_lr (float, optional): Start learning rate within warmup. Defaults to 0.0. + last_epoch (int, optional): Last epoch. Defaults to -1. + by_epoch (bool, optional): Learning rate decays by epoch when by_epoch is True, + else by iter. Defaults to False. + + Examples: + >>> import ppsci + >>> lr = ppsci.optimizer.lr_scheduler.Cosine(10, 2, 1e-3)() + """ + + def __init__( + self, + epochs: int, + iters_per_epoch: int, + learning_rate: float, + eta_min: float = 0.0, + warmup_epoch: int = 0, + warmup_start_lr: float = 0.0, + last_epoch: int = -1, + by_epoch: bool = False, + ): + super().__init__( + epochs, + iters_per_epoch, + learning_rate, + warmup_epoch, + warmup_start_lr, + last_epoch, + by_epoch, + ) + self.T_max = (self.epochs - self.warmup_epoch) * self.iters_per_epoch + self.eta_min = eta_min + if self.by_epoch: + self.T_max = self.epochs - self.warmup_epoch + + def __call__(self): + learning_rate = ( + lr.CosineAnnealingDecay( + learning_rate=self.learning_rate, + T_max=self.T_max, + eta_min=self.eta_min, + last_epoch=self.last_epoch, + ) + if self.T_max > 0 + else Constant(self.learning_rate) + ) + + if self.warmup_steps > 0: + learning_rate = self.linear_warmup(learning_rate) + + setattr(learning_rate, "by_epoch", self.by_epoch) + return learning_rate + + +class Step(LRBase): + """Step learning rate decay. + + Args: + epochs (int): Total epoch(s). + iters_per_epoch (int): Number of iterations within an epoch. + learning_rate (float): Learning rate. + step_size (int): The interval to update. + gamma (float, optional): The Ratio that the learning rate will be reduced. + ``new_lr = origin_lr * gamma``. It should be less than 1.0. Default: 0.1. + warmup_epoch (int, optional): The epoch numbers for LinearWarmup. Defaults to 0. + warmup_start_lr (float, optional): Start learning rate within warmup. Defaults to 0.0. + last_epoch (int, optional): Last epoch. Defaults to -1. + by_epoch (bool, optional): Learning rate decays by epoch when by_epoch is True, + else by iter. Defaults to False. + + Examples: + >>> import ppsci + >>> lr = ppsci.optimizer.lr_scheduler.Step(10, 1, 1e-3, 2, 0.95)() + """ + + def __init__( + self, + epochs: int, + iters_per_epoch: int, + learning_rate: float, + step_size: int, + gamma: float, + warmup_epoch: int = 0, + warmup_start_lr: float = 0.0, + last_epoch: int = -1, + by_epoch: bool = False, + ): + super().__init__( + epochs, + iters_per_epoch, + learning_rate, + warmup_epoch, + warmup_start_lr, + last_epoch, + by_epoch, + ) + self.step_size = step_size * iters_per_epoch + self.gamma = gamma + if self.by_epoch: + self.step_size = step_size + + def __call__(self): + learning_rate = lr.StepDecay( + learning_rate=self.learning_rate, + step_size=self.step_size, + gamma=self.gamma, + last_epoch=self.last_epoch, + ) + + if self.warmup_steps > 0: + learning_rate = self.linear_warmup(learning_rate) + + setattr(learning_rate, "by_epoch", self.by_epoch) + return learning_rate + + +class Piecewise(LRBase): + """Piecewise learning rate decay + + Args: + epochs (int): Total epoch(s) + iters_per_epoch (int): Number of iterations within an epoch + decay_epochs (Tuple[int, ...]): A list of steps numbers. The type of element in the + list is python int. + values (Tuple[float, ...]): Tuple of learning rate values that will be picked during + different epoch boundaries. + warmup_epoch (int, optional): The epoch numbers for LinearWarmup. Defaults to 0. + warmup_start_lr (float, optional): Start learning rate within warmup. Defaults to 0.0. + last_epoch (int, optional): Last epoch. Defaults to -1. + by_epoch (bool, optional): Learning rate decays by epoch when by_epoch is True, + else by iter. Defaults to False. + + Examples: + >>> import ppsci + >>> lr = ppsci.optimizer.lr_scheduler.Piecewise( + ... 10, 1, [2, 4], (1e-3, 1e-4, 1e-5) + ... )() + """ + + def __init__( + self, + epochs: int, + iters_per_epoch: int, + decay_epochs: Tuple[int, ...], + values: Tuple[float, ...], + warmup_epoch: int = 0, + warmup_start_lr: float = 0.0, + last_epoch: int = -1, + by_epoch: bool = False, + ): + super().__init__( + epochs, + iters_per_epoch, + values[0], + warmup_epoch, + warmup_start_lr, + last_epoch, + by_epoch, + ) + self.values = values + self.boundaries_steps = [e * iters_per_epoch for e in decay_epochs] + if self.by_epoch is True: + self.boundaries_steps = decay_epochs + + def __call__(self): + learning_rate = lr.PiecewiseDecay( + boundaries=self.boundaries_steps, + values=self.values, + last_epoch=self.last_epoch, + ) + + if self.warmup_steps > 0: + learning_rate = self.linear_warmup(learning_rate) + + setattr(learning_rate, "by_epoch", self.by_epoch) + return learning_rate + + +class MultiStepDecay(LRBase): + """MultiStepDecay learning rate decay + + Args: + epochs (int): Total epoch(s) + iters_per_epoch (int): Number of iterations within an epoch + learning_rate (float): Learning rate + milestones (Tuple[int, ...]): Tuple of each boundaries. should be increasing. + gamma (float, optional): The Ratio that the learning rate will be reduced. + `new_lr = origin_lr * gamma`. It should be less than 1.0. Defaults to 0.1. + warmup_epoch (int, optional): The epoch numbers for LinearWarmup. Defaults to 0. + warmup_start_lr (float, optional): Start learning rate within warmup. Defaults to 0.0. + last_epoch (int, optional): Last epoch. Defaults to -1. + by_epoch (bool, optional): Learning rate decays by epoch when by_epoch is True, + else by iter. Defaults to False. + + Examples: + >>> import ppsci + >>> lr = ppsci.optimizer.lr_scheduler.MultiStepDecay(10, 1, 1e-3, (4, 5))() + """ + + def __init__( + self, + epochs: int, + iters_per_epoch: int, + learning_rate: float, + milestones: Tuple[int, ...], + gamma: float = 0.1, + warmup_epoch: int = 0, + warmup_start_lr: float = 0.0, + last_epoch: int = -1, + by_epoch: bool = False, + ): + super().__init__( + epochs, + iters_per_epoch, + learning_rate, + warmup_epoch, + warmup_start_lr, + last_epoch, + by_epoch, + ) + self.milestones = [x * iters_per_epoch for x in milestones] + self.gamma = gamma + if self.by_epoch: + self.milestones = milestones + + def __call__(self): + learning_rate = lr.MultiStepDecay( + learning_rate=self.learning_rate, + milestones=self.milestones, + gamma=self.gamma, + last_epoch=self.last_epoch, + ) + + if self.warmup_steps > 0: + learning_rate = self.linear_warmup(learning_rate) + + setattr(learning_rate, "by_epoch", self.by_epoch) + return learning_rate + + +class CosineAnnealingWarmRestarts(lr.LRScheduler): + """The implementation of cosine annealing schedule with warm restarts. + + Args: + learning_rate (float): Learning rate + T_0 (int): Number of iterations for the first restart. + T_mult (int, optional): A factor increases T_i after a restart. Defaults to 1. + eta_min (float, optional): Minimum learning rate. Defaults to 0. + last_epoch (int, optional): The index of last epoch. Defaults to -1. + verbose (bool, optional): If `True`, prints a message to stdout for each update. Defaults to False. + """ + + def __init__( + self, + learning_rate: float, + T_0: int, + T_mult: int = 1, + eta_min: float = 0.0, + last_epoch: int = -1, + verbose: bool = False, + ): + if T_0 <= 0 or not isinstance(T_0, int): + raise ValueError(f"Expected positive integer T_0, but got {T_0}") + if T_mult < 1 or not isinstance(T_mult, int): + raise ValueError(f"Expected integer T_mult >= 1, but got {T_mult}") + self.T_0 = T_0 + self.T_i = T_0 + self.T_mult = T_mult + self.eta_min = eta_min + self.T_cur = last_epoch + super().__init__(learning_rate, last_epoch, verbose) + + def get_lr(self): + return ( + self.eta_min + + (self.base_lr - self.eta_min) + * (1 + math.cos(math.pi * self.T_cur / self.T_i)) + / 2 + ) + + def step(self, epoch=None): + if epoch is None and self.last_epoch < 0: + epoch = 0 + + if epoch is None: + epoch = self.last_epoch + 1 + self.T_cur = self.T_cur + 1 + if self.T_cur >= self.T_i: + self.T_cur = self.T_cur - self.T_i + self.T_i = self.T_i * self.T_mult + else: + if epoch < 0: + raise ValueError(f"Expected non-negative epoch, but got {epoch}") + if epoch >= self.T_0: + if self.T_mult == 1: + self.T_cur = epoch % self.T_0 + else: + n = int( + math.log( + (epoch / self.T_0 * (self.T_mult - 1) + 1), self.T_mult + ) + ) + self.T_cur = epoch - self.T_0 * (self.T_mult**n - 1) / ( + self.T_mult - 1 + ) + self.T_i = self.T_0 * self.T_mult ** (n) + else: + self.T_i = self.T_0 + self.T_cur = epoch + self.last_epoch = math.floor(epoch) + self.last_lr = self.get_lr() + + +class CosineWarmRestarts(LRBase): + """Set the learning rate using a cosine annealing schedule with warm restarts. + + Args: + epochs (int): Total epoch(s) + iters_per_epoch (int): Number of iterations within an epoch + learning_rate (float): Learning rate + T_0 (int): Number of iterations for the first restart. + T_mult (int): A factor increases T_i after a restart + eta_min (float, optional): Minimum learning rate. Defaults to 0.0. + warmup_epoch (int, optional): The epoch numbers for LinearWarmup. Defaults to 0. + warmup_start_lr (float, optional): Start learning rate within warmup. Defaults to 0.0. + last_epoch (int, optional): Last epoch. Defaults to -1. + by_epoch (bool, optional): Learning rate decays by epoch when by_epoch is True, else by iter. Defaults to False. + + Examples: + >>> import ppsci + >>> lr = ppsci.optimizer.lr_scheduler.CosineWarmRestarts(20, 1, 1e-3, 14, 2)() + """ + + def __init__( + self, + epochs: int, + iters_per_epoch: int, + learning_rate: float, + T_0: int, + T_mult: int, + eta_min: float = 0.0, + warmup_epoch: int = 0, + warmup_start_lr: float = 0.0, + last_epoch: int = -1, + by_epoch: bool = False, + ): + super().__init__( + epochs, + iters_per_epoch, + learning_rate, + warmup_epoch, + warmup_start_lr, + last_epoch, + by_epoch, + ) + self.T_0 = T_0 + self.T_mult = T_mult + self.eta_min = eta_min + if self.by_epoch is False: + self.T_0 = T_0 * iters_per_epoch + + def __call__(self): + learning_rate = CosineAnnealingWarmRestarts( + learning_rate=self.learning_rate, + T_0=self.T_0, + T_mult=self.T_mult, + eta_min=self.eta_min, + last_epoch=self.last_epoch, + verbose=self.verbose, + ) + + if self.warmup_steps > 0: + learning_rate = self.linear_warmup(learning_rate) + + setattr(learning_rate, "by_epoch", self.by_epoch) + return learning_rate + + +class OneCycleLR(LRBase): + """Sets the learning rate according to the one cycle learning rate scheduler. + The scheduler adjusts the learning rate from an initial learning rate to the maximum learning rate and then + from that maximum learning rate to the minimum learning rate, which is much less than the initial learning rate. + + It has been proposed in [Super-Convergence: Very Fast Training of Neural Networks Using Large Learning Rates](https://arxiv.org/abs/1708.07120). + + Please note that the default behavior of this scheduler follows the fastai implementation of one cycle, + which claims that **"unpublished work has shown even better results by using only two phases"**. + If you want the behavior of this scheduler to be consistent with the paper, please set `three_phase=True`. + + Args: + epochs (int): Total epoch(s). + iters_per_epoch (int): Number of iterations within an epoch. + max_learning_rate (float): The maximum learning rate. It is a python float number. Functionally, it defines the initial learning rate by `divide_factor` . + divide_factor (float, optional): Initial learning rate will be determined by initial_learning_rate = max_learning_rate / divide_factor. Defaults to 25.0. + end_learning_rate (float, optional): The minimum learning rate during training, it should be much less than initial learning rate. Defaults to 0.0001. + phase_pct (float): The percentage of total steps which used to increasing learning rate. Defaults to 0.3. + anneal_strategy (str, optional): Strategy of adjusting learning rate. "cos" for cosine annealing, "linear" for linear annealing. Defaults to "cos". + three_phase (bool, optional): Whether to use three phase. Defaults to False. + warmup_epoch (int, optional): The epoch numbers for LinearWarmup. Defaults to 0. + warmup_start_lr (float, optional): Start learning rate within warmup. Defaults to 0.0. + last_epoch (int, optional): Last epoch. Defaults to -1. + by_epoch (bool, optional): Learning rate decays by epoch when by_epoch is True, else by iter. Defaults to False. + + Examples: + >>> import ppsci + >>> lr = ppsci.optimizer.lr_scheduler.OneCycleLR(100, 1, 1e-3)() + """ + + def __init__( + self, + epochs: int, + iters_per_epoch: int, + max_learning_rate: float, + divide_factor: float = 25.0, + end_learning_rate: float = 0.0001, + phase_pct: float = 0.3, + anneal_strategy: str = "cos", + three_phase: bool = False, + warmup_epoch: int = 0, + warmup_start_lr: float = 0.0, + last_epoch: int = -1, + by_epoch: bool = False, + ): + super().__init__( + epochs, + iters_per_epoch, + max_learning_rate, + warmup_epoch, + warmup_start_lr, + last_epoch, + by_epoch, + ) + self.total_steps = epochs + if not by_epoch: + self.total_steps *= iters_per_epoch + self.divide_factor = divide_factor + self.end_learning_rate = end_learning_rate + self.phase_pct = phase_pct + self.anneal_strategy = anneal_strategy + self.three_phase = three_phase + + def __call__(self): + learning_rate = lr.OneCycleLR( + max_learning_rate=self.learning_rate, + total_steps=self.total_steps, + divide_factor=self.divide_factor, + end_learning_rate=self.end_learning_rate, + phase_pct=self.phase_pct, + anneal_strategy=self.anneal_strategy, + three_phase=self.three_phase, + last_epoch=self.last_epoch, + verbose=self.verbose, + ) + + if self.warmup_steps > 0: + learning_rate = self.linear_warmup(learning_rate) + + setattr(learning_rate, "by_epoch", self.by_epoch) + return learning_rate + + +class LambdaDecay(LRBase): + """This interface provides a lambda function to set the learning rate strategy. + + Args: + epochs (int): Total epoch(s). + iters_per_epoch (int): Number of iterations within an epoch. + learning_rate (float): Learning rate. + lr_lambda (Callable): A lambda function that calculates a factor through epoch, which is multiplied by the initial learning rate. + warmup_epoch (int, optional): The epoch numbers for LinearWarmup. Defaults to 0. + warmup_start_lr (float, optional): Start learning rate within warmup. Defaults to 0.0. + last_epoch (int, optional): Last epoch. Defaults to -1. + by_epoch (bool, optional): Learning rate decays by epoch when by_epoch is True, + else by iter. Defaults to False. + verbose (bool, optional): If True, prints a message to stdout for each update. Defaults to False. + + Examples: + >>> import ppsci + >>> lr = ppsci.optimizer.lr_scheduler.LambdaDecay(0.5, lr_lambda=lambda x:0.95**x, verbose=True)() + """ + + def __init__( + self, + epochs: int, + iters_per_epoch: int, + learning_rate: float, + lr_lambda: Callable, + warmup_epoch: int = 0, + warmup_start_lr: float = 0.0, + last_epoch: int = -1, + by_epoch: bool = False, + verbose: bool = False, + ): + super().__init__( + epochs, + iters_per_epoch, + learning_rate, + warmup_epoch, + warmup_start_lr, + last_epoch, + by_epoch, + verbose, + ) + self.learning_rate = learning_rate + self.lr_lambda = lr_lambda + self.last_epoch = last_epoch + self.verbose = verbose + self.by_epoch = by_epoch + + def __call__(self): + learning_rate = lr.LambdaDecay( + learning_rate=self.learning_rate, + lr_lambda=self.lr_lambda, + last_epoch=self.last_epoch, + verbose=self.verbose, + ) + + if self.warmup_steps > 0: + learning_rate = self.linear_warmup(learning_rate) + + setattr(learning_rate, "by_epoch", self.by_epoch) + return learning_rate + + +class ReduceOnPlateau(LRBase): + """This interface provides a learning rate scheduler that reduces the learning rate when a metric has stopped improving. + + Args: + epochs (int): Total epoch(s). + iters_per_epoch (int): Number of iterations within an epoch. + learning_rate (float): Initial learning rate. + last_epoch (int, optional): The index of the last epoch. Defaults to -1. + warmup_epoch (int, optional): The epoch numbers for LinearWarmup. Defaults to 0. + warmup_start_lr (float, optional): Start learning rate within warmup. Defaults to 0.0. + mode (str, optional): One of `min` or `max`. In `min` mode, lr will be reduced when the quantity monitored has stopped decreasing; in `max` mode it will be reduced when the quantity monitored has stopped increasing. Defaults to "min". + patience (int, optional): Number of epochs with no improvement after which learning rate will be reduced. Defaults to 20. + factor (float, optional): Factor by which the learning rate will be reduced. new_lr = lr * factor. Defaults to 1e-4. + verbose (bool, optional): If True, prints a message to stdout for each update. Defaults to True. + by_epoch (bool, optional): Learning rate decays by epoch when by_epoch is True, else by iter. Defaults to True. + + Examples: + >>> import ppsci + >>> lr = ppsci.optimizer.lr_scheduler.ReduceOnPlateau(epochs=50, iters_per_epoch=100, learning_rate=0.1, mode='min', patience=10, factor=0.5, verbose=True)() + """ + + def __init__( + self, + epochs: int, + iters_per_epoch: int, + learning_rate: float, + last_epoch: int = -1, + warmup_epoch: int = 0, + warmup_start_lr: float = 0.0, + mode: str = "min", + patience: int = 20, + factor: float = 1e-4, + verbose: bool = True, + by_epoch: bool = True, + ): + super().__init__( + epochs, + iters_per_epoch, + learning_rate, + warmup_epoch, + warmup_start_lr, + last_epoch, + by_epoch, + ) + self.mode = mode + self.patience = patience + self.factor = factor + self.verbose = verbose + self.learning_rate = learning_rate + self.by_epoch = by_epoch + + def __call__(self): + learning_rate = lr.ReduceOnPlateau( + mode=self.mode, + patience=self.patience, + factor=self.factor, + verbose=self.verbose, + learning_rate=self.learning_rate, + ) + + if self.warmup_steps > 0: + learning_rate = self.linear_warmup(learning_rate) + + setattr(learning_rate, "by_epoch", self.by_epoch) + return learning_rate + + +class SchedulerList: + """SchedulerList which wrap more than one scheduler. + + Args: + scheduler_list (Tuple[lr.LRScheduler, ...]): Schedulers listed in a tuple. + + Examples: + >>> import ppsci + >>> sch1 = ppsci.optimizer.lr_scheduler.Linear(10, 2, 0.001)() + >>> sch2 = ppsci.optimizer.lr_scheduler.ExponentialDecay(10, 2, 1e-3, 0.95, 3)() + >>> sch = ppsci.optimizer.lr_scheduler.SchedulerList((sch1, sch2)) + """ + + def __init__(self, scheduler_list: Tuple[lr.LRScheduler, ...]): + super().__init__() + self._sch_list = scheduler_list + self.by_epoch = False + + def step(self): + for sch in self._sch_list: + sch.step() + + def get_lr(self) -> float: + """Return learning rate of first scheduler""" + return self._sch_list[0].get_lr() + + def _state_keys(self) -> List[str]: + return ["last_epoch", "last_lr"] + + def __len__(self) -> int: + return len(self._sch_list) + + def __getitem__(self, idx): + return self._sch_list[idx] + + def __setitem__(self, idx, sch): + raise NotImplementedError("Can not modify any item in SchedulerList.") diff --git a/examples/smc_reac/ppsci/optimizer/optimizer.py b/examples/smc_reac/ppsci/optimizer/optimizer.py new file mode 100644 index 0000000000..fd4b3d447a --- /dev/null +++ b/examples/smc_reac/ppsci/optimizer/optimizer.py @@ -0,0 +1,649 @@ +# 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. + +from __future__ import annotations + +from typing import TYPE_CHECKING +from typing import Dict +from typing import List +from typing import Optional +from typing import Tuple +from typing import Union + +from paddle import nn +from paddle import optimizer as optim +from paddle import regularizer +from paddle.incubate import optimizer as incubate_optim +from typing_extensions import Literal + +from ppsci.optimizer.soap import SOAP as SOAP_impl +from ppsci.utils import logger +from ppsci.utils import misc + +if TYPE_CHECKING: + import paddle + +__all__ = ["SGD", "Momentum", "Adam", "RMSProp", "AdamW", "LBFGS", "OptimizerList"] + + +class SGD: + """Stochastic Gradient Descent. + + Args: + learning_rate (Union[float, optim.lr.LRScheduler], optional): The learning rate + used to update parameter(s). Defaults to 0.001. + weight_decay (Optional[Union[float, regularizer.L1Decay, regularizer.L2Decay]]): + Regularization strategy. Defaults to None. + grad_clip (Optional[Union[nn.ClipGradByNorm, nn.ClipGradByValue, nn.ClipGradByGlobalNorm]]): + Gradient clipping strategy. Defaults to None. + + Examples: + >>> import ppsci + >>> model = ppsci.arch.MLP(("x",), ("u",), 5, 20) + >>> opt = ppsci.optimizer.SGD(1e-3)(model) + """ + + def __init__( + self, + learning_rate: Union[float, optim.lr.LRScheduler] = 0.001, + weight_decay: Optional[ + Union[float, regularizer.L1Decay, regularizer.L2Decay] + ] = None, + grad_clip: Optional[ + Union[nn.ClipGradByNorm, nn.ClipGradByValue, nn.ClipGradByGlobalNorm] + ] = None, + ): + self.learning_rate = learning_rate + self.weight_decay = weight_decay + self.grad_clip = grad_clip + + def __call__(self, model_list: Union[nn.Layer, Tuple[nn.Layer, ...]]): + # model_list is None in static graph + if not isinstance(model_list, (tuple, list)): + model_list = (model_list,) + parameters = ( + sum([m.parameters() for m in model_list], []) if model_list else None + ) + opt = optim.SGD( + learning_rate=self.learning_rate, + parameters=parameters, + weight_decay=self.weight_decay, + grad_clip=self.grad_clip, + ) + return opt + + +class Momentum: + """Simple Momentum optimizer with velocity state. + + Args: + learning_rate (Union[float, optim.lr.LRScheduler]): The learning rate + used to update parameter(s). + momentum (float): Momentum factor. + weight_decay (Optional[Union[float, regularizer.L1Decay, regularizer.L2Decay]]): + Regularization strategy. Defaults to None. + grad_clip (Optional[Union[nn.ClipGradByNorm, nn.ClipGradByValue, nn.ClipGradByGlobalNorm]]): + Gradient clipping strategy. Defaults to None. + use_nesterov (bool, optional): Whether to use nesterov momentum. Defaults to False. + no_weight_decay_name (Optional[str]): List of names of no weight decay parameters split by white space. Defaults to None. + + Examples: + >>> import ppsci + >>> model = ppsci.arch.MLP(("x",), ("u",), 5, 20) + >>> opt = ppsci.optimizer.Momentum(1e-3, 0.9)(model) + """ + + def __init__( + self, + learning_rate: Union[float, optim.lr.LRScheduler], + momentum: float, + weight_decay: Optional[ + Union[float, regularizer.L1Decay, regularizer.L2Decay] + ] = None, + grad_clip: Optional[ + Union[nn.ClipGradByNorm, nn.ClipGradByValue, nn.ClipGradByGlobalNorm] + ] = None, + use_nesterov: bool = False, + no_weight_decay_name: Optional[str] = None, + ): + super().__init__() + self.learning_rate = learning_rate + self.momentum = momentum + self.weight_decay = weight_decay + self.grad_clip = grad_clip + self.use_nesterov = use_nesterov + self.no_weight_decay_name_list = ( + no_weight_decay_name.split() if no_weight_decay_name else [] + ) + + def __call__(self, model_list: Union[nn.Layer, Tuple[nn.Layer, ...]]): + # model_list is None in static graph + if not isinstance(model_list, (tuple, list)): + model_list = (model_list,) + parameters = None + if len(self.no_weight_decay_name_list) > 0: + params_with_decay = [] + params_without_decay = [] + for m in model_list: + params = [ + p + for n, p in m.named_parameters() + if not any(nd in n for nd in self.no_weight_decay_name_list) + ] + params_with_decay.extend(params) + params = [ + p + for n, p in m.named_parameters() + if any(nd in n for nd in self.no_weight_decay_name_list) + ] + params_without_decay.extend(params) + parameters = [ + {"params": params_with_decay, "weight_decay": self.weight_decay}, + {"params": params_without_decay, "weight_decay": 0.0}, + ] + else: + parameters = ( + sum([m.parameters() for m in model_list], []) if model_list else None + ) + opt = optim.Momentum( + learning_rate=self.learning_rate, + momentum=self.momentum, + weight_decay=self.weight_decay, + grad_clip=self.grad_clip, + use_nesterov=self.use_nesterov, + parameters=parameters, + ) + if hasattr(opt, "_use_multi_tensor"): + opt = optim.Momentum( + learning_rate=self.learning_rate, + momentum=self.momentum, + weight_decay=self.weight_decay, + grad_clip=self.grad_clip, + parameters=parameters, + use_nesterov=self.use_nesterov, + use_multi_tensor=True, + ) + return opt + + +class Adam: + """Adam: A Method for Stochastic Optimization. + + Args: + learning_rate (Union[float, optim.lr.LRScheduler], optional): The learning rate + used to update parameter(s). Defaults to 0.001. + beta1 (float, optional): The exponential decay rate for the 1st moment estimates. Defaults to 0.9. + beta2 (float, optional): The exponential decay rate for the 2nd moment estimates. Defaults to 0.999. + epsilon (float, optional): A small float value for numerical stability. Defaults to 1e-08. + weight_decay (Optional[Union[float, regularizer.L1Decay, regularizer.L2Decay]]): Regularization strategy. Defaults to None. + grad_clip (Optional[Union[nn.ClipGradByNorm, nn.ClipGradByValue, nn.ClipGradByGlobalNorm]]): Gradient clipping strategy. Defaults to None. + lazy_mode (bool, optional): Whether to enable lazy mode for moving-average. Defaults to False. + amsgrad (bool, optional): Whether to use the AMSGrad variant of this algorithm from the paper + `On the Convergence of Adam and Beyond `_. Defaults to False. + + Examples: + >>> import ppsci + >>> model = ppsci.arch.MLP(("x",), ("u",), 5, 20) + >>> opt = ppsci.optimizer.Adam(1e-3)(model) + """ + + def __init__( + self, + learning_rate: Union[float, optim.lr.LRScheduler] = 0.001, + beta1: float = 0.9, + beta2: float = 0.999, + epsilon: float = 1e-08, + weight_decay: Optional[ + Union[float, regularizer.L1Decay, regularizer.L2Decay] + ] = None, + grad_clip: Optional[ + Union[nn.ClipGradByNorm, nn.ClipGradByValue, nn.ClipGradByGlobalNorm] + ] = None, + lazy_mode: bool = False, + amsgrad: bool = False, + ): + self.learning_rate = learning_rate + self.beta1 = beta1 + self.beta2 = beta2 + self.epsilon = epsilon + self.learning_rate = learning_rate + self.weight_decay = weight_decay + self.grad_clip = grad_clip + self.lazy_mode = lazy_mode + self.amsgrad = amsgrad + + def __call__(self, model_list: Union[nn.Layer, Tuple[nn.Layer, ...]]): + # model_list is None in static graph + if not isinstance(model_list, (tuple, list)): + model_list = (model_list,) + parameters = ( + sum([m.parameters() for m in model_list], []) if model_list else None + ) + import inspect + + extra_kwargs = {} + if "amsgrad" in inspect.signature(optim.Adam.__init__).parameters: + extra_kwargs["amsgrad"] = self.amsgrad + opt = optim.Adam( + learning_rate=self.learning_rate, + beta1=self.beta1, + beta2=self.beta2, + epsilon=self.epsilon, + weight_decay=self.weight_decay, + grad_clip=self.grad_clip, + lazy_mode=self.lazy_mode, + parameters=parameters, + **extra_kwargs, + ) + return opt + + +class LBFGS: + """The L-BFGS is a quasi-Newton method for solving an unconstrained optimization + problem over a differentiable function. Closely related is the Newton method for minimization. + + Args: + learning_rate (float, optional): The learning rate + used to update parameter(s). Defaults to 1.0. + max_iter (int, optional): Maximal number of iterations per optimization step. + Defaults to 1. + max_eval (Optional[int]): Maximal number of function evaluations per + optimization step. Defaults to None. + tolerance_grad (float, optional): Termination tolerance on first order optimality. + Defaults to 1e-07. + tolerance_change (float, optional): Termination tolerance on function + value/parameter changes. Defaults to 1e-09. + history_size (int, optional): Update history size. Defaults to 100. + line_search_fn (Optional[Literal["strong_wolfe"]]): Either 'strong_wolfe' or None. + Defaults to "strong_wolfe". + + Examples: + >>> import ppsci + >>> model = ppsci.arch.MLP(("x",), ("u",), 5, 20) + >>> opt = ppsci.optimizer.LBFGS(1e-3)(model) + """ + + def __init__( + self, + learning_rate: float = 1.0, + max_iter: int = 1, + max_eval: Optional[int] = None, + tolerance_grad: float = 1e-07, + tolerance_change: float = 1e-09, + history_size: int = 100, + line_search_fn: Optional[Literal["strong_wolfe"]] = "strong_wolfe", + ): + self.lr = learning_rate + self.max_iter = max_iter + self.max_eval = max_eval + self.tolerance_grad = tolerance_grad + self.tolerance_change = tolerance_change + self.history_size = history_size + self.line_search_fn = line_search_fn + + def __call__(self, model_list: Union[nn.Layer, Tuple[nn.Layer, ...]]): + # model_list is None in static graph + if not isinstance(model_list, (tuple, list)): + model_list = (model_list,) + parameters = ( + sum([m.parameters() for m in model_list], []) if model_list else None + ) + try: + opt = getattr(optim, "LBFGS")( + learning_rate=self.lr, + max_iter=self.max_iter, + max_eval=self.max_eval, + tolerance_grad=self.tolerance_grad, + tolerance_change=self.tolerance_change, + history_size=self.history_size, + line_search_fn=self.line_search_fn, + parameters=parameters, + ) + except AttributeError: + opt = getattr(incubate_optim, "LBFGS")( + learning_rate=self.lr, + max_iter=self.max_iter, + max_eval=self.max_eval, + tolerance_grad=self.tolerance_grad, + tolerance_change=self.tolerance_change, + history_size=self.history_size, + line_search_fn=self.line_search_fn, + parameters=parameters, + ) + return opt + + +class RMSProp: + """Root Mean Squared Propagation (RMSProp) is an unpublished, adaptive learning rate method. + + Args: + learning_rate (Union[float, optim.lr.LRScheduler]): The learning rate + used to update parameter(s) + rho (float, optional): Factor ρ in equation. Defaults to 0.95. + epsilon (float, optional): Factor ϵ in equation as a smoothing term. Defaults to 1e-6. + momentum (float, optional):β in equation is the momentum term. Defaults to 0.0. + weight_decay (Optional[Union[float, regularizer.L1Decay, regularizer.L2Decay]]): + Regularization strategy. Defaults to None. + grad_clip (Optional[Union[nn.ClipGradByNorm, nn.ClipGradByValue, nn.ClipGradByGlobalNorm]]): + Gradient clipping strategy. Defaults to None. + + Examples: + >>> import ppsci + >>> model = ppsci.arch.MLP(("x",), ("u",), 5, 20) + >>> opt = ppsci.optimizer.RMSProp(1e-3)(model) + """ + + def __init__( + self, + learning_rate: Union[float, optim.lr.LRScheduler], + rho: float = 0.95, + epsilon: float = 1e-6, + momentum: float = 0.0, + weight_decay: Optional[ + Union[float, regularizer.L1Decay, regularizer.L2Decay] + ] = None, + grad_clip: Optional[ + Union[nn.ClipGradByNorm, nn.ClipGradByValue, nn.ClipGradByGlobalNorm] + ] = None, + ): + super().__init__() + self.learning_rate = learning_rate + self.momentum = momentum + self.rho = rho + self.epsilon = epsilon + self.weight_decay = weight_decay + self.grad_clip = grad_clip + + def __call__(self, model_list: Union[nn.Layer, Tuple[nn.Layer, ...]]): + # model_list is None in static graph + if not isinstance(model_list, (tuple, list)): + model_list = (model_list,) + parameters = ( + sum([m.parameters() for m in model_list], []) if model_list else None + ) + opt = optim.RMSProp( + learning_rate=self.learning_rate, + momentum=self.momentum, + rho=self.rho, + epsilon=self.epsilon, + weight_decay=self.weight_decay, + grad_clip=self.grad_clip, + parameters=parameters, + ) + return opt + + +class AdamW: + """AdamW is implemented based on DECOUPLED WEIGHT DECAY REGULARIZATION. + + Args: + learning_rate (Union[float, optim.lr.LRScheduler], optional): The learning rate + used to update parameter(s). Defaults to 0.001. + beta1 (float, optional): The exponential decay rate for the 1st moment estimates. Defaults to 0.9. + beta2 (float, optional): The exponential decay rate for the 2nd moment estimates. Defaults to 0.999. + epsilon (float, optional): A small float value for numerical stability. Defaults to 1e-8. + weight_decay (float, optional): Regularization coefficient. Defaults to 0.01. + grad_clip (Optional[Union[nn.ClipGradByNorm, nn.ClipGradByValue, nn.ClipGradByGlobalNorm]]): Gradient clipping strategy. Defaults to None. + no_weight_decay_name (Optional[str]): List of names of no weight decay parameters split by white space. Defaults to None. + one_dim_param_no_weight_decay (bool, optional): Apply no weight decay on 1-D parameter(s). Defaults to False. + amsgrad (bool, optional): Whether to use the AMSGrad variant of this algorithm from the paper + `On the Convergence of Adam and Beyond `_. Defaults to False. + + Examples: + >>> import ppsci + >>> model = ppsci.arch.MLP(("x",), ("u",), 5, 20) + >>> opt = ppsci.optimizer.AdamW(1e-3)(model) + """ + + def __init__( + self, + learning_rate: Union[float, optim.lr.LRScheduler] = 0.001, + beta1: float = 0.9, + beta2: float = 0.999, + epsilon: float = 1e-8, + weight_decay: float = 0.001, + grad_clip: Optional[ + Union[nn.ClipGradByNorm, nn.ClipGradByValue, nn.ClipGradByGlobalNorm] + ] = None, + no_weight_decay_name: Optional[str] = None, + one_dim_param_no_weight_decay: bool = False, + amsgrad: bool = False, + ): + super().__init__() + self.learning_rate = learning_rate + self.beta1 = beta1 + self.beta2 = beta2 + self.epsilon = epsilon + self.grad_clip = grad_clip + self.weight_decay = weight_decay + self.no_weight_decay_name_list = ( + no_weight_decay_name.split() if no_weight_decay_name else [] + ) + self.one_dim_param_no_weight_decay = one_dim_param_no_weight_decay + self.amsgrad = amsgrad + + def __call__(self, model_list: Union[nn.Layer, Tuple[nn.Layer, ...]]): + # model_list is None in static graph + if not isinstance(model_list, (tuple, list)): + model_list = (model_list,) + parameters = ( + sum([m.parameters() for m in model_list], []) if model_list else None + ) + + # TODO(gaotingquan): Model_list is None when in static graph, "no_weight_decay" not work. + if model_list is None: + if ( + self.one_dim_param_no_weight_decay + or len(self.no_weight_decay_name_list) != 0 + ): + msg = '"AdamW" does not support setting "no_weight_decay" in static graph. Please use dynamic graph.' + logger.error(Exception(msg)) + raise Exception(msg) + + self.no_weight_decay_param_name_list = ( + [ + p.name + for model in model_list + for n, p in model.named_parameters() + if any(nd in n for nd in self.no_weight_decay_name_list) + ] + if model_list + else [] + ) + + if self.one_dim_param_no_weight_decay: + self.no_weight_decay_param_name_list += ( + [ + p.name + for model in model_list + for n, p in model.named_parameters() + if len(p.shape) == 1 + ] + if model_list + else [] + ) + import inspect + + extra_kwargs = {} + if "amsgrad" in inspect.signature(optim.AdamW.__init__).parameters: + extra_kwargs["amsgrad"] = self.amsgrad + + opt = optim.AdamW( + learning_rate=self.learning_rate, + beta1=self.beta1, + beta2=self.beta2, + epsilon=self.epsilon, + parameters=parameters, + weight_decay=self.weight_decay, + grad_clip=self.grad_clip, + apply_decay_param_fun=self._apply_decay_param_fun, + **extra_kwargs, + ) + return opt + + def _apply_decay_param_fun(self, name): + return name not in self.no_weight_decay_param_name_list + + +class SOAP: + """ + Improving and Stabilizing Shampoo using Adam. Implements SOAP algorithm (https://arxiv.org/abs/2409.11321). + + Args: + learning_rate (float, optional): + The learning rate to use. defaults to 0.003. + beta1 (float, optional): + Adam's betas parameters beta1. defaults to 0.95. + beta2 (float, optional): + Adam's betas parameters beta2. defaults to 0.95. + shampoo_beta (float, optional): + If >= 0, use this beta for the preconditioner (L and R in paper, state['GG'] below) moving average instead of betas[1]. + defaults to -1. + epsilon (float, optional): + Adam's epsilon for numerical stability. defaults to 1e-08. + weight_decay (float, optional): weight decay coefficient. defaults to 0.01. + precondition_frequency (int, optional): + How often to update the preconditioner. defaults to 10. + max_precond_dim (int, optional): + Maximum dimension of the preconditioner. + Set to 10000, so that we exclude most common vocab sizes while including layers. defaults to 10000. + merge_dims (bool, optional): + Whether or not to merge dimensions of the preconditioner. defaults to `False`. + precondition_1d (bool, optional): + Whether or not to precondition 1D gradients. defaults to `False`. + normalize_grads (bool, optional): + Whether or not to normalize gradients per layer. + Helps at large precondition_frequency (~100 in our experiments), + but hurts performance at small precondition_frequency (~10 in our experiments). defaults to `False`. + data_format (str, optional): + Data format of the input for convolutional layers. + Should be "channels_last" for data_format of NHWC and "channels_first" for NCHW. defaults to `channels_first`. + correct_bias (bool, optional): + Whether or not to use bias correction in Adam. defaults to `True`. + + Examples: + >>> import ppsci + >>> model = ppsci.arch.MLP(("x",), ("u",), 5, 20) + >>> opt = ppsci.optimizer.SOAP(1e-3)(model) + """ + + def __init__( + self, + learning_rate: float = 3e-3, + beta1: float = 0.95, + beta2: float = 0.95, + shampoo_beta: float = -1, + epsilon: float = 1e-8, + weight_decay: float = 0.01, + precondition_frequency: int = 10, + max_precond_dim: int = 10000, # + merge_dims: bool = False, # Merge dimensions till the product of the dimensions is less than or equal to max_precond_dim. + precondition_1d: bool = False, + normalize_grads: bool = False, + data_format: str = "channels_first", + correct_bias: bool = True, + ): + self.learning_rate = learning_rate + self.beta1 = beta1 + self.beta2 = beta2 + self.shampoo_beta = shampoo_beta + self.epsilon = epsilon + self.weight_decay = weight_decay + self.precondition_frequency = precondition_frequency + self.max_precond_dim = max_precond_dim + self.merge_dims = merge_dims + self.precondition_1d = precondition_1d + self.normalize_grads = normalize_grads + self.data_format = data_format + self.correct_bias = correct_bias + + def __call__(self, model_list: Union[nn.Layer, Tuple[nn.Layer, ...]]): + # model_list is None in static graph + if not isinstance(model_list, (tuple, list)): + model_list = (model_list,) + parameters = ( + sum([m.parameters() for m in model_list], []) if model_list else None + ) + opt = SOAP_impl( + parameters=parameters, + learning_rate=self.learning_rate, + beta1=self.beta1, + beta2=self.beta2, + shampoo_beta=self.shampoo_beta, + epsilon=self.epsilon, + weight_decay=self.weight_decay, + precondition_frequency=self.precondition_frequency, + max_precond_dim=self.max_precond_dim, + merge_dims=self.merge_dims, + precondition_1d=self.precondition_1d, + normalize_grads=self.normalize_grads, + data_format=self.data_format, + correct_bias=self.correct_bias, + ) + return opt + + +class OptimizerList: + """OptimizerList which wrap more than one optimizer. + NOTE: LBFGS is not supported yet. + + Args: + optimizer_list (Tuple[optim.Optimizer, ...]): Optimizers listed in a tuple. + + Examples: + >>> import ppsci + >>> model1 = ppsci.arch.MLP(("x",), ("u",), 5, 20) + >>> opt1 = ppsci.optimizer.Adam(1e-3)(model1) + >>> model2 = ppsci.arch.MLP(("y",), ("v",), 5, 20) + >>> opt2 = ppsci.optimizer.Adam(1e-3)(model2) + >>> opt = ppsci.optimizer.OptimizerList((opt1, opt2)) + """ + + def __init__(self, optimizer_list: Tuple[optim.Optimizer, ...]): + super().__init__() + self._opt_list = optimizer_list + if "LBFGS" in set(misc.typename(opt) for opt in optimizer_list): + raise ValueError("LBFGS is not supported in OptimizerList yet.") + + def step(self): + for opt in self._opt_list: + opt.step() + + def clear_grad(self): + for opt in self._opt_list: + opt.clear_grad() + + def get_lr(self) -> float: + """Return learning rate of first optimizer""" + return self._opt_list[0].get_lr() + + def set_state_dict(self, state_dicts: List[Dict[str, "paddle.Tensor"]]): + for i, opt in enumerate(self._opt_list): + opt.set_state_dict(state_dicts[i]) + + def state_dict(self) -> List[Dict[str, "paddle.Tensor"]]: + state_dicts = [opt.state_dict() for opt in self._opt_list] + return state_dicts + + def __len__(self) -> int: + return len(self._opt_list) + + def __getitem__(self, idx): + return self._opt_list[idx] + + def __setitem__(self, idx, opt): + raise NotImplementedError("Can not modify any item in OptimizerList.") + + def __iter__(self): + yield from iter(self._opt_list) diff --git a/examples/smc_reac/ppsci/optimizer/soap.py b/examples/smc_reac/ppsci/optimizer/soap.py new file mode 100644 index 0000000000..de239a978b --- /dev/null +++ b/examples/smc_reac/ppsci/optimizer/soap.py @@ -0,0 +1,558 @@ +# Copyright (c) 2025 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. +# +# refs: https://github.com/nikhilvyas/SOAP + +from collections import defaultdict +from itertools import chain + +import paddle +import paddle.optimizer as optim + + +class SOAP(optim.Optimizer): + """ + Improving and Stabilizing Shampoo using Adam. Implements SOAP algorithm (https://arxiv.org/abs/2409.11321). + + Parameters: + parameters (list|tuple): + List/Tuple of ``Tensor`` names to update to minimize ``loss``. + learning_rate (float, optional): + The learning rate to use. defaults to 0.003. + beta1 (float optional): + Adam's betas parameters b1. defaults to 0.95. + beta2 (float optional): + Adam's betas parameters b1. defaults to 0.95. + shampoo_beta (float, optional): + If >= 0, use this beta for the preconditioner (L and R in paper, state['GG'] below) moving average instead of betas[1]. + defaults to -1. + epsilon (float, optional): + Adam's epsilonilon for numerical stability. defaults to 1e-08. + weight_decay (float, optional): weight decay coefficient. defaults to 0.01. + precondition_frequency (int, optional): + How often to update the preconditioner. defaults to 10. + max_precond_dim (int, optional): + Maximum dimension of the preconditioner. + Set to 10000, so that we exclude most common vocab sizes while including layers. defaults to 10000. + merge_dims (bool, optional): + Whether or not to merge dimensions of the preconditioner. defaults to `False`. + precondition_1d (bool, optional): + Whether or not to precondition 1D gradients. defaults to `False`. + normalize_grads (bool, optional): + Whether or not to normalize gradients per layer. + Helps at large precondition_frequency (~100 in our experiments), + but hurts performance at small precondition_frequency (~10 in our experiments). defaults to `False`. + data_format (str, optional): + Data format of the input for convolutional layers. + Should be "channels_last" for data_format of NHWC and "channels_first" for NCHW. defaults to `channels_first`. + correct_bias (bool, optional): + Whether or not to use bias correction in Adam. defaults to `True`. + name (str, optional): Normally there is no need for user to set this property. + For more information, please refer to :ref:`api_guide_Name`. + The default value is None. + + Return: + loss (Tensor): the final loss of closure. + + Examples: + .. code-block:: python + + >>> import paddle + >>> import ppsci + >>> import numpy as np + + >>> np.random.seed(0) + >>> np_w = np.random.rand(1).astype(np.float32) + >>> np_x = np.random.rand(1).astype(np.float32) + + >>> inputs = [np.random.rand(1).astype(np.float32) for i in range(10)] + >>> # y = 2x + >>> targets = [2 * x for x in inputs] + + >>> class Net(paddle.nn.Layer): + ... def __init__(self): + ... super().__init__() + ... w = paddle.to_tensor(np_w) + ... self.w = paddle.create_parameter(shape=w.shape, dtype=w.dtype, default_initializer=paddle.nn.initializer.Assign(w)) + ... + ... def forward(self, x): + ... return self.w * x + ... + >>> net = Net() + >>> opt = ppsci.optimizer.soap.SOAP(parameters=net.parameters()) + >>> def train_step(inputs, targets): + ... def closure(): + ... outputs = net(inputs) + ... loss = paddle.nn.functional.mse_loss(outputs, targets) + ... print('loss: ', loss.item()) + ... opt.clear_grad() + ... loss.backward() + ... return loss + ... opt.step(closure) + ... + >>> for input, target in zip(inputs, targets): + ... input = paddle.to_tensor(input) + ... target = paddle.to_tensor(target) + ... train_step(input, target) + """ + + def __init__( + self, + parameters, + learning_rate: float = 3e-3, + beta1: float = 0.95, + beta2: float = 0.95, + shampoo_beta: float = -1, + epsilon: float = 1e-8, + weight_decay: float = 0.01, + precondition_frequency: int = 10, + max_precond_dim: int = 10000, # + merge_dims: bool = False, # Merge dimensions till the product of the dimensions is less than or equal to max_precond_dim. + precondition_1d: bool = False, + normalize_grads: bool = False, + data_format: str = "channels_first", + correct_bias: bool = True, + name: str = None, + ): + self._betas = paddle.to_tensor((beta1, beta2)) + self._shampoo_beta = shampoo_beta + self._epsilon = epsilon + self._weight_decay = weight_decay + self._precondition_frequency = precondition_frequency + self._max_precond_dim = max_precond_dim + self._merge_dims = merge_dims + self._precondition_1d = precondition_1d + self._normalize_grads = normalize_grads + self._correct_bias = correct_bias + + self.state = defaultdict(dict) + + super().__init__( + learning_rate=learning_rate, + parameters=parameters, + weight_decay=weight_decay, + name=name, + ) + + if isinstance(self._parameter_list[0], dict): + raise TypeError("The parameter groups is not supported on SOAP optimizer.") + + self._data_format = data_format + + def merge_dims(self, grad, max_precond_dim): + """ + Merges dimensions of the gradient tensor till the product of the dimensions is less than or equal to max_precond_dim. + """ + assert self._data_format in ["channels_first", "channels_last"] + if self._data_format == "channels_last" and grad.ndim == 4: + grad = grad.transpose(0, 3, 1, 2) + shape = grad.shape + new_shape = [] + + curr_shape = 1 + for dim_size in shape: + temp_shape = curr_shape * dim_size + if temp_shape > max_precond_dim: + if curr_shape > 1: + new_shape.append(curr_shape) + curr_shape = dim_size + else: + new_shape.append(dim_size) + curr_shape = 1 + else: + curr_shape = temp_shape + + if curr_shape > 1 or len(new_shape) == 0: + new_shape.append(curr_shape) + + new_grad = grad.reshape(new_shape) + return new_grad + + @paddle.base.framework.non_static_only + def step(self, closure=None): + """ + Performs a single optimization step. + + Arguments: + closure (Optional[Callable]): A closure that reevaluates the model and returns the loss. + """ + with paddle.no_grad(): + if closure is None: + loss = None + else: + closure = paddle.enable_grad()(closure) + loss = closure() + + for p in self._parameter_list: + if p.grad is None: + continue + grad = p.grad + + state = self.state[p] + + if "step" not in state: + state["step"] = 0 + + # State initialization + if "exp_avg" not in state: + # Exponential moving average of gradient values + state["exp_avg"] = paddle.zeros_like(grad) + # Exponential moving average of squared gradient values + state["exp_avg_sq"] = paddle.zeros_like(grad) + + if "Q" not in state: + self.init_preconditioner( + grad, + state, + precondition_frequency=self._precondition_frequency, + precondition_1d=self._precondition_1d, + shampoo_beta=( + self._shampoo_beta + if self._shampoo_beta >= 0 + else self._betas[1] + ), + max_precond_dim=self._max_precond_dim, + merge_dims=self._merge_dims, + ) + self.update_preconditioner( + grad, + state, + max_precond_dim=self._max_precond_dim, + merge_dims=self._merge_dims, + precondition_1d=self._precondition_1d, + ) + continue # first step is skipped so that we never use the current gradients in the projection. + + # Projecting gradients to the eigenbases of Shampoo's preconditioner + # i.e. projecting to the eigenbases of matrices in state['GG'] + grad_projected = self.project( + grad, + state, + merge_dims=self._merge_dims, + max_precond_dim=self._max_precond_dim, + ) + + exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"] + beta1, beta2 = self._betas + + state["step"] += 1 + + # Decay the first and second moment running average coefficient + # In-place operations to update the averages at the same time + exp_avg.multiply_(beta1).add_((1.0 - beta1) * grad_projected) + exp_avg_sq.multiply_(beta2).add_( + (1.0 - beta2) * grad_projected.square() + ) + + denom = exp_avg_sq.sqrt().add_( + paddle.full([], self._epsilon, dtype=exp_avg_sq.dtype) + ) + + # Projecting the exponential moving average of gradients to the eigenbases of Shampoo's preconditioner + # i.e. projecting to the eigenbases of matrices in state['GG'] + # exp_avg_projected = self.project(exp_avg, state, merge_dims=self._merge_dims"], + # max_precond_dim=self._max_precond_dim']) + exp_avg_projected = exp_avg + + step_size = self.get_lr() + if self._correct_bias: + bias_correction1 = 1.0 - beta1 ** (state["step"]) + bias_correction2 = 1.0 - beta2 ** (state["step"]) + step_size = step_size * (bias_correction2**0.5) / bias_correction1 + + # Projecting back the preconditioned (by Adam) exponential moving average of gradients + # to the original space + norm_grad = self.project_back( + exp_avg_projected / denom, + state, + merge_dims=self._merge_dims, + max_precond_dim=self._max_precond_dim, + ) + + if self._normalize_grads: + norm_grad = norm_grad / (1e-30 + paddle.mean(norm_grad**2) ** 0.5) + + p.add_(-step_size * norm_grad) + + # From AdamW code: Just adding the square of the weights to the loss function is *not* + # the correct way of using L2 regularization/weight decay with Adam, + # since that will interact with the m and v parameters in strange ways. + # + # Instead we want to decay the weights in a manner that doesn't interact + # with the m/v parameters. This is equivalent to adding the square + # of the weights to the loss with plain (non-momentum) SGD. + # Add weight decay at the end (fixed version) + if self._weight_decay > 0.0: + p.add_((-self.get_lr() * self._weight_decay) * p) + + # Update is done after the gradient step to avoid using current gradients in the projection. + self.update_preconditioner( + grad, + state, + max_precond_dim=self._max_precond_dim, + merge_dims=self._merge_dims, + precondition_1d=self._precondition_1d, + ) + + return loss + + def init_preconditioner( + self, + grad, + state, + precondition_frequency=10, + shampoo_beta=0.95, + max_precond_dim=10000, + precondition_1d=False, + merge_dims=False, + ): + """ + Initializes the preconditioner matrices (L and R in the paper). + """ + state[ + "GG" + ] = [] # Will hold all the preconditioner matrices (L and R in the paper). + if grad.ndim == 1: + if not precondition_1d or grad.shape[0] > max_precond_dim: + state["GG"].append([]) + else: + state["GG"].append(paddle.zeros([grad.shape[0], grad.shape[0]])) + else: + if merge_dims: + grad = self.merge_dims(grad, max_precond_dim) + + for dim_size in grad.shape: + if dim_size > max_precond_dim: + state["GG"].append([]) + else: + state["GG"].append(paddle.zeros([dim_size, dim_size])) + + state["Q"] = None # Will hold all the eigenbases of the preconditioner. + state["precondition_frequency"] = precondition_frequency + state["shampoo_beta"] = shampoo_beta + + def project(self, grad, state, merge_dims=False, max_precond_dim=10000): + """ + Projects the gradient to the eigenbases of the preconditioner. + """ + original_shape = grad.shape + if merge_dims: + if grad.ndim == 4 and self._data_format == "channels_last": + transposed_shape = grad.transpose(0, 3, 1, 2).shape + grad = self.merge_dims(grad, max_precond_dim) + + for mat in state["Q"]: + if len(mat) > 0: + grad = paddle.tensordot( + grad, + mat, + axes=[[0], [0]], + ) + else: + transpose_order = list(range(1, len(grad.shape))) + [0] + grad = grad.transpose(transpose_order) + + if merge_dims: + if self._data_format == "channels_last" and len(original_shape) == 4: + grad = grad.reshape(transposed_shape).transpose(0, 2, 3, 1) + else: + grad = grad.reshape(original_shape) + return grad + + def update_preconditioner( + self, + grad, + state, + max_precond_dim=10000, + merge_dims=False, + precondition_1d=False, + ): + """ + Updates the preconditioner matrices and the eigenbases (L, R, Q_L, Q_R in the paper). + """ + if state["Q"] is not None: + state["exp_avg"] = self.project_back( + state["exp_avg"], + state, + merge_dims=merge_dims, + max_precond_dim=max_precond_dim, + ) + if grad.ndim == 1: + if precondition_1d and grad.shape[0] <= max_precond_dim: + state["GG"][0].lerp_( + grad.unsqueeze(1) @ grad.unsqueeze(0), 1 - state["shampoo_beta"] + ) + else: + if merge_dims: + new_grad = self.merge_dims(grad, max_precond_dim) + for idx, dim_size in enumerate(new_grad.shape): + if dim_size <= max_precond_dim: + outer_product = paddle.tensordot( + new_grad, + new_grad, + axes=[ + [ + *chain( + range(idx), range(idx + 1, len(new_grad.shape)) + ) + ] + ] + * 2, + ) + state["GG"][idx].lerp_(outer_product, 1 - state["shampoo_beta"]) + else: + for idx, dim_size in enumerate(grad.shape): + if dim_size <= max_precond_dim: + outer_product = paddle.tensordot( + grad, + grad, + # Contracts across all dimensions except for k. + axes=[[*chain(range(idx), range(idx + 1, len(grad.shape)))]] + * 2, + ) + state["GG"][idx].lerp_(outer_product, 1 - state["shampoo_beta"]) + + if state["Q"] is None: + state["Q"] = self.get_orthogonal_matrix(state["GG"]) + if state["step"] > 0 and state["step"] % state["precondition_frequency"] == 0: + state["Q"] = self.get_orthogonal_matrix_QR( + state, max_precond_dim, merge_dims + ) + # state['Q'] = self.get_fast_QR(state, max_precond_dim, merge_dims) + + if state["step"] > 0: + state["exp_avg"] = self.project( + state["exp_avg"], + state, + merge_dims=merge_dims, + max_precond_dim=max_precond_dim, + ) + + def project_back(self, grad, state, merge_dims=False, max_precond_dim=10000): + """ + Projects the gradient back to the original space. + """ + original_shape = grad.shape + if merge_dims: + if self._data_format == "channels_last" and grad.ndim == 4: + transposed_shape = grad.transpose(0, 3, 1, 2).shape + grad = self.merge_dims(grad, max_precond_dim) + for mat in state["Q"]: + if len(mat) > 0: + grad = paddle.tensordot( + grad, + mat, + axes=[[0], [1]], + ) + else: + transpose_order = list(range(1, len(grad.shape))) + [0] + grad = grad.transpose(transpose_order) + + if merge_dims: + if self._data_format == "channels_last" and len(original_shape) == 4: + grad = grad.reshape(transposed_shape).transpose(0, 2, 3, 1) + else: + grad = grad.reshape(original_shape) + return grad + + def get_orthogonal_matrix(self, mat): + """ + Computes the eigenbases of the preconditioner using paddle.linalg.eigh decomposition. + """ + matrix = [] + for m in mat: + if len(m) == 0: + matrix.append([]) + continue + if m.dtype != paddle.float32: + float_data = False + original_type = m.dtype + original_device = m.place + matrix.append(m.to(paddle.float32)) + else: + float_data = True + matrix.append(m) + + final = [] + for m in matrix: + if len(m) == 0: + final.append([]) + continue + _, Q = paddle.linalg.eigh(m + 1e-30 * paddle.eye(m.shape[0])) + Q = paddle.flip(Q, [1]) + + if not float_data: + Q = Q.to(original_device, dtype=original_type) + final.append(Q) + return final + + def get_orthogonal_matrix_QR(self, state, max_precond_dim=10000, merge_dims=False): + """ + Computes the eigenbases of the preconditioner using one round of power iteration + followed by paddle.linalg.qr decomposition. + """ + precond_list = state["GG"] + orth_list = state["Q"] + + matrix = [] + orth_matrix = [] + for m, o in zip(precond_list, orth_list): + if len(m) == 0: + matrix.append([]) + orth_matrix.append([]) + continue + if m.dtype != paddle.float32: + float_data = False + original_type = m.dtype + original_device = m.place + matrix.append(m.to(paddle.float32)) + orth_matrix.append(o.to(paddle.float32)) + else: + float_data = True + matrix.append(m.to(paddle.float32)) + orth_matrix.append(o.to(paddle.float32)) + + orig_shape = state["exp_avg_sq"].shape + if self._data_format == "channels_last" and len(orig_shape) == 4: + transposed_shape = state["exp_avg_sq"].transpose(0, 3, 1, 2).shape + if merge_dims: + exp_avg_sq = self.merge_dims(state["exp_avg_sq"], max_precond_dim) + else: + exp_avg_sq = state["exp_avg_sq"] + + final = [] + for ind, (m, o) in enumerate(zip(matrix, orth_matrix)): + if len(m) == 0: + final.append([]) + continue + est_eig = paddle.diag(o.T @ m @ o) + sort_idx = paddle.argsort(est_eig, descending=True) + exp_avg_sq = exp_avg_sq.index_select(sort_idx, ind) + o = o[:, sort_idx] + power_iter = m @ o + Q, _ = paddle.linalg.qr(power_iter) + + if not float_data: + Q = Q.to(original_device, dtype=original_type) + final.append(Q) + + if merge_dims: + if self._data_format == "channels_last" and len(orig_shape) == 4: + exp_avg_sq = exp_avg_sq.reshape(transposed_shape).transpose(0, 2, 3, 1) + else: + exp_avg_sq = exp_avg_sq.reshape(orig_shape) + + state["exp_avg_sq"] = exp_avg_sq + + return final diff --git a/examples/smc_reac/ppsci/probability/__init__.py b/examples/smc_reac/ppsci/probability/__init__.py new file mode 100644 index 0000000000..1068ada02f --- /dev/null +++ b/examples/smc_reac/ppsci/probability/__init__.py @@ -0,0 +1,19 @@ +# 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. + +from ppsci.probability.hmc import HamiltonianMonteCarlo + +__all__ = [ + "HamiltonianMonteCarlo", +] diff --git a/examples/smc_reac/ppsci/probability/hmc.py b/examples/smc_reac/ppsci/probability/hmc.py new file mode 100644 index 0000000000..76a4847f77 --- /dev/null +++ b/examples/smc_reac/ppsci/probability/hmc.py @@ -0,0 +1,175 @@ +# 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. + +from __future__ import annotations + +from typing import Callable +from typing import Dict + +import paddle +from paddle import distribution + +from ppsci import utils + + +class EnableGradient: + """ + This class is for enabling a dict of tensor for autodiff + """ + + def __init__(self, tensor_dict: Dict[str, paddle.Tensor]): + self.tensor_dict = tensor_dict + + def __enter__(self): + for t in self.tensor_dict.values(): + t.stop_gradient = False + t.clear_grad() + + def __exit__(self, exec_type, exec_val, exec_tb): + for t in self.tensor_dict.values(): + t.stop_gradient = True + + +class HamiltonianMonteCarlo: + """ + Using the HamiltonianMonteCarlo(HMC) to sample from the desired probability distribution. The HMC combine the Hamiltonian Dynamics and Markov Chain Monte Carlo sampling algorithm which is a more efficient way compared to the Metropolis Hasting (MH) method. + + Args: + distribution_fn (paddle.distribution.Distribution): The Log (Posterior) Distribution function that of the parameters needed to be sampled. + path_len (float): The total path length. + step_size (float): Every step size. + num_warmup_steps (int): The number of warm-up steps for the MCMC run. + random_seed (int): Random seed number. + + Examples: + >>> import paddle + >>> from ppsci.probability.hmc import HamiltonianMonteCarlo + >>> def log_posterior(**kwargs): + ... dist = paddle.distribution.Normal(loc=0, scale=1) + ... return dist.log_prob(kwargs['x']) + >>> HMC = HamiltonianMonteCarlo(log_posterior, path_len=1.5, step_size=0.25) + >>> trial = HMC.run_chain(1000, {'x': paddle.to_tensor(0.0)}) + """ + + def __init__( + self, + distribution_fn: Callable, + path_len: float = 1.0, + step_size: float = 0.25, + num_warmup_steps: int = 0, + random_seed: int = 1024, + ): + self.distribution_fn = distribution_fn + self.steps = int(path_len / step_size) + self.step_size = step_size + self.path_len = path_len + self.num_warmup_steps = num_warmup_steps + utils.set_random_seed(random_seed) + self._rv_unif = distribution.Uniform(0, 1) + + def sample( + self, last_position: Dict[str, paddle.Tensor] + ) -> Dict[str, paddle.Tensor]: + """ + Single step for sample + """ + q0 = q1 = last_position + p0 = p1 = self._sample_r(q0) + + for s in range(self.steps): + grad = self._potential_energy_gradient(q1) + for site_name in p1.keys(): + p1[site_name] -= self.step_size * grad[site_name] / 2 + for site_name in q1.keys(): + q1[site_name] += self.step_size * p1[site_name] + + grad = self._potential_energy_gradient(q1) + for site_name in p1.keys(): + p1[site_name] -= self.step_size * grad[site_name] / 2 + + # set the next state in the Markov chain + return q1 if self._check_acceptance(q0, q1, p0, p1) else q0 + + def run_chain( + self, epochs: int, initial_position: Dict[str, paddle.Tensor] + ) -> Dict[str, paddle.Tensor]: + sampling_result: Dict[str, paddle.Tensor] = {} + for k in initial_position.keys(): + sampling_result[k] = [] + pos = initial_position + + # warmup + for _ in range(self.num_warmup_steps): + pos = self.sample(pos) + + # begin collecting sampling result + for e in range(epochs): + pos = self.sample(pos) + for k in pos.keys(): + sampling_result[k].append(pos[k].numpy()) + + for k in initial_position.keys(): + sampling_result[k] = paddle.to_tensor(sampling_result[k]) + + return sampling_result + + def _potential_energy_gradient( + self, pos: Dict[str, paddle.Tensor] + ) -> Dict[str, paddle.Tensor]: + """ + Calculate the gradient of potential energy + """ + grads = {} + with EnableGradient(pos): + (-self.distribution_fn(**pos)).backward() + for k, v in pos.items(): + grads[k] = v.grad.detach() + return grads + + def _k_energy_fn(self, r: Dict[str, paddle.Tensor]) -> paddle.Tensor: + energy = 0.0 + for v in r.values(): + energy = energy + v.dot(v) + return 0.5 * energy + + def _sample_r( + self, params_dict: Dict[str, paddle.Tensor] + ) -> Dict[str, paddle.Tensor]: + # sample r for params + r = {} + for k, v in params_dict.items(): + rv_r = distribution.Normal(paddle.zeros_like(v), paddle.ones_like(v)) + r[k] = rv_r.sample([1]) + if not (v.shape == [] or v.shape == 1): + r[k] = r[k].squeeze() + return r + + def _check_acceptance( + self, + q0: Dict[str, paddle.Tensor], + q1: Dict[str, paddle.Tensor], + p0: Dict[str, paddle.Tensor], + p1: Dict[str, paddle.Tensor], + ) -> bool: + # calculate the Metropolis acceptance probability + energy_current = -self.distribution_fn(**q0) + self._k_energy_fn(p0) + energy_proposed = -self.distribution_fn(**q1) + self._k_energy_fn(p1) + + acceptance = paddle.minimum( + paddle.to_tensor(1.0), paddle.exp(energy_current - energy_proposed) + ) + + # whether accept the proposed state position + event = self._rv_unif.sample([]) + return event <= acceptance diff --git a/examples/smc_reac/ppsci/solver/__init__.py b/examples/smc_reac/ppsci/solver/__init__.py new file mode 100644 index 0000000000..03f97bc2d9 --- /dev/null +++ b/examples/smc_reac/ppsci/solver/__init__.py @@ -0,0 +1,25 @@ +# 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. + +from ppsci.solver import eval +from ppsci.solver import train +from ppsci.solver import visu +from ppsci.solver.solver import Solver + +__all__ = [ + "eval", + "train", + "visu", + "Solver", +] diff --git a/examples/smc_reac/ppsci/solver/eval.py b/examples/smc_reac/ppsci/solver/eval.py new file mode 100644 index 0000000000..1af7655901 --- /dev/null +++ b/examples/smc_reac/ppsci/solver/eval.py @@ -0,0 +1,316 @@ +# 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. + +from __future__ import annotations + +import time +from typing import TYPE_CHECKING +from typing import Dict +from typing import Optional +from typing import Tuple +from typing import Union + +import paddle +from paddle import io + +from ppsci.solver import printer +from ppsci.solver.train import _compute_batch_size +from ppsci.utils import misc + +if TYPE_CHECKING: + from pgl.utils import data as pgl_data + + from ppsci import solver + + +def _get_dataset_length( + data_loader: Union["io.DataLoader", "pgl_data.Dataloader", "io.IterableDataset"] +) -> int: + """Get full dataset length of given dataloader. + + Args: + data_loader (Union[io.DataLoader, pgl_data.Dataloader, io.IterableDataset]): + Given dataloader. + + Returns: + int: Length of full dataset. + """ + if isinstance(data_loader, io.DataLoader): + num_samples = len(data_loader.dataset) + elif isinstance(data_loader, io.IterableDataset): + num_samples = data_loader.num_samples + elif str(type(data_loader)) == "": + num_samples = len(data_loader.dataset) + else: + raise NotImplementedError( + f"Can not fetch the length of given dataset({type(data_loader)})." + ) + + return num_samples + + +def _eval_by_dataset( + solver: "solver.Solver", epoch_id: Optional[int], log_freq: int +) -> Tuple[float, Dict[str, Dict[str, float]]]: + """Evaluate with computing metric on total samples(default process). + + NOTE: This is the default evaluation method as general for most cases, but may not + memory-efficiency for large dataset or large output. + + Args: + solver (solver.Solver): Main Solver. + epoch_id (Optional[int]): Epoch id. + log_freq (int): Log evaluation information every `log_freq` steps. + + Returns: + Tuple[float, Dict[str, Dict[str, float]]]: Target metric and all metric dicts + computed during evaluation. + """ + target_metric: float = float("inf") + metric_dict_group: Dict[str, Dict[str, float]] = misc.PrettyOrderedDict() + for _, _validator in solver.validator.items(): + all_output = misc.Prettydefaultdict(list) + all_label = misc.Prettydefaultdict(list) + num_samples = _get_dataset_length(_validator.data_loader) + + loss_dict = misc.Prettydefaultdict(float) + reader_tic = time.perf_counter() + batch_tic = time.perf_counter() + for iter_id, batch in enumerate(_validator.data_loader, start=1): + input_dict, label_dict, weight_dict = batch + reader_cost = time.perf_counter() - reader_tic + + for v in input_dict.values(): + if hasattr(v, "stop_gradient"): + v.stop_gradient = False + + # forward + with solver.autocast_context_manager( + solver.use_amp, solver.amp_level + ), solver.no_grad_context_manager(solver.eval_with_no_grad): + output_dict, validator_loss = solver.forward_helper.eval_forward( + _validator.output_expr, + input_dict, + solver.model, + _validator, + label_dict, + weight_dict, + ) + + loss_dict[f"{_validator.name}/loss"] = float( + sum(list(validator_loss.values())) + ) + + for key, output in output_dict.items(): + all_output[key].append( + (output.detach() if hasattr(output, "detach") else output) + if solver.world_size == 1 + else misc.all_gather(output.detach()) + ) + + for key, label in label_dict.items(): + all_label[key].append( + (label.detach() if hasattr(label, "detach") else label) + if solver.world_size == 1 + else misc.all_gather(label.detach()) + ) + + batch_cost = time.perf_counter() - batch_tic + solver.eval_time_info["reader_cost"].update(reader_cost) + solver.eval_time_info["batch_cost"].update(batch_cost) + batch_size = _compute_batch_size(input_dict) + printer.update_eval_loss(solver, loss_dict, batch_size) + if ( + iter_id == 1 + or iter_id % log_freq == 0 + or iter_id == len(_validator.data_loader) + ): + printer.log_eval_info( + solver, + batch_size, + epoch_id, + len(_validator.data_loader), + iter_id, + ) + + reader_tic = time.perf_counter() + batch_tic = time.perf_counter() + + # concatenate all data and discard padded sample(s) + for key in all_output: + if paddle.is_tensor(all_output[key][0]): + all_output[key] = paddle.concat(all_output[key]) + if len(all_output[key]) > num_samples: + all_output[key] = all_output[key][:num_samples] + + for key in all_label: + if paddle.is_tensor(all_label[key][0]): + all_label[key] = paddle.concat(all_label[key]) + if len(all_label[key]) > num_samples: + all_label[key] = all_label[key][:num_samples] + + for metric_name, metric_func in _validator.metric.items(): + # NOTE: compute metric with entire output and label + metric_dict = metric_func(all_output, all_label) + metric_dict_group[metric_name] = { + k: float(v) for k, v in metric_dict.items() + } + for var_name, metric_value in metric_dict.items(): + metric_str = f"{_validator.name}/{metric_name}.{var_name}" + if metric_str not in solver.eval_output_info: + solver.eval_output_info[metric_str] = misc.AverageMeter( + metric_str, ".5f" + ) + solver.eval_output_info[metric_str].update( + float(metric_value), num_samples + ) + + # use the first metric for return value + tmp = metric_dict_group + while isinstance(tmp, dict): + tmp = next(iter(tmp.values())) + # avoid that none of metric is set + if isinstance(tmp, float): + target_metric = float(tmp) + + return target_metric, metric_dict_group + + +def _eval_by_batch( + solver: "solver.Solver", epoch_id: Optional[int], log_freq: int +) -> Tuple[float, Dict[str, Dict[str, float]]]: + """Evaluate with computing metric by batch, which is memory-efficient. + + NOTE: This is a evaluation function for large dataset or large output, as is more + memory-efficiency than evaluating by dataset, but less general because some metric + is not independent among samples, e.g. L2 relative error. + + Args: + solver (solver.Solver): Main Solver. + epoch_id (Optional[int]): Epoch id. + log_freq (int): Log evaluation information every `log_freq` steps. + + Returns: + Tuple[float, Dict[str, Dict[str, float]]]: Target metric and all metric dicts + computed during evaluation. + """ + target_metric: float = float("inf") + metric_dict_group: Dict[str, Dict[str, float]] = misc.PrettyOrderedDict() + for _, _validator in solver.validator.items(): + num_samples = _get_dataset_length(_validator.data_loader) + + loss_dict = misc.Prettydefaultdict(float) + reader_tic = time.perf_counter() + batch_tic = time.perf_counter() + for iter_id, batch in enumerate(_validator.data_loader, start=1): + input_dict, label_dict, weight_dict = batch + reader_cost = time.perf_counter() - reader_tic + + batch_size = _compute_batch_size(input_dict) + for v in input_dict.values(): + if hasattr(v, "stop_gradient"): + v.stop_gradient = False + + # forward + with solver.autocast_context_manager( + solver.use_amp, solver.amp_level + ), solver.no_grad_context_manager(solver.eval_with_no_grad): + output_dict, validator_loss = solver.forward_helper.eval_forward( + _validator.output_expr, + input_dict, + solver.model, + _validator, + label_dict, + weight_dict, + ) + + loss_dict[f"{_validator.name}/loss"] = float( + sum(list(validator_loss.values())) + ) + + # collect batch metric + for metric_name, metric_func in _validator.metric.items(): + metric_dict_group[metric_name] = misc.Prettydefaultdict(list) + metric_dict = metric_func(output_dict, label_dict) + for var_name, metric_value in metric_dict.items(): + metric_dict_group[metric_name][var_name].append( + metric_value + if solver.world_size == 1 + else misc.all_gather(metric_value) + ) + + batch_cost = time.perf_counter() - batch_tic + solver.eval_time_info["reader_cost"].update(reader_cost) + solver.eval_time_info["batch_cost"].update(batch_cost) + printer.update_eval_loss(solver, loss_dict, batch_size) + if ( + iter_id == 1 + or iter_id % log_freq == 0 + or iter_id == len(_validator.data_loader) + ): + printer.log_eval_info( + solver, + batch_size, + epoch_id, + len(_validator.data_loader), + iter_id, + ) + + reader_tic = time.perf_counter() + batch_tic = time.perf_counter() + + # concatenate all metric and discard metric of padded sample(s) + for metric_name, metric_dict in metric_dict_group.items(): + for var_name, metric_value in metric_dict.items(): + # NOTE: concat single metric(scalar) list into metric vector + metric_value = paddle.concat(metric_value)[:num_samples] + # NOTE: compute metric via averaging metric over all samples, + # this might be not general for certain evaluation case + metric_value = float(metric_value.mean()) + metric_dict_group[metric_name][var_name] = metric_value + metric_str = f"{_validator.name}/{metric_name}.{var_name}" + if metric_str not in solver.eval_output_info: + solver.eval_output_info[metric_str] = misc.AverageMeter( + metric_str, ".5f" + ) + solver.eval_output_info[metric_str].update(metric_value, num_samples) + + # use the first metric for return value + tmp = metric_dict_group + while isinstance(tmp, dict): + tmp = next(iter(tmp.values())) + # avoid that none of metric is set + if isinstance(tmp, float): + target_metric = tmp + + return target_metric, metric_dict_group + + +def eval_func( + solver: "solver.Solver", epoch_id: Optional[int], log_freq: int +) -> Tuple[float, Dict[str, Dict[str, float]]]: + """Evaluation function. + + Args: + solver (solver.Solver): Main Solver. + epoch_id (Optional[int]): Epoch id. + log_freq (int): Log evaluation information every `log_freq` steps. + + Returns: + Tuple[float, Dict[str, Dict[str, float]]]: Target metric and all metric dicts + computed during evaluation. + """ + if solver.compute_metric_by_batch: + return _eval_by_batch(solver, epoch_id, log_freq) + return _eval_by_dataset(solver, epoch_id, log_freq) diff --git a/examples/smc_reac/ppsci/solver/printer.py b/examples/smc_reac/ppsci/solver/printer.py new file mode 100644 index 0000000000..cedaeab7cc --- /dev/null +++ b/examples/smc_reac/ppsci/solver/printer.py @@ -0,0 +1,161 @@ +# 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. + +import datetime +from typing import TYPE_CHECKING +from typing import Dict +from typing import Optional + +from paddle import device + +from ppsci.utils import logger +from ppsci.utils import misc + +if TYPE_CHECKING: + from ppsci import solver + + +def update_train_loss( + solver: "solver.Solver", loss_dict: Dict[str, float], batch_size: int +): + for key in loss_dict: + if key not in solver.train_output_info: + solver.train_output_info[key] = misc.AverageMeter(key, "7.5f") + solver.train_output_info[key].update(float(loss_dict[key]), batch_size) + if key not in solver.train_loss_info: + solver.train_loss_info[key] = misc.AverageMeter(key, ".5f") + solver.train_loss_info[key].update(float(loss_dict[key])) + + +def update_eval_loss( + solver: "solver.Solver", loss_dict: Dict[str, float], batch_size: int +): + for key in loss_dict: + if key not in solver.eval_output_info: + solver.eval_output_info[key] = misc.AverageMeter(key, "7.5f") + solver.eval_output_info[key].update(float(loss_dict[key]), batch_size) + + +def log_train_info( + solver: "solver.Solver", batch_size: int, epoch_id: int, iter_id: int +): + lr_msg = f"lr: {solver.optimizer.get_lr():.5f}" + + metric_msg = ", ".join( + [ + f"{key}: {solver.train_output_info[key].avg:.5f}" + for key in solver.train_output_info + ] + ) + + time_msg = ", ".join( + [solver.train_time_info[key].mean for key in solver.train_time_info] + ) + + ips_msg = f"ips: {batch_size / solver.train_time_info['batch_cost'].avg:.2f}" + if solver.benchmark_flag: + ips_msg += " samples/s" + + eta_sec = ( + (solver.epochs - epoch_id + 1) * solver.iters_per_epoch - iter_id + ) * solver.train_time_info["batch_cost"].avg + eta_msg = f"eta: {str(datetime.timedelta(seconds=int(eta_sec)))}" + + epoch_width = len(str(solver.epochs)) + iters_width = len(str(solver.iters_per_epoch)) + log_str = ( + f"[Train][Epoch {epoch_id:>{epoch_width}}/{solver.epochs}]" + f"[Iter {iter_id:>{iters_width}}/{solver.iters_per_epoch}] {lr_msg}, " + f"{metric_msg}, {time_msg}, {ips_msg}, {eta_msg}" + ) + if solver.benchmark_flag: + max_mem_reserved_msg = ( + f"max_mem_reserved: {device.cuda.max_memory_reserved() // (1 << 20)} MB" + ) + max_mem_allocated_msg = ( + f"max_mem_allocated: {device.cuda.max_memory_allocated() // (1 << 20)} MB" + ) + log_str += f", {max_mem_reserved_msg}, {max_mem_allocated_msg}" + logger.info(log_str) + + # reset time information after printing + for key in solver.train_time_info: + solver.train_time_info[key].reset() + + logger.scalar( + { + "train/lr": solver.optimizer.get_lr(), + **{ + f"train/{key}": solver.train_output_info[key].avg + for key in solver.train_output_info + }, + }, + step=solver.global_step, + vdl_writer=solver.vdl_writer, + wandb_writer=solver.wandb_writer, + tbd_writer=solver.tbd_writer, + ) + + +def log_eval_info( + solver: "solver.Solver", + batch_size: int, + epoch_id: Optional[int], + iters_per_epoch: int, + iter_id: int, +): + metric_msg = ", ".join( + [ + f"{key}: {solver.eval_output_info[key].avg:.5f}" + for key in solver.eval_output_info + ] + ) + + time_msg = ", ".join( + [solver.eval_time_info[key].mean for key in solver.eval_time_info] + ) + + ips_msg = f"ips: {batch_size / solver.eval_time_info['batch_cost'].avg:.2f}" + + eta_sec = (iters_per_epoch - iter_id) * solver.eval_time_info["batch_cost"].avg + eta_msg = f"eta: {str(datetime.timedelta(seconds=int(eta_sec)))}" + + epoch_width = len(str(solver.epochs)) + iters_width = len(str(iters_per_epoch)) + if isinstance(epoch_id, int): + logger.info( + f"[Eval][Epoch {epoch_id:>{epoch_width}}/{solver.epochs}]" + f"[Iter {iter_id:>{iters_width}}/{iters_per_epoch}] " + f"{metric_msg}, {time_msg}, {ips_msg}, {eta_msg}" + ) + else: + logger.info( + f"[Eval][Iter {iter_id:>{iters_width}}/{iters_per_epoch}] " + f"{metric_msg}, {time_msg}, {ips_msg}, {eta_msg}" + ) + + # reset time information after printing + for key in solver.eval_time_info: + solver.eval_time_info[key].reset() + + # logger.scalar( + # { + # f"eval/{key}": solver.eval_output_info[key].avg + # for key in solver.eval_output_info + # }, + # step=solver.global_step, + # vdl_writer=solver.vdl_writer, + # wandb_writer=solver.wandb_writer, + # tbd_writer=solver.tbd_writer, + # ) diff --git a/examples/smc_reac/ppsci/solver/solver.py b/examples/smc_reac/ppsci/solver/solver.py new file mode 100644 index 0000000000..390211c023 --- /dev/null +++ b/examples/smc_reac/ppsci/solver/solver.py @@ -0,0 +1,1219 @@ +# 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. + +from __future__ import annotations + +import contextlib +import functools +import importlib +import itertools +import os +import sys +from os import path as osp +from typing import TYPE_CHECKING +from typing import Callable +from typing import Dict +from typing import List +from typing import Mapping +from typing import Optional +from typing import Tuple +from typing import Union + +import numpy as np +import paddle +import paddle.distributed as dist +import sympy as sp +from omegaconf import DictConfig +from omegaconf import OmegaConf +from packaging import version +from paddle import amp +from paddle import jit +from paddle import nn +from paddle import optimizer as optim +from paddle.distributed import fleet +from paddle.framework import core +from typing_extensions import Literal + +import ppsci +from ppsci.loss import mtl +from ppsci.utils import ema +from ppsci.utils import expression +from ppsci.utils import logger +from ppsci.utils import misc +from ppsci.utils import save_load + +if TYPE_CHECKING: + from types import ModuleType + + from paddle.static import InputSpec + + +class Solver: + """Class for solver. + + Args: + model (nn.Layer): Model. + constraint (Optional[Dict[str, ppsci.constraint.Constraint]]): Constraint(s) applied on model. Defaults to None. + output_dir (Optional[str]): Output directory. Defaults to "./output/". + optimizer (Optional[optimizer.Optimizer]): Optimizer object. Defaults to None. + lr_scheduler (Optional[optimizer.lr.LRScheduler]): Learning rate scheduler. Defaults to None. + epochs (int, optional): Training epoch(s). Defaults to 5. + iters_per_epoch (int, optional): Number of iterations within an epoch. If set to -1, + than will be automatically set to the length of dataloader of given constraint. + Defaults to 20. + update_freq (int, optional): Update frequency of parameters. Defaults to 1. + save_freq (int, optional): Saving frequency for checkpoint. Defaults to 0. + log_freq (int, optional): Logging frequency. Defaults to 10. + eval_during_train (bool, optional): Whether evaluate model during training. Defaults to False. + start_eval_epoch (int, optional): Epoch number evaluation applied begin after. Defaults to 1. + eval_freq (int, optional): Evaluation frequency. Defaults to 1. + seed (int, optional): Random seed. Defaults to 42. + use_vdl (Optional[bool]): Whether use VisualDL to log scalars. Defaults to False. + use_wandb (Optional[bool]): Whether use wandb to log data. Defaults to False. + use_tbd (Optional[bool]): Whether use tensorboardX to log data. Defaults to False. + wandb_config (Optional[Dict[str, str]]): Config dict of WandB. Defaults to None. + device (Literal["cpu", "gpu", "xpu", "sdaa", None], optional): Runtime device. Defaults to None, which means use default device on current platform. + equation (Optional[Dict[str, ppsci.equation.PDE]]): Equation dict. Defaults to None. + geom (Optional[Dict[str, ppsci.geometry.Geometry]]): Geometry dict. Defaults to None. + validator (Optional[Dict[str, ppsci.validate.Validator]]): Validator dict. Defaults to None. + visualizer (Optional[Dict[str, ppsci.visualize.Visualizer]]): Visualizer dict. Defaults to None. + use_amp (bool, optional): Whether use AMP. Defaults to False. + amp_level (Literal["O0", "O1", "O2", "OD"], optional): AMP level. Defaults to "O1". + pretrained_model_path (Optional[str]): Pretrained model path. Defaults to None. + checkpoint_path (Optional[str]): Checkpoint path. Defaults to None. + compute_metric_by_batch (bool, optional): Whether calculate metrics after each batch during evaluation. Defaults to False. + eval_with_no_grad (bool, optional): Whether set `stop_gradient=True` for every Tensor if no differentiation + involved during computation, generally for save GPU memory and accelerate computing. Defaults to False. + to_static (bool, optional): Whether enable to_static for forward pass. Defaults to False. + loss_aggregator (Optional[mtl.LossAggregator]): Loss aggregator, such as a multi-task learning loss aggregator. Defaults to None. + cfg: (Optional[DictConfig]): Running config dict. Defaults to None. NOTE: This will be required in the future. + + Examples: + >>> import ppsci + >>> model = ppsci.arch.MLP(("x",), ("u",), 5, 20) + >>> opt = ppsci.optimizer.AdamW(1e-3)(model) + >>> geom = ppsci.geometry.Rectangle((0, 0), (1, 1)) + >>> pde_constraint = ppsci.constraint.InteriorConstraint( + ... {"u": lambda out: out["u"]}, + ... {"u": 0}, + ... geom, + ... { + ... "dataset": "IterableNamedArrayDataset", + ... "iters_per_epoch": 1, + ... "batch_size": 16, + ... }, + ... ppsci.loss.MSELoss("mean"), + ... name="EQ", + ... ) # doctest: +SKIP + >>> solver = ppsci.solver.Solver( + ... model, + ... {"EQ": pde_constraint}, + ... "./output", + ... opt, + ... None, + ... ) # doctest: +SKIP + """ + + def __init__( + self, + model: nn.Layer, + constraint: Optional[Dict[str, ppsci.constraint.Constraint]] = None, + output_dir: Optional[str] = "./output/", + optimizer: Optional[optim.Optimizer] = None, + lr_scheduler: Optional[optim.lr.LRScheduler] = None, + epochs: int = 5, + iters_per_epoch: int = 20, + update_freq: int = 1, + save_freq: int = 0, + log_freq: int = 10, + eval_during_train: bool = False, + start_eval_epoch: int = 1, + eval_freq: int = 1, + seed: int = 42, + use_vdl: bool = False, + use_wandb: bool = False, + use_tbd: bool = False, + wandb_config: Optional[Mapping] = None, + device: Literal["cpu", "gpu", "xpu", "sdaa", None] = None, + equation: Optional[Dict[str, ppsci.equation.PDE]] = None, + geom: Optional[Dict[str, ppsci.geometry.Geometry]] = None, + validator: Optional[Dict[str, ppsci.validate.Validator]] = None, + visualizer: Optional[Dict[str, ppsci.visualize.Visualizer]] = None, + use_amp: bool = False, + amp_level: Literal["O0", "O1", "O2", "OD"] = "O1", + pretrained_model_path: Optional[str] = None, + checkpoint_path: Optional[str] = None, + compute_metric_by_batch: bool = False, + eval_with_no_grad: bool = False, + to_static: bool = False, + loss_aggregator: Optional[mtl.LossAggregator] = None, + *, + cfg: Optional[DictConfig] = None, + ): + self.cfg = cfg + if isinstance(cfg, DictConfig): + # (Recommended)Params can be passed within cfg + # rather than passed to 'Solver.__init__' one-by-one. + self._parse_params_from_cfg(cfg) + + # set model + self.model = model + # set constraint + self.constraint = constraint + # set output directory + if not cfg: + self.output_dir = output_dir + + # set optimizer + self.optimizer = optimizer + # set learning rate scheduler + if lr_scheduler is not None: + logger.warning( + "The argument: 'lr_scheduler' now automatically retrieves from " + "'optimizer._learning_rate' when 'optimizer' is given, so it is " + "recommended to remove it from the Solver's initialization arguments." + ) + self.lr_scheduler = ( + optimizer._learning_rate + if ( + isinstance(optimizer, optim.Optimizer) + and isinstance(optimizer._learning_rate, optim.lr.LRScheduler) + ) + else None + ) + if isinstance(self.optimizer, ppsci.optimizer.OptimizerList): + self.lr_scheduler = ppsci.optimizer.lr_scheduler.SchedulerList( + tuple( + opt._learning_rate + for opt in self.optimizer + if isinstance(opt._learning_rate, optim.lr.LRScheduler) + ) + ) + + # set training hyper-parameter + if not cfg: + self.epochs = epochs + self.iters_per_epoch = iters_per_epoch + # set update_freq for gradient accumulation + self.update_freq = update_freq + # set checkpoint saving frequency + self.save_freq = save_freq + # set logging frequency + self.log_freq = log_freq + + # set evaluation hyper-parameter + self.eval_during_train = eval_during_train + self.start_eval_epoch = start_eval_epoch + self.eval_freq = eval_freq + + if self.iters_per_epoch == -1 and self.constraint is not None: + if len(self.constraint) != 1: + raise NotImplementedError( + f"Multiple({len(self.constraint)}) constraints are detected, " + "which is not supported yet, when 'iters_per_epoch' is set to -1." + ) + self.iters_per_epoch = len(next(iter(self.constraint.values())).data_loader) + logger.message( + "Detected 'iters_per_epoch' is set to -1, 'iters_per_epoch' is now " + f"reset to the length of dataloader({self.iters_per_epoch}) of given constraint." + ) + + # initialize training log(training loss, time cost, etc.) recorder during one epoch + self.train_output_info: Dict[str, misc.AverageMeter] = {} + self.train_time_info = { + "batch_cost": misc.AverageMeter("batch_cost", ".5f", postfix="s"), + "reader_cost": misc.AverageMeter("reader_cost", ".5f", postfix="s"), + } + self.train_loss_info: Dict[str, misc.AverageMeter] = {} + + # initialize evaluation log(evaluation loss, metric, etc.) recorder. + self.eval_output_info: Dict[str, misc.AverageMeter] = {} + self.eval_time_info = { + "batch_cost": misc.AverageMeter("batch_cost", ".5f", postfix="s"), + "reader_cost": misc.AverageMeter("reader_cost", ".5f", postfix="s"), + } + + # set running device + if not cfg: + self.device = device + if self.device is None: + # set to default device if not specified + self.device: str = paddle.device.get_device() + + if self.device != "cpu" and paddle.device.get_device() == "cpu": + # fall back to cpu if no other device available + logger.warning(f"Set device({device}) to 'cpu' for only cpu available.") + self.device = "cpu" + self.device = paddle.device.set_device(self.device) + + # set equations for physics-driven or data-physics hybrid driven task, such as PINN + self.equation = equation + + # set validator + self.validator = validator + + # set visualizer + self.visualizer = visualizer + + # set automatic mixed precision(AMP) configuration + if not cfg: + self.use_amp = use_amp + self.amp_level = amp_level + self.scaler = amp.GradScaler(True) if self.use_amp else None + + # whether calculate metrics by each batch during evaluation, mainly for memory efficiency + if not cfg: + self.compute_metric_by_batch = compute_metric_by_batch + if validator is not None: + for metric in itertools.chain( + *[_v.metric.values() for _v in self.validator.values()] + ): + if metric.keep_batch ^ self.compute_metric_by_batch: + raise ValueError( + f"{misc.typename(metric)}.keep_batch should be " + f"{self.compute_metric_by_batch} when compute_metric_by_batch=" + f"{self.compute_metric_by_batch}." + ) + # check metric name uniqueness over all validators + _count = {} + for _validator in validator.values(): + for metric_name in _validator.metric: + if metric_name in _count: + logger.warning( + f"Metric name({metric_name}) is duplicated, please ensure " + "all metric names are unique over all given validators." + ) + _count[metric_name] = 1 + del _count + + # whether set `stop_gradient=True` for every Tensor if no differentiation involved during evaluation + if not cfg: + self.eval_with_no_grad = eval_with_no_grad + + self.rank = dist.get_rank() + self.world_size = dist.get_world_size() + # initialize distributed environment + if self.world_size > 1: + # TODO(sensen): Support different kind of DistributedStrategy + fleet.init(is_collective=True) + logger.warning( + f"Detected 'world_size'({self.world_size}) > 1, it is recommended to " + "scale up the learning rate and reduce the 'epochs' or " + "'iters_per_epoch' according to the 'world_size' both linearly if you " + "are training model." + ) + + # set moving average model(optional) + self.ema_model = None + if self.cfg and any(key in self.cfg.TRAIN for key in ["ema", "swa"]): + if "ema" in self.cfg.TRAIN and cfg.TRAIN.ema.get("use_ema", False): + self.ema_model = ema.ExponentialMovingAverage( + self.model, self.cfg.TRAIN.ema.decay + ) + elif "swa" in self.cfg.TRAIN and cfg.TRAIN.swa.get("use_swa", False): + self.ema_model = ema.StochasticWeightAverage(self.model) + + # load pretrained model, usually used for transfer learning + if not cfg: + self.pretrained_model_path = pretrained_model_path + if self.pretrained_model_path is not None: + save_load.load_pretrain( + self.model, self.pretrained_model_path, self.equation + ) + + self.cur_metric = float("inf") + # initialize an dict for tracking best metric during training + self.best_metric = { + "metric": float("inf"), + "epoch": 0, + } + + # use loss aggregator, use Sum if None + if isinstance(loss_aggregator, (mtl.AGDA, mtl.PCGrad)) and self.use_amp: + raise ValueError( + "Auto Mix Precision do not support AGDA, PCGrad loss aggregator yet, " + "please set use_amp=False." + ) + self.loss_aggregator = loss_aggregator or mtl.Sum() + + # load model checkpoint, usually used for resume training + if not cfg: + self.checkpoint_path = checkpoint_path + if self.checkpoint_path is not None: + if self.pretrained_model_path is not None: + logger.warning( + "Detected 'pretrained_model_path' is given, weights in which might be" + "overridden by weights loaded from given 'checkpoint_path'." + ) + loaded_metric = save_load.load_checkpoint( + self.checkpoint_path, + self.model, + self.optimizer, + self.scaler, + self.equation, + self.ema_model, + self.loss_aggregator, + ) + if isinstance(loaded_metric, dict): + self.best_metric.update(loaded_metric) + + # decorate model(s) and optimizer(s) for AMP + if self.use_amp: + self.model, self.optimizer = amp.decorate( + self.model, + self.optimizer, + self.amp_level, + save_dtype="float32", + ) + + # choosing an appropriate training function for different optimizers + if misc.typename(self.optimizer) == "LBFGS": + if self.use_amp: + raise ValueError( + "Auto Mix Precision is not supported for L-BFGS optimizer." + ) + self.train_epoch_func = ppsci.solver.train.train_LBFGS_epoch_func + if self.update_freq != 1: + self.update_freq = 1 + logger.warning("Set 'update_freq' to to 1 when using L-BFGS optimizer.") + else: + self.train_epoch_func = ppsci.solver.train.train_epoch_func + + # wrap model and optimizer to parallel object + if self.world_size > 1: + if isinstance(self.model, paddle.DataParallel): + raise ValueError( + "Given model is already wrapped by paddle.DataParallel." + "Please do not wrap your model with DataParallel " + "before 'Solver.__init__' and keep it's type as 'nn.Layer'." + ) + + def dist_wrapper(model: nn.Layer) -> paddle.DataParallel: + dist_model = fleet.distributed_model(model) + if hasattr(model, "input_keys"): + dist_model.input_keys = dist_model._layers.input_keys + if hasattr(model, "output_keys"): + dist_model.output_keys = dist_model._layers.output_keys + return dist_model + + if isinstance(self.model, ppsci.arch.ModelList): + for i in range(len(self.model.model_list)): + # NOTE: Convert each model in model_list to DataParallel + self.model.model_list[i] = dist_wrapper(self.model.model_list[i]) + else: + self.model = dist_wrapper(self.model) + + if self.optimizer is not None: + self.optimizer = fleet.distributed_optimizer(self.optimizer) + + # set VisualDL tool + self.vdl_writer = None + if not cfg: + self.use_vdl = use_vdl + if self.use_vdl: + try: + import visualdl as vdl + except ModuleNotFoundError: + raise ModuleNotFoundError( + "Please install 'visualdl' with `pip install visualdl` first." + ) + with misc.RankZeroOnly(self.rank) as is_master: + if is_master: + self.vdl_writer = vdl.LogWriter(osp.join(self.output_dir, "vdl")) + logger.info( + "VisualDL is enabled for logging, you can view it by " + f"running:\nvisualdl --logdir {self.vdl_writer._logdir} --port 8080" + ) + + # set WandB tool + self.wandb_writer = None + if not cfg: + self.use_wandb = use_wandb + self.wandb_config = {} + if self.use_wandb: + try: + import wandb + except ModuleNotFoundError: + raise ModuleNotFoundError( + "Please install 'wandb' with `pip install wandb` first." + ) + with misc.RankZeroOnly(self.rank) as is_master: + if is_master: + self.wandb_writer = wandb.init(**self.wandb_config) + + # set TensorBoardX tool + self.tbd_writer = None + if not cfg: + self.use_tbd = use_tbd + if self.use_tbd: + try: + import tensorboardX + except ModuleNotFoundError: + raise ModuleNotFoundError( + "Please install 'tensorboardX' with `pip install tensorboardX` first." + ) + with misc.RankZeroOnly(self.rank) as is_master: + if is_master: + self.tbd_writer = tensorboardX.SummaryWriter( + osp.join(self.output_dir, "tensorboard") + ) + logger.message( + "TensorboardX is enabled for logging, you can view it by " + f"running:\ntensorboard --logdir {self.tbd_writer.logdir}" + ) + + self.global_step = 0 + + # log paddlepaddle's version + if version.Version(paddle.__version__) != version.Version("0.0.0"): + paddle_version = paddle.__version__ + if version.Version(paddle.__version__) < version.Version("2.6.0"): + logger.warning( + f"Detected paddlepaddle version is '{paddle_version}', " + "currently it is recommended to use paddlepaddle >= 2.6 or develop version." + ) + else: + paddle_version = f"develop({paddle.version.commit[:7]})" + + logger.info(f"Using paddlepaddle {paddle_version} on device {self.device}") + + self.forward_helper = expression.ExpressionSolver() + + # whether enable static for forward pass. Defaults to False + if not cfg: + self.to_static = to_static + jit.enable_to_static(self.to_static) + logger.message( + f"Set to_static={self.to_static} for computational optimization." + ) + + # convert sympy to callable object if exist + extra_parameters = [] + if self.equation: + for equation in self.equation.values(): + extra_parameters += list(equation.learnable_parameters) + + def convert_expr( + container_dict: Union[ + Dict[str, ppsci.constraint.Constraint], + Dict[str, ppsci.validate.Validator], + Dict[str, ppsci.visualize.Visualizer], + ] + ) -> None: + for container in container_dict.values(): + exprs = [ + expr + for expr in container.output_expr.values() + if isinstance(expr, sp.Basic) + ] + if len(exprs) > 0: + funcs = ppsci.lambdify( + exprs, + self.model, + extra_parameters=extra_parameters, + # graph_filename=osp.join(self.output_dir, "symbolic_graph_visual"), # HACK: Activate it for DEBUG. + fuse_derivative=True, + ) + ind = 0 + for name in container.output_expr: + if isinstance(container.output_expr[name], sp.Basic): + container.output_expr[name] = funcs[ind] + # FIXME: Equation with parameter not support yet. + # if self.world_size > 1: + # container.output_expr[name] = dist_wrapper( + # container.output_expr[name] + # ) + ind += 1 + + if self.constraint: + convert_expr(self.constraint) + + if self.validator: + convert_expr(self.validator) + + if self.visualizer: + convert_expr(self.visualizer) + + # set up benchmark flag, will print memory stat if enabled + self.benchmark_flag: bool = os.getenv("BENCHMARK_ROOT", None) is not None + + # set up nvtx flag for nsight analysis + self.nvtx_flag: bool = os.getenv("NVTX", None) is not None + self.forward_helper.nvtx_flag = self.nvtx_flag + + # for callbacks + self.callbacks_on_epoch_begin: List[Callable[[Solver]]] = [] + self.callbacks_on_epoch_end: List[Callable[[Solver]]] = [] + self.callbacks_on_iter_begin: List[Callable[[Solver]]] = [] + self.callbacks_on_iter_end: List[Callable[[Solver]]] = [] + + def train(self) -> None: + """Training.""" + self.global_step = self.best_metric["epoch"] * self.iters_per_epoch + self.max_steps = self.epochs * self.iters_per_epoch + + start_epoch = self.best_metric["epoch"] + 1 + + if self.use_tbd and isinstance(self.cfg, DictConfig): + self.tbd_writer.add_text( + "config", f"
{str(OmegaConf.to_yaml(self.cfg))}
" + ) + + if self.nvtx_flag: + core.nvprof_start() + core.nvprof_enable_record_event() + + for epoch_id in range(start_epoch, self.epochs + 1): + self._invoke_callbacks_on_epoch_begin() # [optional] + self.train_epoch_func(self, epoch_id, self.log_freq) + self._invoke_callbacks_on_epoch_end() # [optional] + + self.train_output_info.clear() + + # update average model if exist + if self.ema_model and epoch_id % self.avg_freq == 0: + self.ema_model.update() + + # evaluate during training + if ( + self.eval_during_train + and epoch_id % self.eval_freq == 0 + and epoch_id >= self.start_eval_epoch + ): + self.cur_metric, metric_dict_group = self.eval(epoch_id) + if self.cur_metric < self.best_metric["metric"]: + self.best_metric["metric"] = self.cur_metric + self.best_metric["epoch"] = epoch_id + save_load.save_checkpoint( + self.model, + self.optimizer, + self.best_metric, + self.scaler, + self.output_dir, + "best_model", + self.equation, + aggregator=self.loss_aggregator, + ) + logger.info( + f"[Eval][Epoch {epoch_id}]" + f"[best metric: {self.best_metric['metric']}]" + ) + for metric_name, metric_dict in metric_dict_group.items(): + logger.scalar( + {f"eval/{metric_name}/{k}": v for k, v in metric_dict.items()}, + epoch_id, + self.vdl_writer, + self.wandb_writer, + self.tbd_writer, + ) + + # visualize after evaluation + if self.visualizer is not None: + self.visualize(epoch_id) + + # evaluation for moving average evaluation(almost same procedure) + if self.ema_model and epoch_id % self.avg_freq == 0: + self.ema_model.apply_shadow() + logger.info("Evaluating metric of averaging model...") + cur_metric_ema, metric_dict_group_ema = self.eval(epoch_id) + self.ema_model.restore() + + if cur_metric_ema < self.best_metric["metric"]: + self.best_metric["metric"] = cur_metric_ema + self.best_metric["epoch"] = epoch_id + save_load.save_checkpoint( + self.ema_model, + None, + metric=self.best_metric, + output_dir=self.output_dir, + prefix="best_model_ema", + ) + logger.info( + f"[Eval][Epoch {epoch_id}]" + f"[best metric: {self.best_metric['metric']}]" + ) + for metric_name, metric_dict in metric_dict_group_ema.items(): + logger.scalar( + { + f"eval_ema/{metric_name}/{k}": v + for k, v in metric_dict.items() + }, + epoch_id, + self.vdl_writer, + self.wandb_writer, + self.tbd_writer, + ) + + # update learning rate by epoch + if self.lr_scheduler is not None and self.lr_scheduler.by_epoch: + self.lr_scheduler.step() + + # save epoch model every save_freq epochs + if self.save_freq > 0 and epoch_id % self.save_freq == 0: + save_load.save_checkpoint( + self.model, + self.optimizer, + {"metric": self.cur_metric, "epoch": epoch_id}, + self.scaler, + self.output_dir, + f"epoch_{epoch_id}", + self.equation, + ema_model=self.ema_model, + aggregator=self.loss_aggregator, + ) + + # save the latest model for convenient resume training + save_load.save_checkpoint( + self.model, + self.optimizer, + {"metric": self.cur_metric, "epoch": epoch_id}, + self.scaler, + self.output_dir, + "latest", + self.equation, + print_log=(epoch_id == start_epoch), + ema_model=self.ema_model, + aggregator=self.loss_aggregator, + ) + + def finetune(self, pretrained_model_path: str) -> None: + """Finetune model based on given pretrained model path. + + Args: + pretrained_model_path (str): Pretrained model path or url. + """ + # load pretrained model + save_load.load_pretrain(self.model, pretrained_model_path, self.equation) + + # call train program + self.train() + + @misc.run_on_eval_mode + def eval( + self, epoch_id: Optional[int] = None + ) -> Tuple[float, Dict[str, Dict[str, float]]]: + """Evaluation. + + Args: + epoch_id (Optional[int]): Epoch id. Defaults to None. + + Returns: + Tuple[float, Dict[str, Dict[str, float]]]: A targe metric value(float) and + all metric(s)(dict) of evaluation, used to judge the quality of the model. + """ + # set eval func + self.eval_func = ppsci.solver.eval.eval_func + + result = self.eval_func(self, epoch_id, self.log_freq) + metric_msg = ", ".join( + [self.eval_output_info[key].avg_info for key in self.eval_output_info] + ) + + if isinstance(epoch_id, int): + logger.info(f"[Eval][Epoch {epoch_id}][Avg] {metric_msg}") + else: + logger.info(f"[Eval][Avg] {metric_msg}") + self.eval_output_info.clear() + + return result + + @misc.run_on_eval_mode + def visualize(self, epoch_id: Optional[int] = None): + """Visualization. + + Args: + epoch_id (Optional[int]): Epoch id. Defaults to None. + """ + # set visualize func + self.visu_func = ppsci.solver.visu.visualize_func + + self.visu_func(self, epoch_id) + if isinstance(epoch_id, int): + logger.info(f"[Visualize][Epoch {epoch_id}] Finish visualization") + else: + logger.info("[Visualize] Finish visualization") + + @misc.run_on_eval_mode + def predict( + self, + input_dict: Dict[str, Union[np.ndarray, paddle.Tensor]], + expr_dict: Optional[Dict[str, Callable]] = None, + batch_size: Optional[int] = 64, + no_grad: bool = True, + return_numpy: bool = False, + ) -> Dict[str, Union[paddle.Tensor, np.ndarray]]: + """Pure prediction using model.forward(...) and expression(optional, if given). + + Args: + input_dict (Dict[str, Union[np.ndarray, paddle.Tensor]]): Input data in dict. + expr_dict (Optional[Dict[str, Callable]]): Expression dict, which guide to + compute equation variable with callable function. Defaults to None. + batch_size (Optional[int]): Predicting by batch size. If None, data in + `input_dict` will be used directly for inference without any batch slicing. + Defaults to 64. + no_grad (bool): Whether set stop_gradient=True for entire prediction, mainly + for memory-efficiency. Defaults to True. + return_numpy (bool): Whether convert result from Tensor to numpy ndarray. + Defaults to False. + + Returns: + Dict[str, Union[paddle.Tensor, np.ndarray]]: Prediction in dict. + + Examples: + >>> import paddle + >>> import ppsci + >>> model = ppsci.arch.MLP(('x', 'y'), ('u', 'v'), num_layers=None, hidden_size=[32, 8]) + >>> solver = ppsci.solver.Solver(model) # doctest: +SKIP + >>> input_dict = {'x': paddle.rand([32, 1]), + ... 'y': paddle.rand([32, 1])} + >>> pred = solver.predict(input_dict) # doctest: +SKIP + >>> for k, v in pred.items(): # doctest: +SKIP + ... print(k, v.shape) # doctest: +SKIP + u [32, 1] + v [32, 1] + """ + num_samples = len(next(iter(input_dict.values()))) + num_pad = (self.world_size - num_samples % self.world_size) % self.world_size + # pad with last element if `num_samples` is not divisible by `world_size` + # ensuring every device get same number of data. + if num_pad > 0: + for k, v in input_dict.items(): + repeat_times = (num_pad, *(1 for _ in range(v.ndim - 1))) + if isinstance(v, np.ndarray): + input_dict[k] = np.concatenate( + ( + v, + np.tile(v[num_samples - 1 : num_samples], repeat_times), + ), + ) + elif isinstance(v, paddle.Tensor): + input_dict[k] = paddle.concat( + ( + v, + paddle.tile(v[num_samples - 1 : num_samples], repeat_times), + ), + ) + else: + raise ValueError(f"Unsupported data type {type(v)}.") + + num_samples_pad = num_samples + num_pad + local_num_samples_pad = num_samples_pad // self.world_size + local_input_dict = ( + {k: v[self.rank :: self.world_size] for k, v in input_dict.items()} + if self.world_size > 1 + else input_dict + ) + local_batch_num = ( + (local_num_samples_pad + (batch_size - 1)) // batch_size + if batch_size is not None + else 1 + ) + + pred_dict = misc.Prettydefaultdict(list) + with self.no_grad_context_manager(no_grad), self.no_sync_context_manager( + self.world_size > 1, self.model + ): + for batch_id in range(local_batch_num): + # prepare local batch input + if batch_size is not None: + st = batch_id * batch_size + ed = min(local_num_samples_pad, (batch_id + 1) * batch_size) + batch_input_dict = { + k: v[st:ed] for k, v in local_input_dict.items() + } + else: + batch_input_dict = {**local_input_dict} + # Keep dtype unchanged as all dtype be correct when given into predict function + for key in batch_input_dict: + if not paddle.is_tensor(batch_input_dict[key]): + batch_input_dict[key] = paddle.to_tensor( + batch_input_dict[key], stop_gradient=no_grad + ) + + # forward + with self.autocast_context_manager(self.use_amp, self.amp_level): + batch_output_dict = self.forward_helper.visu_forward( + expr_dict, batch_input_dict, self.model + ) + + # collect local batch output + for key, batch_output in batch_output_dict.items(): + pred_dict[key].append( + batch_output.detach() if no_grad else batch_output + ) + + # concatenate local output + pred_dict = {key: paddle.concat(value) for key, value in pred_dict.items()} + + if self.world_size > 1: + # gather global output from all devices if world_size > 1 + pred_dict = { + key: misc.all_gather(value) for key, value in pred_dict.items() + } + # rearrange output as the same order of input_dict according + # to inverse permutation + perm = np.arange(num_samples_pad, dtype="int64") + perm = np.concatenate( + [perm[rank :: self.world_size] for rank in range(self.world_size)], + axis=0, + ) + perm_inv = np.empty_like(perm) + perm_inv[perm] = np.arange(num_samples_pad, dtype="int64") + perm_inv = paddle.to_tensor(perm_inv) + pred_dict = {key: value[perm_inv] for key, value in pred_dict.items()} + # then discard output of padding data at the end if num_pad > 0 + if num_pad > 0: + pred_dict = { + key: value[:num_samples] for key, value in pred_dict.items() + } + # NOTE: Discard padding data in input_dict for consistency + for k in input_dict: + input_dict[k] = input_dict[k][:num_samples] + + # convert to numpy ndarray if specified + if return_numpy: + pred_dict = { + k: (v.numpy() if paddle.is_tensor(v) else v) + for k, v in pred_dict.items() + } + + return pred_dict + + @misc.run_on_eval_mode + def export( + self, + input_spec: List[Dict[str, InputSpec]], + export_path: str, + with_onnx: bool = False, + skip_prune_program: bool = False, + *, + full_graph: bool = True, + ignore_modules: Optional[List[ModuleType]] = None, + ): + """ + Convert model to static graph model and export to files. + + Args: + input_spec (List[Dict[str, InputSpec]]): InputSpec describes the signature + information of the model input. + export_path (str): The path prefix to save model. + with_onnx (bool, optional): Whether to export model into onnx after + paddle inference models are exported. Defaults to False. + skip_prune_program (bool, optional): Whether prune program, pruning program + may cause unexpectable result, e.g. llm-inference. Defaults to False. + full_graph (bool, optional): Symbolic OpCode Translator(SOT) will be used + when set to True, where otherwise use Abstract Syntax Tree(AST) if False. + Defaults to True. + ignore_modules (List[ModuleType]): Adds modules that should be ignored during + conversion. Builtin modules that have been ignored are collections, pdb, + copy, inspect, re, numpy, logging, six. For example, einops can be added + here. Defaults to None. + """ + if ignore_modules is not None: + jit.ignore_module(ignore_modules) + + jit.enable_to_static(True) + + if self.pretrained_model_path is None: + logger.warning( + "'INFER.pretrained_model_path' is not given, so the weights of exported " + "model will be random initialized." + ) + + # convert model to static graph model + static_model = jit.to_static( + self.model, + input_spec=input_spec, + full_graph=full_graph, + ) + + # save static graph model to disk + if len(osp.dirname(export_path)): + os.makedirs(osp.dirname(export_path), exist_ok=True) + try: + jit.save(static_model, export_path, skip_prune_program=skip_prune_program) + except Exception as e: + raise e + logger.message( + f"Inference model has been exported to: {export_path}, including " + + ( + "*.json, *.pdiparams files." + if paddle.framework.use_pir_api() + else "*.pdmodel, *.pdiparams and *.pdiparams.info files." + ) + ) + jit.enable_to_static(False) + + if with_onnx: + # TODO: support pir + onnx + if not importlib.util.find_spec("paddle2onnx"): + raise ModuleNotFoundError( + "Please install paddle2onnx with `pip install paddle2onnx`" + " before exporting onnx model." + ) + import paddle2onnx + + DEFAULT_OPSET_VERSION = 19 + + paddle2onnx.export( + model_filename=export_path + ".json" + if paddle.framework.use_pir_api() + else ".pdmodel", + params_filename=export_path + ".pdiparams", + save_file=export_path + ".onnx", + opset_version=DEFAULT_OPSET_VERSION, + enable_onnx_checker=True, + ) + logger.message(f"ONNX model has been exported to: {export_path}.onnx") + + def autocast_context_manager( + self, enable: bool, level: Literal["O0", "O1", "O2", "OD"] = "O1" + ) -> contextlib.AbstractContextManager: + """Smart autocast context manager for Auto Mix Precision. + + Args: + enable (bool): Enable autocast. + level (Literal["O0", "O1", "O2", "OD"]): Autocast level. + + Returns: + contextlib.AbstractContextManager: Smart autocast context manager. + """ + if enable: + ctx_manager = amp.auto_cast(level=level) + else: + ctx_manager = ( + contextlib.nullcontext() + if sys.version_info >= (3, 7) + else contextlib.suppress() + ) + return ctx_manager + + @functools.lru_cache() + def no_grad_context_manager( + self, enable: bool + ) -> contextlib.AbstractContextManager: + """Smart no_grad context manager. + + Args: + enable (bool): Enable no_grad. + + Returns: + contextlib.AbstractContextManager: Smart no_grad context manager. + """ + if enable: + ctx_manager = paddle.no_grad() + else: + ctx_manager = ( + contextlib.nullcontext() + if sys.version_info >= (3, 7) + else contextlib.suppress() + ) + return ctx_manager + + def no_sync_context_manager( + self, + enable: bool, + ddp_model: paddle.DataParallel, + ) -> contextlib.AbstractContextManager: + """Smart no_sync context manager for given model. + NOTE: Only `paddle.DataParallel` object has `no_sync` interface. + + Args: + enable (bool): Enable no_sync. + + Returns: + contextlib.AbstractContextManager: Smart no_sync context manager. + """ + if enable: + if isinstance(self.model, ppsci.arch.ModelList): + for model in self.model.model_list: + if not isinstance(model, paddle.DataParallel): + raise TypeError( + "no_sync interface is only for model with type " + "paddle.DataParallel, but got type " + f"{misc.typename(model)}" + ) + ctx_manager = contextlib.ExitStack() + for model in self.model.model_list: + ctx_manager.enter_context(model.no_sync()) + else: + if not isinstance(self.model, paddle.DataParallel): + raise TypeError( + "no_sync interface is only for model with type " + f"paddle.DataParallel, but got type {misc.typename(ddp_model)}" + ) + ctx_manager = ddp_model.no_sync() + else: + ctx_manager = ( + contextlib.nullcontext() + if sys.version_info >= (3, 7) + else contextlib.suppress() + ) + return ctx_manager + + def plot_loss_history( + self, + by_epoch: bool = False, + smooth_step: int = 1, + use_semilogy: bool = True, + ) -> None: + """Plotting iteration/epoch-loss curve. + + Args: + by_epoch (bool, optional): Whether the abscissa axis of the curve is epoch or iteration. Defaults to False. + smooth_step (int, optional): How many steps of loss are squeezed to one point to smooth the curve. Defaults to 1. + use_semilogy (bool, optional): Whether to set non-uniform coordinates for the y-axis. Defaults to True. + """ + loss_dict = {} + for key in self.train_loss_info: + loss_arr = np.asarray(self.train_loss_info[key].history) + if by_epoch: + loss_arr = np.mean( + np.reshape(loss_arr, (-1, self.iters_per_epoch)), + axis=1, + ) + loss_dict[key] = list(loss_arr) + + misc.plot_curve( + data=loss_dict, + xlabel="Epoch" if by_epoch else "Iteration", + ylabel="Loss", + output_dir=self.output_dir, + smooth_step=smooth_step, + use_semilogy=use_semilogy, + ) + + def _parse_params_from_cfg(self, cfg: DictConfig): + """ + Parse hyper-parameters from DictConfig. + """ + self.output_dir = cfg.output_dir + self.log_freq = cfg.log_freq + self.use_tbd = cfg.use_tbd + self.use_vdl = cfg.use_vdl + self.wandb_config = cfg.wandb_config + self.use_wandb = cfg.use_wandb + self.device = cfg.device + self.to_static = cfg.to_static + + self.use_amp = cfg.use_amp + self.amp_level = cfg.amp_level + + self.epochs = cfg.TRAIN.epochs + self.iters_per_epoch = cfg.TRAIN.iters_per_epoch + self.update_freq = cfg.TRAIN.update_freq + self.save_freq = cfg.TRAIN.save_freq + self.eval_during_train = cfg.TRAIN.eval_during_train + self.start_eval_epoch = cfg.TRAIN.start_eval_epoch + self.eval_freq = cfg.TRAIN.eval_freq + self.checkpoint_path = cfg.TRAIN.checkpoint_path + + if "ema" in cfg.TRAIN and cfg.TRAIN.ema.get("use_ema", False): + self.avg_freq = cfg.TRAIN.ema.avg_freq + elif "swa" in cfg.TRAIN and cfg.TRAIN.swa.get("use_swa", False): + self.avg_freq = cfg.TRAIN.swa.avg_freq + + self.compute_metric_by_batch = cfg.EVAL.compute_metric_by_batch + self.eval_with_no_grad = cfg.EVAL.eval_with_no_grad + + if cfg.mode == "train": + self.pretrained_model_path = cfg.TRAIN.pretrained_model_path + elif cfg.mode == "eval": + self.pretrained_model_path = cfg.EVAL.pretrained_model_path + elif cfg.mode in ["export", "infer"]: + self.pretrained_model_path = cfg.INFER.pretrained_model_path + + def register_callback_on_epoch_begin( + self: Solver, callback_fn: Callable[[Solver]] + ) -> None: + """ + Registers a callback function to be executed at the beginning of each training epoch. + + Args: + callback_fn : Callable[[Solver]] + A function that takes a Solver instance as an argument. This function + will be called at the start of every epoch. + """ + self.callbacks_on_epoch_begin.append(callback_fn) + + def register_callback_on_epoch_end( + self: Solver, callback_fn: Callable[[Solver]] + ) -> None: + """ + Registers a callback function to be executed at the end of each training epoch. + + Args: + callback_fn : Callable[[Solver]] + A function that takes a Solver instance as an argument. This function + will be called at the end of every epoch. + """ + self.callbacks_on_epoch_end.append(callback_fn) + + def register_callback_on_iter_begin( + self: Solver, callback_fn: Callable[[Solver]] + ) -> None: + """ + Registers a callback function to be executed at the beginning of each training iteration. + + Args: + callback_fn : Callable[[Solver]] + A function that takes a Solver instance as an argument. This function + will be called at the start of every iteration. + """ + self.callbacks_on_iter_begin.append(callback_fn) + + def register_callback_on_iter_end( + self: Solver, callback_fn: Callable[[Solver]] + ) -> None: + """ + Registers a callback function to be executed at the end of each training iteration. + + Args: + callback_fn : Callable[[Solver]] + A function that takes a Solver instance as an argument. This function + will be called at the end of every iteration. + + Returns: + ------- + None + """ + self.callbacks_on_iter_end.append(callback_fn) + + def _invoke_callbacks_on_epoch_begin(self: Solver) -> None: + """ + Invokes all registered callbacks at the beginning of an epoch. + """ + for callback in self.callbacks_on_epoch_begin: + callback(self) + + def _invoke_callbacks_on_epoch_end(self: Solver) -> None: + """ + Invokes all registered callbacks at the end of an epoch. + """ + for callback in self.callbacks_on_epoch_end: + callback(self) + + def _invoke_callbacks_on_iter_begin(self: Solver) -> None: + """ + Invokes all registered callbacks at the beginning of an iteration. + """ + for callback in self.callbacks_on_iter_begin: + callback(self) + + def _invoke_callbacks_on_iter_end(self: Solver) -> None: + """ + Invokes all registered callbacks at the end of an iteration. + """ + for callback in self.callbacks_on_iter_end: + callback(self) diff --git a/examples/smc_reac/ppsci/solver/train.py b/examples/smc_reac/ppsci/solver/train.py new file mode 100644 index 0000000000..aec95a0dc4 --- /dev/null +++ b/examples/smc_reac/ppsci/solver/train.py @@ -0,0 +1,324 @@ +# 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. + +from __future__ import annotations + +import sys +import time +from typing import TYPE_CHECKING +from typing import Dict +from typing import Sequence +from typing import Union + +import paddle +from paddle.distributed.fleet.utils import hybrid_parallel_util as hpu +from paddle.framework import core + +from ppsci.solver import printer +from ppsci.utils import misc + +if TYPE_CHECKING: + from ppsci import solver + + +def _compute_batch_size( + input_dict: Dict[str, Union[paddle.Tensor, Sequence[paddle.Tensor]]] +) -> int: + """Compute batch size from given input dict. + + NOTE: Returned `batch_size` might be inaccurate, but it won't affect the correctness + of the training results because `batch_size` is now only used for timing. + + Args: + input_dict (Dict[str, Union[paddle.Tensor, Sequence[paddle.Tensor]]]): Given input dict. + + Returns: + int: Batch size of input dict. + """ + sample = next(iter(input_dict.values())) + if hasattr(sample, "shape"): + return sample.shape[0] + elif hasattr(sample, "__len__"): # Might be inaccurate here. + return len(sample) + else: + raise ValueError("Unsupported type of input dict value.") + + +def train_epoch_func(solver: "solver.Solver", epoch_id: int, log_freq: int): + """Train program for one epoch. + + Args: + solver (solver.Solver): Main solver. + epoch_id (int): Epoch id. + log_freq (int): Log training information every `log_freq` steps. + """ + batch_tic = time.perf_counter() + + for iter_id in range(1, solver.iters_per_epoch + 1): + solver._invoke_callbacks_on_iter_begin() + if solver.nvtx_flag: # only for nsight analysis + core.nvprof_nvtx_push( + f"Training iteration {solver.global_step + 1}" + ) # Training iteration + + total_batch_size = 0 + reader_cost = 0.0 + batch_cost = 0.0 + reader_tic = time.perf_counter() + + input_dicts = [] + label_dicts = [] + weight_dicts = [] + for _, _constraint in solver.constraint.items(): + # fetch data from data loader + if solver.nvtx_flag: # only for nsight analysis + core.nvprof_nvtx_push("Data load") + + try: + input_dict, label_dict, weight_dict = next(_constraint.data_iter) + except StopIteration: + _constraint.data_iter = iter(_constraint.data_loader) + input_dict, label_dict, weight_dict = next(_constraint.data_iter) + + if solver.nvtx_flag: # only for nsight analysis + core.nvprof_nvtx_pop() + + reader_cost += time.perf_counter() - reader_tic + + for v in input_dict.values(): + if hasattr(v, "stop_gradient"): + v.stop_gradient = False + + # gather each constraint's input, label, weight to a list + input_dicts.append(input_dict) + label_dicts.append(label_dict) + weight_dicts.append(weight_dict) + total_batch_size += _compute_batch_size(input_dict) + reader_tic = time.perf_counter() + + loss_dict = misc.Prettydefaultdict(float) + loss_dict["loss"] = 0.0 + # forward for every constraint, including model and equation expression + with solver.no_sync_context_manager(solver.world_size > 1, solver.model): + with solver.autocast_context_manager(solver.use_amp, solver.amp_level): + if solver.nvtx_flag: # only for nsight analysis + core.nvprof_nvtx_push("Loss computation") + + losses_all, losses_constraint = solver.forward_helper.train_forward( + tuple( + _constraint.output_expr + for _constraint in solver.constraint.values() + ), + input_dicts, + solver.model, + solver.constraint, + label_dicts, + weight_dicts, + ) + assert "loss" not in losses_all, ( + "Key 'loss' is not allowed in loss_dict for it is an preserved key" + " representing total loss, please use other name instead." + ) + + if solver.nvtx_flag: # only for nsight analysis + core.nvprof_nvtx_pop() # Loss computation + + # accumulate all losses + if solver.nvtx_flag: # only for nsight analysis + core.nvprof_nvtx_push("Loss aggregator") + + total_loss = solver.loss_aggregator(losses_all, solver.global_step) + if solver.update_freq > 1: + total_loss = total_loss / solver.update_freq + + loss_dict.update(losses_constraint) + loss_dict["loss"] = float(total_loss) + + if solver.nvtx_flag: # only for nsight analysis + core.nvprof_nvtx_pop() # Loss aggregator + + # backward + if solver.nvtx_flag: # only for nsight analysis + core.nvprof_nvtx_push("Loss backward") + + if solver.use_amp: + total_loss_scaled = solver.scaler.scale(total_loss) + total_loss_scaled.backward() + else: + total_loss.backward() + + if solver.nvtx_flag: # only for nsight analysis + core.nvprof_nvtx_pop() # Loss backward + + # update parameters + if iter_id % solver.update_freq == 0 or iter_id == solver.iters_per_epoch: + if solver.nvtx_flag: # only for nsight analysis + core.nvprof_nvtx_push("Optimizer update") + + if solver.world_size > 1: + # fuse + allreduce manually before optimization if use DDP + no_sync + # details in https://github.com/PaddlePaddle/Paddle/issues/48898#issuecomment-1343838622 + hpu.fused_allreduce_gradients(list(solver.model.parameters()), None) + if solver.use_amp: + solver.scaler.minimize(solver.optimizer, total_loss_scaled) + else: + solver.optimizer.step() + + if solver.nvtx_flag: # only for nsight analysis + core.nvprof_nvtx_pop() # Optimizer update + + solver.optimizer.clear_grad() + + # update learning rate by step + if solver.lr_scheduler is not None and not solver.lr_scheduler.by_epoch: + solver.lr_scheduler.step() + + if solver.benchmark_flag: + paddle.device.synchronize() + batch_cost += time.perf_counter() - batch_tic + + # update and log training information + solver.global_step += 1 + solver.train_time_info["reader_cost"].update(reader_cost) + solver.train_time_info["batch_cost"].update(batch_cost) + printer.update_train_loss(solver, loss_dict, total_batch_size) + if ( + solver.global_step % log_freq == 0 + or solver.global_step == 1 + or solver.global_step == solver.max_steps + ): + printer.log_train_info(solver, total_batch_size, epoch_id, iter_id) + + batch_tic = time.perf_counter() + + if solver.nvtx_flag: # only for nsight analysis + core.nvprof_nvtx_pop() # Training iteration + NVTX_STOP_ITER = 25 + if solver.global_step >= NVTX_STOP_ITER: + print( + f"Only run {NVTX_STOP_ITER} steps when 'NVTX' is set in environment" + " for nsight analysis. Exit now ......\n" + ) + core.nvprof_stop() + sys.exit(0) + + solver._invoke_callbacks_on_iter_end() + + +def train_LBFGS_epoch_func(solver: "solver.Solver", epoch_id: int, log_freq: int): + """Train function for one epoch with L-BFGS optimizer. + + NOTE: L-BFGS training program do not support AMP now. + + Args: + solver (solver.Solver): Main solver. + epoch_id (int): Epoch id. + log_freq (int): Log training information every `log_freq` steps. + """ + batch_tic = time.perf_counter() + + for iter_id in range(1, solver.iters_per_epoch + 1): + solver._invoke_callbacks_on_iter_begin() + loss_dict = misc.Prettydefaultdict(float) + loss_dict["loss"] = 0.0 + total_batch_size = 0 + reader_cost = 0.0 + batch_cost = 0.0 + reader_tic = time.perf_counter() + + input_dicts = [] + label_dicts = [] + weight_dicts = [] + for _, _constraint in solver.constraint.items(): + # fetch data from data loader + try: + input_dict, label_dict, weight_dict = next(_constraint.data_iter) + except StopIteration: + _constraint.data_iter = iter(_constraint.data_loader) + input_dict, label_dict, weight_dict = next(_constraint.data_iter) + reader_cost += time.perf_counter() - reader_tic + + for v in input_dict.values(): + if hasattr(v, "stop_gradient"): + v.stop_gradient = False + + # gather each constraint's input, label, weight to a list + input_dicts.append(input_dict) + label_dicts.append(label_dict) + weight_dicts.append(weight_dict) + total_batch_size += _compute_batch_size(input_dict) + reader_tic = time.perf_counter() + + def closure() -> paddle.Tensor: + """Forward-backward closure function for LBFGS optimizer. + + Returns: + paddle.Tensor: Computed loss scalar. + """ + with solver.no_sync_context_manager(solver.world_size > 1, solver.model): + with solver.autocast_context_manager(solver.use_amp, solver.amp_level): + # forward for every constraint, including model and equation expression + losses_all, losses_constraint = solver.forward_helper.train_forward( + tuple( + _constraint.output_expr + for _constraint in solver.constraint.values() + ), + input_dicts, + solver.model, + solver.constraint, + label_dicts, + weight_dicts, + ) + + # accumulate all losses + total_loss = solver.loss_aggregator(losses_all, solver.global_step) + loss_dict.update(losses_constraint) + loss_dict["loss"] = float(total_loss) + + # backward + solver.optimizer.clear_grad() + total_loss.backward() + + if solver.world_size > 1: + # fuse + allreduce manually before optimization if use DDP model + # details in https://github.com/PaddlePaddle/Paddle/issues/48898#issuecomment-1343838622 + hpu.fused_allreduce_gradients(list(solver.model.parameters()), None) + + return total_loss + + # update parameters + solver.optimizer.step(closure) + + # update learning rate by step + if solver.lr_scheduler is not None and not solver.lr_scheduler.by_epoch: + solver.lr_scheduler.step() + + if solver.benchmark_flag: + paddle.device.synchronize() + batch_cost += time.perf_counter() - batch_tic + + # update and log training information + solver.global_step += 1 + solver.train_time_info["reader_cost"].update(reader_cost) + solver.train_time_info["batch_cost"].update(batch_cost) + printer.update_train_loss(solver, loss_dict, total_batch_size) + if ( + solver.global_step % log_freq == 0 + or solver.global_step == 1 + or solver.global_step == solver.max_steps + ): + printer.log_train_info(solver, total_batch_size, epoch_id, iter_id) + + batch_tic = time.perf_counter() + solver._invoke_callbacks_on_iter_end() diff --git a/examples/smc_reac/ppsci/solver/visu.py b/examples/smc_reac/ppsci/solver/visu.py new file mode 100644 index 0000000000..80e11abe75 --- /dev/null +++ b/examples/smc_reac/ppsci/solver/visu.py @@ -0,0 +1,98 @@ +# 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. + +from __future__ import annotations + +import os +import os.path as osp +from typing import TYPE_CHECKING +from typing import Optional + +import paddle + +from ppsci.utils import misc + +if TYPE_CHECKING: + from ppsci import solver + + +def visualize_func(solver: "solver.Solver", epoch_id: Optional[int]): + """Visualization program. + + Args: + solver (solver.Solver): Main Solver. + epoch_id (Optional[int]): Epoch id. + """ + for _, _visualizer in solver.visualizer.items(): + all_input = misc.Prettydefaultdict(list) + all_output = misc.Prettydefaultdict(list) + + # NOTE: 'visualize_func' now do not apply data sharding(different from 'Solver.predict'), + # where every rank receive same input data and compute same output data + # (which will cause computational redundancy), + # but only the 0-rank(master) device save the visualization result into disk. + # TODO(HydrogenSulfate): This will be optimized in the future. + + input_dict = _visualizer.input_dict + batch_size = _visualizer.batch_size + num_samples = len(next(iter(input_dict.values()))) + batch_num = (num_samples + (batch_size - 1)) // batch_size + + for batch_id in range(batch_num): + batch_input_dict = {} + st = batch_id * batch_size + ed = min(num_samples, (batch_id + 1) * batch_size) + + # prepare batch input dict + for key in input_dict: + if not paddle.is_tensor(input_dict[key]): + batch_input_dict[key] = paddle.to_tensor( + input_dict[key][st:ed], paddle.get_default_dtype() + ) + else: + batch_input_dict[key] = input_dict[key][st:ed] + batch_input_dict[key].stop_gradient = False + + # forward + with solver.autocast_context_manager( + solver.use_amp, solver.amp_level + ), solver.no_grad_context_manager(solver.eval_with_no_grad): + batch_output_dict = solver.forward_helper.visu_forward( + _visualizer.output_expr, batch_input_dict, solver.model + ) + + # collect batch data with dtype fixed to float32 regardless of the dtypes of + # paddle runtime, which is most compatible with almost visualization tools. + for key, batch_input in batch_input_dict.items(): + all_input[key].append(batch_input.detach().astype("float32")) + for key, batch_output in batch_output_dict.items(): + all_output[key].append(batch_output.detach().astype("float32")) + + # concatenate all data + for key in all_input: + all_input[key] = paddle.concat(all_input[key]) + for key in all_output: + all_output[key] = paddle.concat(all_output[key]) + + # save visualization + with misc.RankZeroOnly(solver.rank) as is_master: + if is_master: + visual_dir = osp.join(solver.output_dir, "visual") + if epoch_id: + visual_dir = osp.join(visual_dir, f"epoch_{epoch_id}") + os.makedirs(visual_dir, exist_ok=True) + _visualizer.save( + osp.join(visual_dir, _visualizer.prefix), + {**all_input, **all_output}, + ) diff --git a/examples/smc_reac/ppsci/utils/__init__.py b/examples/smc_reac/ppsci/utils/__init__.py new file mode 100644 index 0000000000..3382eee856 --- /dev/null +++ b/examples/smc_reac/ppsci/utils/__init__.py @@ -0,0 +1,66 @@ +# 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. + +# NOTE: Put config module import at the top level for register default config(s) in +# ConfigStore at the beginning of ppsci +from ppsci.utils import config # isort:skip # noqa: F401 +from ppsci.utils import ema +from ppsci.utils import initializer +from ppsci.utils import logger +from ppsci.utils import misc +from ppsci.utils import reader +from ppsci.utils import writer +from ppsci.utils.checker import dynamic_import_to_globals +from ppsci.utils.checker import run_check +from ppsci.utils.checker import run_check_mesh +from ppsci.utils.expression import ExpressionSolver +from ppsci.utils.misc import AverageMeter +from ppsci.utils.misc import set_random_seed +from ppsci.utils.reader import load_csv_file +from ppsci.utils.reader import load_mat_file +from ppsci.utils.reader import load_npz_file +from ppsci.utils.reader import load_vtk_file +from ppsci.utils.reader import load_vtk_with_time_file +from ppsci.utils.save_load import load_checkpoint +from ppsci.utils.save_load import load_pretrain +from ppsci.utils.save_load import save_checkpoint +from ppsci.utils.symbolic import lambdify +from ppsci.utils.writer import save_csv_file +from ppsci.utils.writer import save_tecplot_file + +__all__ = [ + "AverageMeter", + "ExpressionSolver", + "initializer", + "logger", + "misc", + "ema", + "reader", + "writer", + "load_csv_file", + "load_mat_file", + "load_npz_file", + "load_vtk_file", + "load_vtk_with_time_file", + "save_csv_file", + "save_tecplot_file", + "dynamic_import_to_globals", + "run_check", + "run_check_mesh", + "set_random_seed", + "load_checkpoint", + "load_pretrain", + "save_checkpoint", + "lambdify", +] diff --git a/examples/smc_reac/ppsci/utils/callbacks.py b/examples/smc_reac/ppsci/utils/callbacks.py new file mode 100644 index 0000000000..a753432b55 --- /dev/null +++ b/examples/smc_reac/ppsci/utils/callbacks.py @@ -0,0 +1,136 @@ +# 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. + +import importlib.util +import inspect +import sys +import traceback +from os import path as osp +from typing import Any + +from hydra.core.hydra_config import HydraConfig +from hydra.experimental.callback import Callback +from omegaconf import DictConfig + +from ppsci.utils import config as config_module +from ppsci.utils import logger +from ppsci.utils import misc + +RUNTIME_EXIT_CODE = 1 # for other errors +VALIDATION_ERROR_EXIT_CODE = 2 # for invalid argument detected in config file + + +class InitCallback(Callback): + """Callback class for: + 1. Parse config dict from given yaml file and check its validity. + 2. Fixing random seed to 'config.seed'. + 3. Initialize logger while creating output directory(if not exist). + 4. Enable prim mode if specified. + + NOTE: This callback is mainly for reducing unnecessary duplicate code in each + examples code when runing with hydra. + + This callback should be added to hydra config file as follows: + + ``` yaml hl_lines="7-11" + # content of example.yaml below + hydra: + run: + ... + job: + ... + callbacks: + init_callback: + _target_: ppsci.utils.callbacks.InitCallback # <-- add callback at here + xxx_callback: + _target_: ppsci.utils.callbacks.XxxCallback # <-- add more callback here + sweep: + ... + ... + ... + ``` + """ + + def on_job_start(self, config: DictConfig, **kwargs: Any) -> None: + if importlib.util.find_spec("pydantic") is not None: + from pydantic import ValidationError + else: + logger.error( + f"ModuleNotFoundError at {__file__}:{inspect.currentframe().f_lineno}\n" + "Please install pydantic with `pip install pydantic` when set callbacks" + " in your config yaml." + ) + sys.exit(RUNTIME_EXIT_CODE) + + # check given cfg using pre-defined pydantic schema in 'SolverConfig', + # error(s) will be printed and exit program if any checking failed at this step + try: + _model_pydantic = config_module.SolverConfig(**dict(config)) + full_cfg = DictConfig(_model_pydantic.model_dump()) + except ValidationError as e: + print(e) + sys.exit(VALIDATION_ERROR_EXIT_CODE) + except Exception as e: + print(e) + sys.exit(RUNTIME_EXIT_CODE) + + # fix random seed for reproducibility + misc.set_random_seed(full_cfg.seed) + + # initialize logger while creating output directory + logger.init_logger( + "ppsci", + osp.join(full_cfg.output_dir, f"{full_cfg.mode}.log") + if full_cfg.output_dir and full_cfg.mode not in ["export", "infer"] + else None, + full_cfg.log_level, + ) + + # set device before running into example function + if "device" in full_cfg: + import paddle + + if isinstance(full_cfg.device, str): + paddle.device.set_device(full_cfg.device) + + try: + if "num" in HydraConfig.get().job: + jobs_id = HydraConfig.get().job.num + else: + jobs_id = None + if "n_jobs" in HydraConfig.get().launcher: + parallel_jobs_num = HydraConfig.get().launcher.n_jobs + else: + parallel_jobs_num = None + + if jobs_id and parallel_jobs_num: + job_device_id = jobs_id % parallel_jobs_num + device_type = paddle.get_device().split(":")[0] + logger.message( + f"Running job {jobs_id} on device {device_type}:{job_device_id}(logical device id)" + ) + paddle.set_device(f"{device_type}:{job_device_id}") + except Exception as e: + print(e) + traceback.print_exc() + sys.exit(RUNTIME_EXIT_CODE) + + # enable prim if specified + if "prim" in full_cfg and bool(full_cfg.prim): + # Mostly for compiler running with dy2st. + from paddle.framework import core + + core.set_prim_eager_enabled(True) + core._set_prim_all_enabled(True) + logger.message("Prim mode is enabled.") diff --git a/examples/smc_reac/ppsci/utils/checker.py b/examples/smc_reac/ppsci/utils/checker.py new file mode 100644 index 0000000000..8991a89e2c --- /dev/null +++ b/examples/smc_reac/ppsci/utils/checker.py @@ -0,0 +1,287 @@ +# 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. + +from __future__ import annotations + +import importlib.util +import traceback +from typing import Dict +from typing import Sequence +from typing import Union + +import paddle + +from ppsci.utils import logger + +__all__ = [ + "run_check", + "run_check_mesh", + "dynamic_import_to_globals", +] + + +def run_check() -> None: + """Check whether PaddleScience is installed correctly and running successfully on + your system. + + Examples: + >>> import ppsci + >>> ppsci.utils.run_check() # doctest: +SKIP + """ + # test demo code below. + import ppsci + + try: + ppsci.utils.set_random_seed(42) + ppsci.utils.logger.init_logger() + model = ppsci.arch.MLP(("x", "y"), ("u", "v", "p"), 3, 16, "tanh") + + equation = {"NavierStokes": ppsci.equation.NavierStokes(0.01, 1.0, 2, False)} + + geom = {"rect": ppsci.geometry.Rectangle((-0.05, -0.05), (0.05, 0.05))} + + ITERS_PER_EPOCH = 5 + train_dataloader_cfg = { + "dataset": "IterableNamedArrayDataset", + "iters_per_epoch": ITERS_PER_EPOCH, + } + + NPOINT_PDE = 8**2 + pde_constraint = 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={ + "continuity": 0.0001, + "momentum_x": 0.0001, + "momentum_y": 0.0001, + }, + name="EQ", + ) + constraint = {pde_constraint.name: pde_constraint} + + residual_validator = ppsci.validate.GeometryValidator( + equation["NavierStokes"].equations, + {"continuity": 0, "momentum_x": 0, "momentum_y": 0}, + geom["rect"], + { + "dataset": "NamedArrayDataset", + "total_size": 8**2, + "batch_size": 32, + "sampler": {"name": "BatchSampler"}, + }, + ppsci.loss.MSELoss("sum"), + evenly=True, + metric={"MSE": ppsci.metric.MSE(False)}, + name="Residual", + ) + validator = {residual_validator.name: residual_validator} + + EPOCHS = 2 + optimizer = ppsci.optimizer.Adam(0.001)(model) + solver = ppsci.solver.Solver( + model, + constraint, + None, + optimizer, + None, + EPOCHS, + ITERS_PER_EPOCH, + device=paddle.device.get_device(), + equation=equation, + validator=validator, + ) + solver.train() + solver.eval(EPOCHS) + except Exception as e: + traceback.print_exc() + logger.error( + f"PaddleScience meets some problem with \n {repr(e)} \nplease check whether " + "Paddle's version and PaddleScience's version are both correct." + ) + else: + logger.message("PaddleScience is installed successfully.✨ 🍰 ✨") + + +def run_check_mesh() -> None: + """Check whether geometry packages is installed correctly and `ppsci.geometry.Mesh` + can running successfully on your system. + + Examples: + >>> import ppsci + >>> ppsci.utils.run_check_mesh() # doctest: +SKIP + """ + # test demo code below. + if importlib.util.find_spec("open3d") is None: + raise ModuleNotFoundError( + "Please install open3d first with: " "`pip install open3d`" + ) + if importlib.util.find_spec("pysdf") is None: + raise ModuleNotFoundError( + "Please install pysdf first with: `pip install pysdf`" + ) + if importlib.util.find_spec("pymesh") is None: + raise ModuleNotFoundError( + "Please install pymesh first as " + "https://paddlescience-docs.readthedocs.io/zh/latest/zh/install_setup/#__tabbed_4_4" + ) + + import numpy as np + import pymesh + + import ppsci + + try: + ppsci.utils.set_random_seed(42) + ppsci.utils.logger.init_logger() + model = ppsci.arch.MLP(("x", "y"), ("u", "v", "p"), 3, 16, "tanh") + + equation = {"NavierStokes": ppsci.equation.NavierStokes(0.01, 1.0, 2, False)} + + # create a 1x1x1 simple cube geometry + vertices = np.array( + [ + [0.0, 0.0, 0.0], + [1.0, 0.0, 0.0], + [0.0, 0.0, 1.0], + [1.0, 0.0, 1.0], + [0.0, 1.0, 0.0], + [1.0, 1.0, 0.0], + [0.0, 1.0, 1.0], + [1.0, 1.0, 1.0], + ] + ) # 8 vertices for mesh + faces = np.array( + [ + [4, 7, 5], + [4, 6, 7], + [0, 2, 4], + [2, 6, 4], + [0, 1, 2], + [1, 3, 2], + [1, 5, 7], + [1, 7, 3], + [2, 3, 7], + [2, 7, 6], + [0, 4, 1], + [1, 4, 5], + ] + ) # 12 triangle faces for mesh + box_mesh = pymesh.form_mesh(vertices, faces) + geom = {"rect": ppsci.geometry.Mesh(box_mesh)} + + ITERS_PER_EPOCH = 5 + train_dataloader_cfg = { + "dataset": "IterableNamedArrayDataset", + "iters_per_epoch": ITERS_PER_EPOCH, + } + + NPOINT_PDE = 8**2 + pde_constraint = 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"), + weight_dict={ + "continuity": "sdf", + "momentum_x": "sdf", + "momentum_y": "sdf", + }, + name="EQ", + ) + constraint = {pde_constraint.name: pde_constraint} + + residual_validator = ppsci.validate.GeometryValidator( + equation["NavierStokes"].equations, + {"continuity": 0, "momentum_x": 0, "momentum_y": 0}, + geom["rect"], + { + "dataset": "NamedArrayDataset", + "total_size": 8**2, + "batch_size": 32, + "sampler": {"name": "BatchSampler"}, + }, + ppsci.loss.MSELoss("sum"), + metric={"MSE": ppsci.metric.MSE(False)}, + name="Residual", + ) + validator = {residual_validator.name: residual_validator} + + EPOCHS = 2 + optimizer = ppsci.optimizer.Adam(0.001)(model) + solver = ppsci.solver.Solver( + model, + constraint, + None, + optimizer, + None, + EPOCHS, + ITERS_PER_EPOCH, + device=paddle.device.get_device(), + equation=equation, + validator=validator, + ) + solver.train() + solver.eval(EPOCHS) + except Exception as e: + traceback.print_exc() + logger.error( + f"PaddleScience meets some problem with \n {repr(e)} \nplease check whether " + "open3d, pysdf, pybind11, PyMesh are all installed correctly." + ) + else: + logger.message("ppsci.geometry.Mesh module running successfully.✨ 🍰 ✨") + + +def dynamic_import_to_globals( + names: Union[str, Sequence[str]], alias: Dict[str, str] = None +) -> bool: + """Import module and add it to globals() by given names dynamically. + + Args: + names (Union[str, Sequence[str]]): Module name or sequence of module names. + alias (Dict[str, str]): Alias name of module when imported into globals(). + + Returns: + bool: Whether given names all exist. + """ + if isinstance(names, str): + names = (names,) + + if alias is None: + alias = {} + + for name in names: + # find module in environment by it's name and alias(if given) + module_spec = importlib.util.find_spec(name) + if module_spec is None and name in alias: + module_spec = importlib.util.find_spec(alias[name]) + + # log error and return False if module do not exist + if not module_spec: + logger.error(f"Module {name} should be installed first.") + return False + + # module exist, add to globals() if not in globals() + add_name = name + if add_name in alias: + add_name = alias[add_name] + if add_name not in globals(): + globals()[add_name] = importlib.import_module(name) + + return True diff --git a/examples/smc_reac/ppsci/utils/config.py b/examples/smc_reac/ppsci/utils/config.py new file mode 100644 index 0000000000..d887a75d3d --- /dev/null +++ b/examples/smc_reac/ppsci/utils/config.py @@ -0,0 +1,457 @@ +# 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. + +from __future__ import annotations + +import importlib.util +from typing import Mapping +from typing import Optional +from typing import Tuple + +from typing_extensions import Literal + +__all__ = [] + +if importlib.util.find_spec("pydantic") is not None: + try: + from hydra.core.config_store import ConfigStore + from omegaconf import OmegaConf + from pydantic import BaseModel + from pydantic import field_validator + from pydantic import model_validator + from pydantic_core.core_schema import ValidationInfo + + __all__.append("SolverConfig") + + class EMAConfig(BaseModel): + use_ema: bool = False + decay: float = 0.9 + avg_freq: int = 1 + + @field_validator("decay") + def decay_check(cls, v): + if v <= 0 or v >= 1: + raise ValueError( + f"'ema.decay' should be in (0, 1) when is type of float, but got {v}" + ) + return v + + @field_validator("avg_freq") + def avg_freq_check(cls, v): + if v <= 0: + raise ValueError( + "'ema.avg_freq' should be a positive integer when is type of int, " + f"but got {v}" + ) + return v + + class SWAConfig(BaseModel): + use_swa: bool = False + avg_freq: int = 1 + avg_range: Optional[Tuple[int, int]] = None + + @field_validator("avg_range") + def avg_range_check(cls, v, info: ValidationInfo): + if isinstance(v, tuple) and v[0] > v[1]: + raise ValueError( + f"'swa.avg_range' should be a valid range, but got {v}." + ) + if isinstance(v, tuple) and v[0] < 0: + raise ValueError( + "The start epoch of 'swa.avg_range' should be a non-negtive integer" + f" , but got {v[0]}." + ) + return v + + @field_validator("avg_freq") + def avg_freq_check(cls, v): + if v <= 0: + raise ValueError( + "'swa.avg_freq' should be a positive integer when is type of int, " + f"but got {v}" + ) + return v + + class TrainConfig(BaseModel): + """ + Schema of training config for pydantic validation. + """ + + epochs: int = 1 + iters_per_epoch: int = 20 + update_freq: int = 1 + save_freq: int = 0 + eval_during_train: bool = False + start_eval_epoch: int = 1 + eval_freq: int = 1 + checkpoint_path: Optional[str] = None + pretrained_model_path: Optional[str] = None + ema: Optional[EMAConfig] = None + swa: Optional[SWAConfig] = None + + # Fine-grained validator(s) below + @field_validator("epochs") + def epochs_check(cls, v): + if v <= 0: + raise ValueError( + "'TRAIN.epochs' should be a positive integer when is type of int, " + f"but got {v}" + ) + return v + + @field_validator("iters_per_epoch") + def iters_per_epoch_check(cls, v): + if v <= 0 and v != -1: + raise ValueError( + f"'TRAIN.iters_per_epoch' received an invalid value({v}), " + "but is expected one of: \n" + "* A positive integer, to manually specify the number of iterations per epoch, " + "which is commonly used in PINN training.\n" + "* -1, to automatically set the number of iterations per epoch to " + "the length of dataloader of given constraint, which is commonly " + f"used in data-driven training.\n" + ) + return v + + @field_validator("update_freq") + def update_freq_check(cls, v): + if v <= 0: + raise ValueError( + "'TRAIN.update_freq' should be a positive integer when is type of int" + f", but got {v}" + ) + return v + + @field_validator("save_freq") + def save_freq_check(cls, v): + if v < 0: + raise ValueError( + "'TRAIN.save_freq' should be a non-negtive integer when is type of int" + f", but got {v}" + ) + return v + + @field_validator("start_eval_epoch") + def start_eval_epoch_check(cls, v, info: ValidationInfo): + if info.data["eval_during_train"]: + if v <= 0: + raise ValueError( + f"'TRAIN.start_eval_epoch' should be a positive integer when " + f"'TRAIN.eval_during_train' is True, but got {v}" + ) + return v + + @field_validator("eval_freq") + def eval_freq_check(cls, v, info: ValidationInfo): + if info.data["eval_during_train"]: + if v <= 0: + raise ValueError( + f"'TRAIN.eval_freq' should be a positive integer when " + f"'TRAIN.eval_during_train' is True, but got {v}" + ) + return v + + @model_validator(mode="after") + def ema_swa_checker(self): + if (self.ema and self.swa) and (self.ema.use_ema and self.swa.use_swa): + raise ValueError( + "Cannot enable both EMA and SWA at the same time, " + "please disable at least one of them." + ) + return self + + @model_validator(mode="after") + def swa_avg_range_checker(self): + if ( + self.swa + and self.swa.use_swa + and self.swa.avg_range[1] > self.epochs + ): + raise ValueError( + "The end epoch of 'swa.avg_range' should not be lager than " + f"'epochs'({self.epochs}), but got {self.swa.avg_range[1]}." + ) + return self + + class EvalConfig(BaseModel): + """ + Schema of evaluation config for pydantic validation. + """ + + pretrained_model_path: Optional[str] = None + eval_with_no_grad: bool = False + compute_metric_by_batch: bool = False + batch_size: Optional[int] = 256 + + @field_validator("batch_size") + def batch_size_check(cls, v): + if isinstance(v, int) and v <= 0: + raise ValueError( + f"'EVAL.batch_size' should be greater than 0 or None, but got {v}" + ) + return v + + class InferConfig(BaseModel): + """ + Schema of inference config for pydantic validation. + """ + + pretrained_model_path: Optional[str] = None + export_path: str = "./inference" + pdmodel_path: Optional[str] = None + pdiparams_path: Optional[str] = None + onnx_path: Optional[str] = None + device: Literal["cpu", "gpu", "npu", "xpu", "sdaa"] = "cpu" + engine: Literal["native", "tensorrt", "onnx", "mkldnn"] = "native" + precision: Literal["fp32", "fp16", "int8"] = "fp32" + ir_optim: bool = True + min_subgraph_size: int = 30 + gpu_mem: int = 2000 + gpu_id: int = 0 + max_batch_size: int = 1024 + num_cpu_threads: int = 10 + batch_size: Optional[int] = 256 + + # Fine-grained validator(s) below + @field_validator("engine") + def engine_check(cls, v, info: ValidationInfo): + if v == "tensorrt" and info.data["device"] != "gpu": + raise ValueError( + "'INFER.device' should be 'gpu' when 'INFER.engine' is 'tensorrt', " + f"but got '{info.data['device']}'" + ) + if v == "mkldnn" and info.data["device"] != "cpu": + raise ValueError( + "'INFER.device' should be 'cpu' when 'INFER.engine' is 'mkldnn', " + f"but got '{info.data['device']}'" + ) + + return v + + @field_validator("min_subgraph_size") + def min_subgraph_size_check(cls, v): + if v <= 0: + raise ValueError( + "'INFER.min_subgraph_size' should be greater than 0, " + f"but got {v}" + ) + return v + + @field_validator("gpu_mem") + def gpu_mem_check(cls, v): + if v <= 0: + raise ValueError( + "'INFER.gpu_mem' should be greater than 0, " f"but got {v}" + ) + return v + + @field_validator("gpu_id") + def gpu_id_check(cls, v): + if v < 0: + raise ValueError( + "'INFER.gpu_id' should be greater than or equal to 0, " + f"but got {v}" + ) + return v + + @field_validator("max_batch_size") + def max_batch_size_check(cls, v): + if v <= 0: + raise ValueError( + "'INFER.max_batch_size' should be greater than 0, " + f"but got {v}" + ) + return v + + @field_validator("num_cpu_threads") + def num_cpu_threads_check(cls, v): + if v < 0: + raise ValueError( + "'INFER.num_cpu_threads' should be greater than or equal to 0, " + f"but got {v}" + ) + return v + + @field_validator("batch_size") + def batch_size_check(cls, v): + if isinstance(v, int) and v <= 0: + raise ValueError( + f"'INFER.batch_size' should be greater than 0 or None, but got {v}" + ) + return v + + class SolverConfig(BaseModel): + """ + Schema of global config for pydantic validation. + """ + + # Global settings config + mode: Literal["train", "eval", "export", "infer"] = "train" + output_dir: Optional[str] = None + log_freq: int = 20 + seed: int = 42 + use_vdl: bool = False + use_tbd: bool = False + wandb_config: Mapping = {} + use_wandb: bool = False + device: Literal["cpu", "gpu", "xpu", "sdaa", None] = None + use_amp: bool = False + amp_level: Literal["O0", "O1", "O2", "OD"] = "O1" + to_static: bool = False + prim: bool = False + log_level: Literal["debug", "info", "warning", "error"] = "info" + + # Training related config + TRAIN: Optional[TrainConfig] = None + + # Evaluation related config + EVAL: Optional[EvalConfig] = None + + # Inference related config + INFER: Optional[InferConfig] = None + + # Fine-grained validator(s) below + @field_validator("log_freq") + def log_freq_check(cls, v): + if v <= 0: + raise ValueError( + "'log_freq' should be a non-negtive integer when is type of int" + f", but got {v}" + ) + return v + + @field_validator("seed") + def seed_check(cls, v): + if v < 0: + raise ValueError( + f"'seed' should be a non-negtive integer, but got {v}" + ) + return v + + @field_validator("use_wandb") + def use_wandb_check(cls, v, info: ValidationInfo): + if v and not isinstance(info.data["wandb_config"], dict): + raise ValueError( + "'wandb_config' should be a dict when 'use_wandb' is True, " + f"but got {info.data['wandb_config'].__class__.__name__}" + ) + return v + + # Register 'XXXConfig' as default node, so as to be used as default config in *.yaml + """ + #### xxx.yaml #### + defaults: + - ppsci_default <-- 'ppsci_default' used here + - TRAIN: train_default <-- 'train_default' used here + - TRAIN/ema: ema_default <-- 'ema_default' used here + - TRAIN/swa: swa_default <-- 'swa_default' used here + - EVAL: eval_default <-- 'eval_default' used here + - INFER: infer_default <-- 'infer_default' used here + - _self_ <-- config defined in current yaml + + mode: train + seed: 42 + ... + ... + ################## + """ + + cs = ConfigStore.instance() + + global_default_cfg = SolverConfig().model_dump() + omegaconf_dict_config = OmegaConf.create(global_default_cfg) + cs.store(name="ppsci_default", node=omegaconf_dict_config) + + train_default_cfg = TrainConfig().model_dump() + train_omegaconf_dict_config = OmegaConf.create(train_default_cfg) + cs.store(group="TRAIN", name="train_default", node=train_omegaconf_dict_config) + + ema_default_cfg = EMAConfig().model_dump() + ema_omegaconf_dict_config = OmegaConf.create(ema_default_cfg) + cs.store(group="TRAIN/ema", name="ema_default", node=ema_omegaconf_dict_config) + + swa_default_cfg = SWAConfig().model_dump() + swa_omegaconf_dict_config = OmegaConf.create(swa_default_cfg) + cs.store(group="TRAIN/swa", name="swa_default", node=swa_omegaconf_dict_config) + + eval_default_cfg = EvalConfig().model_dump() + eval_omegaconf_dict_config = OmegaConf.create(eval_default_cfg) + cs.store(group="EVAL", name="eval_default", node=eval_omegaconf_dict_config) + + infer_default_cfg = InferConfig().model_dump() + infer_omegaconf_dict_config = OmegaConf.create(infer_default_cfg) + cs.store(group="INFER", name="infer_default", node=infer_omegaconf_dict_config) + + exclude_keys_default = [ + "mode", + "output_dir", + "log_freq", + "seed", + "use_vdl", + "use_tbd", + "wandb_config", + "use_wandb", + "device", + "use_amp", + "amp_level", + "to_static", + "prim", + "log_level", + "TRAIN.save_freq", + "TRAIN.eval_during_train", + "TRAIN.start_eval_epoch", + "TRAIN.eval_freq", + "TRAIN.checkpoint_path", + "TRAIN.pretrained_model_path", + "EVAL.pretrained_model_path", + "EVAL.eval_with_no_grad", + "EVAL.compute_metric_by_batch", + "EVAL.batch_size", + "INFER.pretrained_model_path", + "INFER.export_path", + "INFER.pdmodel_path", + "INFER.pdiparams_path", + "INFER.onnx_path", + "INFER.device", + "INFER.engine", + "INFER.precision", + "INFER.ir_optim", + "INFER.min_subgraph_size", + "INFER.gpu_mem", + "INFER.gpu_id", + "INFER.max_batch_size", + "INFER.num_cpu_threads", + "INFER.batch_size", + ] + cs.store( + group="hydra/job/config/override_dirname/exclude_keys", + name="exclude_keys_default", + node=exclude_keys_default, + ) + except ImportError as e: + from ppsci.utils import logger + + logger.error(e) + logger.error( + "paddlesci requires pydantic>=2.5.0; otherwise, built-in examples may not run properly." + ) + except Exception as e: + raise e + +else: + from ppsci.utils import logger + + logger.error( + "paddlesci requires pydantic>=2.5.0; otherwise, built-in examples may not run properly." + ) diff --git a/examples/smc_reac/ppsci/utils/download.py b/examples/smc_reac/ppsci/utils/download.py new file mode 100644 index 0000000000..291703e2d2 --- /dev/null +++ b/examples/smc_reac/ppsci/utils/download.py @@ -0,0 +1,285 @@ +# 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. + +from __future__ import annotations + +import hashlib +import os +import os.path as osp +import shutil +import tarfile +import time +import zipfile + +import requests +import tqdm + +from ppsci.utils import logger +from ppsci.utils import misc + +__all__ = ["get_weights_path_from_url"] + +WEIGHTS_HOME = osp.expanduser("~/.paddlesci/weights") + +DOWNLOAD_RETRY_LIMIT = 3 + + +def is_url(path): + """ + Whether path is URL. + + Args: + path (str): URL string or not. + """ + return path.startswith("http://") or path.startswith("https://") + + +def get_weights_path_from_url(url, md5sum=None): + """Get weights path from WEIGHT_HOME, if not exists, + download it from url. + + Args: + url (str): Download url + md5sum (str): md5 sum of download package + + Returns: + str: a local path to save downloaded weights. + """ + path = get_path_from_url(url, WEIGHTS_HOME, md5sum) + return path + + +def _map_path(url, root_dir): + # parse path after download under root_dir + fname = osp.split(url)[-1] + fpath = fname + return osp.join(root_dir, fpath) + + +def get_path_from_url(url, root_dir, md5sum=None, check_exist=True, decompress=True): + """Download from given url to root_dir. + if file or directory specified by url is exists under + root_dir, return the path directly, otherwise download + from url and decompress it, return the path. + + Args: + url (str): Download url + root_dir (str): Root dir for downloading, it should be + WEIGHTS_HOME or DATASET_HOME + md5sum (str): md5 sum of download package + + Returns: + str: a local path to save downloaded models & weights & datasets. + """ + if not is_url(url): + raise ValueError(f"Given url({url}) is not valid") + # parse path after download to decompress under root_dir + fullpath = _map_path(url, root_dir) + # Mainly used to solve the problem of downloading data from different + # machines in the case of multiple machines. Different nodes will download + # data, and the same node will only download data once. + rank_id_curr_node = int(os.environ.get("PADDLE_RANK_IN_NODE", 0)) + + if osp.exists(fullpath) and check_exist and _md5check(fullpath, md5sum): + logger.message(f"Found {fullpath} already in {WEIGHTS_HOME}, skip downloading.") + else: + with misc.RankZeroOnly(rank_id_curr_node) as is_master: + if is_master: + fullpath = _download(url, root_dir, md5sum) + + if decompress and (tarfile.is_tarfile(fullpath) or zipfile.is_zipfile(fullpath)): + with misc.RankZeroOnly(rank_id_curr_node) as is_master: + if is_master: + fullpath = _decompress(fullpath) + + return fullpath + + +def _download(url, path, md5sum=None): + """ + Download from url, save to path. + + url (str): Download url + path (str): Download to given path + """ + if not osp.exists(path): + os.makedirs(path) + + fname = osp.split(url)[-1] + fullname = osp.join(path, fname) + retry_cnt = 0 + + while not (osp.exists(fullname) and _md5check(fullname, md5sum)): + if retry_cnt < DOWNLOAD_RETRY_LIMIT: + retry_cnt += 1 + else: + raise RuntimeError(f"Download from {url} failed. " "Retry limit reached") + + logger.message(f"Downloading {fname} from {url}") + + try: + req = requests.get(url, stream=True) + except Exception as e: # requests.exceptions.ConnectionError + logger.warning( + f"Downloading {fname} from {url} failed {retry_cnt + 1} times with exception {str(e)}" + ) + time.sleep(1) + continue + + if req.status_code != 200: + raise RuntimeError( + f"Downloading from {url} failed with code " f"{req.status_code}!" + ) + + # For protecting download interrupted, download to + # tmp_fullname firstly, move tmp_fullname to fullname + # after download finished + tmp_fullname = fullname + "_tmp" + total_size = req.headers.get("content-length") + with open(tmp_fullname, "wb") as f: + if total_size: + with tqdm.tqdm(total=(int(total_size) + 1023) // 1024) as pbar: + for chunk in req.iter_content(chunk_size=1024): + f.write(chunk) + pbar.update(1) + else: + for chunk in req.iter_content(chunk_size=1024): + if chunk: + f.write(chunk) + shutil.move(tmp_fullname, fullname) + logger.message(f"Finish downloading pretrained model and saved to {fullname}") + + return fullname + + +def _md5check(fullname, md5sum=None): + if md5sum is None: + return True + + logger.message(f"File {fullname} md5 checking...") + md5 = hashlib.md5() + with open(fullname, "rb") as f: + for chunk in iter(lambda: f.read(4096), b""): + md5.update(chunk) + calc_md5sum = md5.hexdigest() + + if calc_md5sum != md5sum: + logger.error( + f"File {fullname} md5 check failed, {calc_md5sum}(calc) != " + f"{md5sum}(base)" + ) + return False + return True + + +def _decompress(fname): + """ + Decompress for zip and tar file + """ + logger.message(f"Decompressing {fname}...") + + # For protecting decompressing interrupted, + # decompress to fpath_tmp directory firstly, if decompress + # succeed, move decompress files to fpath and delete + # fpath_tmp and remove download compress file. + + if tarfile.is_tarfile(fname): + uncompressed_path = _uncompress_file_tar(fname) + elif zipfile.is_zipfile(fname): + uncompressed_path = _uncompress_file_zip(fname) + else: + raise TypeError(f"Unsupported compress file type {fname}") + + return uncompressed_path + + +def _uncompress_file_zip(filepath): + with zipfile.ZipFile(filepath, "r") as files: + file_list = files.namelist() + + file_dir = os.path.dirname(filepath) + + if _is_a_single_file(file_list): + rootpath = file_list[0] + uncompressed_path = os.path.join(file_dir, rootpath) + + for item in file_list: + files.extract(item, file_dir) + + elif _is_a_single_dir(file_list): + rootpath = os.path.splitext(file_list[0])[0].split(os.sep)[-1] + uncompressed_path = os.path.join(file_dir, rootpath) + + for item in file_list: + files.extract(item, file_dir) + + else: + rootpath = os.path.splitext(filepath)[0].split(os.sep)[-1] + uncompressed_path = os.path.join(file_dir, rootpath) + if not os.path.exists(uncompressed_path): + os.makedirs(uncompressed_path) + for item in file_list: + files.extract(item, os.path.join(file_dir, rootpath)) + + return uncompressed_path + + +def _uncompress_file_tar(filepath, mode="r:*"): + with tarfile.open(filepath, mode) as files: + file_list = files.getnames() + + file_dir = os.path.dirname(filepath) + + if _is_a_single_file(file_list): + rootpath = file_list[0] + uncompressed_path = os.path.join(file_dir, rootpath) + for item in file_list: + files.extract(item, file_dir) + elif _is_a_single_dir(file_list): + rootpath = os.path.splitext(file_list[0])[0].split(os.sep)[-1] + uncompressed_path = os.path.join(file_dir, rootpath) + for item in file_list: + files.extract(item, file_dir) + else: + rootpath = os.path.splitext(filepath)[0].split(os.sep)[-1] + uncompressed_path = os.path.join(file_dir, rootpath) + if not os.path.exists(uncompressed_path): + os.makedirs(uncompressed_path) + + for item in file_list: + files.extract(item, os.path.join(file_dir, rootpath)) + + return uncompressed_path + + +def _is_a_single_file(file_list): + if len(file_list) == 1 and file_list[0].find(os.sep) < -1: + return True + return False + + +def _is_a_single_dir(file_list): + new_file_list = [] + for file_path in file_list: + if "/" in file_path: + file_path = file_path.replace("/", os.sep) + elif "\\" in file_path: + file_path = file_path.replace("\\", os.sep) + new_file_list.append(file_path) + + file_name = new_file_list[0].split(os.sep)[0] + for i in range(1, len(new_file_list)): + if file_name != new_file_list[i].split(os.sep)[0]: + return False + return True diff --git a/examples/smc_reac/ppsci/utils/ema.py b/examples/smc_reac/ppsci/utils/ema.py new file mode 100644 index 0000000000..690ee6fda8 --- /dev/null +++ b/examples/smc_reac/ppsci/utils/ema.py @@ -0,0 +1,172 @@ +# Copyright (c) 2024 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. + +from __future__ import annotations + +import itertools +from typing import Dict +from typing import Optional + +import paddle +from paddle import nn + +__all__ = [ + "AveragedModel", + "ExponentialMovingAverage", + "StochasticWeightAverage", +] + + +class AveragedModel(nn.Layer): + """Base class for Averaged Model. + + Args: + model (nn.Layer): The model to be averaged. + decay (float): The decay rate for averaging. + """ + + def __init__(self, model: nn.Layer, decay: Optional[float] = None): + super().__init__() + self.model = model # As a quick reference to online model + self.decay = decay + + self.params_shadow: Dict[str, paddle.Tensor] = {} # ema param or buffer + self.params_backup: Dict[str, paddle.Tensor] = {} # used for apply and restore + for name, param_or_buffer in itertools.chain( + self.model.named_parameters(), self.model.named_buffers() + ): + self.params_shadow[name] = param_or_buffer.clone().detach() + + self.register_buffer("n_avg", paddle.to_tensor(0, "int64"), True) + + def _update_fn_( + self, + shadow_param: paddle.Tensor, + model_param: paddle.Tensor, + step: paddle.Tensor, + ): + raise NotImplementedError("AveragedModel._update_fn_ should be implemented.") + + def update(self): + for name, param_or_buffer in itertools.chain( + self.model.named_parameters(), self.model.named_buffers() + ): + if not param_or_buffer.stop_gradient: + assert ( + name in self.params_shadow + ), f"Parameter: {name} should be in params_shadow dict, but not found." + + # only update floating and complex data + if paddle.is_floating_point(param_or_buffer) or paddle.is_complex( + param_or_buffer + ): + with paddle.no_grad(): + self._update_fn_( + self.params_shadow[name], + param_or_buffer, + self.n_avg, + ) + self.n_avg += 1 + + def apply_shadow(self): + """Set averaged model parameters to online model.""" + for name, param_or_buffer in itertools.chain( + self.model.named_parameters(), self.model.named_buffers() + ): + if name in self.params_shadow: + stop_gradient = param_or_buffer.stop_gradient + with paddle.no_grad(): + self.params_backup[name] = paddle.assign(param_or_buffer) + paddle.assign(self.params_shadow[name], param_or_buffer) + param_or_buffer.stop_gradient = stop_gradient + + def restore(self): + """Restore online model parameters from backup parameter dict.""" + assert self.params_backup, ( + "params_backup should not be empty, may be caused by calling 'restore' " + "before 'apply_shadow'." + ) + for name, param_or_buffer in itertools.chain( + self.model.named_parameters(), self.model.named_buffers() + ): + if name in self.params_backup: + assert name in self.params_shadow + stop_gradient = param_or_buffer.stop_gradient + with paddle.no_grad(): + paddle.assign(self.params_backup[name], param_or_buffer) + param_or_buffer.stop_gradient = stop_gradient + + self.params_backup = {} + + def set_state_dict(self, state_dict: Dict[str, paddle.Tensor]): + assert ( + "n_avg" in state_dict + ), "state_dict should contain 'n_avg' key, but not found." + self.n_avg.set_value(state_dict.pop("n_avg")) + self.params_shadow.update(state_dict) + + def state_dict(self) -> Dict[str, paddle.Tensor]: + return { + **self.params_shadow, + "n_avg": self.n_avg, + } + + +class ExponentialMovingAverage(AveragedModel): + r"""Implements the exponential moving average (EMA) of the model. + + All parameters are updated by the formula as below: + + $$ + \mathbf{\theta}_{EMA}^{t+1} = \alpha \mathbf{\theta}_{EMA}^{t} + (1 - \alpha) \mathbf{\theta}^{t} + $$ + + Where $\alpha$ is the decay rate, $\theta_{EMA}^{t}$ is the moving average parameters and $\theta^{t}$ is the online parameters at step $t$. + + Args: + model (nn.Layer): The model to be averaged. + decay (float): The decay rate for averaging. + """ + + def __init__(self, model: nn.Layer, decay: float = 0.9): + super().__init__(model, decay) + + def _update_fn_(self, shadow_param, model_param, step): + shadow_param.lerp_(model_param, 1.0 - self.decay) + + +class StochasticWeightAverage(AveragedModel): + r"""Implements the stochastic weight averaging (SWA) of the model. + + Stochastic Weight Averaging was proposed in [Averaging Weights Leads to Wider Optima and Better Generalization](https://arxiv.org/abs/1803.05407), + + All parameters are updated by the formula as below: + + $$ + \mathbf{\theta}_{SWA}^{t} = \frac{1}{t-t_0+1}\sum_{i=t_0}^t{\mathbf{\theta}^{i}} + $$ + + Where $\theta_{SWA}^{t}$ is the average parameters between step $t_0$ and $t$, $\theta^{i}$ is the online parameters at step $i$. + + Args: + model (nn.Layer): The model to be averaged. + """ + + def __init__(self, model: nn.Layer): + super().__init__(model, None) + self.n_avg += 1 # Set to 1 for model already initialized + + def _update_fn_(self, shadow_param, model_param, step): + dynamic_decay = step / (step + 1) + shadow_param.lerp_(model_param, 1.0 - dynamic_decay) diff --git a/examples/smc_reac/ppsci/utils/expression.py b/examples/smc_reac/ppsci/utils/expression.py new file mode 100644 index 0000000000..6bfcddd214 --- /dev/null +++ b/examples/smc_reac/ppsci/utils/expression.py @@ -0,0 +1,212 @@ +# 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. + +from __future__ import annotations + +from typing import TYPE_CHECKING +from typing import Callable +from typing import Dict +from typing import Optional +from typing import Tuple + +from paddle import jit +from paddle import nn +from paddle.framework import core + +if TYPE_CHECKING: + import paddle + from ppsci import constraint + from ppsci import validate + from ppsci import arch + +from ppsci.autodiff import clear + +__all__ = [ + "ExpressionSolver", +] + + +class ExpressionSolver(nn.Layer): + """Expression computing helper, which compute named result according to corresponding + function and related inputs. + + Examples: + >>> import ppsci + >>> model = ppsci.arch.MLP(("x", "y"), ("u", "v"), 5, 128) + >>> expr_solver = ExpressionSolver() + """ + + nvtx_flag: bool # only for nsight analysis + + def __init__(self): + super().__init__() + + def forward(self, *args, **kwargs): + raise NotImplementedError( + "Use train_forward/eval_forward/visu_forward instead of forward." + ) + + @jit.to_static + def train_forward( + self, + expr_dicts: Tuple[Dict[str, Callable], ...], + input_dicts: Tuple[Dict[str, "paddle.Tensor"], ...], + model: arch.Arch, + constraint: Dict[str, "constraint.Constraint"], + label_dicts: Tuple[Dict[str, "paddle.Tensor"], ...], + weight_dicts: Tuple[Dict[str, "paddle.Tensor"], ...], + ) -> Tuple[Dict[str, "paddle.Tensor"], Dict[str, float]]: + """Forward computation for training, including model forward and equation + forward. + + Args: + expr_dicts (Tuple[Dict[str, Callable], ...]): Tuple of expression dicts. + input_dicts (Tuple[Dict[str, paddle.Tensor], ...]): Tuple of input dicts. + model (arch.Arch): NN model. + constraint (Dict[str, "constraint.Constraint"]): Constraint dict. + label_dicts (Tuple[Dict[str, paddle.Tensor], ...]): Tuple of label dicts. + weight_dicts (Tuple[Dict[str, paddle.Tensor], ...]): Tuple of weight dicts. + + Returns: + Tuple[Dict[str, "paddle.Tensor"], Dict[str, float]]: + all_losses: A loss dictionary containing the output terms of all constraints, + constraint_losses: The loss values of all constraints. + """ + losses_all: Dict[str, "paddle.Tensor"] = {} + losses_constraint: Dict[str, float] = {} + + for i, cst_name in enumerate(constraint): + cst_obj = constraint[cst_name] + + # model forward + if self.nvtx_flag: # only for nsight analysis + core.nvprof_nvtx_push(f"Constraint {cst_name}") + + output_dict = model(input_dicts[i]) + + # equation forward + data_dict = {k: v for k, v in input_dicts[i].items()} + data_dict.update(output_dict) + for name, expr in expr_dicts[i].items(): + output_dict[name] = expr(data_dict) + + # put field 'area' into output_dict + if "area" in input_dicts[i]: + output_dict["area"] = input_dicts[i]["area"] + + # clear differentiation cache + clear() + + # compute loss for each constraint according to its' own output, label and weight + losses: Dict[str, "paddle.Tensor"] = cst_obj.loss( + output_dict, + label_dicts[i], + weight_dicts[i], + ) + # update losses into 'losses_all' and 'losses_constraint' + # 'losses_all': Will be send to loss aggregator for further computing final loss(scalar) + # 'losses_constraint': Will be used in logging + losses_constraint[cst_name] = 0.0 + for key in losses: + losses_constraint[cst_name] += losses[key].item() + if key in losses_all: + losses_all[key] += losses[key] + else: + losses_all[key] = losses[key] + + if self.nvtx_flag: # only for nsight analysis + core.nvprof_nvtx_pop() + + return losses_all, losses_constraint + + @jit.to_static + def eval_forward( + self, + expr_dict: Dict[str, Callable], + input_dict: Dict[str, "paddle.Tensor"], + model: arch.Arch, + validator: "validate.Validator", + label_dict: Dict[str, "paddle.Tensor"], + weight_dict: Dict[str, "paddle.Tensor"], + ) -> Tuple[Dict[str, "paddle.Tensor"], Dict[str, "paddle.Tensor"]]: + """Forward computation for evaluation, including model forward and equation + forward. + + Args: + expr_dict (Dict[str, Callable]): Expression dict. + input_dict (Dict[str, paddle.Tensor]): Input dict. + model (arch.Arch): NN model. + validator (validate.Validator): Validator. + label_dict (Dict[str, paddle.Tensor]): Label dict. + weight_dict (Dict[str, paddle.Tensor]): Weight dict. + + Returns: + Tuple[Dict[str, paddle.Tensor], Dict[str, paddle.Tensor]]: Result dict and loss for + given validator. + """ + # model forward + output_dict = model(input_dict) + + # equation forward + data_dict = {k: v for k, v in input_dict.items()} + data_dict.update(output_dict) + for name, expr in expr_dict.items(): + output_dict[name] = expr(data_dict) + + # put field 'area' into output_dict + if "area" in input_dict: + output_dict["area"] = input_dict["area"] + + # clear differentiation cache + clear() + + # compute loss for each validator according to its' own output, label and weight + validator_losses = validator.loss( + output_dict, + label_dict, + weight_dict, + ) + return output_dict, validator_losses + + def visu_forward( + self, + expr_dict: Optional[Dict[str, Callable]], + input_dict: Dict[str, "paddle.Tensor"], + model: arch.Arch, + ) -> Dict[str, "paddle.Tensor"]: + """Forward computation for visualization, including model forward and equation + forward. + + Args: + expr_dict (Optional[Dict[str, Callable]]): Expression dict. + input_dict (Dict[str, paddle.Tensor]): Input dict. + model (arch.Arch): NN model. + + Returns: + Dict[str, paddle.Tensor]: Result dict for given expression dict. + """ + # model forward + output_dict = model(input_dict) + + if isinstance(expr_dict, dict): + # equation forward + data_dict = {k: v for k, v in input_dict.items()} + data_dict.update(output_dict) + for name, expr in expr_dict.items(): + output_dict[name] = expr(data_dict) + + # clear differentiation cache + clear() + + return output_dict diff --git a/examples/smc_reac/ppsci/utils/initializer.py b/examples/smc_reac/ppsci/utils/initializer.py new file mode 100644 index 0000000000..0a5ececf84 --- /dev/null +++ b/examples/smc_reac/ppsci/utils/initializer.py @@ -0,0 +1,498 @@ +# 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. + +""" +The initialization method under this module is aligned with pytorch initialization. +If you need to use the initialization method of PaddlePaddle, please refer to +[paddle.nn.initializer](https://github.com/PaddlePaddle/Paddle/tree/develop/python/paddle/nn/initializer) + +This code is based on [torch.nn.init](https://github.com/pytorch/pytorch/blob/main/torch/nn/init.py) +The copyright of pytorch/pytorch is a BSD-style license, as found in the LICENSE file. +""" + +from __future__ import annotations + +import math + +import numpy as np +import paddle +from paddle import nn +from typing_extensions import Literal + +from ppsci.utils import logger + +__all__ = [ + "uniform_", + "normal_", + "trunc_normal_", + "glorot_normal_", + "constant_", + "ones_", + "zeros_", + "xavier_uniform_", + "xavier_normal_", + "kaiming_uniform_", + "kaiming_normal_", + "linear_init_", + "conv_init_", +] + + +def _no_grad_uniform_(tensor, a, b): + with paddle.no_grad(): + tensor.set_value( + paddle.uniform(shape=tensor.shape, dtype=tensor.dtype, min=a, max=b) + ) + return tensor + + +def _no_grad_normal_(tensor, mean=0.0, std=1.0): + with paddle.no_grad(): + tensor.set_value(paddle.normal(mean=mean, std=std, shape=tensor.shape)) + return tensor + + +def _no_grad_trunc_normal_(tensor, mean=0.0, std=1.0, a=-2.0, b=2.0): + # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf + def norm_cdf(x): + # Computes standard normal cumulative distribution function + return (1.0 + math.erf(x / math.sqrt(2.0))) / 2.0 + + if (mean < a - 2 * std) or (mean > b + 2 * std): + logger.warning( + f"mean({mean}) is more than 2 std({std}) from [a, b]([{a}, {b}]) in _no_grad_trunc_normal_. " + "The distribution of values may be incorrect." + ) + with paddle.no_grad(): + # Values are generated by using a truncated uniform distribution and + # then using the inverse CDF for the normal distribution. + # Get upper and lower cdf values + l = norm_cdf((a - mean) / std) + u = norm_cdf((b - mean) / std) + + # Uniformly fill tensor with values from [l, u], then translate to + # [2l-1, 2u-1]. + _tensor = paddle.uniform( + shape=tensor.shape, dtype=tensor.dtype, min=2 * l - 1, max=2 * u - 1 + ) + + # Use inverse cdf transform for normal distribution to get truncated + # standard normal + _tensor.erfinv_() + + # Transform to proper mean, std + _tensor = paddle.multiply( + _tensor, paddle.to_tensor(std * math.sqrt(2.0), tensor.dtype) + ) + _tensor = paddle.add(_tensor, paddle.to_tensor(mean, tensor.dtype)) + + # Clamp to ensure it"s in the proper range + _tensor = paddle.clip(_tensor, min=a, max=b) + tensor.set_value(_tensor) + return tensor + + +def _no_grad_fill_(tensor, value=0.0): + with paddle.no_grad(): + tensor.set_value(paddle.full_like(tensor, value, dtype=tensor.dtype)) + return tensor + + +def uniform_(tensor: paddle.Tensor, a: float, b: float) -> paddle.Tensor: + """Modify tensor inplace using uniform_. + + Args: + tensor (paddle.Tensor): Paddle Tensor. + a (float): Min value. + b (float): Max value. + + Returns: + paddle.Tensor: Initialized tensor. + + Examples: + >>> import paddle + >>> import ppsci + >>> param = paddle.empty((128, 256), "float32") + >>> param = ppsci.utils.initializer.uniform_(param, -1, 1) + """ + return _no_grad_uniform_(tensor, a, b) + + +def normal_( + tensor: paddle.Tensor, mean: float = 0.0, std: float = 1.0 +) -> paddle.Tensor: + """Modify tensor inplace using normal_. + + Args: + tensor (paddle.Tensor): Paddle Tensor. + mean (float, optional): Mean value. Defaults to 0.0. + std (float, optional): Std value. Defaults to 1.0. + + Returns: + paddle.Tensor: Initialized tensor. + + Examples: + >>> import paddle + >>> import ppsci + >>> param = paddle.empty((128, 256), "float32") + >>> param = ppsci.utils.initializer.normal_(param, 0, 1) + """ + return _no_grad_normal_(tensor, mean, std) + + +def trunc_normal_( + tensor: paddle.Tensor, + mean: float = 0.0, + std: float = 1.0, + a: float = -2.0, + b: float = 2.0, +) -> paddle.Tensor: + """Modify tensor inplace using trunc_normal_. + + Args: + tensor (paddle.Tensor): Paddle Tensor. + mean (float, optional): The mean of the normal distribution. Defaults to 0.0. + std (float, optional): The standard deviation of the normal distribution. Defaults to 1.0. + a (float, optional): The minimum cutoff value. Defaults to -2.0. + b (float, optional): The maximum cutoff value. Defaults to 2.0. + + Returns: + paddle.Tensor: Initialized tensor. + + Examples: + >>> import paddle + >>> import ppsci + >>> param = paddle.empty((128, 256), "float32") + >>> param = ppsci.utils.initializer.trunc_normal_(param, 0.0, 1.0) + """ + return _no_grad_trunc_normal_(tensor, mean, std, a, b) + + +def constant_(tensor: paddle.Tensor, value: float = 0.0) -> paddle.Tensor: + """Modify tensor inplace using constant_. + + Args: + tensor (paddle.Tensor): Paddle Tensor. + value (float, optional): Value to fill tensor. Defaults to 0.0. + + Returns: + paddle.Tensor: Initialized tensor. + + Examples: + >>> import paddle + >>> import ppsci + >>> param = paddle.empty((128, 256), "float32") + >>> param = ppsci.utils.initializer.constant_(param, 2) + """ + return _no_grad_fill_(tensor, value) + + +def ones_(tensor: paddle.Tensor) -> paddle.Tensor: + """Modify tensor inplace using ones_. + + Args: + tensor (paddle.Tensor): Paddle Tensor. + + Returns: + paddle.Tensor: Initialized tensor. + + Examples: + >>> import paddle + >>> import ppsci + >>> param = paddle.empty((128, 256), "float32") + >>> param = ppsci.utils.initializer.ones_(param) + """ + return _no_grad_fill_(tensor, 1) + + +def zeros_(tensor: paddle.Tensor) -> paddle.Tensor: + """Modify tensor inplace using zeros_. + + Args: + tensor (paddle.Tensor): Paddle Tensor. + + Returns: + paddle.Tensor: Initialized tensor. + + Examples: + >>> import paddle + >>> import ppsci + >>> param = paddle.empty((128, 256), "float32") + >>> param = ppsci.utils.initializer.zeros_(param) + """ + return _no_grad_fill_(tensor, 0) + + +def _calculate_fan_in_and_fan_out(tensor, reverse=False): + """ + Calculate (fan_in, _fan_out) for tensor. + + Args: + tensor (paddle.Tensor): paddle.Tensor. + reverse (bool): Tensor data format order, False by default as [fout, fin, ...]. + e.g. : conv.weight [cout, cin, kh, kw] is False; linear.weight [cin, cout] + is True. + + Return: + Tuple[float, float]: (fan_in, fan_out). + """ + if tensor.ndim < 2: + raise ValueError( + f"tensor.ndim should be no less than 2, but got {tensor.ndim}." + ) + + if reverse: + num_input_fmaps, num_output_fmaps = tensor.shape[0], tensor.shape[1] + else: + num_input_fmaps, num_output_fmaps = tensor.shape[1], tensor.shape[0] + + receptive_field_size = 1 + if tensor.ndim > 2: + receptive_field_size = np.prod(tensor.shape[2:]) + + fan_in = num_input_fmaps * receptive_field_size + fan_out = num_output_fmaps * receptive_field_size + + return fan_in, fan_out + + +def xavier_uniform_( + tensor: paddle.Tensor, gain: float = 1.0, reverse: bool = False +) -> paddle.Tensor: + """Modify tensor inplace using xavier_uniform_. + + Args: + tensor (paddle.Tensor): Paddle Tensor. + gain (float, optional): Hyperparameter. Defaults to 1.0. + reverse (bool, optional): Tensor data format order, False by default as + [fout, fin, ...].. Defaults to False. + + Returns: + paddle.Tensor: Initialized tensor. + + Examples: + >>> import paddle + >>> import ppsci + >>> param = paddle.empty((128, 256), "float32") + >>> param = ppsci.utils.initializer.xavier_uniform_(param) + """ + fan_in, fan_out = _calculate_fan_in_and_fan_out(tensor, reverse=reverse) + std = gain * math.sqrt(2.0 / float(fan_in + fan_out)) + k = math.sqrt(3.0) * std + return _no_grad_uniform_(tensor, -k, k) + + +def xavier_normal_( + tensor: paddle.Tensor, gain: float = 1.0, reverse: bool = False +) -> paddle.Tensor: + """Modify tensor inplace using xavier_normal_. + + Args: + tensor (paddle.Tensor): Paddle Tensor. + gain (float, optional): Hyperparameter. Defaults to 1.0. + reverse (bool, optional): Tensor data format order, False by + default as [fout, fin, ...]. Defaults to False. + + Returns: + paddle.Tensor: Initialized tensor. + + Examples: + >>> import paddle + >>> import ppsci + >>> param = paddle.empty((128, 256), "float32") + >>> param = ppsci.utils.initializer.xavier_normal_(param) + """ + fan_in, fan_out = _calculate_fan_in_and_fan_out(tensor, reverse=reverse) + std = gain * math.sqrt(2.0 / float(fan_in + fan_out)) + return _no_grad_normal_(tensor, 0, std) + + +# reference: https://pytorch.org/docs/stable/_modules/torch/nn/init.html +def _calculate_correct_fan(tensor, mode, reverse=False): + mode = mode.lower() + valid_modes = ["fan_in", "fan_out"] + if mode not in valid_modes: + raise ValueError(f"Mode {mode} not supported, please use one of {valid_modes}") + + fan_in, fan_out = _calculate_fan_in_and_fan_out(tensor, reverse) + + return fan_in if mode == "fan_in" else fan_out + + +def _calculate_gain(nonlinearity, param=None): + linear_fns = [ + "linear", + "conv1d", + "conv2d", + "conv3d", + "conv_transpose1d", + "conv_transpose2d", + "conv_transpose3d", + ] + if nonlinearity in linear_fns or nonlinearity == "sigmoid": + return 1 + elif nonlinearity == "tanh": + return 5.0 / 3 + elif nonlinearity == "relu": + return math.sqrt(2.0) + elif nonlinearity == "leaky_relu": + if param is None: + negative_slope = 0.01 + elif ( + not isinstance(param, bool) + and isinstance(param, int) + or isinstance(param, float) + ): + # True/False are instances of int, hence check above + negative_slope = param + else: + raise ValueError(f"negative_slope {param} not a valid number") + return math.sqrt(2.0 / (1 + negative_slope**2)) + elif nonlinearity == "selu": + return 3.0 / 4 + else: + raise ValueError(f"Unsupported nonlinearity {nonlinearity}") + + +def kaiming_uniform_( + tensor: paddle.Tensor, + a: float = 0, + mode: Literal["fan_in", "fan_out"] = "fan_in", + nonlinearity: str = "leaky_relu", + reverse: bool = False, +) -> paddle.Tensor: + """Modify tensor inplace using kaiming_uniform method. + + Args: + tensor (paddle.Tensor): Paddle Tensor. + a (float, optional): The negative slope of the rectifier used after this layer. + Defaults to 0. + mode (Literal["fan_in", "fan_out"], optional): + ["fan_in", "fan_out"]. Defaults to "fan_in". + nonlinearity (str, optional): Nonlinearity method name. Defaults to "leaky_relu". + reverse (bool, optional): Tensor data format order, False by default as + [fout, fin, ...].. Defaults to False. + + Returns: + paddle.Tensor: Initialized tensor. + + Examples: + >>> import paddle + >>> import ppsci + >>> param = paddle.empty((128, 256), "float32") + >>> param = ppsci.utils.initializer.kaiming_uniform_(param) + """ + fan = _calculate_correct_fan(tensor, mode, reverse) + gain = _calculate_gain(nonlinearity, a) + std = gain / math.sqrt(fan) + k = math.sqrt(3.0) * std + return _no_grad_uniform_(tensor, -k, k) + + +def kaiming_normal_( + tensor: paddle.Tensor, + a: float = 0, + mode: Literal["fan_in", "fan_out"] = "fan_in", + nonlinearity: str = "leaky_relu", + reverse: bool = False, +) -> paddle.Tensor: + """Modify tensor inplace using kaiming_normal_. + + Args: + tensor (paddle.Tensor): Paddle Tensor. + a (float, optional): The negative slope of the rectifier used after this layer. + Defaults to 0. + mode (Literal["fan_in", "fan_out"], optional): Either + 'fan_in' (default) or 'fan_out'. Defaults to "fan_in". + nonlinearity (str, optional): Nonlinearity method name. Defaults to "leaky_relu". + reverse (bool, optional): Tensor data format order. Defaults to False. + + Returns: + paddle.Tensor: Initialized tensor. + + Examples: + >>> import paddle + >>> import ppsci + >>> param = paddle.empty((128, 256), "float32") + >>> param = ppsci.utils.initializer.kaiming_normal_(param) + """ + fan = _calculate_correct_fan(tensor, mode, reverse) + gain = _calculate_gain(nonlinearity, a) + std = gain / math.sqrt(fan) + return _no_grad_normal_(tensor, 0, std) + + +def linear_init_(module: nn.Layer) -> None: + """Initialize module's weight and bias as it is a linear layer. + + Args: + module (nn.Layer): Linear Layer to be initialized. + + Examples: + >>> import paddle + >>> import ppsci + >>> layer = paddle.nn.Linear(128, 256) + >>> ppsci.utils.initializer.linear_init_(layer) + """ + kaiming_uniform_(module.weight, a=math.sqrt(5), reverse=True) + if module.bias is not None: + fan_in, _ = _calculate_fan_in_and_fan_out(module.weight, reverse=True) + bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0 + uniform_(module.bias, -bound, bound) + + +def conv_init_(module: nn.Layer) -> None: + """Initialize module's weight and bias as it is a conv layer. + + Args: + module (nn.Layer): Convolution Layer to be initialized. + + Examples: + >>> import paddle + >>> import ppsci + >>> layer = paddle.nn.Conv2D(4, 16, 2) + >>> ppsci.utils.initializer.conv_init_(layer) + """ + kaiming_uniform_(module.weight, a=math.sqrt(5)) + if module.bias is not None: + fan_in, _ = _calculate_fan_in_and_fan_out(module.weight, reverse=False) + if fan_in != 0: + bound = 1 / math.sqrt(fan_in) + uniform_(module.bias, -bound, bound) + + +def glorot_normal_(tensor: paddle.Tensor) -> paddle.Tensor: + """Modify tensor inplace using jax-style glorot_normal. + + Args: + tensor (paddle.Tensor): Paddle Tensor/Parameter. + + Returns: + paddle.Tensor: Initialized tensor. + + Examples: + >>> import paddle + >>> import ppsci + >>> param = paddle.empty((128, 256), "float32") + >>> param = ppsci.utils.initializer.glorot_normal_(param) + """ + assert ( + tensor.ndim == 2 + ), f"glorot_normal_ only support 2D tensor now, but got ndim={tensor.ndim}" + fin, fout = tensor.shape + var = 2.0 / (fin + fout) + stddev = math.sqrt(var) * 0.87962566103423978 + trunc_normal_(tensor) + tensor.set_value(tensor * stddev) + return tensor diff --git a/examples/smc_reac/ppsci/utils/logger.py b/examples/smc_reac/ppsci/utils/logger.py new file mode 100644 index 0000000000..46ca57bd46 --- /dev/null +++ b/examples/smc_reac/ppsci/utils/logger.py @@ -0,0 +1,264 @@ +# 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. + +from __future__ import annotations + +import functools +import logging +import os +import sys +from typing import TYPE_CHECKING +from typing import Callable +from typing import Dict +from typing import Optional + +import colorlog +import paddle.distributed as dist + +from ppsci.utils import misc + +if TYPE_CHECKING: + import visualdl # isort:skip + import wandb # isort:skip + import tensorboardX as tbd + +_logger: logging.Logger = None + +# INFO(20) is white(no color) +# use custom log level `MESSAGE` for printing message in color +_MESSAGE_LEVEL = 25 + +_COLORLOG_CONFIG = { + "DEBUG": "green", + "WARNING": "yellow", + "ERROR": "red", + "MESSAGE": "bold_cyan", +} + +__all__ = [ + "init_logger", + "set_log_level", + "info", + "message", + "debug", + "warning", + "error", + "scalar", +] + + +def init_logger( + name: str = "ppsci", + log_file: Optional[str] = None, + log_level: int = logging.INFO, +) -> None: + """Initialize and get a logger by name. + + If the logger has not been initialized, this method will initialize the logger by + adding one or two handlers, otherwise the initialized logger will be directly + returned. During initialization, a StreamHandler will always be added. If `log_file` + is specified a FileHandler will also be added. + + Args: + name (str, optional): Logger name. Defaults to "ppsci". + log_file (Optional[str]): The log filename. If specified, a FileHandler + will be added to the logger. Defaults to None. + log_level (int, optional): The logger level. Note that only the process of + rank 0 is affected, and other processes will set the level to + "Error" thus be silent most of the time. Defaults to logging.INFO. + """ + # Add custom log level MESSAGE(25), between WARNING(30) and INFO(20) + logging.addLevelName(_MESSAGE_LEVEL, "MESSAGE") + + if isinstance(log_level, str): + log_level = getattr(logging, log_level.upper()) + + global _logger + + # get a clean logger + _logger = logging.getLogger(name) + _logger.handlers.clear() + + # add stream_handler, output to stdout such as terminal + stream_formatter = colorlog.ColoredFormatter( + "%(log_color)s[%(asctime)s] %(name)s %(levelname)s: %(message)s", + datefmt="%Y/%m/%d %H:%M:%S", + log_colors=_COLORLOG_CONFIG, + ) + stream_handler = logging.StreamHandler(stream=sys.stdout) + stream_handler.setFormatter(stream_formatter) + stream_handler._name = "stream_handler" + _logger.addHandler(stream_handler) + + # add file_handler, output to log_file(if specified), only for rank 0 device + if log_file is not None and dist.get_rank() == 0: + log_file_folder = os.path.dirname(log_file) + if len(log_file_folder): + os.makedirs(log_file_folder, exist_ok=True) + file_formatter = logging.Formatter( + "[%(asctime)s] %(name)s %(levelname)s: %(message)s", + datefmt="%Y/%m/%d %H:%M:%S", + ) + file_handler = logging.FileHandler(log_file, "a") # append mode + file_handler.setFormatter(file_formatter) + file_handler._name = "file_handler" + _logger.addHandler(file_handler) + + if dist.get_rank() == 0: + _logger.setLevel(log_level) + else: + _logger.setLevel(logging.ERROR) + + _logger.propagate = False + + +def set_log_level(log_level: int): + """Set logger level, only message of level >= `log_level` will be printed. + + Built-in log level are below: + + CRITICAL = 50, + FATAL = 50, + ERROR = 40, + WARNING = 30, + WARN = 30, + INFO = 20, + DEBUG = 10, + NOTSET = 0. + + Args: + log_level (int): Log level. + """ + if dist.get_rank() == 0: + _logger.setLevel(log_level) + else: + _logger.setLevel(logging.ERROR) + + +def ensure_logger(log_func: Callable) -> Callable: + """ + A decorator which automatically initialize `logger` by default arguments + when init_logger() is not called manually. + """ + + @functools.wraps(log_func) + def wrapped_log_func(msg, *args): + if _logger is None: + init_logger() + _logger.warning( + "Logger has already been automatically initialized as `log_file` is " + "set to None by default, information will only be printed to terminal " + "without writing to any file." + ) + + log_func(msg, *args) + + return wrapped_log_func + + +@ensure_logger +@misc.run_at_rank0 +def info(msg, *args): + _logger.info(msg, *args) + + +@ensure_logger +@misc.run_at_rank0 +def message(msg, *args): + _logger.log(_MESSAGE_LEVEL, msg, *args) + + +@ensure_logger +@misc.run_at_rank0 +def debug(msg, *args): + _logger.debug(msg, *args) + + +@ensure_logger +@misc.run_at_rank0 +def warning(msg, *args): + _logger.warning(msg, *args) + + +@ensure_logger +@misc.run_at_rank0 +def error(msg, *args): + _logger.error(msg, *args) + + +def scalar( + metric_dict: Dict[str, float], + step: int, + vdl_writer: Optional["visualdl.LogWriter"] = None, + wandb_writer: Optional["wandb.run"] = None, + tbd_writer: Optional["tbd.SummaryWriter"] = None, +): + """This function will add scalar data to VisualDL or WandB for plotting curve(s). + + Args: + metric_dict (Dict[str, float]): Metrics dict with metric name and value. + step (int): The step of the metric. + vdl_writer (Optional[visualdl.LogWriter]): VisualDL writer to record metrics. Defaults to None. + wandb_writer (Optional[wandb.run]): Run object of WandB to record metrics. Defaults to None. + tbd_writer (Optional[tbd.SummaryWriter]): Run object of WandB to record metrics. Defaults to None. + """ + if vdl_writer is not None: + with misc.RankZeroOnly() as is_master: + if is_master: + for name, value in metric_dict.items(): + vdl_writer.add_scalar(name, value, step) + + if wandb_writer is not None: + with misc.RankZeroOnly() as is_master: + if is_master: + wandb_writer.log({"step": step, **metric_dict}) + + if tbd_writer is not None: + with misc.RankZeroOnly() as is_master: + if is_master: + for name, value in metric_dict.items(): + tbd_writer.add_scalar(name, value, global_step=step) + + +def advertise(): + """ + Show the advertising message like the following: + + =========================================================== + == PaddleScience is powered by PaddlePaddle ! == + =========================================================== + == == + == For more info please go to the following website. == + == == + == https://github.com/PaddlePaddle/PaddleScience == + =========================================================== + """ + + _copyright = "PaddleScience is powered by PaddlePaddle !" + ad = "Please refer to the following website for more info." + website = "https://github.com/PaddlePaddle/PaddleScience" + AD_LEN = 6 + len(max([_copyright, ad, website], key=len)) + + info( + "\n{0}\n{1}\n{2}\n{3}\n{4}\n{5}\n{6}\n{7}\n".format( + "=" * (AD_LEN + 4), + "=={}==".format(_copyright.center(AD_LEN)), + "=" * (AD_LEN + 4), + "=={}==".format(" " * AD_LEN), + "=={}==".format(ad.center(AD_LEN)), + "=={}==".format(" " * AD_LEN), + "=={}==".format(website.center(AD_LEN)), + "=" * (AD_LEN + 4), + ) + ) diff --git a/examples/smc_reac/ppsci/utils/misc.py b/examples/smc_reac/ppsci/utils/misc.py new file mode 100644 index 0000000000..7874290d4e --- /dev/null +++ b/examples/smc_reac/ppsci/utils/misc.py @@ -0,0 +1,684 @@ +# 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. + +from __future__ import annotations + +import collections +import functools +import os +import random +import time +from contextlib import ContextDecorator +from typing import Callable +from typing import Dict +from typing import List +from typing import Optional +from typing import Sequence +from typing import Tuple +from typing import Union + +import numpy as np +import paddle +from matplotlib import pyplot as plt +from paddle import distributed as dist + +from ppsci.utils import logger + +__all__ = [ + "AverageMeter", + "PrettyOrderedDict", + "Prettydefaultdict", + "RankZeroOnly", + "RankZeroFirst", + "Timer", + "all_gather", + "concat_dict_list", + "convert_to_array", + "convert_to_dict", + "stack_dict_list", + "cartesian_product", + "combine_array_with_time", + "set_random_seed", + "run_on_eval_mode", + "run_at_rank0", + "plot_curve", + "check_flag_enabled", +] + + +class AverageMeter: + """ + Computes and stores the average and current value + Code was based on https://github.com/pytorch/examples/blob/master/imagenet/main.py + """ + + def __init__(self, name="", fmt="f", postfix="", need_avg=True): + self.name = name + self.fmt = fmt + self.postfix = postfix + self.need_avg = need_avg + self.reset() + + def reset(self): + """Reset.""" + self.val = 0 + self.avg = 0 + self.sum = 0 + self.count = 0 + self.history = [] + + def update(self, val, n=1): + """Update.""" + self.val = val + self.sum += val * n + self.count += n + self.avg = self.sum / self.count + self.history.append(val) + + @property + def avg_info(self): + if isinstance(self.avg, paddle.Tensor): + self.avg = float(self.avg) + return f"{self.name}: {self.avg:.5f}" + + @property + def total(self): + return f"{self.name}_sum: {self.sum:{self.fmt}}{self.postfix}" + + @property + def total_minute(self): + return f"{self.name} {self.sum / 60:{self.fmt}}{self.postfix} min" + + @property + def mean(self): + return ( + f"{self.name}: {self.avg:{self.fmt}}{self.postfix}" if self.need_avg else "" + ) + + @property + def value(self): + return f"{self.name}: {self.val:{self.fmt}}{self.postfix}" + + +class PrettyOrderedDict(collections.OrderedDict): + """ + The ordered dict which can be prettily printed. + + Examples: + >>> import ppsci + >>> dic = ppsci.utils.misc.PrettyOrderedDict() + >>> dic.update({'a':1, 'b':2, 'c':3}) + >>> print(dic) + ('a', 1)('b', 2)('c', 3) + """ + + def __str__(self): + return "".join([str((k, v)) for k, v in self.items()]) + + +class Prettydefaultdict(collections.defaultdict): + """ + The default dict which can be prettily printed. + + Examples: + >>> import ppsci + >>> dic = ppsci.utils.misc.Prettydefaultdict() + >>> dic.update({'a':1, 'b':2, 'c':3}) + >>> print(dic) + ('a', 1)('b', 2)('c', 3) + """ + + def __str__(self): + return "".join([str((k, v)) for k, v in self.items()]) + + +class RankZeroOnly: + """ + A context manager that ensures the code inside it is only executed by the process + with rank zero. All rank will be synchronized by `dist.barrier` in + distributed environment. + + NOTE: Always used for time consuming code blocks, such as initialization of log + writer, saving result to disk, etc. + + Args: + rank (Optional[int]): The rank of the current process. If not provided, + it will be obtained from `dist.get_rank()`. + + Examples: + >>> import paddle.distributed as dist + >>> with RankZeroOnly(dist.get_rank()) as is_master: + ... if is_master: + ... # code here which should only be executed in the master process + ... pass + """ + + def __init__(self, rank: Optional[int] = None): + """ + Enter the context and check if the current process is the master. + + Args: + rank (Optional[int]): The rank of the current process. If not provided, + it will be obtained from `dist.get_rank()`. + """ + super().__init__() + self.rank = rank if (rank is not None) else dist.get_rank() + self.is_master = self.rank == 0 + + def __enter__(self) -> bool: + """ + Enter the context and check if the current process is the master. + + Returns: + bool: True if the current process is the master (rank zero), False otherwise. + """ + return self.is_master + + def __exit__(self, exc_type, exc_value, traceback): + if dist.get_world_size() > 1: + dist.barrier() + + +class RankZeroFirst(ContextDecorator): + """ + A context manager that ensures the code inside it is only executed by the process + with rank zero first. All ranks will be synchronized by `dist.barrier()`. + + Args: + rank (Optional[int]): The rank of the current process. If not provided, + it will be obtained from `dist.get_rank()`. + + Examples: + >>> import paddle.distributed as dist + >>> with RankZeroFirst(dist.get_rank()): + ... # code here which should be executed first in the master(rank-0) process + ... pass + """ + + def __init__(self, rank: Optional[int] = None): + if dist.is_initialized(): + self.rank = rank if rank is not None else dist.get_rank() + self.world_size = dist.get_world_size() + else: + self.rank = 0 + self.world_size = 1 + self.is_master = self.rank == 0 + + def __enter__(self): + if self.world_size > 1 and not self.is_master: + dist.barrier() # Non-master processs wait for master to finish + + def __exit__(self, type, value, traceback): + if self.world_size > 1 and self.is_master: + dist.barrier() # Allow others to proceed + + +class Timer(ContextDecorator): + """Count time cost for code block within context. + + Args: + name (str, optional): Name of timer discriminate different code block. + Defaults to "Timer". + auto_print (bool, optional): Whether print time cost when exit context. + Defaults to True. + + Examples: + >>> import paddle + >>> from ppsci.utils import misc + >>> with misc.Timer("test1", auto_print=False) as timer: + ... w = sum(range(0, 10)) + >>> print(f"time cost of 'sum(range(0, 10))' is {timer.interval:.2f}") # doctest: +SKIP + time cost of 'sum(range(0, 10))' is 0.00 + + >>> @misc.Timer("test2", auto_print=True) + ... def func(): + ... w = sum(range(0, 10)) + >>> func() # doctest: +SKIP + + >>> timer = misc.Timer("cost_of_func", auto_print=False) + >>> timer.start() + >>> def func(): + ... w = sum(range(0, 10)) + >>> func() + >>> timer.end() + >>> print(f"time cost of 'cost_of_func' is {timer.interval:.2f}") # doctest: +SKIP + time cost of 'cost_of_func' is 0.00 + """ + + interval: float # Time cost for code within Timer context + + def __init__(self, name: str = "Timer", auto_print: bool = True): + super().__init__() + self.name = name + self.auto_print = auto_print + + def __enter__(self): + paddle.device.synchronize() + self.start_time = time.perf_counter() + return self + + def __exit__(self, type, value, traceback): + paddle.device.synchronize() + self.end_time = time.perf_counter() + self.interval = self.end_time - self.start_time + if self.auto_print: + logger.message(f"{self.name}.time_cost = {self.interval * 1000:.2f} ms") + + def start(self, name: str = "Timer"): + """Push a new timer context. + + Args: + name (str, optional): Name of code block to be clocked. Defaults to "Timer". + """ + paddle.device.synchronize() + self.start_time = time.perf_counter() + + def end(self): + """End current timer context and print time cost.""" + paddle.device.synchronize() + self.end_time = time.perf_counter() + self.interval = self.end_time - self.start_time + if self.auto_print: + logger.message(f"{self.name}.time_cost = {self.interval:.2f} s") + + +def convert_to_dict(array: np.ndarray, keys: Tuple[str, ...]) -> Dict[str, np.ndarray]: + """Split given array into single channel array at axis -1 in order of given keys. + + Args: + array (np.ndarray): Array to be split. + keys (Tuple[str, ...]): Keys used in split. + + Returns: + Dict[str, np.ndarray]: Split dict. + + Examples: + >>> import numpy as np + >>> import ppsci + >>> arr = np.array([[1., 2., 3.], [4., 5., 6.]]) + >>> result = ppsci.utils.misc.convert_to_dict(arr, ("x", "y", "z")) + >>> print(arr.shape) + (2, 3) + >>> for k, v in result.items(): + ... print(k, v.shape) + x (2, 1) + y (2, 1) + z (2, 1) + """ + if array.shape[-1] != len(keys): + raise ValueError( + f"dim of array({array.shape[-1]}) must equal to " f"len(keys)({len(keys)})" + ) + + split_array = np.split(array, len(keys), axis=-1) + return {key: split_array[i] for i, key in enumerate(keys)} + + +def all_gather( + tensor: paddle.Tensor, concat: bool = True, axis: int = 0 +) -> Union[paddle.Tensor, List[paddle.Tensor]]: + """Gather tensor from all devices, concatenate them along given axis if specified. + + Args: + tensor (paddle.Tensor): Tensor to be gathered from all GPUs. + concat (bool, optional): Whether to concatenate gathered Tensors. Defaults to True. + axis (int, optional): Axis which concatenated along. Defaults to 0. + + Returns: + Union[paddle.Tensor, List[paddle.Tensor]]: Gathered Tensors. + + Examples: + >>> import paddle + >>> import ppsci + >>> import paddle.distributed as dist + >>> dist.init_parallel_env() # doctest: +SKIP + >>> if dist.get_rank() == 0: # doctest: +SKIP + ... data = paddle.to_tensor([[1, 2, 3], [4, 5, 6]]) + ... else: + ... data = paddle.to_tensor([[7, 8, 9], [10, 11, 12]]) + >>> result = ppsci.utils.misc.all_gather(data) # doctest: +SKIP + >>> print(result.numpy()) # doctest: +SKIP + [[ 1 2 3] + [ 4 5 6] + [ 7 8 9] + [10 11 12]] + """ + result: List[paddle.Tensor] = [] + + # NOTE: Put tensor to CUDAPlace from CUDAPinnedPlace to use communication. + if tensor.place.is_cuda_pinned_place(): + tensor = tensor.cuda() + + # TODO(HydrogenSulfate): As non-contiguous(strided) tensor is not supported in + # dist.all_gather, manually convert given Tensor to contiguous below. Strided tensor + # will be supported in future. + dist.all_gather(result, tensor.contiguous()) + + if concat: + return paddle.concat(result, axis) + return result + + +def convert_to_array(dict_: Dict[str, np.ndarray], keys: Tuple[str, ...]) -> np.ndarray: + """Concatenate arrays in axis -1 in order of given keys. + + Args: + dict_ (Dict[str, np.ndarray]): Dict contains arrays. + keys (Tuple[str, ...]): Concatenate keys used in concatenation. + + Returns: + np.ndarray: Concatenated array. + + Examples: + >>> import numpy as np + >>> import ppsci + >>> dic = {"x": np.array([[1., 2.], [3., 4.]]), + ... "y": np.array([[5., 6.], [7., 8.]]), + ... "z": np.array([[9., 10.], [11., 12.]])} + >>> result = ppsci.utils.misc.convert_to_array(dic, ("x", "z")) + >>> print(result) + [[ 1. 2. 9. 10.] + [ 3. 4. 11. 12.]] + """ + return np.concatenate([dict_[key] for key in keys], axis=-1) + + +def concat_dict_list( + dict_list: Sequence[Dict[str, np.ndarray]] +) -> Dict[str, np.ndarray]: + """Concatenate arrays in tuple of dicts at axis 0. + + Args: + dict_list (Sequence[Dict[str, np.ndarray]]): Sequence of dicts. + + Returns: + Dict[str, np.ndarray]: A dict with concatenated arrays for each key. + + Examples: + >>> import numpy as np + >>> import ppsci + >>> dic1 = {"x": np.array([[1., 2.], [3., 4.]]), "y": np.array([[5., 6.], [7., 8.]])} + >>> dic2 = {"x": np.array([[1., 2.], [3., 4.]]), "y": np.array([[5., 6.], [7., 8.]])} + >>> result = ppsci.utils.misc.concat_dict_list((dic1, dic2)) + >>> print(result) + {'x': array([[1., 2.], + [3., 4.], + [1., 2.], + [3., 4.]]), 'y': array([[5., 6.], + [7., 8.], + [5., 6.], + [7., 8.]])} + """ + ret = {} + for key in dict_list[0].keys(): + ret[key] = np.concatenate([_dict[key] for _dict in dict_list], axis=0) + return ret + + +def stack_dict_list( + dict_list: Sequence[Dict[str, np.ndarray]] +) -> Dict[str, np.ndarray]: + """Stack arrays in tuple of dicts at axis 0. + + Args: + dict_list (Sequence[Dict[str, np.ndarray]]): Sequence of dicts. + + Returns: + Dict[str, np.ndarray]: A dict with stacked arrays for each key. + + Examples: + >>> import numpy as np + >>> import ppsci + >>> dic1 = {"x": np.array([[1., 2.], [3., 4.]]), "y": np.array([[5., 6.], [7., 8.]])} + >>> dic2 = {"x": np.array([[1., 2.], [3., 4.]]), "y": np.array([[5., 6.], [7., 8.]])} + >>> result = ppsci.utils.misc.stack_dict_list((dic1, dic2)) + >>> for k, v in result.items(): + ... print(k, v.shape) + x (2, 2, 2) + y (2, 2, 2) + """ + ret = {} + for key in dict_list[0].keys(): + ret[key] = np.stack([_dict[key] for _dict in dict_list], axis=0) + return ret + + +def typename(obj: object) -> str: + """Return type name of given object. + + Args: + obj (object): Python object which is instantiated from a class. + + Returns: + str: Class name of given object. + """ + return obj.__class__.__name__ + + +def combine_array_with_time(x: np.ndarray, t: Tuple[int, ...]) -> np.ndarray: + """Combine given data x with time sequence t. + Given x with shape (N, D) and t with shape (T, ), + this function will repeat t_i for N times and will concat it with data x for each t_i in t, + finally return the stacked result, which is of shape (N×T, D+1). + + Args: + x (np.ndarray): Points data with shape (N, D). + t (Tuple[int, ...]): Time sequence with shape (T, ). + + Returns: + np.ndarray: Combined data with shape of (N×T, D+1). + + Examples: + >>> import numpy as np + >>> import ppsci + >>> data_point = np.arange(10).reshape((2, 5)) + >>> time = (1, 2, 3) + >>> result = ppsci.utils.misc.combine_array_with_time(data_point, time) + >>> print(result) + [[1. 0. 1. 2. 3. 4.] + [1. 5. 6. 7. 8. 9.] + [2. 0. 1. 2. 3. 4.] + [2. 5. 6. 7. 8. 9.] + [3. 0. 1. 2. 3. 4.] + [3. 5. 6. 7. 8. 9.]] + """ + nx = len(x) + tx = [] + for ti in t: + tx.append( + np.hstack( + (np.full([nx, 1], float(ti), dtype=paddle.get_default_dtype()), x) + ) + ) + tx = np.vstack(tx) + return tx + + +def cartesian_product(*arrays: np.ndarray) -> np.ndarray: + """Cartesian product for input sequence of array(s). + + Reference: https://stackoverflow.com/questions/11144513/cartesian-product-of-x-and-y-array-points-into-single-array-of-2d-points + + Assume shapes of input arrays are: $(N_1,), (N_2,), (N_3,), ..., (N_M,)$, + then the cartesian product result will be shape of $(N_1xN_2xN_3x...xN_M, M)$. + + Args: + arrays (np.ndarray): Input arrays. + + Returns: + np.ndarray: Cartesian product result of shape $(N_1xN_2xN_3x...xN_M, M)$. + + Examples: + >>> t = np.array([1, 2]) + >>> x = np.array([10, 20]) + >>> y = np.array([100, 200]) + >>> txy = cartesian_product(t, x, y) + >>> print(txy) + [[ 1 10 100] + [ 1 10 200] + [ 1 20 100] + [ 1 20 200] + [ 2 10 100] + [ 2 10 200] + [ 2 20 100] + [ 2 20 200]] + """ + la = len(arrays) + dtype = np.result_type(*arrays) + arr = np.empty([len(a) for a in arrays] + [la], dtype=dtype) + for i, a in enumerate(np.ix_(*arrays)): + arr[..., i] = a + return arr.reshape(-1, la) + + +def set_random_seed(seed: int): + """Set numpy, random, paddle random_seed to given seed. + + Args: + seed (int): Random seed. + """ + paddle.seed(seed) + np.random.seed(seed) + random.seed(seed) + + +def run_on_eval_mode(func: Callable) -> Callable: + """A decorator automatically running given class method in eval mode and keep + training state unchanged after function finished. + + Args: + func (Callable): Class method which is expected running in eval mode. + + Returns: + Callable: Decorated class method. + """ + + @functools.wraps(func) + def function_with_eval_state(self, *args, **kwargs): + # log original state + train_state = self.model.training + + # switch to eval mode + if train_state: + self.model.eval() + + # run func in eval mode + result = func(self, *args, **kwargs) + + # restore state + if train_state: + self.model.train() + + return result + + return function_with_eval_state + + +def run_at_rank0(func: Callable) -> Callable: + """A decorator that allow given function run only at rank 0 to avoid + multiple logs or other events. Usually effected in distributed environment. + + Args: + func (Callable): Given function. + + Returns: + Callable: Wrapped function which will only run at at rank 0, + skipped at other rank. + + Examples: + >>> import paddle + >>> from ppsci.utils import misc + >>> @misc.run_at_rank0 + ... def func(): + ... print(f"now_rank is {paddle.distributed.get_rank()}") + >>> func() + now_rank is 0 + """ + + @functools.wraps(func) + def wrapped_func(*args, **kwargs): + if dist.get_rank() == 0: + return func(*args, **kwargs) + + return wrapped_func + + +def plot_curve( + data: Dict[str, List], + xlabel: str = "X", + ylabel: str = "Y", + output_dir: str = "./output/", + smooth_step: int = 1, + use_semilogy: bool = False, +) -> None: + """Plotting curve. + + Args: + data (Dict[str, List]): Dict of all data, keys are curves' name. + xlabel (str, optional): Label of x-axis. Defaults to "X". + ylabel (str, optional): Label of y-axis. Defaults to "Y". + output_dir (str, optional): Output directory of figure. Defaults to "./output/". + smooth_step (int, optional): How many points are squeezed to one point to smooth the curve. Defaults to 1. + use_semilogy (bool, optional): Whether to set non-uniform coordinates for the y-axis. Defaults to False. + """ + data_arr = np.concatenate( + [np.asarray(arr).reshape(-1, 1) for arr in data.values()], axis=1 + ) + + # smooth + if data_arr.shape[0] % smooth_step != 0: + data_arr = np.reshape( + data_arr[: -(data_arr.shape[0] % smooth_step), :], + (-1, smooth_step, data_arr.shape[1]), + ) + else: + data_arr = np.reshape(data_arr, (-1, smooth_step, data_arr.shape[1])) + data_arr = np.mean(data_arr, axis=1) + + # plot + plt.figure() + if use_semilogy: + plt.yscale("log") + plt.xscale("log") + plt.plot(np.arange(data_arr.shape[0]) * smooth_step, data_arr) + plt.legend( + list(data.keys()), + loc="upper left", + bbox_to_anchor=(1, 1), + ) + plt.xlabel(xlabel) + plt.ylabel(ylabel) + plt.grid() + plt.yticks(size=10) + plt.xticks(size=10) + plt.tight_layout() + + plt.savefig(os.path.join(output_dir, f"{xlabel}-{ylabel}_curve.jpg"), dpi=200) + plt.clf() + plt.close() + + +def check_flag_enabled(flag_name: str) -> bool: + """Check whether the flag is enabled. + + Args: + flag_name (str): Flag name to be checked whether enabled or disabled. + + Returns: + bool: Whether given flag name is enabled in environment. + """ + value = os.getenv(flag_name, False) + if isinstance(value, str): + return value.lower() in ["true", "1"] + return False diff --git a/examples/smc_reac/ppsci/utils/reader.py b/examples/smc_reac/ppsci/utils/reader.py new file mode 100644 index 0000000000..ef0fb8f191 --- /dev/null +++ b/examples/smc_reac/ppsci/utils/reader.py @@ -0,0 +1,266 @@ +# 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. + +from __future__ import annotations + +import collections +import csv +import pickle +from typing import Dict +from typing import Optional +from typing import Tuple + +import meshio +import numpy as np +import paddle +import scipy.io as sio + +__all__ = [ + "load_csv_file", + "load_mat_file", + "load_npz_file", + "load_vtk_file", + "load_vtk_with_time_file", + "load_dat_file", +] + + +def load_csv_file( + file_path: str, + keys: Tuple[str, ...], + alias_dict: Optional[Dict[str, str]] = None, + delimiter: str = ",", + encoding: str = "utf-8", +) -> Dict[str, np.ndarray]: + """Load *.csv file and fetch data as given keys. + + Args: + file_path (str): CSV file path. + keys (Tuple[str, ...]): Required fetching keys. + alias_dict (Optional[Dict[str, str]]): Alias for keys, + i.e. {inner_key: outer_key}. Defaults to None. + encoding (str, optional): Encoding code when open file. Defaults to "utf-8". + + Returns: + Dict[str, np.ndarray]: Loaded data in dict. + """ + if alias_dict is None: + alias_dict = {} + + try: + # read all data from csv file + with open(file_path, "r", encoding=encoding) as csv_file: + reader = csv.DictReader(csv_file, delimiter=delimiter) + raw_data = collections.defaultdict(list) + for _, line_dict in enumerate(reader): + for key, value in line_dict.items(): + raw_data[key].append(value) + except FileNotFoundError as e: + raise e + + # convert to numpy array + data_dict = {} + for key in keys: + fetch_key = alias_dict[key] if key in alias_dict else key + if fetch_key not in raw_data: + raise KeyError(f"fetch_key({fetch_key}) do not exist in raw_data.") + data_dict[key] = np.asarray(raw_data[fetch_key]) + if not np.issubdtype(data_dict[key].dtype, np.integer): + data_dict[key] = data_dict[key].astype(paddle.get_default_dtype()) + data_dict[key] = data_dict[key].reshape([-1, 1]) + + return data_dict + + +def load_mat_file( + file_path: str, keys: Tuple[str, ...], alias_dict: Optional[Dict[str, str]] = None +) -> Dict[str, np.ndarray]: + """Load *.mat file and fetch data as given keys. + + Args: + file_path (str): Mat file path. + keys (Tuple[str, ...]): Required fetching keys. + alias_dict (Optional[Dict[str, str]]): Alias for keys, + i.e. {original_key: original_key}. Defaults to None. + + Returns: + Dict[str, np.ndarray]: Loaded data in dict. + """ + + if alias_dict is None: + alias_dict = {} + + try: + # read all data from mat file + raw_data = sio.loadmat(file_path) + except FileNotFoundError as e: + raise e + + # convert to numpy array + data_dict = {} + for key in keys: + fetch_key = alias_dict[key] if key in alias_dict else key + if fetch_key not in raw_data: + raise KeyError(f"fetch_key({fetch_key}) do not exist in raw_data.") + data_dict[key] = np.asarray(raw_data[fetch_key]) + if not np.issubdtype(data_dict[key].dtype, np.integer): + data_dict[key] = data_dict[key].astype(paddle.get_default_dtype()) + data_dict[key] = data_dict[key].reshape([-1, 1]) + + return data_dict + + +def load_npz_file( + file_path: str, keys: Tuple[str, ...], alias_dict: Optional[Dict[str, str]] = None +) -> Dict[str, np.ndarray]: + """Load *.npz file and fetch data as given keys. + + Args: + file_path (str): Npz file path. + keys (Tuple[str, ...]): Required fetching keys. + alias_dict (Optional[Dict[str, str]]): Alias for keys, + i.e. {original_key: original_key}. Defaults to None. + + Returns: + Dict[str, np.ndarray]: Loaded data in dict. + """ + + if alias_dict is None: + alias_dict = {} + + try: + # read all data from npz file + raw_data = np.load(file_path, allow_pickle=True) + except FileNotFoundError as e: + raise e + + # convert to numpy array + data_dict = {} + for key in keys: + fetch_key = alias_dict[key] if key in alias_dict else key + if fetch_key not in raw_data: + raise KeyError(f"fetch_key({fetch_key}) do not exist in raw_data.") + data_dict[key] = np.asarray(raw_data[fetch_key]) + if data_dict[key].dtype in (np.float16, np.float32, np.float64): + data_dict[key] = data_dict[key].astype(paddle.get_default_dtype()) + + return data_dict + + +def load_vtk_file( + filename_without_timeid: str, + time_step: float, + time_index: Tuple[int, ...], + input_keys: Tuple[str, ...], + label_keys: Optional[Tuple[str, ...]], +) -> Dict[str, np.ndarray]: + """Load coordinates and attached label from the *.vtu file. + + Args: + filename_without_timeid (str): File name without time id. + time_step (float): Physical time step. + time_index (Tuple[int, ...]): Physical time indexes. + input_keys (Tuple[str, ...]): Input coordinates name keys. + label_keys (Optional[Tuple[str, ...]]): Input label name keys. + + Returns: + Dict[str, np.ndarray]: Input coordinates dict, label coordinates dict + """ + input_dict = {var: [] for var in input_keys} + label_dict = {var: [] for var in label_keys} + for index in time_index: + file = filename_without_timeid + f"{index}.vtu" + mesh = meshio.read(file) + n = mesh.points.shape[0] + i = 0 + for key in input_dict: + if key == "t": + input_dict[key].append( + np.full((n, 1), index * time_step, paddle.get_default_dtype()) + ) + else: + input_dict[key].append( + mesh.points[:, i].reshape(n, 1).astype(paddle.get_default_dtype()) + ) + i += 1 + for i, key in enumerate(label_dict): + label_dict[key].append( + np.array(mesh.point_data[key], paddle.get_default_dtype()) + ) + for key in input_dict: + input_dict[key] = np.concatenate(input_dict[key]) + for key in label_dict: + label_dict[key] = np.concatenate(label_dict[key]) + + return input_dict, label_dict + + +def load_vtk_with_time_file(file: str) -> Dict[str, np.ndarray]: + """Temporary interface for points cloud, will be banished sooner. + + Args: + file (str): Input file name. + + Returns: + Dict[str, np.ndarray]: Input coordinates dict. + """ + mesh = meshio.read(file) + n = mesh.points.shape[0] + t = np.array(mesh.point_data["time"]) + x = mesh.points[:, 0].reshape(n, 1) + y = mesh.points[:, 1].reshape(n, 1) + z = mesh.points[:, 2].reshape(n, 1) + input_dict = {"t": t, "x": x, "y": y, "z": z} + return input_dict + + +def load_dat_file( + file_path: str, + keys: Tuple[str, ...] = None, + alias_dict: Optional[Dict[str, str]] = None, +) -> Dict[str, np.ndarray]: + """Load *.dat file and fetch data as given keys. + + Args: + file_path (str): Dat file path. + keys (Tuple[str, ...]): Required fetching keys. + alias_dict (Optional[Dict[str, str]]): Alias for keys, + i.e. {original_key: original_key}. Defaults to None. + + Returns: + Dict[str, np.ndarray]: Loaded data in dict. + """ + + if alias_dict is None: + alias_dict = {} + + try: + # read all data from .dat file + raw_data = pickle.load(open(file_path, "rb")) + except FileNotFoundError as e: + raise e + + # convert to numpy array + data_dict = {} + if keys is None: + keys = raw_data.keys() + for key in keys: + fetch_key = alias_dict[key] if key in alias_dict else key + if fetch_key not in raw_data: + raise KeyError(f"fetch_key({fetch_key}) do not exist in raw_data.") + data_dict[key] = np.asarray(raw_data[fetch_key]) + if data_dict[key].dtype in (np.float16, np.float32, np.float64): + data_dict[key] = data_dict[key].astype(paddle.get_default_dtype()) + + return data_dict diff --git a/examples/smc_reac/ppsci/utils/save_load.py b/examples/smc_reac/ppsci/utils/save_load.py new file mode 100644 index 0000000000..cdf14cce87 --- /dev/null +++ b/examples/smc_reac/ppsci/utils/save_load.py @@ -0,0 +1,300 @@ +# 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. + +from __future__ import annotations + +import os +from typing import TYPE_CHECKING +from typing import Any +from typing import Dict +from typing import Optional + +import paddle + +from ppsci.utils import download +from ppsci.utils import logger + +if TYPE_CHECKING: + from paddle import amp + from paddle import nn + from paddle import optimizer + + from ppsci import equation + from ppsci.loss import mtl + from ppsci.utils import ema + + +__all__ = [ + "load_checkpoint", + "save_checkpoint", + "load_pretrain", +] + + +def _load_pretrain_from_path( + path: str, + model: nn.Layer, + equation: Optional[Dict[str, equation.PDE]] = None, +): + """Load pretrained model from given path. + + Args: + path (str): File path of pretrained model, i.e. `/path/to/model.pdparams`. + model (nn.Layer): Model with parameters. + equation (Optional[Dict[str, equation.PDE]]): Equations. Defaults to None. + """ + if not (os.path.isdir(path) or os.path.exists(f"{path}.pdparams")): + raise FileNotFoundError( + f"Pretrained model path {path}.pdparams does not exists." + ) + + param_state_dict = paddle.load(f"{path}.pdparams") + model.set_state_dict(param_state_dict) + logger.message(f"Finish loading pretrained model from: {path}.pdparams") + if equation is not None: + if not os.path.exists(f"{path}.pdeqn"): + num_learnable_params = sum( + [len(eq.learnable_parameters) for eq in equation.values()] + ) + if num_learnable_params > 0: + logger.warning( + f"There are a total of {num_learnable_params} learnable parameters" + f" in the equation, but {path}.pdeqn not found." + ) + else: + equation_dict = paddle.load(f"{path}.pdeqn") + for name, _equation in equation.items(): + _equation.set_state_dict(equation_dict[name]) + logger.message( + f"Finish loading pretrained equation parameters from: {path}.pdeqn" + ) + + +def load_pretrain( + model: nn.Layer, + path: str, + equation: Optional[Dict[str, equation.PDE]] = None, +): + """ + Load pretrained model from given path or url. + + Args: + model (nn.Layer): Model with parameters. + path (str): File path or url of pretrained model, i.e. `/path/to/model.pdparams` + or `http://xxx.com/model.pdparams`. + equation (Optional[Dict[str, equation.PDE]]): Equations. Defaults to None. + + Examples: + >>> import ppsci + >>> from ppsci.utils import save_load + >>> model = ppsci.arch.MLP(("x", "y"), ("u", "v", "p"), 9, 50, "tanh") + >>> save_load.load_pretrain( + ... model=model, + ... path="path/to/pretrain_model") # doctest: +SKIP + """ + if path.startswith("http"): + # download from path(url) and get its' physical path + eqn_path = path.replace(".pdparams", ".pdeqn", 1) + path = download.get_weights_path_from_url(path) + + # automatically download additional equation weights if available + def is_url_accessible(url: str): + try: + import requests + + response = requests.head(url, timeout=5) + return response.status_code == requests.codes.ok + except requests.RequestException: + return False + except Exception: + return False + + if is_url_accessible(eqn_path): + download.get_weights_path_from_url(eqn_path) + + # remove ".pdparams" in suffix of path for convenient + if path.endswith(".pdparams"): + path = path[:-9] + _load_pretrain_from_path(path, model, equation) + + +def load_checkpoint( + path: str, + model: nn.Layer, + optimizer: optimizer.Optimizer, + grad_scaler: Optional[amp.GradScaler] = None, + equation: Optional[Dict[str, equation.PDE]] = None, + ema_model: Optional[ema.AveragedModel] = None, + aggregator: Optional[mtl.LossAggregator] = None, +) -> Dict[str, Any]: + """Load from checkpoint. + + Args: + path (str): Path for checkpoint. + model (nn.Layer): Model with parameters. + optimizer (optimizer.Optimizer): Optimizer for model. + grad_scaler (Optional[amp.GradScaler]): GradScaler for AMP. Defaults to None. + equation (Optional[Dict[str, equation.PDE]]): Equations. Defaults to None. + ema_model: Optional[ema.AveragedModel]: Average model. Defaults to None. + aggregator: Optional[mtl.LossAggregator]: Loss aggregator. Defaults to None. + + Returns: + Dict[str, Any]: Loaded metric information. + """ + if not os.path.exists(f"{path}.pdparams"): + raise FileNotFoundError(f"{path}.pdparams not exist.") + if not os.path.exists(f"{path}.pdopt"): + raise FileNotFoundError(f"{path}.pdopt not exist.") + if grad_scaler is not None and not os.path.exists(f"{path}.pdscaler"): + raise FileNotFoundError(f"{path}.scaler not exist.") + + # load state dict + model_dict = paddle.load(f"{path}.pdparams") + optim_dict = paddle.load(f"{path}.pdopt") + metric_dict = {} + if os.path.exists(f"{path}.pdstates"): + metric_dict = paddle.load(f"{path}.pdstates") + if grad_scaler is not None: + scaler_dict = paddle.load(f"{path}.pdscaler") + if equation is not None: + if not os.path.exists(f"{path}.pdeqn"): + logger.warning(f"{path}.pdeqn not found.") + equation_dict = None + else: + equation_dict = paddle.load(f"{path}.pdeqn") + + # set model state dict + logger.message(f"* Loading model checkpoint from {path}.pdparams") + missing_keys, unexpected_keys = model.set_state_dict(model_dict) + if missing_keys: + logger.warning( + f"There are missing keys when loading checkpoint: {missing_keys}, " + "and corresponding parameters will be initialized by default." + ) + if unexpected_keys: + logger.warning( + f"There are redundant keys: {unexpected_keys}, " + "and corresponding weights will be ignored." + ) + + # set optimizer state dict + logger.message(f"* Loading optimizer checkpoint from {path}.pdopt") + optimizer.set_state_dict(optim_dict) + + if grad_scaler is not None: + logger.message(f"* Loading grad scaler checkpoint from {path}.pdscaler") + grad_scaler.load_state_dict(scaler_dict) + + if equation is not None and equation_dict is not None: + logger.message(f"* Loading equation checkpoint from {path}.pdeqn") + for name, _equation in equation.items(): + _equation.set_state_dict(equation_dict[name]) + + if ema_model is not None: + logger.message(f"* Loading EMA checkpoint from {path}_ema.pdparams") + avg_model_dict = paddle.load(f"{path}_ema.pdparams") + ema_model.set_state_dict(avg_model_dict) + + if aggregator is not None and aggregator.should_persist: + logger.message(f"* Loading loss aggregator checkpoint from {path}.pdagg") + aggregator_dict = paddle.load(f"{path}.pdagg") + aggregator.set_state_dict(aggregator_dict) + + logger.message(f"Finish loading checkpoint from {path}") + return metric_dict + + +def save_checkpoint( + model: nn.Layer, + optimizer: Optional[optimizer.Optimizer], + metric: Optional[Dict[str, float]] = None, + grad_scaler: Optional[amp.GradScaler] = None, + output_dir: Optional[str] = None, + prefix: str = "model", + equation: Optional[Dict[str, equation.PDE]] = None, + print_log: bool = True, + ema_model: Optional[ema.AveragedModel] = None, + aggregator: Optional[mtl.LossAggregator] = None, +): + """ + Save checkpoint, including model params, optimizer params, metric information. + + Args: + model (nn.Layer): Model with parameters. + optimizer (Optional[optimizer.Optimizer]): Optimizer for model. + metric (Optional[Dict[str, float]]): Metric information, such as {"RMSE": 0.1, "MAE": 0.2}. Defaults to None. + grad_scaler (Optional[amp.GradScaler]): GradScaler for AMP. Defaults to None. + output_dir (Optional[str]): Directory for checkpoint storage. + prefix (str, optional): Prefix for storage. Defaults to "model". + equation (Optional[Dict[str, equation.PDE]]): Equations. Defaults to None. + print_log (bool, optional): Whether print saving log information, mainly for + keeping log tidy without duplicate 'Finish saving checkpoint ...' log strings. + Defaults to True. + ema_model: Optional[ema.AveragedModel]: Average model. Defaults to None. + aggregator: Optional[mtl.LossAggregator]: Loss aggregator. Defaults to None. + + Examples: + >>> import ppsci + >>> import paddle + >>> from ppsci.utils import save_load + >>> model = ppsci.arch.MLP(("x", "y", "z"), ("u", "v", "w"), 5, 64, "tanh") + >>> optimizer = ppsci.optimizer.Adam(0.001)(model) + >>> save_load.save_checkpoint(model, optimizer, {"RMSE": 0.1}, output_dir="path/to/output/dir") # doctest: +SKIP + """ + if paddle.distributed.get_rank() != 0: + return + + if output_dir is None: + logger.warning("output_dir is None, skip save_checkpoint") + return + + ckpt_dir = os.path.join(output_dir, "checkpoints") + ckpt_path = os.path.join(ckpt_dir, prefix) + os.makedirs(ckpt_dir, exist_ok=True) + + paddle.save(model.state_dict(), f"{ckpt_path}.pdparams") + + if optimizer is not None: + paddle.save(optimizer.state_dict(), f"{ckpt_path}.pdopt") + + if metric is not None and len(metric) > 0: + paddle.save(metric, f"{ckpt_path}.pdstates") + + if grad_scaler is not None: + paddle.save(grad_scaler.state_dict(), f"{ckpt_path}.pdscaler") + + if equation is not None: + num_learnable_params = sum( + [len(eq.learnable_parameters) for eq in equation.values()] + ) + if num_learnable_params > 0: + paddle.save( + {key: eq.state_dict() for key, eq in equation.items()}, + f"{ckpt_path}.pdeqn", + ) + + if ema_model is not None: + paddle.save(ema_model.state_dict(), f"{ckpt_path}_ema.pdparams") + + if aggregator is not None and aggregator.should_persist: + paddle.save(aggregator.state_dict(), f"{ckpt_path}.pdagg") + + if print_log: + log_str = f"Finish saving checkpoint to: {ckpt_path}" + if prefix == "latest": + log_str += ( + "(latest checkpoint will be saved every epoch as expected, " + "but this log will be printed only once for tidy logging)" + ) + logger.message(log_str) diff --git a/examples/smc_reac/ppsci/utils/symbolic.py b/examples/smc_reac/ppsci/utils/symbolic.py new file mode 100644 index 0000000000..dcd089d017 --- /dev/null +++ b/examples/smc_reac/ppsci/utils/symbolic.py @@ -0,0 +1,981 @@ +# 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. + +""" +Sympy to python function conversion module +""" + +from __future__ import annotations + +import functools +import os +from typing import Dict +from typing import List +from typing import Optional +from typing import Sequence +from typing import Tuple +from typing import Union + +import paddle +import sympy as sp +from paddle import nn +from typing_extensions import TypeAlias + +from ppsci import arch +from ppsci import equation +from ppsci.autodiff import hessian +from ppsci.autodiff import jacobian +from ppsci.utils import logger + +__all__ = [ + "lambdify", + "_cvt_to_key", +] + + +DATA_DICT: TypeAlias = Dict[str, paddle.Tensor] + +SYMPY_BUILTIN_FUNC: TypeAlias = Union[ + sp.sin, + sp.sinh, + sp.asin, + sp.cos, + sp.acos, + sp.cosh, + sp.tan, + sp.atan, + sp.atan2, + sp.acosh, + sp.asinh, + sp.tanh, + sp.atanh, + sp.erf, + sp.loggamma, + sp.exp, + sp.Pow, + sp.log, + sp.Max, + sp.Min, + sp.Abs, + sp.Heaviside, + sp.sign, + sp.ceiling, + sp.floor, + sp.Add, + sp.Mul, +] + +SYMPY_TO_PADDLE = { + sp.sin: paddle.sin, + sp.sinh: paddle.sinh, + sp.asin: paddle.asin, + sp.cos: paddle.cos, + sp.acos: paddle.acos, + sp.cosh: paddle.cosh, + sp.tan: paddle.tan, + sp.atan: paddle.atan, + sp.atan2: paddle.atan2, + sp.acosh: paddle.acosh, + sp.asinh: paddle.asinh, + sp.tanh: paddle.tanh, + sp.atanh: paddle.atanh, + sp.erf: paddle.erf, + sp.loggamma: paddle.lgamma, + sp.exp: paddle.exp, + sp.Pow: paddle.pow, + sp.log: paddle.log, + sp.Max: paddle.maximum, + sp.Min: paddle.minimum, + sp.Abs: paddle.abs, + sp.Heaviside: paddle.heaviside, + sp.sign: paddle.sign, + sp.ceiling: paddle.ceil, + sp.floor: paddle.floor, + # NOTE: sp.Add and sp.Mul is not included here for un-alignment with paddle + # and are implemented manually in 'OperatorNode._add_operator_func' and + # 'OperatorNode._mul_operator_func' +} + + +def _cvt_to_key(expr: sp.Basic) -> str: + """Convert sympy expression to a string key, mainly as retrieval key in dict. + + Args: + expr (sp.Basic): Sympy expression. + + Returns: + str: Converted string key. + """ + if isinstance(expr, sp.Function) and str(expr.func) == equation.DETACH_FUNC_NAME: + return f"{_cvt_to_key(expr.args[0])}_{equation.DETACH_FUNC_NAME}" + + if isinstance(expr, (sp.Symbol, sp.core.function.UndefinedFunction, sp.Function)): + # use name of custom function(e.g. "f") instead of itself(e.g. "f(x, y)") + # for simplicity. + if hasattr(expr, "name"): + return expr.name + else: + return str(expr) + elif isinstance(expr, sp.Derivative): + # convert "Derivative(u(x,y),(x,2),(y,2))" to "u__x__x__y__y" + expr_str = expr.args[0].name + for symbol, order in expr.args[1:]: + expr_str += f"__{symbol}" * order + return expr_str + else: + return str(expr) + + +class Node(nn.Layer): + """The base class of the node in expression tree. + + Args: + expr (sp.Basic): Sympy expression. + """ + + def __init__(self, expr: sp.Basic): + super().__init__() + self.expr = expr + self.key = _cvt_to_key(self.expr) + + def forward(self, **kwargs): + raise NotImplementedError("Node.forward is not implemented") + + def __str__(self): + return ( + f"{self.__class__.__name__}(expr: {self.expr}, " + f"expr_type: {type(self.expr)})" + ) + + def __repr__(self): + return f"{self.__class__.__name__}(expr: {self.expr})" + + +class DetachNode(Node): + """Class for detach operation in converted expression tree. + + Args: + expr (sp.Basic): Sympy expression. + """ + + def __init__(self, expr: sp.Basic): + super().__init__(expr) + self.child = _cvt_to_key(self.expr.args[0]) + + def forward(self, data_dict: DATA_DICT): + if self.key in data_dict: + return data_dict + + data_dict[self.key] = data_dict[self.child].detach() + return data_dict + + +class OperatorNode(Node): + """Class for operator node in converted expression tree. + + Args: + expr (SYMPY_BUILTIN_FUNC): Sympy expression. + """ + + def __init__( + self, + expr: SYMPY_BUILTIN_FUNC, + ): + super().__init__(expr) + # preprocess children's key instead of processing at run-time in forward + # which can reduce considerable overhead of time for calling "_cvt_to_key" + self.childs = [_cvt_to_key(arg) for arg in self.expr.args] + + if self.expr.func == sp.Add: + self._apply_func = self._add_operator_func + elif self.expr.func == sp.Mul: + self._apply_func = self._mul_operator_func + elif self.expr.func == sp.Heaviside: + self._apply_func = self._heaviside_operator_func + self._auxiliary_func = SYMPY_TO_PADDLE[sp.Heaviside] + self._auxiliary_func = functools.partial( + self._auxiliary_func, y=paddle.zeros([]) + ) + elif self.expr.func == sp.Min: + self._apply_func = self._minimum_operator_func + elif self.expr.func == sp.Max: + self._apply_func = self._maximum_operator_func + else: + self._apply_func = self._vanilla_operator_func + self._auxiliary_func = SYMPY_TO_PADDLE[self.expr.func] + + def forward(self, data_dict: DATA_DICT): + # use cache + if self.key in data_dict: + return data_dict + + return self._apply_func(data_dict) + + def _add_operator_func(self, data_dict: DATA_DICT) -> DATA_DICT: + data_dict[self.key] = data_dict[self.childs[0]] + for p in self.childs[1:]: + data_dict[self.key] += data_dict[p] + return data_dict + + def _mul_operator_func(self, data_dict: DATA_DICT) -> DATA_DICT: + data_dict[self.key] = data_dict[self.childs[0]] + for child in self.childs[1:]: + data_dict[self.key] *= data_dict[child] + return data_dict + + def _heaviside_operator_func(self, data_dict: DATA_DICT) -> DATA_DICT: + data_dict[self.key] = self._auxiliary_func(data_dict[self.childs[0]]) + return data_dict + + def _minimum_operator_func(self, data_dict: DATA_DICT) -> DATA_DICT: + data_dict[self.key] = paddle.minimum( + data_dict[self.childs[0]], data_dict[self.childs[1]] + ) + for i in range(2, len(self.childs)): + data_dict[self.key] = paddle.minimum( + data_dict[self.key], + data_dict[self.childs[i]], + ) + return data_dict + + def _maximum_operator_func(self, data_dict: DATA_DICT) -> DATA_DICT: + data_dict[self.key] = paddle.maximum( + data_dict[self.childs[0]], data_dict[self.childs[1]] + ) + for i in range(2, len(self.childs)): + data_dict[self.key] = paddle.maximum( + data_dict[self.key], + data_dict[self.childs[i]], + ) + return data_dict + + def _vanilla_operator_func(self, data_dict: DATA_DICT) -> DATA_DICT: + data_dict[self.key] = self._auxiliary_func( + *tuple(data_dict[child] for child in self.childs) + ) + return data_dict + + +class DerivativeNode(Node): + """Class for operator node in converted expression tree. + + Args: + expr (sp.Derivative): Sympy derivative expression. + create_graph (bool, optional): Whether to create the gradient graphs of + the computing process. When it is True, higher order derivatives are + supported to compute; when it is False, the gradient graphs of the + computing process would be discarded. Defaults to True. + retain_graph (Optional[bool]): Whether to retain the forward graph which + is used to calculate the gradient. When it is True, the graph would + be retained, in which way users can calculate backward twice for the + same graph. When it is False, the graph would be freed. Defaults to None, + which means it is equal to `create_graph`. + """ + + def __init__( + self, + expr: sp.Derivative, + create_graph: bool = True, + retain_graph: Optional[bool] = None, + ): + super().__init__(expr) + # preprocess children's key instead of processing at run-time in forward + # which can reduce considerable overhead of time for calling "_cvt_to_key" + self.childs = [_cvt_to_key(self.expr.args[0])] + [ + (_cvt_to_key(arg), int(order)) for (arg, order) in self.expr.args[1:] + ] + self.create_graph = create_graph + self.retain_graph = retain_graph + self._apply_func = self._derivate_operator_func + self.merged = False + + def forward(self, data_dict: DATA_DICT): + # use cache + if self.key in data_dict: + return data_dict + + return self._apply_func(data_dict) + + def _derivate_operator_func(self, data_dict: DATA_DICT) -> DATA_DICT: + # NOTE: Derivative of 'sdf' function will not be executed here, which is already + # generated in 'data_dict' during points sampling using discrete difference + # method(see also: ppsci/geometry/geometry.py: Geometry.sdf_derivatives), + # such as 'sdf__x', 'sdf__y'. + data_dict[self.key] = data_dict[self.childs[0]] + for child, order in self.childs[1:]: + if order & 1: + data_dict[self.key] = jacobian( + data_dict[self.key], + data_dict[child], + create_graph=self.create_graph, + retain_graph=self.retain_graph, + ) + order -= 1 + for _ in range(0, order, 2): + data_dict[self.key] = hessian( + data_dict[self.key], + data_dict[child], + create_graph=self.create_graph, + retain_graph=self.retain_graph, + ) + order -= 2 + return data_dict + + +class FusedDerivativeNode(nn.Layer): + """Class for fused DerivativeNode. + + Args: + f_x_tuples (List[Tuple[Union[sp.Function, sp.Derivative], sp.Symbol]]): + indicate all derivatives of a function in list of tuples. e.g. + [(func1, var1), (func1, var2), (func1, var3), ...]. + create_graph (bool, optional): Whether to create the gradient graphs of + the computing process. When it is True, higher order derivatives are + supported to compute; when it is False, the gradient graphs of the + computing process would be discarded. Defaults to True. + retain_graph (Optional[bool]): Whether to retain the forward graph which + is used to calculate the gradient. When it is True, the graph would + be retained, in which way users can calculate backward twice for the + same graph. When it is False, the graph would be freed. Defaults to None, + which means it is equal to `create_graph`. + """ + + def __init__( + self, + f_x_tuples: List[Tuple[Union[sp.Function, sp.Derivative], sp.Symbol]], + create_graph: bool = True, + retain_graph: Optional[bool] = None, + ): + super().__init__() + self.expr: List[sp.Derivative] = [f.diff(x) for f, x in f_x_tuples] + self.key: List[str] = [_cvt_to_key(expr) for expr in self.expr] + self.create_graph = create_graph + self.retain_graph = retain_graph + + # preprocess children's key instead of processing at run-time in forward + # which can reduce considerable overhead of time for calling "_cvt_to_key" + self.y_key: str = _cvt_to_key(f_x_tuples[0][0]) + self.childs: List[str] = [_cvt_to_key(x) for _, x in f_x_tuples] + self._apply_func = self._parallel_derivate_operator_func + + def forward(self, data_dict: DATA_DICT): + # use cache + if all([key in data_dict for key in self.key]): + return data_dict + + return self._apply_func(data_dict) + + def _parallel_derivate_operator_func(self, data_dict: DATA_DICT) -> DATA_DICT: + # NOTE: Derivative of 'sdf' function will not be executed here, which is already + # generated in 'data_dict' during points sampling using discrete difference + # method(see also: ppsci/geometry/geometry.py: Geometry.sdf_derivatives), + # such as 'sdf__x', 'sdf__y'. + y_data: paddle.Tensor = data_dict[self.y_key] + xs_data: List[paddle.Tensor] = [data_dict[x_key] for x_key in self.childs] + y_wrt_xs_grad: List[paddle.Tensor] = jacobian( + y_data, + xs_data, + create_graph=self.create_graph, + retain_graph=self.retain_graph, + ) + for i, key in enumerate(self.key): + data_dict[key] = y_wrt_xs_grad[i] + return data_dict + + def __str__(self): + return ( + f"{self.__class__.__name__}(expr: {self.expr}, " + f"expr_type: {type(self.expr)})" + ) + + def __repr__(self): + return f"{self.__class__.__name__}(expr: {self.expr})" + + +class LayerNode(Node): + """Class for layer node in converted expression tree. + + Args: + expr (sp.core.function.UndefinedFunction): Sympy expression. + model (arch.Arch): NN model for computing forward result in this node. + """ + + def __init__( + self, + expr: sp.core.function.UndefinedFunction, + model: arch.Arch, + ): + super().__init__(expr) + self.model = model + + def forward(self, data_dict: DATA_DICT) -> DATA_DICT: + # use cache + if self.key in data_dict: + return data_dict + + output_dict = self.model(data_dict) + data_dict.update(output_dict) + + return data_dict + + +class ConstantNode(Node): + """Class for constant variable node in converted expression tree. + + Args: + expr (Union[sp.Number, sp.NumberSymbol]): Number expression. + """ + + def __init__(self, expr: Union[sp.Number, sp.NumberSymbol]): + super().__init__(expr) + if ( + self.expr.is_Float + or self.expr.is_Integer + or self.expr.is_Boolean + or self.expr.is_Rational + ): + self.expr = float(self.expr) + else: + raise TypeError( + "expr({expr}) should be Float/Integer/Boolean/Rational, " + f"but got {type(self.expr)}" + ) + self.expr = paddle.to_tensor(self.expr) + + def forward(self, data_dict: DATA_DICT) -> DATA_DICT: + # use cache + if self.key in data_dict: + return data_dict + + data_dict[self.key] = self.expr + return data_dict + + def __str__(self): + return ( + f"{self.__class__.__name__}(expr: {float(self.expr)}, " + f"expr_type: {type(self.expr)})" + ) + + +class ParameterNode(Node): + """Class for constant variable node in converted expression tree. + + Args: + expr (sp.Symbol): Parameter expression. + parameter (paddle.framework.io.EagerParamBase): Parameter tensor. + """ + + def __init__(self, expr: sp.Symbol, parameter: paddle.framework.io.EagerParamBase): + super().__init__(expr) + self.parameter = parameter + + def forward(self, data_dict: DATA_DICT) -> DATA_DICT: + data_dict[self.key] = self.parameter + return data_dict + + +class ComposedNode(nn.Layer): + """ + Compose list of several callable objects together. + """ + + def __init__(self, callable_nodes: List[Node]): + super().__init__() + assert len(callable_nodes) + self.callable_nodes = nn.LayerList(callable_nodes) + + def forward(self, data_dict: DATA_DICT) -> paddle.Tensor: + # call all callable_nodes in order + for i, func in enumerate(self.callable_nodes): + data_dict = func(data_dict) + + # return result of last node(root node) for target + return data_dict[self.callable_nodes[-1].key] + + +def _post_traverse(cur_node: sp.Basic, nodes: List[sp.Basic]) -> List[sp.Basic]: + """Traverse sympy expression tree in post-order. + + Args: + cur_node (sp.Basic): Sympy expression of current node. + nodes (List[sp.Basic]): Node list storing all tree nodes in post-order. + + Returns: + List[sp.Basic]: Node list storing all tree nodes in post-order. + """ + # traverse into sub-nodes + if isinstance(cur_node, sp.Function): + for arg in cur_node.args: + nodes = _post_traverse(arg, nodes) + nodes.append(cur_node) + elif isinstance(cur_node, sp.Derivative): + nodes = _post_traverse(cur_node.args[0], nodes) + nodes.append(cur_node) + elif isinstance(cur_node, sp.Symbol): + nodes.append(cur_node) + return nodes + elif isinstance(cur_node, sp.Number): + nodes.append(cur_node) + else: + for arg in cur_node.args: + nodes = _post_traverse(arg, nodes) + nodes.append(cur_node) + return nodes + + +def _visualize_graph(nodes: List[sp.Basic], graph_filename: str): + try: + import pygraphviz + except ModuleNotFoundError: + raise ModuleNotFoundError( + "Please install pygraphviz by steps below:\n" + "1. apt-get install graphviz graphviz-dev\n" + "2. python -m pip install pygraphviz" + ) + + SYMPY_BUILTIN_NAME = { + sp.sin: "sin", + sp.sinh: "sinh", + sp.asin: "asin", + sp.cos: "cos", + sp.acos: "acos", + sp.cosh: "cosh", + sp.tan: "tan", + sp.atan: "atan", + sp.atan2: "atan2", + sp.acosh: "acosh", + sp.asinh: "asinh", + sp.tanh: "tanh", + sp.atanh: "atanh", + sp.erf: "erf", + sp.loggamma: "loggamma", + sp.exp: "exp", + sp.Pow: "Pow", + sp.log: "log", + sp.Max: "Max", + sp.Min: "Min", + sp.Abs: "Abs", + sp.Heaviside: "Heaviside", + sp.sign: "sign", + sp.ceiling: "ceiling", + sp.floor: "floor", + sp.Add: "Add", + sp.Mul: "Mul", + } + naming_counter = {k: 0 for k in SYMPY_BUILTIN_NAME} + + def get_operator_name(node: sp.Function): + ret = f"{SYMPY_BUILTIN_NAME[node.func]}_{naming_counter[node.func]}" + naming_counter[node.func] += 1 + return ret + + graph = pygraphviz.AGraph(directed=True, rankdir="TB") + C_FUNC = "#9196f1" # purple color function node + C_DATA = "#feb64d" # orange color for data node + C_EDGE = "#000000" # black color for edge + + def add_edge(u: str, v: str, u_color: str = C_DATA, v_color: str = C_DATA): + """Add an edge from `u` to `v`. + + Args: + u (str): Name of begin node u. + v (str): Name of end node v. + u_color (str, optional): Color of node u. Defaults to '#feb64d'. + v_color (str, optional): Color of node v. Defaults to '#feb64d'. + """ + graph.add_node(u, style="filled", shape="ellipse", color=u_color) + graph.add_node(v, style="filled", shape="ellipse", color=v_color) + graph.add_edge(u, v, color=C_EDGE, style="solid", penwidth=0.5, arrowsize=0.5) + + for node in nodes: + if isinstance(node, tuple(SYMPY_BUILTIN_NAME.keys())): + operator_str = get_operator_name(node) + for arg in node.args: + add_edge(_cvt_to_key(arg), operator_str, v_color=C_FUNC) + add_edge(operator_str, _cvt_to_key(node), u_color=C_FUNC) + elif isinstance(node, sp.Function): + for arg in node.args: + add_edge(_cvt_to_key(arg), str(node), v_color=C_FUNC) + add_edge(str(node), _cvt_to_key(node), u_color=C_FUNC) + elif isinstance(node, sp.Derivative): + add_edge(str(node), _cvt_to_key(node), u_color=C_FUNC) + add_edge(_cvt_to_key(node.args[0]), str(node), v_color=C_FUNC) + for arg in node.args[1:]: + add_edge(_cvt_to_key(arg[0]), str(node), v_color=C_FUNC) + + # export graph to image + graph.layout() + image_path = f"{graph_filename}.png" + dot_path = f"{graph_filename}.dot" + if len(os.path.dirname(image_path)): + os.makedirs(os.path.dirname(image_path), exist_ok=True) + graph.draw(image_path, prog="dot") + graph.write(dot_path) + logger.message( + f"Computational graph has been written to: {image_path} and {dot_path}, " + "which can be visualized at: https://dreampuf.github.io/GraphvizOnline/" + ) + + +def _fuse_derivative_nodes( + derivative_exprs: List[sp.Derivative], +) -> List[FusedDerivativeNode]: + """Merge derivative nodes and return in list of FusedDerivativeNode after merger. + + Args: + derivative_exprs (List[sp.Derivative]): Derivatives sympy expression of same + function, e.g. [Derivative(u(x,y), x), Derivative(u(x,y), y)] + + Returns: + List[FusedDerivativeNode]: List of FusedDerivativeNode converting from mergeable + derivatives. + """ + + class DerivativeTrie: + """Trie for unrolling derivative.""" + + def __init__(self, expr: sp.Basic): + self.expr: sp.Basic = expr + self.next: Dict["sp.Symbol", "DerivativeTrie"] = {} + + # unroll derivative expressions into a trie structure + trie_root = DerivativeTrie(derivative_exprs[0].args[0]) + for derivative_expr in derivative_exprs: + cur_node = trie_root + for (child, order) in derivative_expr.args[1:]: + for _ in range(order): + if child not in cur_node.next: + cur_node.next[child] = DerivativeTrie(cur_node.expr.diff(child)) + cur_node = cur_node.next[child] + + def dfs_trie( + node: DerivativeTrie, fused_derivative_nodes: List[FusedDerivativeNode] + ) -> None: + if node.next: + fused_derivative_nodes.append( + FusedDerivativeNode( + [(node.expr, name) for name in node.next], + ) + ) + for child in node.next: + dfs_trie(node.next[child], fused_derivative_nodes) + + # walk on derivative trie in pre-order and log fusable nodes + fused_derivative_nodes: List[FusedDerivativeNode] = [] + dfs_trie(trie_root, fused_derivative_nodes) + + return fused_derivative_nodes + + +def lambdify( + expr: Union[sp.Basic, List[sp.Basic]], + models: Optional[Union[arch.Arch, Tuple[arch.Arch, ...]]] = None, + extra_parameters: Optional[Sequence[paddle.Tensor]] = None, + graph_filename: Optional[str] = None, + create_graph: bool = True, + retain_graph: Optional[bool] = None, + fuse_derivative: bool = False, +) -> Union[ComposedNode, List[ComposedNode]]: + """Convert sympy expression to callable function. + + Args: + expr (Union[sp.Basic, List[sp.Basic]]): Sympy expression(s) to be converted. + Will return callable functions in list if multiple expressions are given, + else return one single callable function. + models (Optional[Union[arch.Arch, Tuple[arch.Arch, ...]]]): Model(s) for + computing forward result in `LayerNode`. + extra_parameters (Optional[nn.ParameterList]): Extra learnable parameters. + Defaults to None. + graph_filename (Optional[str]): Save computational graph to `graph_filename.png` + for given `expr`, if `graph_filename` is not None and a valid string, + such as 'momentum_x'. Defaults to None. + create_graph (bool, optional): Whether to create the gradient graphs of + the computing process. When it is True, higher order derivatives are + supported to compute. When it is False, the gradient graphs of the + computing process would be discarded. Defaults to True. + retain_graph (Optional[bool]): Whether to retain the forward graph which + is used to calculate the gradient. When it is True, the graph would + be retained, in which way users can calculate backward twice for the + same graph. When it is False, the graph would be freed. Defaults to None, + which means it is equal to `create_graph`. + fuse_derivative (bool, optional): Whether to fuse the derivative nodes. + For example, if `expr` is 'Derivative(u, x) + Derivative(u, y)' + It will compute grad(u, x) + grad(u, y) if fuse_derivative=False, + else will compute sum(grad(u, [x, y])) if fuse_derivative=True as is more + efficient in backward-graph. Defaults to False, as it is experimental so not + enabled by default if used independently. + + Returns: + Union[ComposedNode, List[ComposedNode]]: Callable object(s) for computing expr + with necessary input(s) data in dict given. + + Examples: + >>> import paddle + >>> import ppsci + >>> import sympy as sp + + >>> a, b, c, x, y = sp.symbols("a b c x y") + >>> u = sp.Function("u")(x, y) + >>> v = sp.Function("v")(x, y) + >>> z = -a + b * (c ** 2) + u * v + 2.3 + + >>> model = ppsci.arch.MLP(("x", "y"), ("u", "v"), 4, 16) + + >>> batch_size = 13 + >>> a_tensor = paddle.randn([batch_size, 1]) + >>> b_tensor = paddle.randn([batch_size, 1]) + >>> c_tensor = paddle.randn([batch_size, 1]) + >>> x_tensor = paddle.randn([batch_size, 1]) + >>> y_tensor = paddle.randn([batch_size, 1]) + + >>> model_output_dict = model({"x": x_tensor, "y": y_tensor}) + >>> u_tensor, v_tensor = model_output_dict["u"], model_output_dict["v"] + + >>> z_tensor_manually = ( + ... -a_tensor + b_tensor * (c_tensor ** 2) + ... + u_tensor * v_tensor + 2.3 + ... ) + >>> z_tensor_sympy = ppsci.lambdify(z, model)( + ... { + ... "a": a_tensor, + ... "b": b_tensor, + ... "c": c_tensor, + ... "x": x_tensor, + ... "y": y_tensor, + ... } + ... ) + + >>> paddle.allclose(z_tensor_manually, z_tensor_sympy).item() + True + """ + if not extra_parameters: + extra_parameters = () + + if isinstance(models, arch.ModelList): + models = tuple(models.model_list[i] for i in range(len(models.model_list))) + if not isinstance(models, (tuple, list)): + models = (models,) + + def _expr_to_callable_nodes( + single_expr: sp.Basic, graph_filename_: Optional[str] = None + ) -> List[Node]: + """Convert sympy expression to a sequence of nodes in topologic order. + + Args: + single_expr (sp.Basic): Single sympy expression, such as "a+b*c". + graph_filename_ (Optional[str]): Save computational graph to + `/path/to/graph_filename.png` for given `expr`, if `graph_filename` is not + None and a valid string, such as 'momentum_x'. Defaults to None. + + Returns: + List[Node]: Sequence of callable nodes. + """ + # NOTE: Those simplify methods may complicate given expr instead, so not use here + # simplify expression to reduce nodes in tree + # expr = sp.nsimplify(expr) + # expr = sp.expand(expr) + # expr = sp.simplify(expr) + + # remove 1.0 from sympy expression tree + single_expr = single_expr.subs(1.0, 1) + + # convert sympy expression tree to list of nodes in post-order + sympy_nodes: List[sp.Basic] = [] + sympy_nodes = _post_traverse(single_expr, sympy_nodes) + + # remove unnecessary symbol nodes already in input dict(except for parameter symbol) + _parameter_names = tuple(param.name for param in extra_parameters) + sympy_nodes = [ + node + for node in sympy_nodes + if (not node.is_Symbol) or (_cvt_to_key(node) in _parameter_names) + ] + + # remove duplicated node(s) with topological order kept + sympy_nodes = list(dict.fromkeys(sympy_nodes)) + + # convert sympy node to callable node + callable_nodes = [] + for i, node in enumerate(sympy_nodes): + if isinstance( + node, tuple(SYMPY_TO_PADDLE.keys()) + (sp.Add, sp.Mul, sp.Derivative) + ): + if isinstance(node, sp.Derivative): + callable_nodes.append( + DerivativeNode(node, create_graph, retain_graph) + ) + else: + callable_nodes.append(OperatorNode(node)) + elif isinstance(node, sp.Function): + if str(node.func) == equation.DETACH_FUNC_NAME: + callable_nodes.append(DetachNode(node)) + logger.debug(f"Detected detach node {node}") + else: + match_index = None + for j, model in enumerate(models): + if str(node.func) in model.output_keys: + callable_nodes.append( + LayerNode( + node, + model, + ) + ) + if match_index is not None: + raise ValueError( + f"Name of function: '{node}' should be unique along given" + f" models, but got same output_key: '{str(node.func)}' " + f"in given models[{match_index}] and models[{j}]." + ) + match_index = j + # NOTE: Skip 'sdf' function, which should be already generated in + # given data_dict + if match_index is None and str(node.func) != "sdf": + raise ValueError( + f"Node {node} can not match any model in given model(s)." + ) + elif node.is_Number or node.is_NumberSymbol: + callable_nodes.append(ConstantNode(node)) + elif isinstance(node, sp.Symbol): + callable_nodes.append( + ParameterNode( + node, + *[ + param + for param in extra_parameters + if param.name == node.name + ], + ) + ) + else: + raise NotImplementedError( + f"The node {node} is not supported in lambdify." + ) + + # NOTE: visualize computational graph using 'pygraphviz' + if isinstance(graph_filename, str): + _visualize_graph(sympy_nodes, os.path.join(graph_filename, graph_filename_)) + + return callable_nodes + + if isinstance(expr, sp.Basic): + callable_nodes_group = [_expr_to_callable_nodes(expr, "expr")] + else: + callable_nodes_group = [ + _expr_to_callable_nodes(expr_i, f"expr_{i}") + for i, expr_i in enumerate(expr) + ] + + # [Optional] Fused derivatives nodes that with same function to be differentiated + while fuse_derivative: + candidate_pos: List[Tuple[int, int]] = [] # [(group_id, node_id), ...] + + # use 4-nested for-loop to find all potential mergeable derivative nodes + for i in range(len(callable_nodes_group)): + for j in range(len(callable_nodes_group[i])): + # skip non-derivative node + if not isinstance(callable_nodes_group[i][j], DerivativeNode): + continue + # skip sdf function since it is always already given in data_dict + if callable_nodes_group[i][j].expr.args[0].name == "sdf": + continue + # skip merged node + if callable_nodes_group[i][j].merged: + continue + + candidate_pos = [[i, j]] + for ii in range(len(callable_nodes_group)): + for jj in range(len(callable_nodes_group[ii])): + # skip non-derivative node + if not isinstance(callable_nodes_group[ii][jj], DerivativeNode): + continue + + # skip same node + if i == ii and j == jj: + continue + # skip merged node + if callable_nodes_group[ii][jj].merged: + continue + + # has same function item + if ( + callable_nodes_group[i][j].expr.args[0] + == callable_nodes_group[ii][jj].expr.args[0] + ): + candidate_pos.append([ii, jj]) + + if len(candidate_pos) > 1: + break + if len(candidate_pos) > 1: + break + + # merge all candidate nodes into one or more FusedDerivativeNode node + if len(candidate_pos) > 1: + fused_node_seq = _fuse_derivative_nodes( + [callable_nodes_group[gid][nid].expr for gid, nid in candidate_pos] + ) + assert isinstance( + fused_node_seq, list + ), "'fused_node_seq' should be list of 'FusedDerivativeNode'" + gid0, nid0 = candidate_pos[0] + logger.debug( + f"Fused {len(candidate_pos)} derivatives nodes: " + f"{[callable_nodes_group[i][j].expr for i, j in candidate_pos]} into" + f" {len(fused_node_seq)} fuse node sequence: {fused_node_seq} at position: ([{gid0}][{nid0}])" + ) + + # mark merged node + for i, (gid, nid) in enumerate(candidate_pos): + assert isinstance(callable_nodes_group[gid][nid], DerivativeNode) + callable_nodes_group[gid][nid].merged = True + + # replace first mergeable node with fused node sequence(packed in list) + # then mask the rest merged node to None(except [gid0, nid0]) + for i, (gid, nid) in enumerate(candidate_pos[1:]): + # keep the end node of each group to avoid generating empty callable + # node sequence, this will not effect performance since cache strategy + # in Node.forward + if nid != len(callable_nodes_group[gid]) - 1: + callable_nodes_group[gid][nid] = None + + if nid0 == len(callable_nodes_group[gid0]) - 1: + callable_nodes_group[gid0].insert(nid0, fused_node_seq) + else: + callable_nodes_group[gid0][nid0] = fused_node_seq + + # re-organize callable_nodes_group, remove None element and unpack list + for i in range(len(callable_nodes_group)): + tmp = [] + for j in range(len(callable_nodes_group[i])): + if isinstance( + callable_nodes_group[i][j], (Node, FusedDerivativeNode) + ): + tmp.append(callable_nodes_group[i][j]) + elif isinstance(callable_nodes_group[i][j], list) and isinstance( + callable_nodes_group[i][j][0], FusedDerivativeNode + ): + tmp.extend(callable_nodes_group[i][j]) + else: + assert ( + callable_nodes_group[i][j] is None + ), f"Unexpected element: {callable_nodes_group[i][j]}" + callable_nodes_group[i] = tmp + else: + # exit while loop if no more fused + break + + # Compose callable nodes into one callable object + if isinstance(expr, sp.Basic): + return ComposedNode(callable_nodes_group[0]) + else: + return [ComposedNode(callable_nodes) for callable_nodes in callable_nodes_group] diff --git a/examples/smc_reac/ppsci/utils/writer.py b/examples/smc_reac/ppsci/utils/writer.py new file mode 100644 index 0000000000..d1b4c503f3 --- /dev/null +++ b/examples/smc_reac/ppsci/utils/writer.py @@ -0,0 +1,225 @@ +# 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. + +from __future__ import annotations + +import csv +import os +from typing import Dict +from typing import Optional +from typing import Tuple +from typing import Union + +import numpy as np +import paddle + +from ppsci.utils import logger + +__all__ = [ + "save_csv_file", +] + + +def save_csv_file( + filename: str, + data_dict: Dict[str, Union[np.ndarray, "paddle.Tensor"]], + keys: Tuple[str, ...], + alias_dict: Optional[Dict[str, str]] = None, + use_header: bool = True, + delimiter: str = ",", + encoding: str = "utf-8", +): + """Write numpy or tensor data into csv file. + + Args: + filename (str): Dump file path. + data_dict (Dict[str, Union[np.ndarray, paddle.Tensor]]): Numpy or tensor data in dict. + keys (Tuple[str, ...]): Keys for data_dict to be fetched. + alias_dict (Optional[Dict[str, str]], optional): Alias dict for keys, + i.e. {dump_key: dict_key}. Defaults to None. + use_header (bool, optional): Whether save csv with header. Defaults to True. + delimiter (str, optional): Delimiter for splitting different data field. Defaults to ",". + encoding (str, optional): Encoding. Defaults to "utf-8". + + Examples: + >>> import numpy as np + >>> from ppsci.utils import save_csv_file + >>> data_dict = { + ... "a": np.array([[1], [2], [3]]).astype("int64"), # [3, 1] + ... "b": np.array([[4.12], [5.25], [6.3370]]).astype("float32"), # [3, 1] + ... } + >>> save_csv_file( + ... "test.csv", + ... data_dict, + ... ("A", "B"), + ... alias_dict={"A": "a", "B": "b"}, + ... use_header=True, + ... delimiter=",", + ... encoding="utf-8", + ... ) # doctest: +SKIP + + >>> # == test.csv == + >>> # A,B + >>> # 1,4.12 + >>> # 2,5.25 + >>> # 3,6.337 + """ + if alias_dict is None: + alias_dict = {} + + # convert to numpy array + data_fields = [] + header = [] + for key in keys: + fetch_key = alias_dict.get(key, key) + data = data_dict[fetch_key] + if isinstance(data, paddle.Tensor): + data = data.numpy() # [num_of_samples, ] + + if isinstance(data, np.ndarray): + data = data.flatten() + data_fields.append(data) + + header.append(key) + + assert len(header) == len(data_fields) + + if len(os.path.dirname(filename)): + os.makedirs(os.path.dirname(filename), exist_ok=True) + + data_fields = zip(*data_fields) # transpose col data to row data + with open(filename, "w", newline="", encoding=encoding) as file: + writer = csv.writer(file, delimiter=delimiter) + + if use_header: + writer.writerow(header) + + writer.writerows(data_fields) + + logger.message(f"csv file has been dumped to: {filename}") + + +def save_tecplot_file( + filename: str, + data_dict: Dict[str, Union[np.ndarray, "paddle.Tensor"]], + keys: Tuple[str, ...], + num_x: int, + num_y: int, + alias_dict: Optional[Dict[str, str]] = None, + delimiter: str = " ", + encoding: str = "utf-8", + num_timestamps: int = 1, +): + """Write numpy or tensor data into tecplot file(s). + + Args: + filename (str): Tecplot file path. + data_dict (Dict[str, Union[np.ndarray, paddle.Tensor]]): Numpy or Tensor data in dict. + keys (Tuple[str, ...]): Target keys to be dumped. + num_x (int): The number of discrete points of the grid in the X-axis. Assuming + the discrete grid size is 20 x 30, then num_x=20. + num_y (int): The number of discrete points of the grid in the Y-axis. Assuming + the discrete grid size is 20 x 30, then num_y=30. + alias_dict (Optional[Dict[str, str]], optional): Alias dict for keys, + i.e. {dump_key: dict_key}. Defaults to None. + delimiter (str, optional): Delimiter for splitting different data field. Defaults to " ". + encoding (str, optional): Encoding. Defaults to "utf-8". + num_timestamps (int, optional): Number of timestamp over coord and value. Defaults to 1. + + Examples: + >>> import numpy as np + >>> from ppsci.utils import save_tecplot_file + >>> data_dict = { + ... "x": np.array([[-1.0], [-1.0], [-1.0], [-1.0], [-1.0], [-1.0]]), # [6, 1] + ... "y": np.array([[1.0], [2.0], [3.0], [1.0], [2.0], [3.0]]), # [6, 1] + ... "value": np.array([[3], [33], [333], [3333], [33333], [333333]]), # [6, 1] + ... } + >>> save_tecplot_file( + ... "./test.dat", + ... data_dict, + ... ("X", "Y", "value"), + ... num_x=1, + ... num_y=3, + ... alias_dict={"X": "x", "Y": "y"}, + ... num_timestamps=2, + ... ) # doctest: +SKIP + >>> # == test_t-0.dat == + >>> # title = "./test_t-0.dat" + >>> # variables = "X", "Y" + >>> # Zone I = 3, J = 1, F = POINT + >>> # -1.0 1.0 3.0 + >>> # -1.0 2.0 33.0 + >>> # -1.0 3.0 333.0 + + + >>> # == test_t-1.dat == + >>> # title = "./test_t-1.dat" + >>> # variables = "X", "Y" + >>> # Zone I = 3, J = 1, F = POINT + >>> # -1.0 1.0 3333.0 + >>> # -1.0 2.0 33333.0 + >>> # -1.0 3.0 333333.0 + """ + if alias_dict is None: + alias_dict = {} + + ntxy = len(next(iter(data_dict.values()))) + if ntxy % num_timestamps != 0: + raise ValueError( + f"num_points({ntxy}) must be a multiple of " + f"num_timestamps({num_timestamps})." + ) + nxy = ntxy // num_timestamps + + nx, ny = num_x, num_y + assert nx * ny == nxy, f"nx({nx}) * ny({ny}) != nxy({nxy})" + + if len(os.path.dirname(filename)): + os.makedirs(os.path.dirname(filename), exist_ok=True) + + if filename.endswith(".dat"): + filename = filename[:-4] + + for t in range(num_timestamps): + # write 1 tecplot file for each timestep + if num_timestamps > 1: + dump_filename = f"{filename}_t-{t}.dat" + else: + dump_filename = f"{filename}.dat" + + fetch_keys = [alias_dict.get(key, key) for key in keys] + with open(dump_filename, "w", encoding=encoding) as f: + # write meta information of tec + f.write(f'title = "{dump_filename}"\n') + header = ", ".join([f'"{key}"' for key in keys]) + f.write(f"variables = {header}\n") + + # NOTE: Tecplot is column-major, so we need to specify I = ny, J = nx, + # which is in contrast to our habits. + f.write(f"Zone I = {ny}, J = {nx}, F = POINT\n") + + # write points data into file + data_cur_time_step = [ + data_dict[key][t * nxy : (t + 1) * nxy] for key in fetch_keys + ] + + for items in zip(*data_cur_time_step): + f.write(delimiter.join([str(float(x)) for x in items]) + "\n") + + if num_timestamps > 1: + logger.message( + f"tecplot files are saved to: {filename}_t-0.dat ~ {filename}_t-{num_timestamps - 1}.dat" + ) + else: + logger.message(f"tecplot file is saved to: {filename}.dat") diff --git a/examples/smc_reac/ppsci/validate/__init__.py b/examples/smc_reac/ppsci/validate/__init__.py new file mode 100644 index 0000000000..3bc1c9ae4d --- /dev/null +++ b/examples/smc_reac/ppsci/validate/__init__.py @@ -0,0 +1,81 @@ +# 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. + +import copy + +from ppsci.loss import build_loss +from ppsci.metric import build_metric +from ppsci.utils import logger +from ppsci.utils import misc +from ppsci.validate.base import Validator +from ppsci.validate.geo_validator import GeometryValidator +from ppsci.validate.sup_validator import SupervisedValidator + +__all__ = [ + "Validator", + "GeometryValidator", + "SupervisedValidator", +] + + +def build_validator(cfg, equation_dict, geom_dict): + """Build validator(s). + + Args: + cfg (List[DictConfig]): Validator(s) config list. + geom_dict (Dct[str, Geometry]): Geometry(ies) in dict. + equation_dict (Dct[str, Equation]): Equation(s) in dict. + + Returns: + Dict[str, Validator]: Validator(s) in dict. + """ + if cfg is None: + return None + cfg = copy.deepcopy(cfg) + global_dataloader_cfg = cfg["dataloader"] + validator_cfg = cfg["content"] + + validator_dict = misc.PrettyOrderedDict() + for _item in validator_cfg: + validator_cls = next(iter(_item.keys())) + _validator_cfg = _item[validator_cls] + validator_name = _validator_cfg.get("name", validator_cls) + # select geometry + geom_name = _validator_cfg.pop("geom") + _validator_cfg["geom"] = geom_dict[geom_name] + + # update complete dataloader config + local_dataloader_cfg = _validator_cfg["dataloader"] + local_dataloader_cfg.update(global_dataloader_cfg) + + # select equation + for name, expr in _validator_cfg["output_expr"].items(): + if isinstance(expr, str) and expr in equation_dict: + _validator_cfg["output_expr"][name] = equation_dict[expr].equations[ + name + ] + + # build loss + _validator_cfg["loss"] = build_loss(_validator_cfg["loss"]) + + # build metric + _validator_cfg["metric"] = build_metric(_validator_cfg["metric"]) + + # instantiate validator + _validator_cfg["dataloader_cfg"] = _validator_cfg.pop("dataloader") + validator_dict[validator_name] = eval(validator_cls)(**_validator_cfg) + + logger.debug(str(validator_dict[validator_name])) + + return validator_dict diff --git a/examples/smc_reac/ppsci/validate/base.py b/examples/smc_reac/ppsci/validate/base.py new file mode 100644 index 0000000000..84760f25d7 --- /dev/null +++ b/examples/smc_reac/ppsci/validate/base.py @@ -0,0 +1,69 @@ +# 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. + +from __future__ import annotations + +from typing import TYPE_CHECKING +from typing import Any +from typing import Dict +from typing import Optional + +from paddle import io + +from ppsci import data + +if TYPE_CHECKING: + from ppsci import loss + from ppsci import metric + + +class Validator: + """Base class for validators. + + Args: + dataset (io.Dataset): Dataset for validator. + dataloader_cfg (Dict[str, Any]): Dataloader config. + loss (loss.Loss): Loss functor. + metric (Optional[Dict[str, metric.Metric]]): Named metric functors in dict. + name (str): Name of validator. + """ + + def __init__( + self, + dataset: io.Dataset, + dataloader_cfg: Dict[str, Any], + loss: "loss.Loss", + metric: Optional[Dict[str, "metric.Metric"]], + name: str, + ): + self.data_loader = data.build_dataloader(dataset, dataloader_cfg) + self.data_iter = iter(self.data_loader) + self.loss = loss + self.metric = metric + self.name = name + + def __str__(self): + return ", ".join( + [ + self.__class__.__name__, + f"name = {self.name}", + f"input_keys = {self.input_keys}", + f"output_keys = {self.output_keys}", + f"output_expr = {self.output_expr}", + f"label_dict = {self.label_dict}", + f"len(dataloader) = {len(self.data_loader)}", + f"loss = {self.loss}", + f"metric = {list(self.metric.keys())}", + ] + ) diff --git a/examples/smc_reac/ppsci/validate/geo_validator.py b/examples/smc_reac/ppsci/validate/geo_validator.py new file mode 100644 index 0000000000..08f2e663a4 --- /dev/null +++ b/examples/smc_reac/ppsci/validate/geo_validator.py @@ -0,0 +1,161 @@ +# 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. + +from __future__ import annotations + +from typing import Any +from typing import Callable +from typing import Dict +from typing import Optional +from typing import Union + +import numpy as np +import paddle +import sympy +from typing_extensions import Literal + +from ppsci import geometry +from ppsci import loss +from ppsci import metric +from ppsci.data import dataset +from ppsci.validate import base + + +class GeometryValidator(base.Validator): + """Validator for geometry. + + Args: + output_expr (Dict[str, Callable]): Function in dict for computing output. + e.g. {"u_mul_v": lambda out: out["u"] * out["v"]} means the model output u + will be multiplied by model output v and the result will be named "u_mul_v". + label_dict (Dict[str, Union[float, Callable]]): Function in dict for computing + label, which will be a reference value to participate in the loss calculation. + geom (geometry.Geometry): Geometry where data sampled from. + dataloader_cfg (Dict[str, Any]): Dataloader config. + loss (loss.Loss): Loss functor. + random (Literal["pseudo", "Halton", "LHS"], optional): Random method for sampling data in + geometry. Defaults to "pseudo". + criteria (Optional[Callable]): Criteria for refining specified domain. Defaults to None. + evenly (bool, optional): Whether to use evenly distribution sampling. Defaults to False. + metric (Optional[Dict[str, metric.Metric]]): Named metric functors in dict. Defaults to None. + with_initial (bool, optional): Whether the data contains time t0. Defaults to False. + name (Optional[str]): Name of validator. Defaults to None. + + Examples: + >>> import ppsci + >>> rect = ppsci.geometry.Rectangle((0, 0), (1, 1)) + >>> geom_validator = ppsci.validate.GeometryValidator( + ... {"u": lambda out: out["u"]}, + ... {"u": 0}, + ... rect, + ... { + ... "dataset": "IterableNamedArrayDataset", + ... "iters_per_epoch": 1, + ... "total_size": 32, + ... "batch_size": 16, + ... }, + ... ppsci.loss.MSELoss("mean"), + ... ) + """ + + def __init__( + self, + output_expr: Dict[str, Callable], + label_dict: Dict[str, Union[float, Callable]], + geom: geometry.Geometry, + dataloader_cfg: Dict[str, Any], + loss: loss.Loss, + random: Literal["pseudo", "Halton", "LHS"] = "pseudo", + criteria: Optional[Callable] = None, + evenly: bool = False, + metric: Optional[Dict[str, metric.Metric]] = None, + with_initial: bool = False, + name: Optional[str] = None, + ): + self.output_expr = output_expr + self.label_dict = label_dict + self.input_keys = geom.dim_keys + self.output_keys = tuple(label_dict.keys()) + + nx = dataloader_cfg["total_size"] + self.num_timestamps = 1 + # TODO(sensen): Simplify code below + if isinstance(geom, geometry.TimeXGeometry): + if geom.timedomain.num_timestamps is not None: + if with_initial: + # include t0 + self.num_timestamps = geom.timedomain.num_timestamps + assert ( + nx % self.num_timestamps == 0 + ), f"{nx} % {self.num_timestamps} != 0" + nx //= self.num_timestamps + input = geom.sample_interior( + nx * (geom.timedomain.num_timestamps - 1), + random, + criteria, + evenly, + ) + initial = geom.sample_initial_interior(nx, random, criteria, evenly) + input = { + key: np.vstack((initial[key], input[key])) for key in input + } + else: + # exclude t0 + self.num_timestamps = geom.timedomain.num_timestamps - 1 + assert ( + nx % self.num_timestamps == 0 + ), f"{nx} % {self.num_timestamps} != 0" + nx //= self.num_timestamps + input = geom.sample_interior( + nx * (geom.timedomain.num_timestamps - 1), + random, + criteria, + evenly, + ) + else: + raise NotImplementedError( + "TimeXGeometry with random timestamp not implemented yet." + ) + else: + input = geom.sample_interior(nx, random, criteria, evenly) + + label = {} + for key, value in label_dict.items(): + if isinstance(value, (int, float)): + label[key] = np.full_like(next(iter(input.values())), value) + elif isinstance(value, sympy.Basic): + func = sympy.lambdify( + sympy.symbols(geom.dim_keys), + value, + [{"amax": lambda xy, _: np.maximum(xy[0], xy[1])}, "numpy"], + ) + label[key] = func( + **{k: v for k, v in input.items() if k in geom.dim_keys} + ) + elif callable(value): + func = value + label[key] = func(input) + if isinstance(label[key], (int, float)): + label[key] = np.full( + (next(iter(input.values())).shape[0], 1), + label[key], + paddle.get_default_dtype(), + ) + else: + raise NotImplementedError(f"type of {type(value)} is invalid yet.") + + weight = {key: np.ones_like(next(iter(label.values()))) for key in label} + + _dataset = getattr(dataset, dataloader_cfg["dataset"])(input, label, weight) + super().__init__(_dataset, dataloader_cfg, loss, metric, name) diff --git a/examples/smc_reac/ppsci/validate/sup_validator.py b/examples/smc_reac/ppsci/validate/sup_validator.py new file mode 100644 index 0000000000..a88a02af5e --- /dev/null +++ b/examples/smc_reac/ppsci/validate/sup_validator.py @@ -0,0 +1,103 @@ +# 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. + +from __future__ import annotations + +from typing import Any +from typing import Callable +from typing import Dict +from typing import Optional + +from ppsci import loss +from ppsci import metric +from ppsci.data import dataset +from ppsci.validate import base + + +class SupervisedValidator(base.Validator): + """Validator for supervised models. + + Args: + dataloader_cfg (Dict[str, Any]): Config of building a dataloader. + loss (loss.Loss): Loss functor. + output_expr (Optional[Dict[str, Callable]]): List of label expression. + metric (Optional[Dict[str, metric.Metric]]): Named metric functors in dict. Defaults to None. + name (Optional[str]): Name of validator. Defaults to None. + + Examples: + >>> import ppsci + >>> valid_dataloader_cfg = { + ... "dataset": { + ... "name": "MatDataset", + ... "file_path": "/path/to/file.mat", + ... "input_keys": ("t_f",), + ... "label_keys": ("eta", "f"), + ... }, + ... "batch_size": 32, + ... "sampler": { + ... "name": "BatchSampler", + ... "drop_last": False, + ... "shuffle": False, + ... }, + ... } # doctest: +SKIP + >>> eta_mse_validator = ppsci.validate.SupervisedValidator( + ... valid_dataloader_cfg, + ... ppsci.loss.MSELoss("mean"), + ... {"eta": lambda out: out["eta"]}, + ... metric={"MSE": ppsci.metric.MSE()}, + ... name="eta_mse", + ... ) # doctest: +SKIP + """ + + def __init__( + self, + dataloader_cfg: Dict[str, Any], + loss: loss.Loss, + output_expr: Optional[Dict[str, Callable]] = None, + metric: Optional[Dict[str, metric.Metric]] = None, + name: Optional[str] = None, + ): + self.output_expr = output_expr + + # build dataset + _dataset = dataset.build_dataset(dataloader_cfg["dataset"]) + + self.input_keys = _dataset.input_keys + self.output_keys = ( + tuple(output_expr.keys()) + if output_expr is not None + else _dataset.label_keys + ) + + if self.output_expr is None: + self.output_expr = { + key: lambda out, k=key: out[k] for key in self.output_keys + } + + # construct dataloader with dataset and dataloader_cfg + super().__init__(_dataset, dataloader_cfg, loss, metric, name) + + def __str__(self): + return ", ".join( + [ + self.__class__.__name__, + f"name = {self.name}", + f"input_keys = {self.input_keys}", + f"output_keys = {self.output_keys}", + f"output_expr = {self.output_expr}", + f"len(dataloader) = {len(self.data_loader)}", + f"loss = {self.loss}", + f"metric = {list(self.metric.keys())}", + ] + ) diff --git a/examples/smc_reac/ppsci/visualize/__init__.py b/examples/smc_reac/ppsci/visualize/__init__.py new file mode 100644 index 0000000000..e6a5d49186 --- /dev/null +++ b/examples/smc_reac/ppsci/visualize/__init__.py @@ -0,0 +1,82 @@ +# 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. + +from __future__ import annotations + +import copy + +from ppsci.visualize.vtu import save_vtu_to_mesh + +from ppsci.visualize.base import Visualizer # isort:skip +from ppsci.visualize.visualizer import VisualizerScatter1D # isort:skip +from ppsci.visualize.visualizer import VisualizerScatter3D # isort:skip +from ppsci.visualize.visualizer import VisualizerVtu # isort:skip +from ppsci.visualize.visualizer import Visualizer2D # isort:skip +from ppsci.visualize.visualizer import Visualizer2DPlot # isort:skip +from ppsci.visualize.visualizer import Visualizer3D # isort:skip +from ppsci.visualize.visualizer import VisualizerWeather # isort:skip +from ppsci.visualize.radar import VisualizerRadar # isort:skip +from ppsci.visualize.vtu import save_vtu_from_dict # isort:skip +from ppsci.visualize.vtu import save_vtp_from_dict # isort:skip +from ppsci.visualize.plot import save_plot_from_1d_dict # isort:skip +from ppsci.visualize.plot import save_plot_from_3d_dict # isort:skip +from ppsci.visualize.plot import save_plot_weather_from_dict # isort:skip + + +__all__ = [ + "Visualizer", + "VisualizerScatter1D", + "VisualizerScatter3D", + "VisualizerVtu", + "Visualizer2D", + "Visualizer2DPlot", + "Visualizer3D", + "VisualizerWeather", + "VisualizerRadar", + "save_vtu_from_dict", + "save_vtp_from_dict", + "save_vtu_to_mesh", + "save_plot_from_1d_dict", + "save_plot_from_3d_dict", + "save_plot_weather_from_dict", +] + + +def build_visualizer(cfg): + """Build visualizer(s). + + Args: + cfg (List[DictConfig]): Visualizer(s) config list. + geom_dict (Dct[str, Geometry]): Geometry(ies) in dict. + equation_dict (Dct[str, Equation]): Equation(s) in dict. + + Returns: + Dict[str, Visualizer]: Visualizer(s) in dict. + """ + if cfg is None: + return None + cfg = copy.deepcopy(cfg) + + visualizer_dict = {} + for _item in cfg: + visualizer_cls = next(iter(_item.keys())) + visualizer_cfg = _item[visualizer_cls] + visualizer = eval(visualizer_cls)(**visualizer_cfg) + + visualizer_name = visualizer_cfg.get("name", visualizer_cls) + if visualizer_name in visualizer_dict: + raise ValueError(f"Name of visualizer({visualizer_name}) should be unique") + visualizer_dict[visualizer_name] = visualizer + + return visualizer_dict diff --git a/examples/smc_reac/ppsci/visualize/base.py b/examples/smc_reac/ppsci/visualize/base.py new file mode 100644 index 0000000000..b249efcabc --- /dev/null +++ b/examples/smc_reac/ppsci/visualize/base.py @@ -0,0 +1,65 @@ +# 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. + +from __future__ import annotations + +import abc +from typing import Callable +from typing import Dict + +import numpy as np + + +class Visualizer: + """Base class for visualizer. + + Args: + input_dict (Dict[str, np.ndarray]): Input dict. + output_expr (Dict[str, Callable]): Output expression. + batch_size (int): Batch size of data when computing result in visu.py. + num_timestamps (int): Number of timestamps. + prefix (str): Prefix for output file. + """ + + def __init__( + self, + input_dict: Dict[str, np.ndarray], + output_expr: Dict[str, Callable], + batch_size: int, + num_timestamps: int, + prefix: str, + ): + self.input_dict = input_dict + self.input_keys = tuple(input_dict.keys()) + self.output_expr = output_expr + self.output_keys = tuple(output_expr.keys()) + self.batch_size = batch_size + self.num_timestamps = num_timestamps + self.prefix = prefix + + @abc.abstractmethod + def save(self, data_dict): + """Visualize result from data_dict and save as files""" + + def __str__(self): + return ", ".join( + [ + f"input_keys: {self.input_keys}", + f"output_keys: {self.output_keys}", + f"output_expr: {self.output_expr}", + f"batch_size: {self.batch_size}", + f"num_timestamps: {self.num_timestamps}", + f"output file prefix: {self.prefix}", + ] + ) diff --git a/examples/smc_reac/ppsci/visualize/plot.py b/examples/smc_reac/ppsci/visualize/plot.py new file mode 100644 index 0000000000..8d41ac17c6 --- /dev/null +++ b/examples/smc_reac/ppsci/visualize/plot.py @@ -0,0 +1,580 @@ +# 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. + +from __future__ import annotations + +import os +from typing import Dict +from typing import Optional +from typing import Tuple +from typing import Union + +import imageio +import matplotlib +import numpy as np +import paddle +from matplotlib import pyplot as plt +from matplotlib.legend_handler import HandlerBase +from matplotlib.patches import Rectangle +from mpl_toolkits.mplot3d.art3d import Line3DCollection + +from ppsci.utils import logger + +cnames = [ + "bisque", + "black", + "blanchedalmond", + "blue", + "blueviolet", + "brown", + "burlywood", + "cadetblue", + "chartreuse", + "orangered", + "orchid", + "palegoldenrod", + "palegreen", +] + +CMAPS = [ + "Reds", + "Blues", + "Greys", + "Purples", + "Greens", + "Oranges", + "YlOrBr", + "YlOrRd", + "OrRd", + "PuRd", + "RdPu", + "BuPu", + "GnBu", + "PuBu", + "YlGnBu", + "PuBuGn", + "BuGn", + "YlGn", +] + + +def _save_plot_from_1d_array(filename, coord, value, value_keys, num_timestamps=1): + """Save plot from given 1D data. + + Args: + filename (str): Filename. + coord (np.ndarray): Coordinate array. + value (Dict[str, np.ndarray]): Dict of value array. + value_keys (Tuple[str, ...]): Value keys. + num_timestamps (int, optional): Number of timestamps coord/value contains. Defaults to 1. + """ + if len(os.path.dirname(filename)): + os.makedirs(os.path.dirname(filename), exist_ok=True) + fig, a = plt.subplots(len(value_keys), num_timestamps, squeeze=False) + fig.subplots_adjust(hspace=0.8) + + len_ts = len(coord) // num_timestamps + for t in range(num_timestamps): + st = t * len_ts + ed = (t + 1) * len_ts + coord_t = coord[st:ed] + + for i, key in enumerate(value_keys): + _value_t: np.ndarray = value[st:ed, i] + a[i][t].scatter( + coord_t, + _value_t, + color=cnames[i], + label=key, + s=2, + ) + if num_timestamps > 1: + a[i][t].set_title(f"{key}(t={t})") + else: + a[i][t].set_title(f"{key}") + a[i][t].grid(color="#c2ccd0", linestyle="--", linewidth=0.5) + a[i][t].legend() + + if num_timestamps == 1: + fig.savefig(filename, dpi=300) + else: + fig.savefig(f"{filename}_{t}", dpi=300) + + if num_timestamps == 1: + logger.message(f"1D result is saved to: {filename}.png") + else: + logger.message( + f"1D result is saved to: {filename}_0.png" + f" ~ {filename}_{num_timestamps - 1}.png" + ) + + +def save_plot_from_1d_dict( + filename, data_dict, coord_keys, value_keys, num_timestamps=1 +): + """Plot dict data as file. + + Args: + filename (str): Output filename. + data_dict (Dict[str, Union[np.ndarray, paddle.Tensor]]): Data in dict. + coord_keys (Tuple[str, ...]): Tuple of coord key. such as ("x", "y"). + value_keys (Tuple[str, ...]): Tuple of value key. such as ("u", "v"). + num_timestamps (int, optional): Number of timestamp in data_dict. Defaults to 1. + + Examples: + >>> import ppsci + >>> import numpy as np + >>> filename = "path/to/file" + >>> data_dict = { + ... "x": np.array([[1], [2], [3],[4]]), + ... "u": np.array([[4], [5], [6],[4]]), + ... } + >>> coord_keys = ("x",) + >>> value_keys = ("u",) + >>> ppsci.visualize.save_plot_from_1d_dict(filename, data_dict, coord_keys, value_keys) # doctest: +SKIP + """ + space_ndim = len(coord_keys) - int("t" in coord_keys) + if space_ndim not in [1, 2, 3]: + raise ValueError(f"ndim of space coord ({space_ndim}) should be 1, 2 or 3") + + coord = [data_dict[k] for k in coord_keys if k != "t"] + value = [data_dict[k] for k in value_keys] if value_keys else None + + if isinstance(coord[0], paddle.Tensor): + coord = [x.numpy() for x in coord] + else: + coord = [x for x in coord] + coord = np.concatenate(coord, axis=1) + + if value is not None: + if isinstance(value[0], paddle.Tensor): + value = [x.numpy() for x in value] + else: + value = [x for x in value] + value = np.concatenate(value, axis=1) + + _save_plot_from_1d_array(filename, coord, value, value_keys, num_timestamps) + + +def _save_plot_from_2d_array( + filename: str, + visu_data: Tuple[np.ndarray, ...], + visu_keys: Tuple[str, ...], + num_timestamps: int = 1, + stride: int = 1, + xticks: Optional[Tuple[float, ...]] = None, + yticks: Optional[Tuple[float, ...]] = None, +): + """Save plot from given 2D data. + + Args: + filename (str): Filename. + visu_data (Tuple[np.ndarray, ...]): Data that requires visualization. + visu_keys (Tuple[str, ...]): Keys for visualizing data. such as ("u", "v"). + num_timestamps (int, optional): Number of timestamps coord/value contains. Defaults to 1. + stride (int, optional): The time stride of visualization. Defaults to 1. + xticks (Optional[Tuple[float, ...]]): Tuple of xtick locations. Defaults to None. + yticks (Optional[Tuple[float, ...]]): Tuple of ytick locations. Defaults to None. + """ + if len(os.path.dirname(filename)): + os.makedirs(os.path.dirname(filename), exist_ok=True) + + plt.close("all") + matplotlib.rcParams["xtick.labelsize"] = 5 + matplotlib.rcParams["ytick.labelsize"] = 5 + + fig, ax = plt.subplots( + len(visu_keys), + num_timestamps, + squeeze=False, + sharey=True, + figsize=(num_timestamps, len(visu_keys)), + ) + fig.subplots_adjust(hspace=0.3) + target_flag = any("target" in key for key in visu_keys) + for i, data in enumerate(visu_data): + if target_flag is False or "target" in visu_keys[i]: + c_max = np.amax(data) + c_min = np.amin(data) + + for t_idx in range(num_timestamps): + t = t_idx * stride + ax[i, t_idx].imshow( + data[t, :, :], + extent=[xticks.min(), xticks.max(), yticks.min(), yticks.max()], + cmap="inferno", + origin="lower", + vmax=c_max, + vmin=c_min, + ) + if xticks is not None: + ax[i, t_idx].set_xticks(xticks) + if yticks is not None: + ax[i, t_idx].set_yticks(yticks) + + ax[i, t_idx].set_title(f"t={t}", fontsize=8) + if t_idx == 0: + ax[i, 0].set_ylabel(visu_keys[i], fontsize=8) + + p0 = ax[i, -1].get_position().get_points().flatten() + ax_cbar = fig.add_axes([p0[2] + 0.005, p0[1], 0.0075, p0[3] - p0[1]]) + ticks = np.linspace(0, 1, 5) + tickLabels = np.linspace(c_min, c_max, 5) + tickLabels = [f"{t0:02.2f}" for t0 in tickLabels] + cbar = matplotlib.colorbar.ColorbarBase( + ax_cbar, cmap=plt.get_cmap("inferno"), orientation="vertical", ticks=ticks + ) + cbar.set_ticklabels(tickLabels, fontsize=5) + plt.savefig(f"{filename}", dpi=300) + + +def save_plot_from_2d_dict( + filename: str, + data_dict: Dict[str, Union[np.ndarray, paddle.Tensor]], + visu_keys: Tuple[str, ...], + num_timestamps: int = 1, + stride: int = 1, + xticks: Optional[Tuple[float, ...]] = None, + yticks: Optional[Tuple[float, ...]] = None, +): + """Plot 2d dict data as file. + + Args: + filename (str): Output filename. + data_dict (Dict[str, Union[np.ndarray, paddle.Tensor]]): Data in dict. + visu_keys (Tuple[str, ...]): Keys for visualizing data. such as ("u", "v"). + num_timestamps (int, optional): Number of timestamp in data_dict. Defaults to 1. + stride (int, optional): The time stride of visualization. Defaults to 1. + xticks (Optional[Tuple[float,...]]): The list of xtick locations. Defaults to None. + yticks (Optional[Tuple[float,...]]): The list of ytick locations. Defaults to None. + """ + visu_data = [data_dict[k] for k in visu_keys] + if isinstance(visu_data[0], paddle.Tensor): + visu_data = [x.numpy() for x in visu_data] + _save_plot_from_2d_array( + filename, visu_data, visu_keys, num_timestamps, stride, xticks, yticks + ) + + +# Interface to LineCollection: +def _colorline3d( + x, y, z, t=None, cmap=plt.get_cmap("viridis"), linewidth=1, alpha=1.0, ax=None +): + """ + Plot a colored line with coordinates x and y + Optionally specify colors in the array z + Optionally specify a colormap, a norm function and a line width + https://stackoverflow.com/questions/52884221/how-to-plot-a-matplotlib-line-plot-using-colormap + """ + # Default colors equally spaced on [0, 1]: + if t is None: + t = np.linspace(0.25, 1.0, len(x)) + if ax is None: + ax = plt.gca() + + points = np.array([x, y, z]).T.reshape(-1, 1, 3) + segments = np.concatenate([points[:-1], points[1:]], axis=1) + + colors = np.array([cmap(i) for i in t]) + lc = Line3DCollection(segments, colors=colors, linewidth=linewidth, alpha=alpha) + ax.add_collection(lc) + ax.scatter(x, y, z, c=colors, marker="*", alpha=alpha) # Adding line markers + + +class HandlerColormap(HandlerBase): + """Class for creating colormap legend rectangles. + + Args: + cmap (matplotlib.cm): Matplotlib colormap. + num_stripes (int, optional): Number of contour levels (strips) in rectangle. Defaults to 8. + """ + + def __init__(self, cmap: matplotlib.cm, num_stripes: int = 8, **kw): + HandlerBase.__init__(self, **kw) + self.cmap = cmap + self.num_stripes = num_stripes + + def create_artists( + self, legend, orig_handle, xdescent, ydescent, width, height, fontsize, trans + ): + stripes = [] + for i in range(self.num_stripes): + s = Rectangle( + [xdescent + i * width / self.num_stripes, ydescent], + width / self.num_stripes, + height, + fc=self.cmap((2 * i + 1) / (2 * self.num_stripes)), + transform=trans, + ) + stripes.append(s) + return stripes + + +def _save_plot_from_3d_array( + filename: str, + visu_data: Tuple[np.ndarray, ...], + visu_keys: Tuple[str, ...], + num_timestamps: int = 1, +): + """Save plot from given 3D data. + + Args: + filename (str): Filename. + visu_data (Tuple[np.ndarray, ...]): Data that requires visualization. + visu_keys (Tuple[str, ...]): Keys for visualizing data. such as ("u", "v"). + num_timestamps (int, optional): Number of timestamps coord/value contains. Defaults to 1. + """ + if len(os.path.dirname(filename)): + os.makedirs(os.path.dirname(filename), exist_ok=True) + + fig = plt.figure(figsize=(10, 10)) + len_ts = len(visu_data[0]) // num_timestamps + for t in range(num_timestamps): + ax = fig.add_subplot(1, num_timestamps, t + 1, projection="3d") + st = t * len_ts + ed = (t + 1) * len_ts + visu_data_t = [data[st:ed] for data in visu_data] + cmaps = [] + for i, data in enumerate(visu_data_t): + cmap = plt.get_cmap(CMAPS[i % len(CMAPS)]) + _colorline3d(data[:, 0], data[:, 1], data[:, 2], cmap=cmap, ax=ax) + cmaps.append(cmap) + cmap_handles = [Rectangle((0, 0), 1, 1) for _ in visu_keys] + handler_map = dict( + zip(cmap_handles, [HandlerColormap(cm, num_stripes=8) for cm in cmaps]) + ) + # Create custom legend with color map rectangles + ax.legend( + handles=cmap_handles, + labels=visu_keys, + handler_map=handler_map, + loc="upper right", + framealpha=0.95, + ) + if num_timestamps == 1: + fig.savefig(filename, dpi=300) + else: + fig.savefig(f"{filename}_{t}", dpi=300) + + if num_timestamps == 1: + logger.message(f"3D result is saved to: {filename}.png") + else: + logger.message( + f"3D result is saved to: {filename}_0.png" + f" ~ {filename}_{num_timestamps - 1}.png" + ) + + +def save_plot_from_3d_dict( + filename: str, + data_dict: Dict[str, Union[np.ndarray, paddle.Tensor]], + visu_keys: Tuple[str, ...], + num_timestamps: int = 1, +): + """Plot dict data as file. + + Args: + filename (str): Output filename. + data_dict (Dict[str, Union[np.ndarray, paddle.Tensor]]): Data in dict. + visu_keys (Tuple[str, ...]): Keys for visualizing data. such as ("u", "v"). + num_timestamps (int, optional): Number of timestamp in data_dict. Defaults to 1. + + Examples: + >>> import numpy as np + >>> import ppsci + + >>> data_dict = { + ... "u": np.array([[[10], [20], [30], [40], [50]]]), + ... "v": np.array([[[5], [15], [25], [35], [45]]]), + ... } + + >>> ppsci.visualize.save_plot_from_3d_dict( + ... "path/to/file", + ... data_dict, + ... ("u", "v"), + ... 1, + ... ) # doctest: +SKIP + """ + visu_data = [data_dict[k] for k in visu_keys] + if isinstance(visu_data[0], paddle.Tensor): + visu_data = [x.numpy() for x in visu_data] + + _save_plot_from_3d_array(filename, visu_data, visu_keys, num_timestamps) + + +def _save_plot_weather_from_array( + filename: str, + pred: np.ndarray, + target: np.ndarray, + pred_key: str, + target_key: str, + xticks: Tuple[float, ...], + xticklabels: Tuple[str, ...], + yticks: Tuple[float, ...], + yticklabels: Tuple[str, ...], + vmin: float, + vmax: float, + colorbar_label: str = "", + log_norm: bool = False, +): + """Plot weather result as file from array data. + + Args: + filename (str): Output file name. + pred (np.ndarray): The predict data. + target (np.ndarray): The target data. + pred_key (str): The key of predict data. + target_key (str): The key of target data. + xticks (Tuple[float, ...]): The list of xtick locations. + xticklabels (Tuple[str, ...]): The x-axis' tick labels. + yticks (Tuple[float, ...]): The list of ytick locations. + yticklabels (Tuple[str, ...]): The y-axis' tick labels. + vmin (float): Minimum value that the colormap covers. + vmax (float): Maximal value that the colormap covers. + colorbar_label (str, optional): The color-bar label. Defaults to "". + log_norm (bool, optional): Whether use log norm. Defaults to False. + """ + + def plot_weather( + ax, + data, + title_text, + xticks, + xticklabels, + yticks, + yticklabels, + vmin, + vmax, + log_norm, + cmap=plt.get_cmap("turbo", 1000), + ): + ax.title.set_text(title_text) + ax.set_yticks(yticks) + ax.set_yticklabels(yticklabels) + ax.set_xticks(xticks) + ax.set_xticklabels(xticklabels) + if not log_norm: + map_ = ax.imshow( + data, + interpolation="nearest", + cmap=cmap, + aspect="auto", + vmin=vmin, + vmax=vmax, + ) + else: + norm = matplotlib.colors.LogNorm(vmin=vmin, vmax=vmax, clip=True) + map_ = ax.imshow( + data, interpolation="nearest", cmap=cmap, aspect="auto", norm=norm + ) + plt.colorbar(mappable=map_, cax=None, ax=None, shrink=0.5, label=colorbar_label) + + if len(os.path.dirname(filename)): + os.makedirs(os.path.dirname(filename), exist_ok=True) + fig = plt.figure(facecolor="w", figsize=(7, 7)) + ax = fig.add_subplot(2, 1, 1) + plot_weather( + ax, + pred, + pred_key, + xticks, + xticklabels, + yticks, + yticklabels, + vmin, + vmax, + log_norm, + ) + bx = fig.add_subplot(2, 1, 2) + plot_weather( + bx, + target, + target_key, + xticks, + xticklabels, + yticks, + yticklabels, + vmin, + vmax, + log_norm, + ) + fig.savefig(filename, dpi=300) + plt.close() + + +def save_plot_weather_from_dict( + foldername: str, + data_dict: Dict[str, Union[np.ndarray, paddle.Tensor]], + visu_keys: Tuple[str, ...], + xticks: Tuple[float, ...], + xticklabels: Tuple[str, ...], + yticks: Tuple[float, ...], + yticklabels: Tuple[str, ...], + vmin: float, + vmax: float, + colorbar_label: str = "", + log_norm: bool = False, + num_timestamps: int = 1, +): + """Plot weather result as file from dict data. + + Args: + foldername (str): Output folder name. + data_dict (Dict[str, Union[np.ndarray, paddle.Tensor]]): Data in dict. + visu_keys (Tuple[str, ...]): Keys for visualizing data. such as ("output_6h", "target_6h"). + xticks (Tuple[float, ...]): The list of xtick locations. + xticklabels (Tuple[str, ...]): The x-axis' tick labels. + yticks (Tuple[float, ...]): The list of ytick locations, + yticklabels (Tuple[str, ...]): The y-axis' tick labels. + vmin (float): Minimum value that the colormap covers. + vmax (float): Maximal value that the colormap covers. + colorbar_label (str, optional): The colorbar label. Defaults to "". + log_norm (bool, optional): Whether use log norm. Defaults to False. + num_timestamps (int): Number of timestamp in data_dict. Defaults to 1. + """ + os.makedirs(foldername, exist_ok=True) + + visu_data = [data_dict[k] for k in visu_keys] + if isinstance(visu_data[0], paddle.Tensor): + visu_data = [x.numpy() for x in visu_data] + + frames = [] + for t in range(num_timestamps): + pred_key, target_key = visu_keys[2 * t], visu_keys[2 * t + 1] + pred_data = visu_data[2 * t] + target_data = visu_data[2 * t + 1] + filename_t = os.path.join(foldername, f"{t}.png") + _save_plot_weather_from_array( + filename_t, + pred_data, + target_data, + pred_key, + target_key, + xticks, + xticklabels, + yticks, + yticklabels, + vmin=vmin, + vmax=vmax, + colorbar_label=colorbar_label, + log_norm=log_norm, + ) + frames.append(imageio.imread(filename_t)) + filename = os.path.join(foldername, "result.gif") + imageio.mimsave( + filename, + frames, + "GIF", + duration=1, + ) diff --git a/examples/smc_reac/ppsci/visualize/radar.py b/examples/smc_reac/ppsci/visualize/radar.py new file mode 100644 index 0000000000..abde75b775 --- /dev/null +++ b/examples/smc_reac/ppsci/visualize/radar.py @@ -0,0 +1,124 @@ +# 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. + +from __future__ import annotations + +import os +from typing import Callable +from typing import Dict + +import matplotlib.pyplot as plt +import numpy as np + +from ppsci.visualize import base + + +class VisualizerRadar(base.Visualizer): + """Visualizer for NowcastNet Radar Dataset. + + Args: + input_dict (Dict[str, np.ndarray]): Input dict. + output_expr (Dict[str, Callable]): Output expression. + batch_size (int, optional): Batch size of data when computing result in visu.py. Defaults to 64. + num_timestamps (int, optional): Number of timestamps + prefix (str, optional): Prefix for output file. + case_type (str, optional): Case type. + total_length (str, optional): Total length. + + Examples: + >>> import ppsci + >>> import paddle + >>> frames_tensor = paddle.randn([1, 29, 512, 512, 2]) + >>> visualizer = ppsci.visualize.VisualizerRadar( + ... {"input": frames_tensor}, + ... {"output": lambda out: out["output"]}, + ... num_timestamps=1, + ... prefix="v_nowcastnet", + ... ) + """ + + def __init__( + self, + input_dict: Dict[str, np.ndarray], + output_expr: Dict[str, Callable], + batch_size: int = 64, + num_timestamps: int = 1, + prefix: str = "vtu", + case_type: str = "normal", + total_length: int = 29, + ): + super().__init__(input_dict, output_expr, batch_size, num_timestamps, prefix) + self.case_type = case_type + self.total_length = total_length + self.input_dict = input_dict + + def save(self, path, data_dict): + if not os.path.exists(path): + os.makedirs(path) + test_ims = self.input_dict[list(self.input_dict.keys())[0]] + # keys: {"input", "output"} + img_gen = data_dict[list(data_dict.keys())[1]] + vis_info = {"vmin": 1, "vmax": 40} + if self.case_type == "normal": + test_ims_plot = test_ims[0][ + :-2, 256 - 192 : 256 + 192, 256 - 192 : 256 + 192 + ] + img_gen_plot = img_gen[0][:-2, 256 - 192 : 256 + 192, 256 - 192 : 256 + 192] + else: + test_ims_plot = test_ims[0][:-2] + img_gen_plot = img_gen[0][:-2] + save_plots( + test_ims_plot, + labels=[f"gt{i + 1}" for i in range(self.total_length)], + res_path=path, + vmin=vis_info["vmin"], + vmax=vis_info["vmax"], + ) + save_plots( + img_gen_plot, + labels=[f"pd{i + 1}" for i in range(9, self.total_length)], + res_path=path, + vmin=vis_info["vmin"], + vmax=vis_info["vmax"], + ) + + +def save_plots( + field, + labels, + res_path, + figsize=None, + vmin=0, + vmax=10, + cmap="viridis", + npy=False, + **imshow_args, +): + for i, data in enumerate(field): + if i >= len(labels): + break + plt.figure(figsize=figsize) + ax = plt.axes() + ax.set_axis_off() + alpha = data[..., 0] / 1 + alpha[alpha < 1] = 0 + alpha[alpha > 1] = 1 + ax.imshow( + data[..., 0], alpha=alpha, vmin=vmin, vmax=vmax, cmap=cmap, **imshow_args + ) + plt.savefig(os.path.join(res_path, labels[i] + ".png")) + plt.close() + if npy: + with open(os.path.join(res_path, labels[i] + ".npy"), "wb") as f: + np.save(f, data[..., 0]) diff --git a/examples/smc_reac/ppsci/visualize/visualizer.py b/examples/smc_reac/ppsci/visualize/visualizer.py new file mode 100644 index 0000000000..e3e602daa5 --- /dev/null +++ b/examples/smc_reac/ppsci/visualize/visualizer.py @@ -0,0 +1,409 @@ +# 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. + +from __future__ import annotations + +import os.path as osp +from typing import Callable +from typing import Dict +from typing import Optional +from typing import Tuple + +import numpy as np + +from ppsci.visualize import base +from ppsci.visualize import plot +from ppsci.visualize import vtu + + +class VisualizerScatter1D(base.Visualizer): + """Visualizer for 1d scatter data. + + Args: + input_dict (Dict[str, np.ndarray]): Input dict. + coord_keys (Tuple[str, ...]): Coordinate keys, such as ("x", "y"). + output_expr (Dict[str, Callable]): Output expression. + batch_size (int, optional): Batch size of data when computing result in visu.py. Defaults to 64. + num_timestamps (int, optional): Number of timestamps. Defaults to 1. + prefix (str, optional): Prefix for output file. Defaults to "plot". + + Examples: + >>> import ppsci + >>> visu_mat = {"t_f": np.random.randn(16, 1), "eta": np.random.randn(16, 1)} + >>> visualizer_eta = ppsci.visualize.VisualizerScatter1D( + ... visu_mat, + ... ("t_f",), + ... {"eta": lambda d: d["eta"]}, + ... num_timestamps=1, + ... prefix="viv_pred", + ... ) + """ + + def __init__( + self, + input_dict: Dict[str, np.ndarray], + coord_keys: Tuple[str, ...], + output_expr: Dict[str, Callable], + batch_size: int = 64, + num_timestamps: int = 1, + prefix: str = "plot", + ): + super().__init__(input_dict, output_expr, batch_size, num_timestamps, prefix) + self.coord_keys = coord_keys + + def save(self, filename, data_dict): + plot.save_plot_from_1d_dict( + filename, data_dict, self.coord_keys, self.output_keys, self.num_timestamps + ) + + +class VisualizerScatter3D(base.Visualizer): + """Visualizer for 3d scatter data. + + Args: + input_dict (Dict[str, np.ndarray]): Input dict. + output_expr (Dict[str, Callable]): Output expression. + batch_size (int, optional): Batch size of data when computing result in visu.py. Defaults to 64. + num_timestamps (int, optional): Number of timestamps. Defaults to 1. + prefix (str, optional): Prefix for output file. Defaults to "plot3d_scatter". + + Examples: + >>> import ppsci + >>> vis_data = {"states": np.random.randn(16, 1)} + >>> visualizer = ppsci.visualize.VisualizerScatter3D( + ... vis_data, + ... {"states": lambda d: d["states"]}, + ... num_timestamps=1, + ... prefix="result_states", + ... ) + """ + + def __init__( + self, + input_dict: Dict[str, np.ndarray], + output_expr: Dict[str, Callable], + batch_size: int = 64, + num_timestamps: int = 1, + prefix: str = "plot3d_scatter", + ): + super().__init__(input_dict, output_expr, batch_size, num_timestamps, prefix) + + def save(self, filename, data_dict): + data_dict = { + key: value for key, value in data_dict.items() if key in self.output_keys + } + value = data_dict[self.output_keys[0]] + dim = len(value.shape) + if dim == 3: + # value.shape=(B, T, 3) + for i in range(value.shape[0]): + cur_data_dict = {key: value[i] for key, value in data_dict.items()} + plot.save_plot_from_3d_dict( + filename + str(i), + cur_data_dict, + self.output_keys, + self.num_timestamps, + ) + else: + # value.shape=(T, 3) + plot.save_plot_from_3d_dict( + filename, data_dict, self.output_keys, self.num_timestamps + ) + + +class VisualizerVtu(base.Visualizer): + """Visualizer for 2D points data. + + Args: + input_dict (Dict[str, np.ndarray]): Input dict. + output_expr (Dict[str, Callable]): Output expression. + batch_size (int, optional): Batch size of data when computing result in visu.py. Defaults to 64. + num_timestamps (int, optional): Number of timestamps + prefix (str, optional): Prefix for output file. + + Examples: + >>> import ppsci + >>> vis_points = { + ... "x": np.random.randn(128, 1), + ... "y": np.random.randn(128, 1), + ... "u": np.random.randn(128, 1), + ... "v": np.random.randn(128, 1), + ... } + >>> visualizer_u_v = ppsci.visualize.VisualizerVtu( + ... vis_points, + ... {"u": lambda d: d["u"], "v": lambda d: d["v"]}, + ... num_timestamps=1, + ... prefix="result_u_v", + ... ) + """ + + def __init__( + self, + input_dict: Dict[str, np.ndarray], + output_expr: Dict[str, Callable], + batch_size: int = 64, + num_timestamps: int = 1, + prefix: str = "vtu", + ): + super().__init__(input_dict, output_expr, batch_size, num_timestamps, prefix) + + def save(self, filename, data_dict): + vtu.save_vtu_from_dict( + filename, data_dict, self.input_keys, self.output_keys, self.num_timestamps + ) + + +class Visualizer2D(base.Visualizer): + """Visualizer for 2D data. + + Args: + input_dict (Dict[str, np.ndarray]): Input dict. + output_expr (Dict[str, Callable]): Output expression. + batch_size (int, optional): Batch size of data when computing result in visu.py. Defaults to 64. + num_timestamps (int, optional): Number of timestamps. Defaults to 1. + prefix (str, optional): Prefix for output file. Defaults to "plot2d". + + Examples: + >>> import ppsci + >>> vis_points = { + ... "x": np.random.randn(128, 1), + ... "y": np.random.randn(128, 1), + ... "u": np.random.randn(128, 1), + ... "v": np.random.randn(128, 1), + ... } + >>> visualizer_u_v = ppsci.visualize.Visualizer2D( + ... vis_points, + ... {"u": lambda d: d["u"], "v": lambda d: d["v"]}, + ... num_timestamps=1, + ... prefix="result_u_v", + ... ) + """ + + def __init__( + self, + input_dict: Dict[str, np.ndarray], + output_expr: Dict[str, Callable], + batch_size: int = 64, + num_timestamps: int = 1, + prefix: str = "plot2d", + ): + super().__init__(input_dict, output_expr, batch_size, num_timestamps, prefix) + + +class Visualizer2DPlot(Visualizer2D): + """Visualizer for 2D data use matplotlib. + + Args: + input_dict (Dict[str, np.ndarray]): Input dict. + output_expr (Dict[str, Callable]): Output expression. + batch_size (int, optional): Batch size of data when computing result in visu.py. Defaults to 64. + num_timestamps (int, optional): Number of timestamps. + stride (int, optional): The time stride of visualization. Defaults to 1. + xticks (Optional[Tuple[float,...]]): The list of xtick locations. Defaults to None. + yticks (Optional[Tuple[float,...]]): The list of ytick locations. Defaults to None. + prefix (str, optional): Prefix for output file. Defaults to "plot2d". + + Examples: + >>> import ppsci + >>> vis_data = { + ... "target_ux": np.random.randn(128, 20, 1), + ... "pred_ux": np.random.randn(128, 20, 1), + ... } + >>> visualizer_states = ppsci.visualize.Visualizer2DPlot( + ... vis_data, + ... { + ... "target_ux": lambda d: d["states"][:, :, 0], + ... "pred_ux": lambda d: output_transform(d)[:, :, 0], + ... }, + ... batch_size=1, + ... num_timestamps=10, + ... stride=20, + ... xticks=np.linspace(-2, 14, 9), + ... yticks=np.linspace(-4, 4, 5), + ... prefix="result_states", + ... ) + """ + + def __init__( + self, + input_dict: Dict[str, np.ndarray], + output_expr: Dict[str, Callable], + batch_size: int = 64, + num_timestamps: int = 1, + stride: int = 1, + xticks: Optional[Tuple[float, ...]] = None, + yticks: Optional[Tuple[float, ...]] = None, + prefix: str = "plot2d", + ): + super().__init__(input_dict, output_expr, batch_size, num_timestamps, prefix) + self.stride = stride + self.xticks = xticks + self.yticks = yticks + + def save(self, filename, data_dict): + data_dict = { + key: value for key, value in data_dict.items() if key in self.output_keys + } + value = data_dict[self.output_keys[0]] + dim = len(value.shape) + if dim == 4: + # value.shape=(B, T, H, W) + for i in range(value.shape[0]): + cur_data_dict = {key: value[i] for key, value in data_dict.items()} + plot.save_plot_from_2d_dict( + filename + str(i), + cur_data_dict, + self.output_keys, + self.num_timestamps, + self.stride, + self.xticks, + self.yticks, + ) + else: + # value.shape=(T, H, W) + plot.save_plot_from_2d_dict( + filename, + data_dict, + self.output_keys, + self.num_timestamps, + self.stride, + self.xticks, + self.yticks, + ) + + +class Visualizer3D(base.Visualizer): + """Visualizer for 3D plot data. + + Args: + input_dict (Dict[str, np.ndarray]): Input dict. + output_expr (Dict[str, Callable]): Output expression. + batch_size (int, optional): Batch size of data when computing result in visu.py. Defaults to 64. + label_dict (Dict[str, np.ndarray]): Label dict. + time_list (Optional[Tuple[float, ...]]): Time list. + prefix (str, optional): Prefix for output file. + """ + + def __init__( + self, + input_dict: Dict[str, np.ndarray], + output_expr: Dict[str, Callable], + batch_size: int = 64, + label_dict: Optional[Dict[str, np.ndarray]] = None, + time_list: Optional[Tuple[float, ...]] = None, + prefix: str = "vtu", + ): + self.label = label_dict + self.time_list = time_list + super().__init__(input_dict, output_expr, batch_size, len(time_list), prefix) + + def save(self, filename: str, data_dict: Dict[str, np.ndarray]): + n = int((next(iter(data_dict.values()))).shape[0] / self.num_timestamps) + coord_keys = [x for x in self.input_dict if x != "t"] + for i in range(len(self.time_list)): + vtu.save_vtu_to_mesh( + osp.join(filename, f"predict_{i+1}.vtu"), + {key: (data_dict[key][i * n : (i + 1) * n]) for key in data_dict}, + coord_keys, + self.output_keys, + ) + + +class VisualizerWeather(base.Visualizer): + """Visualizer for weather data use matplotlib. + + Args: + input_dict (Dict[str, np.ndarray]): Input dict. + output_expr (Dict[str, Callable]): Output expression. + xticks (Tuple[float, ...]): The list of xtick locations. + xticklabels (Tuple[str, ...]): The x-axis' tick labels. + yticks (Tuple[float, ...]): The list of ytick locations. + yticklabels (Tuple[str, ...]): The y-axis' tick labels. + vmin (float): Minimum value that the colormap covers. + vmax (float): Maximal value that the colormap covers. + colorbar_label (str, optional): The color-bar label. Defaults to "". + log_norm (bool, optional): Whether use log norm. Defaults to False. + batch_size (int, optional): : Batch size of data when computing result in visu.py. Defaults to 1. + num_timestamps (int, optional): Number of timestamps. Defaults to 1. + prefix (str, optional): Prefix for output file. Defaults to "plot_weather". + + Examples: + >>> import ppsci + >>> import numpy as np + >>> vis_data = { + ... "output_6h": np.random.randn(1, 720, 1440), + ... "target_6h": np.random.randn(1, 720, 1440), + ... } + >>> visualizer_weather = ppsci.visualize.VisualizerWeather( + ... vis_data, + ... { + ... "output_6h": lambda d: d["output_6h"], + ... "target_6h": lambda d: d["target_6h"], + ... }, + ... xticks=np.linspace(0, 1439, 13), + ... xticklabels=[str(i) for i in range(360, -1, -30)], + ... yticks=np.linspace(0, 719, 7), + ... yticklabels=[str(i) for i in range(90, -91, -30)], + ... vmin=0, + ... vmax=25, + ... prefix="result_states", + ... ) + """ + + def __init__( + self, + input_dict: Dict[str, np.ndarray], + output_expr: Dict[str, Callable], + xticks: Tuple[float, ...], + xticklabels: Tuple[str, ...], + yticks: Tuple[float, ...], + yticklabels: Tuple[str, ...], + vmin: float, + vmax: float, + colorbar_label: str = "", + log_norm: bool = False, + batch_size: int = 1, + num_timestamps: int = 1, + prefix: str = "plot_weather", + ): + super().__init__(input_dict, output_expr, batch_size, num_timestamps, prefix) + self.xticks = xticks + self.xticklabels = xticklabels + self.yticks = yticks + self.yticklabels = yticklabels + self.vmin = vmin + self.vmax = vmax + self.colorbar_label = colorbar_label + self.log_norm = log_norm + + def save(self, filename, data_dict): + data_dict = {key: data_dict[key] for key in self.output_keys} + value = data_dict[self.output_keys[0]] + # value.shape=(B, H, W) + for i in range(value.shape[0]): + cur_data_dict = {key: value[i] for key, value in data_dict.items()} + plot.save_plot_weather_from_dict( + filename + str(i), + cur_data_dict, + self.output_keys, + self.xticks, + self.xticklabels, + self.yticks, + self.yticklabels, + self.vmin, + self.vmax, + self.colorbar_label, + self.log_norm, + self.num_timestamps, + ) diff --git a/examples/smc_reac/ppsci/visualize/vtu.py b/examples/smc_reac/ppsci/visualize/vtu.py new file mode 100644 index 0000000000..500c7e2e84 --- /dev/null +++ b/examples/smc_reac/ppsci/visualize/vtu.py @@ -0,0 +1,278 @@ +# 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. + +from __future__ import annotations + +import os +from typing import Dict +from typing import Tuple + +import meshio +import numpy as np +import paddle +from pyevtk import hl + +from ppsci.utils import logger + + +def _save_vtu_from_array(filename, coord, value, value_keys, num_timestamps=1): + """Save data to '*.vtu' file(s). + + Args: + filename (str): Output filename. + coord (np.ndarray): Coordinate points with shape of [N, 2] or [N, 3]. + value (np.ndarray): Value of each coord points with shape of [N, M]. + value_keys (Tuple[str, ...]): Names of each dimension of value, such as ("u", "v"). + num_timestamps (int, optional): Number of timestamp over coord and value. + Defaults to 1. + """ + if not isinstance(coord, np.ndarray): + raise ValueError(f"type of coord({type(coord)}) should be ndarray.") + if value is not None and not isinstance(value, np.ndarray): + raise ValueError(f"type of value({type(value)}) should be ndarray.") + if value is not None and len(coord) != len(value): + raise ValueError( + f"coord length({len(coord)}) should be equal to value length({len(value)})" + ) + if len(coord) % num_timestamps != 0: + raise ValueError( + f"coord length({len(coord)}) should be an integer multiple of " + f"num_timestamps({num_timestamps})" + ) + if coord.shape[1] not in [2, 3]: + raise ValueError(f"ndim of coord({coord.shape[1]}) should be 2 or 3.") + + if len(os.path.dirname(filename)): + os.makedirs(os.path.dirname(filename), exist_ok=True) + + # discard extension name + if filename.endswith(".vtu"): + filename = filename[:-4] + npoint = len(coord) + coord_ndim = coord.shape[1] + + if value is None: + value = np.ones([npoint, 1], dtype=coord.dtype) + value_keys = ["dummy_key"] + + data_ndim = value.shape[1] + nx = npoint // num_timestamps + for t in range(num_timestamps): + # NOTE: each array in data_vtu should be 1-dim, i.e. [N, 1] will occur error. + if coord_ndim == 2: + axis_x = np.ascontiguousarray(coord[t * nx : (t + 1) * nx, 0]) + axis_y = np.ascontiguousarray(coord[t * nx : (t + 1) * nx, 1]) + axis_z = np.zeros([nx], dtype=paddle.get_default_dtype()) + elif coord_ndim == 3: + axis_x = np.ascontiguousarray(coord[t * nx : (t + 1) * nx, 0]) + axis_y = np.ascontiguousarray(coord[t * nx : (t + 1) * nx, 1]) + axis_z = np.ascontiguousarray(coord[t * nx : (t + 1) * nx, 2]) + + data_vtu = {} + for j in range(data_ndim): + data_vtu[value_keys[j]] = np.ascontiguousarray( + value[t * nx : (t + 1) * nx, j] + ) + + if num_timestamps > 1: + width = len(str(num_timestamps - 1)) + hl.pointsToVTK( + f"{filename}_t-{t:0{width}}", axis_x, axis_y, axis_z, data=data_vtu + ) + else: + hl.pointsToVTK(filename, axis_x, axis_y, axis_z, data=data_vtu) + + if num_timestamps > 1: + logger.message( + f"Visualization results are saved to: {filename}_t-{0:0{width}}.vtu ~ " + f"{filename}_t-{num_timestamps - 1:0{width}}.vtu" + ) + else: + logger.message(f"Visualization result is saved to: {filename}.vtu") + + +def save_vtu_from_dict( + filename: str, + data_dict: Dict[str, np.ndarray], + coord_keys: Tuple[str, ...], + value_keys: Tuple[str, ...], + num_timestamps: int = 1, +): + """Save dict data to '*.vtu' file. + + Args: + filename (str): Output filename. + data_dict (Dict[str, np.ndarray]): Data in dict. + coord_keys (Tuple[str, ...]): Tuple of coord key. such as ("x", "y"). + value_keys (Tuple[str, ...]): Tuple of value key. such as ("u", "v"). + num_timestamps (int, optional): Number of timestamp in data_dict. Defaults to 1. + + Examples: + >>> import ppsci + >>> import numpy as np + >>> filename = "path/to/file.vtu" + >>> data_dict = { + ... "x": np.array([[1], [2], [3],[4]]), + ... "y": np.array([[2], [3], [4],[4]]), + ... "z": np.array([[3], [4], [5],[4]]), + ... "u": np.array([[4], [5], [6],[4]]), + ... "v": np.array([[5], [6], [7],[4]]), + ... } + >>> coord_keys = ("x","y","z") + >>> value_keys = ("u","v") + >>> ppsci.visualize.save_vtu_from_dict(filename, data_dict, coord_keys, value_keys) # doctest: +SKIP + """ + if len(coord_keys) not in [2, 3, 4]: + raise ValueError(f"ndim of coord ({len(coord_keys)}) should be 2, 3 or 4") + + coord = [data_dict[k] for k in coord_keys if k not in ("t", "sdf")] + value = [data_dict[k] for k in value_keys] if value_keys else None + + coord = np.concatenate(coord, axis=1) + + if value is not None: + value = np.concatenate(value, axis=1) + + _save_vtu_from_array(filename, coord, value, value_keys, num_timestamps) + + +def save_vtp_from_dict( + filename: str, + data_dict: Dict[str, np.ndarray], + coord_keys: Tuple[str, ...], + value_keys: Tuple[str, ...], + num_timestamps: int = 1, +): + """Save dict data to '*.vtp' file. + + Args: + filename (str): Output filename. + data_dict (Dict[str, np.ndarray]): Data in dict. + coord_keys (Tuple[str, ...]): Tuple of coord key. such as ("x", "y"). + value_keys (Tuple[str, ...]): Tuple of value key. such as ("u", "v"). + num_timestamps (int, optional): Number of timestamp in data_dict. Defaults to 1. + + Examples: + >>> import ppsci + >>> import numpy as np + >>> filename = "path/to/file.vtp" + >>> data_dict = { + ... "x": np.array([[1], [2], [3],[4]]), + ... "y": np.array([[2], [3], [4],[4]]), + ... "z": np.array([[3], [4], [5],[4]]), + ... "u": np.array([[4], [5], [6],[4]]), + ... "v": np.array([[5], [6], [7],[4]]), + ... } + >>> coord_keys = ("x","y","z") + >>> value_keys = ("u","v") + >>> ppsci.visualize.save_vtp_from_dict(filename, data_dict, coord_keys, value_keys) # doctest: +SKIP + """ + import pyvista as pv + + if len(coord_keys) not in [3]: + raise ValueError(f"ndim of coord ({len(coord_keys)}) should be 3 in vtp format") + + coord = [data_dict[k] for k in coord_keys if k not in ("t", "sdf")] + assert all([c.ndim == 2 for c in coord]), "array of each axis should be [*, 1]" + coord = np.concatenate(coord, axis=1) + + if not isinstance(coord, np.ndarray): + raise ValueError(f"type of coord({type(coord)}) should be ndarray.") + if len(coord) % num_timestamps != 0: + raise ValueError( + f"coord length({len(coord)}) should be an integer multiple of " + f"num_timestamps({num_timestamps})" + ) + if coord.shape[1] not in [3]: + raise ValueError(f"ndim of coord({coord.shape[1]}) should be 3 in vtp format.") + + if len(os.path.dirname(filename)): + os.makedirs(os.path.dirname(filename), exist_ok=True) + + npoint = len(coord) + nx = npoint // num_timestamps + if filename.endswith(".vtp"): + filename = filename[:-4] + + for t in range(num_timestamps): + coord_ = coord[t * nx : (t + 1) * nx] + point_cloud = pv.PolyData(coord_) + for k in value_keys: + value_ = data_dict[k][t * nx : (t + 1) * nx] + if value_ is not None and not isinstance(value_, np.ndarray): + raise ValueError(f"type of value({type(value_)}) should be ndarray.") + if value_ is not None and len(coord_) != len(value_): + raise ValueError( + f"coord length({len(coord_)}) should be equal to value length({len(value_)})" + ) + point_cloud[k] = value_ + + if num_timestamps > 1: + width = len(str(num_timestamps - 1)) + point_cloud.save(f"{filename}_t-{t:0{width}}.vtp") + else: + point_cloud.save(f"{filename}.vtp") + + if num_timestamps > 1: + logger.message( + f"Visualization results are saved to: {filename}_t-{0:0{width}}.vtp ~ " + f"{filename}_t-{num_timestamps - 1:0{width}}.vtp" + ) + else: + logger.message(f"Visualization result is saved to: {filename}.vtp") + + +def save_vtu_to_mesh( + filename: str, + data_dict: Dict[str, np.ndarray], + coord_keys: Tuple[str, ...], + value_keys: Tuple[str, ...], +): + """Save data into .vtu format by meshio. + + Args: + filename (str): File name. + data_dict (Dict[str, np.ndarray]): Data in dict. + coord_keys (Tuple[str, ...]): Tuple of coord key. such as ("x", "y"). + value_keys (Tuple[str, ...]): Tuple of value key. such as ("u", "v"). + + Examples: + >>> import ppsci + >>> import numpy as np + >>> filename = "path/to/file.vtu" + >>> data_dict = { + ... "x": np.array([[1], [2], [3],[4]]), + ... "y": np.array([[2], [3], [4],[4]]), + ... "z": np.array([[3], [4], [5],[4]]), + ... "u": np.array([[4], [5], [6],[4]]), + ... "v": np.array([[5], [6], [7],[4]]), + ... } + >>> coord_keys = ("x","y","z") + >>> value_keys = ("u","v") + >>> ppsci.visualize.save_vtu_to_mesh(filename, data_dict, coord_keys, value_keys) # doctest: +SKIP + """ + npoint = len(next(iter(data_dict.values()))) + coord_ndim = len(coord_keys) + + # get the list variable transposed + points = np.stack(tuple(data_dict[key] for key in coord_keys)).reshape( + coord_ndim, npoint + ) + mesh = meshio.Mesh( + points=points.T, cells=[("vertex", np.arange(npoint).reshape(npoint, 1))] + ) + mesh.point_data = {key: data_dict[key] for key in value_keys} + if len(os.path.dirname(filename)): + os.makedirs(os.path.dirname(filename), exist_ok=True) + mesh.write(filename) From 3370e1d49a450cf99ea83d43d25be1043fe444ff Mon Sep 17 00:00:00 2001 From: Dubhe-Chang Date: Wed, 13 Aug 2025 18:07:38 +0800 Subject: [PATCH 32/32] delete unnecessary files --- examples/smc_reac/ppsci/__init__.py | 78 - examples/smc_reac/ppsci/arch/__init__.py | 136 -- examples/smc_reac/ppsci/arch/activation.py | 160 -- examples/smc_reac/ppsci/arch/afno.py | 687 ------ examples/smc_reac/ppsci/arch/amgnet.py | 649 ------ examples/smc_reac/ppsci/arch/base.py | 279 --- examples/smc_reac/ppsci/arch/cfdgcn.py | 350 --- .../smc_reac/ppsci/arch/chip_deeponets.py | 214 -- .../ppsci/arch/crystalgraphconvnet.py | 167 -- .../smc_reac/ppsci/arch/cuboid_transformer.py | 958 -------- .../ppsci/arch/cuboid_transformer_decoder.py | 1245 ----------- .../ppsci/arch/cuboid_transformer_encoder.py | 1515 ------------- .../ppsci/arch/cuboid_transformer_utils.py | 347 --- examples/smc_reac/ppsci/arch/cvit.py | 1095 --------- examples/smc_reac/ppsci/arch/deeponet.py | 154 -- examples/smc_reac/ppsci/arch/dgmr.py | 1151 ---------- .../smc_reac/ppsci/arch/embedding_koopman.py | 544 ----- examples/smc_reac/ppsci/arch/epnn.py | 126 -- .../ppsci/arch/extformer_moe_cuboid.py | 996 --------- .../arch/extformer_moe_cuboid_decoder.py | 1475 ------------ .../arch/extformer_moe_cuboid_encoder.py | 1992 ----------------- .../ppsci/arch/extformer_moe_cuboid_utils.py | 350 --- .../ppsci/arch/extformer_moe_utils.py | 563 ----- examples/smc_reac/ppsci/arch/fno_block.py | 1269 ----------- examples/smc_reac/ppsci/arch/gan.py | 400 ---- examples/smc_reac/ppsci/arch/geofno.py | 205 -- examples/smc_reac/ppsci/arch/graphcast.py | 492 ---- examples/smc_reac/ppsci/arch/he_deeponets.py | 197 -- examples/smc_reac/ppsci/arch/ifm_mlp.py | 540 ----- examples/smc_reac/ppsci/arch/kan.py | 385 ---- examples/smc_reac/ppsci/arch/lno.py | 312 --- examples/smc_reac/ppsci/arch/mlp.py | 828 ------- examples/smc_reac/ppsci/arch/model_list.py | 72 - examples/smc_reac/ppsci/arch/moflow_basic.py | 297 --- examples/smc_reac/ppsci/arch/moflow_glow.py | 477 ---- examples/smc_reac/ppsci/arch/moflow_net.py | 335 --- examples/smc_reac/ppsci/arch/nowcastnet.py | 639 ------ .../ppsci/arch/paddle_harmonics/legendre.py | 176 -- .../ppsci/arch/paddle_harmonics/quadrature.py | 156 -- .../arch/paddle_harmonics/random_fields.py | 148 -- .../ppsci/arch/paddle_harmonics/sht.py | 461 ---- examples/smc_reac/ppsci/arch/phycrnet.py | 540 ----- examples/smc_reac/ppsci/arch/phylstm.py | 239 -- .../smc_reac/ppsci/arch/physx_transformer.py | 407 ---- examples/smc_reac/ppsci/arch/regdgcnn.py | 250 --- examples/smc_reac/ppsci/arch/regpointnet.py | 146 -- examples/smc_reac/ppsci/arch/sfnonet.py | 568 ----- examples/smc_reac/ppsci/arch/smc_reac.py | 107 - examples/smc_reac/ppsci/arch/spinn.py | 180 -- examples/smc_reac/ppsci/arch/tfnonet.py | 514 ----- examples/smc_reac/ppsci/arch/tgcn.py | 200 -- examples/smc_reac/ppsci/arch/transformer.py | 417 ---- examples/smc_reac/ppsci/arch/unetex.py | 290 --- examples/smc_reac/ppsci/arch/unonet.py | 289 --- examples/smc_reac/ppsci/arch/uscnn.py | 124 - examples/smc_reac/ppsci/arch/vae.py | 103 - examples/smc_reac/ppsci/arch/velocitygan.py | 354 --- examples/smc_reac/ppsci/autodiff/__init__.py | 17 - examples/smc_reac/ppsci/autodiff/ad.py | 341 --- .../smc_reac/ppsci/constraint/__init__.py | 86 - examples/smc_reac/ppsci/constraint/base.py | 62 - .../ppsci/constraint/boundary_constraint.py | 163 -- .../ppsci/constraint/initial_constraint.py | 172 -- .../ppsci/constraint/integral_constraint.py | 178 -- .../ppsci/constraint/interior_constraint.py | 174 -- .../ppsci/constraint/periodic_constraint.py | 169 -- .../ppsci/constraint/supervised_constraint.py | 92 - examples/smc_reac/ppsci/data/__init__.py | 205 -- examples/smc_reac/ppsci/data/dataloader.py | 47 - .../smc_reac/ppsci/data/dataset/__init__.py | 118 - .../ppsci/data/dataset/airfoil_dataset.py | 241 -- .../ppsci/data/dataset/array_dataset.py | 390 ---- .../ppsci/data/dataset/atmospheric_dataset.py | 1781 --------------- .../ppsci/data/dataset/cgcnn_dataset.py | 312 --- .../ppsci/data/dataset/csv_dataset.py | 287 --- .../ppsci/data/dataset/cylinder_dataset.py | 215 -- .../ppsci/data/dataset/darcyflow_dataset.py | 296 --- .../ppsci/data/dataset/dgmr_dataset.py | 95 - .../ppsci/data/dataset/drivaernet_dataset.py | 316 --- .../dataset/drivaernetplusplus_dataset.py | 321 --- .../ppsci/data/dataset/enso_dataset.py | 405 ---- .../ppsci/data/dataset/era5_dataset.py | 249 --- .../data/dataset/ext_moe_enso_dataset.py | 406 ---- .../ppsci/data/dataset/fwi_dataset.py | 103 - .../ppsci/data/dataset/ifm_moe_dataset.py | 462 ---- .../ppsci/data/dataset/mat_dataset.py | 287 --- .../ppsci/data/dataset/moflow_dataset.py | 437 ---- .../ppsci/data/dataset/mrms_dataset.py | 251 --- .../ppsci/data/dataset/npz_dataset.py | 279 --- .../ppsci/data/dataset/pems_dataset.py | 151 -- .../ppsci/data/dataset/radar_dataset.py | 146 -- .../ppsci/data/dataset/sevir_dataset.py | 814 ------- .../data/dataset/spherical_swe_dataset.py | 104 - .../ppsci/data/dataset/trphysx_dataset.py | 326 --- .../ppsci/data/dataset/vtu_dataset.py | 106 - .../smc_reac/ppsci/data/process/__init__.py | 21 - .../data/process/batch_transform/__init__.py | 135 -- .../process/batch_transform/preprocess.py | 74 - .../ppsci/data/process/transform/__init__.py | 72 - .../data/process/transform/preprocess.py | 331 --- examples/smc_reac/ppsci/equation/__init__.py | 76 - .../smc_reac/ppsci/equation/fpde/__init__.py | 19 - .../ppsci/equation/fpde/fractional_poisson.py | 196 -- .../smc_reac/ppsci/equation/ide/__init__.py | 19 - .../smc_reac/ppsci/equation/ide/volterra.py | 127 -- .../smc_reac/ppsci/equation/pde/__init__.py | 43 - .../smc_reac/ppsci/equation/pde/allen_cahn.py | 64 - examples/smc_reac/ppsci/equation/pde/base.py | 243 -- .../smc_reac/ppsci/equation/pde/biharmonic.py | 74 - .../ppsci/equation/pde/heat_exchanger.py | 94 - .../smc_reac/ppsci/equation/pde/helmholtz.py | 119 - .../smc_reac/ppsci/equation/pde/laplace.py | 55 - .../ppsci/equation/pde/linear_elasticity.py | 184 -- .../ppsci/equation/pde/navier_stokes.py | 151 -- .../smc_reac/ppsci/equation/pde/nls_m_b.py | 101 - .../ppsci/equation/pde/normal_dot_vec.py | 59 - .../smc_reac/ppsci/equation/pde/poisson.py | 53 - examples/smc_reac/ppsci/equation/pde/viv.py | 64 - .../smc_reac/ppsci/experimental/__init__.py | 37 - .../ppsci/experimental/math_module.py | 646 ------ examples/smc_reac/ppsci/externals/__init__.py | 20 - examples/smc_reac/ppsci/geometry/__init__.py | 83 - examples/smc_reac/ppsci/geometry/csg.py | 337 --- examples/smc_reac/ppsci/geometry/geometry.py | 696 ------ .../smc_reac/ppsci/geometry/geometry_1d.py | 119 - .../smc_reac/ppsci/geometry/geometry_2d.py | 706 ------ .../smc_reac/ppsci/geometry/geometry_3d.py | 203 -- .../smc_reac/ppsci/geometry/geometry_nd.py | 196 -- examples/smc_reac/ppsci/geometry/inflation.py | 192 -- examples/smc_reac/ppsci/geometry/mesh.py | 1392 ------------ .../smc_reac/ppsci/geometry/pointcloud.py | 312 --- examples/smc_reac/ppsci/geometry/sampler.py | 92 - examples/smc_reac/ppsci/geometry/sdf.py | 198 -- .../smc_reac/ppsci/geometry/timedomain.py | 793 ------- examples/smc_reac/ppsci/loss/__init__.py | 67 - examples/smc_reac/ppsci/loss/base.py | 38 - examples/smc_reac/ppsci/loss/chamfer.py | 92 - examples/smc_reac/ppsci/loss/func.py | 94 - examples/smc_reac/ppsci/loss/integral.py | 112 - examples/smc_reac/ppsci/loss/kl.py | 51 - examples/smc_reac/ppsci/loss/l1.py | 219 -- examples/smc_reac/ppsci/loss/l2.py | 310 --- examples/smc_reac/ppsci/loss/mae.py | 109 - examples/smc_reac/ppsci/loss/mse.py | 355 --- examples/smc_reac/ppsci/loss/mtl/__init__.py | 49 - examples/smc_reac/ppsci/loss/mtl/agda.py | 161 -- examples/smc_reac/ppsci/loss/mtl/base.py | 68 - examples/smc_reac/ppsci/loss/mtl/grad_norm.py | 145 -- examples/smc_reac/ppsci/loss/mtl/ntk.py | 118 - examples/smc_reac/ppsci/loss/mtl/pcgrad.py | 124 - examples/smc_reac/ppsci/loss/mtl/relobralo.py | 127 -- examples/smc_reac/ppsci/loss/mtl/sum.py | 60 - examples/smc_reac/ppsci/metric/__init__.py | 63 - .../smc_reac/ppsci/metric/anomaly_coef.py | 122 - examples/smc_reac/ppsci/metric/base.py | 25 - examples/smc_reac/ppsci/metric/func.py | 66 - examples/smc_reac/ppsci/metric/l2_rel.py | 139 -- examples/smc_reac/ppsci/metric/mae.py | 73 - examples/smc_reac/ppsci/metric/max_ae.py | 77 - examples/smc_reac/ppsci/metric/mse.py | 73 - examples/smc_reac/ppsci/metric/r2_score.py | 97 - examples/smc_reac/ppsci/metric/rmse.py | 155 -- examples/smc_reac/ppsci/optimizer/__init__.py | 84 - .../smc_reac/ppsci/optimizer/lr_scheduler.py | 911 -------- .../smc_reac/ppsci/optimizer/optimizer.py | 649 ------ examples/smc_reac/ppsci/optimizer/soap.py | 558 ----- .../smc_reac/ppsci/probability/__init__.py | 19 - examples/smc_reac/ppsci/probability/hmc.py | 175 -- examples/smc_reac/ppsci/solver/__init__.py | 25 - examples/smc_reac/ppsci/solver/eval.py | 316 --- examples/smc_reac/ppsci/solver/printer.py | 161 -- examples/smc_reac/ppsci/solver/solver.py | 1219 ---------- examples/smc_reac/ppsci/solver/train.py | 324 --- examples/smc_reac/ppsci/solver/visu.py | 98 - examples/smc_reac/ppsci/utils/__init__.py | 66 - examples/smc_reac/ppsci/utils/callbacks.py | 136 -- examples/smc_reac/ppsci/utils/checker.py | 287 --- examples/smc_reac/ppsci/utils/config.py | 457 ---- examples/smc_reac/ppsci/utils/download.py | 285 --- examples/smc_reac/ppsci/utils/ema.py | 172 -- examples/smc_reac/ppsci/utils/expression.py | 212 -- examples/smc_reac/ppsci/utils/initializer.py | 498 ----- examples/smc_reac/ppsci/utils/logger.py | 264 --- examples/smc_reac/ppsci/utils/misc.py | 684 ------ examples/smc_reac/ppsci/utils/reader.py | 266 --- examples/smc_reac/ppsci/utils/save_load.py | 300 --- examples/smc_reac/ppsci/utils/symbolic.py | 981 -------- examples/smc_reac/ppsci/utils/writer.py | 225 -- examples/smc_reac/ppsci/validate/__init__.py | 81 - examples/smc_reac/ppsci/validate/base.py | 69 - .../smc_reac/ppsci/validate/geo_validator.py | 161 -- .../smc_reac/ppsci/validate/sup_validator.py | 103 - examples/smc_reac/ppsci/visualize/__init__.py | 82 - examples/smc_reac/ppsci/visualize/base.py | 65 - examples/smc_reac/ppsci/visualize/plot.py | 580 ----- examples/smc_reac/ppsci/visualize/radar.py | 124 - .../smc_reac/ppsci/visualize/visualizer.py | 409 ---- examples/smc_reac/ppsci/visualize/vtu.py | 278 --- 198 files changed, 60861 deletions(-) delete mode 100644 examples/smc_reac/ppsci/__init__.py delete mode 100644 examples/smc_reac/ppsci/arch/__init__.py delete mode 100644 examples/smc_reac/ppsci/arch/activation.py delete mode 100644 examples/smc_reac/ppsci/arch/afno.py delete mode 100644 examples/smc_reac/ppsci/arch/amgnet.py delete mode 100644 examples/smc_reac/ppsci/arch/base.py delete mode 100644 examples/smc_reac/ppsci/arch/cfdgcn.py delete mode 100644 examples/smc_reac/ppsci/arch/chip_deeponets.py delete mode 100644 examples/smc_reac/ppsci/arch/crystalgraphconvnet.py delete mode 100644 examples/smc_reac/ppsci/arch/cuboid_transformer.py delete mode 100644 examples/smc_reac/ppsci/arch/cuboid_transformer_decoder.py delete mode 100644 examples/smc_reac/ppsci/arch/cuboid_transformer_encoder.py delete mode 100644 examples/smc_reac/ppsci/arch/cuboid_transformer_utils.py delete mode 100644 examples/smc_reac/ppsci/arch/cvit.py delete mode 100644 examples/smc_reac/ppsci/arch/deeponet.py delete mode 100644 examples/smc_reac/ppsci/arch/dgmr.py delete mode 100644 examples/smc_reac/ppsci/arch/embedding_koopman.py delete mode 100644 examples/smc_reac/ppsci/arch/epnn.py delete mode 100644 examples/smc_reac/ppsci/arch/extformer_moe_cuboid.py delete mode 100644 examples/smc_reac/ppsci/arch/extformer_moe_cuboid_decoder.py delete mode 100644 examples/smc_reac/ppsci/arch/extformer_moe_cuboid_encoder.py delete mode 100644 examples/smc_reac/ppsci/arch/extformer_moe_cuboid_utils.py delete mode 100644 examples/smc_reac/ppsci/arch/extformer_moe_utils.py delete mode 100644 examples/smc_reac/ppsci/arch/fno_block.py delete mode 100644 examples/smc_reac/ppsci/arch/gan.py delete mode 100644 examples/smc_reac/ppsci/arch/geofno.py delete mode 100644 examples/smc_reac/ppsci/arch/graphcast.py delete mode 100644 examples/smc_reac/ppsci/arch/he_deeponets.py delete mode 100644 examples/smc_reac/ppsci/arch/ifm_mlp.py delete mode 100644 examples/smc_reac/ppsci/arch/kan.py delete mode 100644 examples/smc_reac/ppsci/arch/lno.py delete mode 100644 examples/smc_reac/ppsci/arch/mlp.py delete mode 100644 examples/smc_reac/ppsci/arch/model_list.py delete mode 100644 examples/smc_reac/ppsci/arch/moflow_basic.py delete mode 100644 examples/smc_reac/ppsci/arch/moflow_glow.py delete mode 100644 examples/smc_reac/ppsci/arch/moflow_net.py delete mode 100644 examples/smc_reac/ppsci/arch/nowcastnet.py delete mode 100644 examples/smc_reac/ppsci/arch/paddle_harmonics/legendre.py delete mode 100644 examples/smc_reac/ppsci/arch/paddle_harmonics/quadrature.py delete mode 100644 examples/smc_reac/ppsci/arch/paddle_harmonics/random_fields.py delete mode 100644 examples/smc_reac/ppsci/arch/paddle_harmonics/sht.py delete mode 100644 examples/smc_reac/ppsci/arch/phycrnet.py delete mode 100644 examples/smc_reac/ppsci/arch/phylstm.py delete mode 100644 examples/smc_reac/ppsci/arch/physx_transformer.py delete mode 100644 examples/smc_reac/ppsci/arch/regdgcnn.py delete mode 100644 examples/smc_reac/ppsci/arch/regpointnet.py delete mode 100644 examples/smc_reac/ppsci/arch/sfnonet.py delete mode 100644 examples/smc_reac/ppsci/arch/smc_reac.py delete mode 100644 examples/smc_reac/ppsci/arch/spinn.py delete mode 100644 examples/smc_reac/ppsci/arch/tfnonet.py delete mode 100644 examples/smc_reac/ppsci/arch/tgcn.py delete mode 100644 examples/smc_reac/ppsci/arch/transformer.py delete mode 100644 examples/smc_reac/ppsci/arch/unetex.py delete mode 100644 examples/smc_reac/ppsci/arch/unonet.py delete mode 100644 examples/smc_reac/ppsci/arch/uscnn.py delete mode 100644 examples/smc_reac/ppsci/arch/vae.py delete mode 100644 examples/smc_reac/ppsci/arch/velocitygan.py delete mode 100644 examples/smc_reac/ppsci/autodiff/__init__.py delete mode 100644 examples/smc_reac/ppsci/autodiff/ad.py delete mode 100644 examples/smc_reac/ppsci/constraint/__init__.py delete mode 100644 examples/smc_reac/ppsci/constraint/base.py delete mode 100644 examples/smc_reac/ppsci/constraint/boundary_constraint.py delete mode 100644 examples/smc_reac/ppsci/constraint/initial_constraint.py delete mode 100644 examples/smc_reac/ppsci/constraint/integral_constraint.py delete mode 100644 examples/smc_reac/ppsci/constraint/interior_constraint.py delete mode 100644 examples/smc_reac/ppsci/constraint/periodic_constraint.py delete mode 100644 examples/smc_reac/ppsci/constraint/supervised_constraint.py delete mode 100644 examples/smc_reac/ppsci/data/__init__.py delete mode 100644 examples/smc_reac/ppsci/data/dataloader.py delete mode 100644 examples/smc_reac/ppsci/data/dataset/__init__.py delete mode 100644 examples/smc_reac/ppsci/data/dataset/airfoil_dataset.py delete mode 100644 examples/smc_reac/ppsci/data/dataset/array_dataset.py delete mode 100644 examples/smc_reac/ppsci/data/dataset/atmospheric_dataset.py delete mode 100644 examples/smc_reac/ppsci/data/dataset/cgcnn_dataset.py delete mode 100644 examples/smc_reac/ppsci/data/dataset/csv_dataset.py delete mode 100644 examples/smc_reac/ppsci/data/dataset/cylinder_dataset.py delete mode 100644 examples/smc_reac/ppsci/data/dataset/darcyflow_dataset.py delete mode 100644 examples/smc_reac/ppsci/data/dataset/dgmr_dataset.py delete mode 100644 examples/smc_reac/ppsci/data/dataset/drivaernet_dataset.py delete mode 100644 examples/smc_reac/ppsci/data/dataset/drivaernetplusplus_dataset.py delete mode 100644 examples/smc_reac/ppsci/data/dataset/enso_dataset.py delete mode 100644 examples/smc_reac/ppsci/data/dataset/era5_dataset.py delete mode 100644 examples/smc_reac/ppsci/data/dataset/ext_moe_enso_dataset.py delete mode 100644 examples/smc_reac/ppsci/data/dataset/fwi_dataset.py delete mode 100644 examples/smc_reac/ppsci/data/dataset/ifm_moe_dataset.py delete mode 100644 examples/smc_reac/ppsci/data/dataset/mat_dataset.py delete mode 100644 examples/smc_reac/ppsci/data/dataset/moflow_dataset.py delete mode 100644 examples/smc_reac/ppsci/data/dataset/mrms_dataset.py delete mode 100644 examples/smc_reac/ppsci/data/dataset/npz_dataset.py delete mode 100644 examples/smc_reac/ppsci/data/dataset/pems_dataset.py delete mode 100644 examples/smc_reac/ppsci/data/dataset/radar_dataset.py delete mode 100644 examples/smc_reac/ppsci/data/dataset/sevir_dataset.py delete mode 100644 examples/smc_reac/ppsci/data/dataset/spherical_swe_dataset.py delete mode 100644 examples/smc_reac/ppsci/data/dataset/trphysx_dataset.py delete mode 100644 examples/smc_reac/ppsci/data/dataset/vtu_dataset.py delete mode 100644 examples/smc_reac/ppsci/data/process/__init__.py delete mode 100644 examples/smc_reac/ppsci/data/process/batch_transform/__init__.py delete mode 100644 examples/smc_reac/ppsci/data/process/batch_transform/preprocess.py delete mode 100644 examples/smc_reac/ppsci/data/process/transform/__init__.py delete mode 100644 examples/smc_reac/ppsci/data/process/transform/preprocess.py delete mode 100644 examples/smc_reac/ppsci/equation/__init__.py delete mode 100644 examples/smc_reac/ppsci/equation/fpde/__init__.py delete mode 100644 examples/smc_reac/ppsci/equation/fpde/fractional_poisson.py delete mode 100644 examples/smc_reac/ppsci/equation/ide/__init__.py delete mode 100644 examples/smc_reac/ppsci/equation/ide/volterra.py delete mode 100644 examples/smc_reac/ppsci/equation/pde/__init__.py delete mode 100644 examples/smc_reac/ppsci/equation/pde/allen_cahn.py delete mode 100644 examples/smc_reac/ppsci/equation/pde/base.py delete mode 100644 examples/smc_reac/ppsci/equation/pde/biharmonic.py delete mode 100644 examples/smc_reac/ppsci/equation/pde/heat_exchanger.py delete mode 100644 examples/smc_reac/ppsci/equation/pde/helmholtz.py delete mode 100644 examples/smc_reac/ppsci/equation/pde/laplace.py delete mode 100644 examples/smc_reac/ppsci/equation/pde/linear_elasticity.py delete mode 100644 examples/smc_reac/ppsci/equation/pde/navier_stokes.py delete mode 100644 examples/smc_reac/ppsci/equation/pde/nls_m_b.py delete mode 100644 examples/smc_reac/ppsci/equation/pde/normal_dot_vec.py delete mode 100644 examples/smc_reac/ppsci/equation/pde/poisson.py delete mode 100644 examples/smc_reac/ppsci/equation/pde/viv.py delete mode 100644 examples/smc_reac/ppsci/experimental/__init__.py delete mode 100644 examples/smc_reac/ppsci/experimental/math_module.py delete mode 100644 examples/smc_reac/ppsci/externals/__init__.py delete mode 100644 examples/smc_reac/ppsci/geometry/__init__.py delete mode 100644 examples/smc_reac/ppsci/geometry/csg.py delete mode 100644 examples/smc_reac/ppsci/geometry/geometry.py delete mode 100644 examples/smc_reac/ppsci/geometry/geometry_1d.py delete mode 100644 examples/smc_reac/ppsci/geometry/geometry_2d.py delete mode 100644 examples/smc_reac/ppsci/geometry/geometry_3d.py delete mode 100644 examples/smc_reac/ppsci/geometry/geometry_nd.py delete mode 100644 examples/smc_reac/ppsci/geometry/inflation.py delete mode 100644 examples/smc_reac/ppsci/geometry/mesh.py delete mode 100644 examples/smc_reac/ppsci/geometry/pointcloud.py delete mode 100644 examples/smc_reac/ppsci/geometry/sampler.py delete mode 100644 examples/smc_reac/ppsci/geometry/sdf.py delete mode 100644 examples/smc_reac/ppsci/geometry/timedomain.py delete mode 100644 examples/smc_reac/ppsci/loss/__init__.py delete mode 100644 examples/smc_reac/ppsci/loss/base.py delete mode 100644 examples/smc_reac/ppsci/loss/chamfer.py delete mode 100644 examples/smc_reac/ppsci/loss/func.py delete mode 100644 examples/smc_reac/ppsci/loss/integral.py delete mode 100644 examples/smc_reac/ppsci/loss/kl.py delete mode 100644 examples/smc_reac/ppsci/loss/l1.py delete mode 100644 examples/smc_reac/ppsci/loss/l2.py delete mode 100644 examples/smc_reac/ppsci/loss/mae.py delete mode 100644 examples/smc_reac/ppsci/loss/mse.py delete mode 100644 examples/smc_reac/ppsci/loss/mtl/__init__.py delete mode 100644 examples/smc_reac/ppsci/loss/mtl/agda.py delete mode 100644 examples/smc_reac/ppsci/loss/mtl/base.py delete mode 100644 examples/smc_reac/ppsci/loss/mtl/grad_norm.py delete mode 100644 examples/smc_reac/ppsci/loss/mtl/ntk.py delete mode 100644 examples/smc_reac/ppsci/loss/mtl/pcgrad.py delete mode 100644 examples/smc_reac/ppsci/loss/mtl/relobralo.py delete mode 100644 examples/smc_reac/ppsci/loss/mtl/sum.py delete mode 100644 examples/smc_reac/ppsci/metric/__init__.py delete mode 100644 examples/smc_reac/ppsci/metric/anomaly_coef.py delete mode 100644 examples/smc_reac/ppsci/metric/base.py delete mode 100644 examples/smc_reac/ppsci/metric/func.py delete mode 100644 examples/smc_reac/ppsci/metric/l2_rel.py delete mode 100644 examples/smc_reac/ppsci/metric/mae.py delete mode 100644 examples/smc_reac/ppsci/metric/max_ae.py delete mode 100644 examples/smc_reac/ppsci/metric/mse.py delete mode 100644 examples/smc_reac/ppsci/metric/r2_score.py delete mode 100644 examples/smc_reac/ppsci/metric/rmse.py delete mode 100644 examples/smc_reac/ppsci/optimizer/__init__.py delete mode 100644 examples/smc_reac/ppsci/optimizer/lr_scheduler.py delete mode 100644 examples/smc_reac/ppsci/optimizer/optimizer.py delete mode 100644 examples/smc_reac/ppsci/optimizer/soap.py delete mode 100644 examples/smc_reac/ppsci/probability/__init__.py delete mode 100644 examples/smc_reac/ppsci/probability/hmc.py delete mode 100644 examples/smc_reac/ppsci/solver/__init__.py delete mode 100644 examples/smc_reac/ppsci/solver/eval.py delete mode 100644 examples/smc_reac/ppsci/solver/printer.py delete mode 100644 examples/smc_reac/ppsci/solver/solver.py delete mode 100644 examples/smc_reac/ppsci/solver/train.py delete mode 100644 examples/smc_reac/ppsci/solver/visu.py delete mode 100644 examples/smc_reac/ppsci/utils/__init__.py delete mode 100644 examples/smc_reac/ppsci/utils/callbacks.py delete mode 100644 examples/smc_reac/ppsci/utils/checker.py delete mode 100644 examples/smc_reac/ppsci/utils/config.py delete mode 100644 examples/smc_reac/ppsci/utils/download.py delete mode 100644 examples/smc_reac/ppsci/utils/ema.py delete mode 100644 examples/smc_reac/ppsci/utils/expression.py delete mode 100644 examples/smc_reac/ppsci/utils/initializer.py delete mode 100644 examples/smc_reac/ppsci/utils/logger.py delete mode 100644 examples/smc_reac/ppsci/utils/misc.py delete mode 100644 examples/smc_reac/ppsci/utils/reader.py delete mode 100644 examples/smc_reac/ppsci/utils/save_load.py delete mode 100644 examples/smc_reac/ppsci/utils/symbolic.py delete mode 100644 examples/smc_reac/ppsci/utils/writer.py delete mode 100644 examples/smc_reac/ppsci/validate/__init__.py delete mode 100644 examples/smc_reac/ppsci/validate/base.py delete mode 100644 examples/smc_reac/ppsci/validate/geo_validator.py delete mode 100644 examples/smc_reac/ppsci/validate/sup_validator.py delete mode 100644 examples/smc_reac/ppsci/visualize/__init__.py delete mode 100644 examples/smc_reac/ppsci/visualize/base.py delete mode 100644 examples/smc_reac/ppsci/visualize/plot.py delete mode 100644 examples/smc_reac/ppsci/visualize/radar.py delete mode 100644 examples/smc_reac/ppsci/visualize/visualizer.py delete mode 100644 examples/smc_reac/ppsci/visualize/vtu.py diff --git a/examples/smc_reac/ppsci/__init__.py b/examples/smc_reac/ppsci/__init__.py deleted file mode 100644 index bde7aa49a5..0000000000 --- a/examples/smc_reac/ppsci/__init__.py +++ /dev/null @@ -1,78 +0,0 @@ -# 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. - -from ppsci import arch # isort:skip -from ppsci import autodiff # isort:skip -from ppsci import constraint # isort:skip -from ppsci import data # isort:skip -from ppsci import equation # isort:skip -from ppsci import geometry # isort:skip -from ppsci import loss # isort:skip -from ppsci import metric # isort:skip -from ppsci import optimizer # isort:skip -from ppsci import utils # isort:skip -from ppsci import visualize # isort:skip -from ppsci import validate # isort:skip -from ppsci import solver # isort:skip -from ppsci import experimental # isort:skip - -from ppsci.utils.checker import run_check # isort:skip -from ppsci.utils.checker import run_check_mesh # isort:skip -from ppsci.utils import lambdify # isort:skip - - -try: - # import auto-generated version information from '._version' file, using - # setuptools_scm via 'pip install'. Details of versioning rule can be referd to: - # https://peps.python.org/pep-0440/#public-version-identifiers - from ._version import version as __version__ -except ImportError: - __version__ = "unknown version" - -__all__ = [ - "arch", - "autodiff", - "constraint", - "data", - "equation", - "geometry", - "loss", - "metric", - "optimizer", - "utils", - "visualize", - "validate", - "solver", - "experimental", - "run_check", - "run_check_mesh", - "lambdify", -] - - -# NOTE: Register custom solvers for parsing values from omegaconf more flexible -def _register_config_solvers(): - import numpy as np - from omegaconf import OmegaConf - - # register solver for "${numpy:xxx}" item, e.g. pi: "${numpy:pi}" - if not OmegaConf.has_resolver("numpy"): - OmegaConf.register_new_resolver("numpy", lambda x: getattr(np, x)) - - # register solver for "${sum:xxx}" item, e.g. pi: "${sum:[10, 20, 30]}" - if not OmegaConf.has_resolver("sum"): - OmegaConf.register_new_resolver("sum", lambda x: sum(x)) - - -_register_config_solvers() diff --git a/examples/smc_reac/ppsci/arch/__init__.py b/examples/smc_reac/ppsci/arch/__init__.py deleted file mode 100644 index 16387994c7..0000000000 --- a/examples/smc_reac/ppsci/arch/__init__.py +++ /dev/null @@ -1,136 +0,0 @@ -# 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. - -from __future__ import annotations - -import copy - -from ppsci.arch.afno import AFNONet # isort:skip -from ppsci.arch.afno import PrecipNet # isort:skip -from ppsci.arch.amgnet import AMGNet # isort:skip -from ppsci.arch.base import Arch # isort:skip -from ppsci.arch.cfdgcn import CFDGCN # isort:skip -from ppsci.arch.smc_reac import SuzukiMiyauraModel # isort:skip -from ppsci.arch.chip_deeponets import ChipDeepONets # isort:skip -from ppsci.arch.crystalgraphconvnet import CrystalGraphConvNet # isort:skip -from ppsci.arch.cuboid_transformer import CuboidTransformer # isort:skip -from ppsci.arch.cvit import CVit # isort:skip -from ppsci.arch.cvit import CVit1D # isort:skip -from ppsci.arch.deeponet import DeepONet # isort:skip -from ppsci.arch.dgmr import DGMR # isort:skip -from ppsci.arch.embedding_koopman import CylinderEmbedding # isort:skip -from ppsci.arch.embedding_koopman import LorenzEmbedding # isort:skip -from ppsci.arch.embedding_koopman import RosslerEmbedding # isort:skip -from ppsci.arch.epnn import Epnn # isort:skip -from ppsci.arch.extformer_moe_cuboid import ExtFormerMoECuboid # isort:skip -from ppsci.arch.gan import Discriminator # isort:skip -from ppsci.arch.gan import Generator # isort:skip -from ppsci.arch.geofno import FNO1d # isort:skip -from ppsci.arch.graphcast import GraphCastNet # isort:skip -from ppsci.arch.he_deeponets import HEDeepONets # isort:skip -from ppsci.arch.lno import LNO # isort:skip -from ppsci.arch.mlp import MLP # isort:skip -from ppsci.arch.mlp import ModifiedMLP # isort:skip -from ppsci.arch.mlp import PirateNet # isort:skip -from ppsci.arch.model_list import ModelList # isort:skip -from ppsci.arch.nowcastnet import NowcastNet # isort:skip -from ppsci.arch.phycrnet import PhyCRNet # isort:skip -from ppsci.arch.phylstm import DeepPhyLSTM # isort:skip -from ppsci.arch.physx_transformer import PhysformerGPT2 # isort:skip -from ppsci.arch.sfnonet import SFNONet # isort:skip -from ppsci.arch.spinn import SPINN # isort:skip -from ppsci.arch.tfnonet import TFNO1dNet, TFNO2dNet, TFNO3dNet # isort:skip -from ppsci.arch.transformer import Transformer # isort:skip -from ppsci.arch.unetex import UNetEx # isort:skip -from ppsci.arch.unonet import UNONet # isort:skip -from ppsci.arch.uscnn import USCNN # isort:skip -from ppsci.arch.vae import AutoEncoder # isort:skip -from ppsci.arch.velocitygan import VelocityDiscriminator # isort:skip -from ppsci.arch.velocitygan import VelocityGenerator # isort:skip -from ppsci.arch.moflow_net import MoFlowNet, MoFlowProp # isort:skip -from ppsci.utils import logger # isort:skip -from ppsci.arch.regdgcnn import RegDGCNN # isort:skip -from ppsci.arch.regpointnet import RegPointNet # isort:skip -from ppsci.arch.ifm_mlp import IFMMLP # isort:skip - -__all__ = [ - "MoFlowNet", - "MoFlowProp", - "AFNONet", - "AMGNet", - "Arch", - "AutoEncoder", - "build_model", - "CFDGCN", - "SuzukiMiyauraModel", - "ChipDeepONets", - "CrystalGraphConvNet", - "CuboidTransformer", - "CVit", - "CVit1D", - "CylinderEmbedding", - "DeepONet", - "DeepPhyLSTM", - "DGMR", - "Discriminator", - "Epnn", - "ExtFormerMoECuboid", - "FNO1d", - "Generator", - "GraphCastNet", - "HEDeepONets", - "LorenzEmbedding", - "LNO", - "MLP", - "ModelList", - "ModifiedMLP", - "NowcastNet", - "PhyCRNet", - "PhysformerGPT2", - "PirateNet", - "PrecipNet", - "RosslerEmbedding", - "SFNONet", - "SPINN", - "TFNO1dNet", - "TFNO2dNet", - "TFNO3dNet", - "Transformer", - "UNetEx", - "UNONet", - "USCNN", - "VelocityDiscriminator", - "VelocityGenerator", - "RegDGCNN", - "RegPointNet", - "IFMMLP", -] - - -def build_model(cfg): - """Build model - - Args: - cfg (DictConfig): Arch config. - - Returns: - nn.Layer: Model. - """ - cfg = copy.deepcopy(cfg) - arch_cls = cfg.pop("name") - arch = eval(arch_cls)(**cfg) - - logger.debug(str(arch)) - - return arch diff --git a/examples/smc_reac/ppsci/arch/activation.py b/examples/smc_reac/ppsci/arch/activation.py deleted file mode 100644 index 3f78eb1a61..0000000000 --- a/examples/smc_reac/ppsci/arch/activation.py +++ /dev/null @@ -1,160 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import Callable - -import numpy as np -import paddle -import paddle.nn.functional as F -from paddle import nn - -from ppsci.utils import initializer -from ppsci.utils import misc - - -class Stan(nn.Layer): - """Self-scalable Tanh. - paper: https://arxiv.org/abs/2204.12589v1 - - Args: - out_features (int, optional): Output features. Defaults to 1. - """ - - def __init__(self, out_features: int = 1): - super().__init__() - self.beta = self.create_parameter( - shape=(out_features,), - default_initializer=nn.initializer.Constant(1), - ) - - def forward(self, x): - # TODO: manually broadcast beta to x.shape for preventing backward error yet. - return F.tanh(x) * (1 + paddle.broadcast_to(self.beta, x.shape) * x) - # return F.tanh(x) * (1 + self.beta * x) - - -class Swish(nn.Layer): - def __init__(self, beta: float = 1.0): - super().__init__() - self.beta = self.create_parameter( - shape=[], - default_initializer=nn.initializer.Constant(beta), - ) - - def forward(self, x): - return x * F.sigmoid(self.beta * x) - - -class Cos(nn.Layer): - def __init__(self): - super().__init__() - - def forward(self, x): - return paddle.cos(x) - - -class Sin(nn.Layer): - def __init__(self): - super().__init__() - - def forward(self, x): - return paddle.sin(x) - - -class Siren(nn.Layer): - """Implicit Neural Representations with Periodic Activation Functions. - paper link: https://arxiv.org/abs/2006.09661 - code ref: https://github.com/vsitzmann/siren/tree/master - """ - - def __init__(self, w0: float = 30): - super().__init__() - self.w0 = w0 - - def forward(self, x): - return paddle.sin(self.w0 * x) - - @staticmethod - def init_for_first_layer(layer: nn.Linear): - """Initialization only for first hidden layer. - ref: https://github.com/vsitzmann/siren/blob/master/modules.py#L630 - """ - if not isinstance(layer, nn.Linear): - raise TypeError( - "Siren initialization only support Linear layer now, " - f"but got {misc.typename(layer)}" - ) - in_features = layer.weight.shape[0] - with paddle.no_grad(): - initializer.uniform_(layer.weight, -1 / in_features, 1 / in_features) - initializer.zeros_(layer.bias) - - @staticmethod - def init_for_hidden_layer(layer: nn.Linear, w0: float = 30): - """Initialization for hidden layer except first layer. - ref: https://github.com/vsitzmann/siren/blob/master/modules.py#L622 - """ - if not isinstance(layer, nn.Linear): - raise TypeError( - "Siren initialization only support Linear layer now, " - f"but got {misc.typename(layer)}" - ) - in_features = layer.weight.shape[0] - with paddle.no_grad(): - initializer.uniform_( - layer.weight, - -np.sqrt(6 / in_features) / w0, - np.sqrt(6 / in_features) / w0, - ) - initializer.zeros_(layer.bias) - - -act_func_dict = { - "elu": nn.ELU(), - "relu": nn.ReLU(), - "selu": nn.SELU(), - "gelu": nn.GELU(), - "leaky_relu": nn.LeakyReLU(), - "sigmoid": nn.Sigmoid(), - "silu": nn.Silu(), - "sin": Sin(), - "cos": Cos(), - "swish": Swish, - "tanh": nn.Tanh(), - "identity": nn.Identity(), - "siren": Siren(), - "stan": Stan, -} - - -def get_activation(act_name: str) -> Callable: - """Get activation function according to act_name. - - Args: - act_name (str): Name of activation, such as "tanh". - - Returns: - Callable: Paddle activation function. - """ - if act_name.lower() not in act_func_dict: - raise ValueError(f"act_name({act_name}) not found in act_func_dict") - - act_layer = act_func_dict[act_name.lower()] - if isinstance(act_layer, type) and act_name != "stan": - # Is a activation class but not a instance of it, instantiate manually(except for 'Stan') - return act_layer() - - return act_layer diff --git a/examples/smc_reac/ppsci/arch/afno.py b/examples/smc_reac/ppsci/arch/afno.py deleted file mode 100644 index 62ec6fdd7c..0000000000 --- a/examples/smc_reac/ppsci/arch/afno.py +++ /dev/null @@ -1,687 +0,0 @@ -# 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. - -""" -Code below is heavily based on [FourCastNet](https://github.com/NVlabs/FourCastNet) -""" -from __future__ import annotations - -from functools import partial -from typing import Optional -from typing import Tuple - -import paddle -import paddle.fft -import paddle.nn.functional as F -from paddle import nn - -from ppsci.arch import activation as act_mod -from ppsci.arch import base -from ppsci.utils import initializer - - -def drop_path( - x: paddle.Tensor, - drop_prob: float = 0.0, - training: bool = False, - scale_by_keep: bool = True, -) -> paddle.Tensor: - """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks). - the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper... - See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... - - Args: - x (paddle.Tensor): The tensor to apply. - drop_prob (float, optional): Drop paths probability. Defaults to 0.0. - training (bool, optional): Whether at training mode. Defaults to False. - scale_by_keep (bool, optional): Whether upscale the output. Defaults to True. - - Returns: - paddle.Tensor: Output tensor after apply dropout. - """ - if drop_prob == 0.0 or not training: - return x - keep_prob = 1 - drop_prob - shape = (x.shape[0],) + (1,) * (x.ndim - 1) - random_tensor = paddle.full(shape, keep_prob, x.dtype) - random_tensor = paddle.bernoulli(random_tensor) - if keep_prob > 0.0 and scale_by_keep: - random_tensor = random_tensor / keep_prob - return x * random_tensor - - -class DropPath(nn.Layer): - """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks). - - Args: - drop_prob (float, optional): Drop paths probability. Defaults to 0.0. - scale_by_keep (bool, optional): Whether upscale the output. Defaults to True. - """ - - def __init__(self, drop_prob: float = 0.0, scale_by_keep: bool = True): - super().__init__() - self.drop_prob = drop_prob - self.scale_by_keep = scale_by_keep - - def forward(self, x): - return drop_path(x, self.drop_prob, self.training, self.scale_by_keep) - - def extra_repr(self): - return f"drop_prob={round(self.drop_prob,3):0.3f}" - - -class PeriodicPad2d(nn.Layer): - """Pad longitudinal (left-right) circular and pad latitude (top-bottom) with zeros. - - Args: - pad (int): Number of pad. - """ - - def __init__(self, pad: int): - super(PeriodicPad2d, self).__init__() - self.pad = pad - - def forward(self, x): - # pad left and right circular - out = F.pad(x, (self.pad, self.pad, 0, 0), mode="circular") - # pad top and bottom zeros - out = F.pad( - out, - (0, 0, 0, 0, self.pad, self.pad, 0, 0), - mode="constant", - value=0, - ) - return out - - -class MLP(nn.Layer): - """Multi layer perceptron module used in Transformer. - - Args: - in_features (int): Number of the input features. - hidden_features (Optional[int]): Number of the hidden size. Defaults to None. - out_features (Optional[int]): Number of the output features. Defaults to None. - activation (str, optional): Name of activation function. Defaults to "gelu". - drop (float, optional): Probability of dropout the units. Defaults to 0.0. - """ - - def __init__( - self, - in_features: int, - hidden_features: Optional[int] = None, - out_features: Optional[int] = None, - activation: str = "gelu", - drop: float = 0.0, - ): - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features - self.fc1 = nn.Linear(in_features, hidden_features) - self.act = act_mod.get_activation(activation) - self.fc2 = nn.Linear(hidden_features, out_features) - self.drop = nn.Dropout(drop) - - def forward(self, x): - x = self.fc1(x) - x = self.act(x) - x = self.drop(x) - x = self.fc2(x) - x = self.drop(x) - return x - - -class AFNO2D(nn.Layer): - """2D Adaptive Fourier Neural Operators. - - Args: - hidden_size (int): Number of hidden size. - num_blocks (int, optional): Number of blocks. Defaults to 8. - sparsity_threshold (float, optional): The value of threshold for softshrink. Defaults to 0.01. - hard_thresholding_fraction (float, optional): The value of threshold for keep mode. Defaults to 1.0. - hidden_size_factor (int, optional): The factor of hidden size. Defaults to 1. - scale (float, optional): The scale factor of the parameter when initialization. Defaults to 0.02. - """ - - def __init__( - self, - hidden_size: int, - num_blocks: int = 8, - sparsity_threshold: float = 0.01, - hard_thresholding_fraction: float = 1.0, - hidden_size_factor: int = 1, - scale: float = 0.02, - ): - super().__init__() - if hidden_size % num_blocks != 0: - raise ValueError( - f"hidden_size({hidden_size}) should be divisible by num_blocks({num_blocks})." - ) - - self.hidden_size = hidden_size - self.sparsity_threshold = sparsity_threshold - self.num_blocks = num_blocks - self.block_size = self.hidden_size // self.num_blocks - self.hard_thresholding_fraction = hard_thresholding_fraction - self.hidden_size_factor = hidden_size_factor - self.scale = scale - - self.w1 = self.create_parameter( - shape=( - 2, - self.num_blocks, - self.block_size, - self.block_size * self.hidden_size_factor, - ), - default_initializer=nn.initializer.Normal(std=self.scale), - ) - self.b1 = self.create_parameter( - shape=(2, self.num_blocks, self.block_size * self.hidden_size_factor), - default_initializer=nn.initializer.Normal(std=self.scale), - ) - self.w2 = self.create_parameter( - shape=( - 2, - self.num_blocks, - self.block_size * self.hidden_size_factor, - self.block_size, - ), - default_initializer=nn.initializer.Normal(std=self.scale), - ) - self.b2 = self.create_parameter( - shape=(2, self.num_blocks, self.block_size), - default_initializer=nn.initializer.Normal(std=self.scale), - ) - - def forward(self, x): - bias = x - - B, H, W, C = x.shape - - x = paddle.fft.rfft2(x, axes=(1, 2), norm="ortho") - x = x.reshape((B, H, W // 2 + 1, self.num_blocks, self.block_size)) - - o1_shape = ( - B, - H, - W // 2 + 1, - self.num_blocks, - self.block_size * self.hidden_size_factor, - ) - o1_real = paddle.zeros(o1_shape) - o1_imag = paddle.zeros(o1_shape) - o2_real = paddle.zeros(x.shape) - o2_imag = paddle.zeros(x.shape) - - total_modes = H // 2 + 1 - kept_modes = int(total_modes * self.hard_thresholding_fraction) - - st, end = total_modes - kept_modes, total_modes + kept_modes - - o1_real[:, st:end, :kept_modes] = F.relu( - paddle.einsum( - "xyzbi,bio->xyzbo", - x[:, st:end, :kept_modes].real(), - self.w1[0], - ) - - paddle.einsum( - "xyzbi,bio->xyzbo", - x[:, st:end, :kept_modes].imag(), - self.w1[1], - ) - + self.b1[0] - ) - - o1_imag[:, st:end, :kept_modes] = F.relu( - paddle.einsum( - "xyzbi,bio->xyzbo", - x[:, st:end, :kept_modes].imag(), - self.w1[0], - ) - + paddle.einsum( - "xyzbi,bio->xyzbo", - x[:, st:end, :kept_modes].real(), - self.w1[1], - ) - + self.b1[1] - ) - - o2_real[:, st:end, :kept_modes] = ( - paddle.einsum( - "xyzbi,bio->xyzbo", - o1_real[:, st:end, :kept_modes], - self.w2[0], - ) - - paddle.einsum( - "xyzbi,bio->xyzbo", - o1_imag[:, st:end, :kept_modes], - self.w2[1], - ) - + self.b2[0] - ) - - o2_imag[:, st:end, :kept_modes] = ( - paddle.einsum( - "xyzbi,bio->xyzbo", - o1_imag[:, st:end, :kept_modes], - self.w2[0], - ) - + paddle.einsum( - "xyzbi,bio->xyzbo", - o1_real[:, st:end, :kept_modes], - self.w2[1], - ) - + self.b2[1] - ) - - x = paddle.stack([o2_real, o2_imag], axis=-1) - x = F.softshrink(x, threshold=self.sparsity_threshold) - x = paddle.as_complex(x) - x = x.reshape((B, H, W // 2 + 1, C)) - x = paddle.fft.irfft2(x, s=(H, W), axes=(1, 2), norm="ortho") - - return x + bias - - -class Block(nn.Layer): - """AFNO network block. - - Args: - dim (int): The input tensor dimension. - mlp_ratio (float, optional): The ratio used in MLP. Defaults to 4.0. - drop (float, optional): The drop ratio used in MLP. Defaults to 0.0. - drop_path (float, optional): The drop ratio used in DropPath. Defaults to 0.0. - activation (str, optional): Name of activation function. Defaults to "gelu". - norm_layer (nn.Layer, optional): Class of norm layer. Defaults to nn.LayerNorm. - double_skip (bool, optional): Whether use double skip. Defaults to True. - num_blocks (int, optional): The number of blocks. Defaults to 8. - sparsity_threshold (float, optional): The value of threshold for softshrink. Defaults to 0.01. - hard_thresholding_fraction (float, optional): The value of threshold for keep mode. Defaults to 1.0. - """ - - def __init__( - self, - dim: int, - mlp_ratio: float = 4.0, - drop: float = 0.0, - drop_path: float = 0.0, - activation: str = "gelu", - norm_layer: nn.Layer = nn.LayerNorm, - double_skip: bool = True, - num_blocks: int = 8, - sparsity_threshold: float = 0.01, - hard_thresholding_fraction: float = 1.0, - ): - super().__init__() - self.norm1 = norm_layer(dim) - self.filter = AFNO2D( - dim, num_blocks, sparsity_threshold, hard_thresholding_fraction - ) - self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() - - self.norm2 = norm_layer(dim) - mlp_hidden_dim = int(dim * mlp_ratio) - self.mlp = MLP( - in_features=dim, - hidden_features=mlp_hidden_dim, - activation=activation, - drop=drop, - ) - self.double_skip = double_skip - - def forward(self, x): - residual = x - x = self.norm1(x) - x = self.filter(x) - - if self.double_skip: - x = x + residual - residual = x - - x = self.norm2(x) - x = self.mlp(x) - x = self.drop_path(x) - x = x + residual - return x - - -class PatchEmbed(nn.Layer): - """Patch embedding module. - - Args: - img_size (Tuple[int, ...], optional): Image size. Defaults to (224, 224). - patch_size (Tuple[int, ...], optional): Patch size. Defaults to (16, 16). - in_channels (int, optional): The input tensor channels. Defaults to 3. - embed_dim (int, optional): The output tensor channels. Defaults to 768. - """ - - def __init__( - self, - img_size: Tuple[int, ...] = (224, 224), - patch_size: Tuple[int, ...] = (16, 16), - in_channels: int = 3, - embed_dim: int = 768, - ): - super().__init__() - num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0]) - self.img_size = img_size - self.patch_size = patch_size - self.num_patches = num_patches - self.proj = nn.Conv2D( - in_channels, embed_dim, kernel_size=patch_size, stride=patch_size - ) - - def forward(self, x): - _, _, H, W = x.shape - if not (H == self.img_size[0] and W == self.img_size[1]): - raise ValueError( - f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})." - ) - x = self.proj(x).flatten(2).transpose((0, 2, 1)) - return x - - -class AFNONet(base.Arch): - """Adaptive Fourier Neural Network. - - Args: - input_keys (Tuple[str, ...]): Name of input keys, such as ("input",). - output_keys (Tuple[str, ...]): Name of output keys, such as ("output",). - img_size (Tuple[int, ...], optional): Image size. Defaults to (720, 1440). - patch_size (Tuple[int, ...], optional): Path. Defaults to (8, 8). - in_channels (int, optional): The input tensor channels. Defaults to 20. - out_channels (int, optional): The output tensor channels. Defaults to 20. - embed_dim (int, optional): The embedding dimension for PatchEmbed. Defaults to 768. - depth (int, optional): Number of transformer depth. Defaults to 12. - mlp_ratio (float, optional): Number of ratio used in MLP. Defaults to 4.0. - drop_rate (float, optional): The drop ratio used in MLP. Defaults to 0.0. - drop_path_rate (float, optional): The drop ratio used in DropPath. Defaults to 0.0. - num_blocks (int, optional): Number of blocks. Defaults to 8. - sparsity_threshold (float, optional): The value of threshold for softshrink. Defaults to 0.01. - hard_thresholding_fraction (float, optional): The value of threshold for keep mode. Defaults to 1.0. - num_timestamps (int, optional): Number of timestamp. Defaults to 1. - - Examples: - >>> import ppsci - >>> model = ppsci.arch.AFNONet(("input", ), ("output", )) - >>> input_data = {"input": paddle.randn([1, 20, 720, 1440])} - >>> output_data = model(input_data) - >>> for k, v in output_data.items(): - ... print(k, v.shape) - output [1, 20, 720, 1440] - """ - - def __init__( - self, - input_keys: Tuple[str, ...], - output_keys: Tuple[str, ...], - img_size: Tuple[int, ...] = (720, 1440), - patch_size: Tuple[int, ...] = (8, 8), - in_channels: int = 20, - out_channels: int = 20, - embed_dim: int = 768, - depth: int = 12, - mlp_ratio: float = 4.0, - drop_rate: float = 0.0, - drop_path_rate: float = 0.0, - num_blocks: int = 8, - sparsity_threshold: float = 0.01, - hard_thresholding_fraction: float = 1.0, - num_timestamps: int = 1, - ): - super().__init__() - self.input_keys = input_keys - self.output_keys = output_keys - - self.img_size = img_size - self.patch_size = patch_size - self.in_channels = in_channels - self.out_channels = out_channels - self.embed_dim = embed_dim - self.num_blocks = num_blocks - self.num_timestamps = num_timestamps - norm_layer = partial(nn.LayerNorm, epsilon=1e-6) - - self.patch_embed = PatchEmbed( - img_size=img_size, - patch_size=self.patch_size, - in_channels=self.in_channels, - embed_dim=embed_dim, - ) - num_patches = self.patch_embed.num_patches - - data = paddle.zeros((1, num_patches, embed_dim)) - data = initializer.trunc_normal_(data, std=0.02) - self.pos_embed = paddle.create_parameter( - shape=data.shape, - dtype=data.dtype, - default_initializer=nn.initializer.Assign(data), - ) - self.pos_drop = nn.Dropout(p=drop_rate) - - dpr = [x.item() for x in paddle.linspace(0, drop_path_rate, depth)] - - self.h = img_size[0] // self.patch_size[0] - self.w = img_size[1] // self.patch_size[1] - - self.blocks = nn.LayerList( - [ - Block( - dim=embed_dim, - mlp_ratio=mlp_ratio, - drop=drop_rate, - drop_path=dpr[i], - norm_layer=norm_layer, - num_blocks=self.num_blocks, - sparsity_threshold=sparsity_threshold, - hard_thresholding_fraction=hard_thresholding_fraction, - ) - for i in range(depth) - ] - ) - - self.norm = norm_layer(embed_dim) - self.head = nn.Linear( - embed_dim, - self.out_channels * self.patch_size[0] * self.patch_size[1], - bias_attr=False, - ) - - self.apply(self._init_weights) - - def _init_weights(self, m): - if isinstance(m, nn.Linear): - initializer.trunc_normal_(m.weight, std=0.02) - if m.bias is not None: - initializer.zeros_(m.bias) - elif isinstance(m, nn.LayerNorm): - initializer.ones_(m.weight) - initializer.zeros_(m.bias) - elif isinstance(m, nn.Conv2D): - initializer.conv_init_(m) - - def forward_tensor(self, x): - B = x.shape[0] - x = self.patch_embed(x) - x = x + self.pos_embed - x = self.pos_drop(x) - - x = x.reshape((B, self.h, self.w, self.embed_dim)) - for block in self.blocks: - x = block(x) - - x = self.head(x) - - b = x.shape[0] - p1 = self.patch_size[0] - p2 = self.patch_size[1] - h = self.img_size[0] // self.patch_size[0] - w = self.img_size[1] // self.patch_size[1] - c_out = x.shape[3] // (p1 * p2) - x = x.reshape((b, h, w, p1, p2, c_out)) - x = x.transpose((0, 5, 1, 3, 2, 4)) - x = x.reshape((b, c_out, h * p1, w * p2)) - - return x - - @staticmethod - def split_to_dict(data_tensors: Tuple[paddle.Tensor, ...], keys: Tuple[str, ...]): - return {key: data_tensors[i] for i, key in enumerate(keys)} - - def forward(self, x): - if self._input_transform is not None: - x = self._input_transform(x) - - x_tensor = self.concat_to_tensor(x, self.input_keys) - - y = [] - input = x_tensor - for _ in range(self.num_timestamps): - out = self.forward_tensor(input) - y.append(out) - input = out - y = self.split_to_dict(y, self.output_keys) - - if self._output_transform is not None: - y = self._output_transform(x, y) - return y - - -class PrecipNet(base.Arch): - """Precipitation Network. - - Args: - input_keys (Tuple[str, ...]): Name of input keys, such as ("input",). - output_keys (Tuple[str, ...]): Name of output keys, such as ("output",). - wind_model (base.Arch): Wind model. - img_size (Tuple[int, ...], optional): Image size. Defaults to (720, 1440). - patch_size (Tuple[int, ...], optional): Path. Defaults to (8, 8). - in_channels (int, optional): The input tensor channels. Defaults to 20. - out_channels (int, optional): The output tensor channels. Defaults to 1. - embed_dim (int, optional): The embedding dimension for PatchEmbed. Defaults to 768. - depth (int, optional): Number of transformer depth. Defaults to 12. - mlp_ratio (float, optional): Number of ratio used in MLP. Defaults to 4.0. - drop_rate (float, optional): The drop ratio used in MLP. Defaults to 0.0. - drop_path_rate (float, optional): The drop ratio used in DropPath. Defaults to 0.0. - num_blocks (int, optional): Number of blocks. Defaults to 8. - sparsity_threshold (float, optional): The value of threshold for softshrink. Defaults to 0.01. - hard_thresholding_fraction (float, optional): The value of threshold for keep mode. Defaults to 1.0. - num_timestamps (int, optional): Number of timestamp. Defaults to 1. - - Examples: - >>> import ppsci - >>> wind_model = ppsci.arch.AFNONet(("input", ), ("output", )) - >>> model = ppsci.arch.PrecipNet(("input", ), ("output", ), wind_model) - >>> data = paddle.randn([1, 20, 720, 1440]) - >>> data_dict = {"input": data} - >>> output = model.forward(data_dict) - >>> print(output['output'].shape) - [1, 1, 720, 1440] - """ - - def __init__( - self, - input_keys: Tuple[str, ...], - output_keys: Tuple[str, ...], - wind_model: base.Arch, - img_size: Tuple[int, ...] = (720, 1440), - patch_size: Tuple[int, ...] = (8, 8), - in_channels: int = 20, - out_channels: int = 1, - embed_dim: int = 768, - depth: int = 12, - mlp_ratio: float = 4.0, - drop_rate: float = 0.0, - drop_path_rate: float = 0.0, - num_blocks: int = 8, - sparsity_threshold: float = 0.01, - hard_thresholding_fraction: float = 1.0, - num_timestamps=1, - ): - super().__init__() - self.input_keys = input_keys - self.output_keys = output_keys - - self.img_size = img_size - self.patch_size = patch_size - self.in_channels = in_channels - self.out_channels = out_channels - self.embed_dim = embed_dim - self.num_blocks = num_blocks - self.num_timestamps = num_timestamps - self.backbone = AFNONet( - ("input",), - ("output",), - img_size=img_size, - patch_size=patch_size, - in_channels=in_channels, - out_channels=out_channels, - embed_dim=embed_dim, - depth=depth, - mlp_ratio=mlp_ratio, - drop_rate=drop_rate, - drop_path_rate=drop_path_rate, - num_blocks=num_blocks, - sparsity_threshold=sparsity_threshold, - hard_thresholding_fraction=hard_thresholding_fraction, - ) - self.ppad = PeriodicPad2d(1) - self.conv = nn.Conv2D( - self.out_channels, self.out_channels, kernel_size=3, stride=1, padding=0 - ) - self.act = nn.ReLU() - self.apply(self._init_weights) - self.wind_model = wind_model - self.wind_model.eval() - - def _init_weights(self, m): - if isinstance(m, nn.Linear): - initializer.trunc_normal_(m.weight, std=0.02) - if m.bias is not None: - initializer.zeros_(m.bias) - elif isinstance(m, nn.LayerNorm): - initializer.ones_(m.weight) - initializer.zeros_(m.bias) - elif isinstance(m, nn.Conv2D): - initializer.conv_init_(m) - - def forward_tensor(self, x): - x = self.backbone.forward_tensor(x) - x = self.ppad(x) - x = self.conv(x) - x = self.act(x) - return x - - @staticmethod - def split_to_dict(data_tensors: Tuple[paddle.Tensor, ...], keys: Tuple[str, ...]): - return {key: data_tensors[i] for i, key in enumerate(keys)} - - def forward(self, x): - if self._input_transform is not None: - x = self._input_transform(x) - - x_tensor = self.concat_to_tensor(x, self.input_keys) - - input_wind = x_tensor - y = [] - for _ in range(self.num_timestamps): - with paddle.no_grad(): - out_wind = self.wind_model.forward_tensor(input_wind) - out = self.forward_tensor(out_wind) - y.append(out) - input_wind = out_wind - y = self.split_to_dict(y, self.output_keys) - - if self._output_transform is not None: - y = self._output_transform(x, y) - return y diff --git a/examples/smc_reac/ppsci/arch/amgnet.py b/examples/smc_reac/ppsci/arch/amgnet.py deleted file mode 100644 index ce728317d6..0000000000 --- a/examples/smc_reac/ppsci/arch/amgnet.py +++ /dev/null @@ -1,649 +0,0 @@ -# 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. - -from __future__ import annotations - -import functools -from typing import Callable -from typing import Dict -from typing import List -from typing import Optional -from typing import Tuple - -import numpy as np -import paddle -import paddle.nn as nn -from typing_extensions import Literal - -try: - import pgl -except ModuleNotFoundError: - pass - -try: - import pyamg -except ModuleNotFoundError: - pass - -from paddle import sparse as pd_sparse -from scipy import sparse as sci_sparse - - -def _knn_interpolate( - features: paddle.Tensor, coarse_nodes: paddle.Tensor, fine_nodes: paddle.Tensor -) -> paddle.Tensor: - coarse_nodes_input = paddle.repeat_interleave( - coarse_nodes.unsqueeze(0), fine_nodes.shape[0], axis=0 - ) # [6684,352,2] - fine_nodes_input = paddle.repeat_interleave( - fine_nodes.unsqueeze(1), coarse_nodes.shape[0], axis=1 - ) # [6684,352,2] - dist_w = 1.0 / ( - paddle.norm(x=coarse_nodes_input - fine_nodes_input, p=2, axis=-1) + 1e-9 - ) # [6684,352] - knn_value, knn_index = paddle.topk(dist_w, k=3, largest=True) # [6684,3],[6684,3] - weight = knn_value.unsqueeze(-2) - features_input = features[knn_index] - output = paddle.bmm(weight, features_input).squeeze(-2) / paddle.sum( - knn_value, axis=-1, keepdim=True - ) - return output - - -def _get_corse_node(latent_graph: "pgl.Graph") -> paddle.Tensor: - row = latent_graph.edge_index[0].numpy() - col = latent_graph.edge_index[1].numpy() - data = paddle.ones(shape=[row.size]).numpy() - A = sci_sparse.coo_matrix((data, (row, col))).tocsr() - splitting = pyamg.classical.split.RS(A) - index = np.array(np.nonzero(splitting)) - b = paddle.to_tensor(index) - b = paddle.squeeze(b) - return b - - -def StAS( - index_A: paddle.Tensor, - value_A: paddle.Tensor, - index_S: paddle.Tensor, - value_S: paddle.Tensor, - N: int, - kN: int, - norm_layer: nn.Layer, -) -> Tuple[paddle.Tensor, paddle.Tensor]: - """ASAP: Adaptive Structure Aware Pooling for Learning Hierarchical Graph Representations. - Ranjan, E., Sanyal, S., Talukdar, P. (2020, April). AAAI(2020) - - Args: - index_A (paddle.Tensor): Indices of sparse matrix A. - value_A (paddle.Tensor): Values of sparse matrix A. - index_S (paddle.Tensor): Indices of sparse matrix S. - value_S (paddle.Tensor): Values of sparse matrix S. - N (int): Dimension N. - kN (int): Dimension kN. - norm_layer (nn.Layer): Normalization layer. - - Returns: - Tuple[paddle.Tensor, paddle.Tensor]: Indices and values of result matrix E. - """ - sp_x = pd_sparse.sparse_coo_tensor(index_A, value_A) - sp_x = pd_sparse.coalesce(sp_x) - index_A = sp_x.indices() - value_A = sp_x.values() - - sp_s = pd_sparse.sparse_coo_tensor(index_S, value_S) - sp_s = pd_sparse.coalesce(sp_s) - index_S = sp_s.indices() - value_S = sp_s.values() - - indices_A = index_A.numpy() - values_A = value_A.numpy() - coo_A = sci_sparse.coo_matrix( - (values_A, (indices_A[0], indices_A[1])), shape=(N, N) - ) - - indices_S = index_S.numpy() - values_S = value_S.numpy() - coo_S = sci_sparse.coo_matrix( - (values_S, (indices_S[0], indices_S[1])), shape=(N, kN) - ) - - ans = coo_A.dot(coo_S).tocoo() - row = paddle.to_tensor(ans.row) - col = paddle.to_tensor(ans.col) - index_B = paddle.stack([row, col], axis=0) - value_B = paddle.to_tensor(ans.data) - - indices_A = index_S - values_A = value_S - coo_A = pd_sparse.sparse_coo_tensor(indices_A, values_A) - out = pd_sparse.transpose(coo_A, [1, 0]) - index_St = out.indices() - value_St = out.values() - - sp_x = pd_sparse.sparse_coo_tensor(index_B, value_B) - sp_x = pd_sparse.coalesce(sp_x) - index_B = sp_x.indices() - value_B = sp_x.values() - - indices_A = index_St.numpy() - values_A = value_St.numpy() - coo_A = sci_sparse.coo_matrix( - (values_A, (indices_A[0], indices_A[1])), shape=(kN, N) - ) - - indices_S = index_B.numpy() - values_S = value_B.numpy() - coo_S = sci_sparse.coo_matrix( - (values_S, (indices_S[0], indices_S[1])), shape=(N, kN) - ) - - ans = coo_A.dot(coo_S).tocoo() - row = paddle.to_tensor(ans.row) - col = paddle.to_tensor(ans.col) - index_E = paddle.stack([row, col], axis=0) - value_E = paddle.to_tensor(ans.data) - - # index_E排序 - sp_x = pd_sparse.sparse_coo_tensor(index_E, value_E) - sp_x = pd_sparse.coalesce(sp_x) - index_E = sp_x.indices() - value_E = sp_x.values() - - return index_E.astype("int64"), value_E - - -def FillZeros( - index_E: paddle.Tensor, value_E: paddle.Tensor, standard_index, kN: int -) -> Tuple[paddle.Tensor, paddle.Tensor]: - shape = [kN, kN] - row_E = index_E[0] - col_E = index_E[1] - DenseMatrix_E = sci_sparse.coo_matrix( - (paddle.ones_like(value_E), (row_E, col_E)), shape - ).toarray() - - row_S = standard_index[0] - col_S = standard_index[1] - DenseMatrix_S = sci_sparse.coo_matrix( - (paddle.ones([row_S.shape[0]]), (row_S, col_S)), shape - ).toarray() - - diff = DenseMatrix_S - DenseMatrix_E - rows, cols = np.nonzero(diff) - rows = paddle.to_tensor(rows, dtype="int32") - cols = paddle.to_tensor(cols, dtype="int32") - index = paddle.stack([rows, cols], axis=0) - value = paddle.zeros([index.shape[1]]) - index_E = paddle.concat([index_E, index], axis=1) - value_E = paddle.concat([value_E, value], axis=-1) - - sp_x = pd_sparse.sparse_coo_tensor(index_E, value_E) - sp_x = pd_sparse.coalesce(sp_x) - index_E = sp_x.indices() - value_E = sp_x.values() - - return index_E.astype("int64"), value_E - - -def remove_self_loops( - edge_index: paddle.Tensor, edge_attr: Optional[paddle.Tensor] = None -) -> Tuple[paddle.Tensor, Optional[paddle.Tensor]]: - # remove self-loop - mask = edge_index[0] != edge_index[1] - mask = mask.tolist() - edge_index = edge_index.t() - edge_index = edge_index[mask] - edge_index = edge_index.t() - if edge_attr is None: - return edge_index, None - else: - return edge_index, edge_attr[mask] - - -def faster_graph_connectivity(perm, edge_index, edge_weight, score, pos, N, norm_layer): - """ - Adapted from Ranjan, E., Sanyal, S., Talukdar, P. (2020, April). Asap: Adaptive structure aware pooling - for learning hierarchical graph representations. AAAI(2020) - """ - - kN = perm.shape[0] - perm2 = perm.reshape((-1, 1)) - mask = (edge_index[0] == perm2).sum(axis=0).astype("bool") - - S0 = edge_index[1][mask].reshape((1, -1)) - S1 = edge_index[0][mask].reshape((1, -1)) - index_S = paddle.concat([S0, S1], axis=0) - value_S = score[mask].detach().squeeze() - n_idx = paddle.zeros([N], dtype=paddle.int64) - n_idx[perm] = paddle.arange(perm.shape[0]) - index_S = index_S.astype("int64") - index_S[1] = n_idx[index_S[1]] - subgraphnode_pos = pos[perm] - index_A = edge_index.clone() - if edge_weight is None: - value_A = value_S.new_ones(edge_index[0].shape[0]) - else: - value_A = edge_weight.clone() - - value_A = paddle.squeeze(value_A) - model_1 = nn.Sequential( - ("l1", nn.Linear(128, 256)), - ("act1", nn.ReLU()), - ("l2", nn.Linear(256, 256)), - ("act2", nn.ReLU()), - ("l4", nn.Linear(256, 128)), - ("act4", nn.ReLU()), - ("l5", nn.Linear(128, 1)), - ) - model_2 = nn.Sequential( - ("l1", nn.Linear(1, 64)), - ("act1", nn.ReLU()), - ("l2", nn.Linear(64, 128)), - ("act2", nn.ReLU()), - ("l4", nn.Linear(128, 128)), - ) - - val_A = model_1(value_A) - val_A = paddle.squeeze(val_A) - index_E, value_E = StAS(index_A, val_A, index_S, value_S, N, kN, norm_layer) - value_E = paddle.reshape(value_E, shape=[-1, 1]) - edge_weight = model_2(value_E) - - return index_E, edge_weight, subgraphnode_pos - - -def norm_graph_connectivity(perm, edge_index, edge_weight, score, pos, N, norm_layer): - """ - Come from Ranjan, E., Sanyal, S., Talukdar, P. (2020, April). Asap: Adaptive - structure aware pooling for learning hierarchical graph representations. AAAI(2020) - """ - - kN = perm.shape[0] - perm2 = perm.reshape((-1, 1)) - mask = (edge_index[0] == perm2).sum(axis=0).astype("bool") - S0 = edge_index[1][mask].reshape((1, -1)) - S1 = edge_index[0][mask].reshape((1, -1)) - - index_S = paddle.concat([S0, S1], axis=0) - value_S = score[mask].detach().squeeze() - n_idx = paddle.zeros([N], dtype=paddle.int64) - n_idx[perm] = paddle.arange(perm.shape[0]) - - index_S = index_S.astype("int64") - index_S[1] = n_idx[index_S[1]] - subgraphnode_pos = pos[perm] - index_A = edge_index.clone() - - if edge_weight is None: - value_A = value_S.new_ones(edge_index[0].shape[0]) - else: - value_A = edge_weight.clone() - - value_A = paddle.squeeze(value_A) - eps_mask = (value_S == 0).astype(paddle.get_default_dtype()) - value_S = paddle.full_like(value_S, 1e-4) * eps_mask + (1 - eps_mask) * value_S - attrlist = [] - standard_index, _ = StAS( - index_A, - paddle.ones_like(value_A[:, 0]), - index_S, - paddle.ones_like(value_S), - N, - kN, - norm_layer, - ) - for i in range(128): - mask = (value_A[:, i] == 0).astype(paddle.get_default_dtype()) - val_A = paddle.full_like(mask, 1e-4) * mask + (1 - mask) * value_A[:, i] - index_E, value_E = StAS(index_A, val_A, index_S, value_S, N, kN, norm_layer) - - if index_E.shape[1] != standard_index.shape[1]: - index_E, value_E = FillZeros(index_E, value_E, standard_index, kN) - - index_E, value_E = remove_self_loops(edge_index=index_E, edge_attr=value_E) - attrlist.append(value_E) - edge_weight = paddle.stack(attrlist, axis=1) - - return index_E, edge_weight, subgraphnode_pos - - -class GraphNetBlock(nn.Layer): - """Multi-Edge Interaction Network with residual connections.""" - - def __init__( - self, model_fn, output_dim, message_passing_aggregator, attention=False - ): - super().__init__() - self.edge_model = model_fn(output_dim, 384) - self.node_model = model_fn(output_dim, 256) - self.message_passing_aggregator = message_passing_aggregator - - def _update_edge_features(self, graph): - """Aggregates node features, and applies edge function.""" - senders = graph.edge_index[0] - receivers = graph.edge_index[1] - sender_features = paddle.index_select(x=graph.x, index=senders, axis=0) - receiver_features = paddle.index_select(x=graph.x, index=receivers, axis=0) - features = [sender_features, receiver_features, graph.edge_attr] - features = paddle.concat(features, axis=-1) - return self.edge_model(features) - - def unsorted_segment_operation(self, data, segment_ids, num_segments, operation): - """Computes the sum along segments of a tensor. Analogous to tf.unsorted_segment_sum. - - Args: - data (paddle.Tensor): A tensor whose segments are to be summed. - segment_ids (paddle.Tensor): The segment indices tensor. - num_segments (int): The number of segments. - operation (str): _description_ - - Returns: - paddle.Tensor: A tensor of same data type as the data argument. - """ - if not all([i in data.shape for i in segment_ids.shape]): - raise ValueError("segment_ids.shape should be a prefix of data.shape") - - if not (data.shape[0] == segment_ids.shape[0]): - raise ValueError("data.shape and segment_ids.shape should be equal") - - shape = [num_segments] + list(data.shape[1:]) - result_shape = paddle.zeros(shape) - if operation == "sum": - result = paddle.scatter(result_shape, segment_ids, data, overwrite=False) - return result - - def _update_node_features(self, node_features, edge_attr, edge_index): - """Aggregates edge features, and applies node function.""" - num_nodes = node_features.shape[0] - features = [node_features] - features.append( - self.unsorted_segment_operation( - edge_attr, - edge_index[1], - num_nodes, - operation=self.message_passing_aggregator, - ) - ) - features = paddle.concat(features, axis=-1) - return self.node_model(features) - - def forward(self, graph): - """Applies GraphNetBlock and returns updated MultiGraph.""" - new_edge_features = self._update_edge_features(graph) - new_node_features = self._update_node_features( - graph.x, graph.edge_attr, graph.edge_index - ) - - new_node_features += graph.x - new_edge_features += graph.edge_attr - latent_graph = pgl.Graph( - num_nodes=new_node_features.shape[0], edges=graph.edge_index - ) - latent_graph.x = new_node_features - latent_graph.edge_attr = new_edge_features - latent_graph.pos = graph.pos - latent_graph.edge_index = graph.edge_index - return latent_graph - - -class Processor(nn.Layer): - """This class takes the nodes with the most influential feature (sum of square) - The the chosen numbers of nodes in each ripple will establish connection(features and distances) with the most influential nodes and this connection will be learned - Then the result is add to output latent graph of encoder and the modified latent graph will be feed into original processor - - Args: - make_mlp (Callable): Function to make MLP. - output_dim (int): Number of dimension of output. - message_passing_steps (int): Message passing steps. - message_passing_aggregator (str): Message passing aggregator. - attention (bool, optional): Whether use attention. Defaults to False. - use_stochastic_message_passing (bool, optional): Whether use stochastic message passing. Defaults to False. - """ - - # Each mesh can be coarsened to have no fewer points than this value - min_nodes = 2000 - - def __init__( - self, - make_mlp: Callable, - output_dim: int, - message_passing_steps: int, - message_passing_aggregator: str, - attention: bool = False, - use_stochastic_message_passing: bool = False, - ): - super().__init__() - self.use_stochastic_message_passing = use_stochastic_message_passing - self.graphnet_blocks = nn.LayerList() - self.cofe_edge_blocks = nn.LayerList() - self.pool_blocks = nn.LayerList() - self.latent_dim = output_dim - self.normalization = nn.LayerNorm(128) - for index in range(message_passing_steps): - self.graphnet_blocks.append( - GraphNetBlock( - model_fn=make_mlp, - output_dim=output_dim, - message_passing_aggregator=message_passing_aggregator, - attention=attention, - ) - ) - - self.pool_blocks.append( - GraphNetBlock( - model_fn=make_mlp, - output_dim=output_dim, - message_passing_aggregator=message_passing_aggregator, - attention=attention, - ) - ) - - def forward(self, latent_graph, speed, normalized_adj_mat=None): - x = [] - pos = [] - new = [] - for graphnet_block, pool in zip(self.graphnet_blocks, self.pool_blocks): - if latent_graph.x.shape[0] > self.min_nodes: - pre_matrix = graphnet_block(latent_graph) - x.append(pre_matrix) - cofe_graph = pool(pre_matrix) - coarsenodes = _get_corse_node(pre_matrix) - nodesfeatures = cofe_graph.x[coarsenodes] - if speed == "fast": - subedge_index, edge_weight, subpos = faster_graph_connectivity( - perm=coarsenodes, - edge_index=cofe_graph.edge_index, - edge_weight=cofe_graph.edge_attr, - score=cofe_graph.edge_attr[:, 0], - pos=cofe_graph.pos, - N=cofe_graph.x.shape[0], - norm_layer=self.normalization, - ) - elif speed == "norm": - subedge_index, edge_weight, subpos = norm_graph_connectivity( - perm=coarsenodes, - edge_index=cofe_graph.edge_index, - edge_weight=cofe_graph.edge_attr, - score=cofe_graph.edge_attr[:, 0], - pos=cofe_graph.pos, - N=cofe_graph.x.shape[0], - norm_layer=self.normalization, - ) - else: - raise ValueError( - f"Argument 'speed' should be 'sum' or 'fast', bot got {speed}." - ) - edge_weight = self.normalization(edge_weight) - pos.append(subpos) - latent_graph = pgl.Graph( - num_nodes=nodesfeatures.shape[0], edges=subedge_index - ) - latent_graph.x = nodesfeatures - latent_graph.edge_attr = edge_weight - latent_graph.pos = subpos - latent_graph.edge_index = subedge_index - else: - latent_graph = graphnet_block(latent_graph) - new.append(latent_graph) - if len(new): - x.append(new[-1]) - return x, pos - - -class FullyConnectedLayer(nn.Layer): - def __init__(self, input_dim: int, hidden_size: Tuple[int, ...]): - super(FullyConnectedLayer, self).__init__() - num_layers = len(hidden_size) - self._layers_ordered_dict = {} - self.in_dim = input_dim - for index, output_dim in enumerate(hidden_size): - self._layers_ordered_dict["linear_" + str(index)] = nn.Linear( - self.in_dim, output_dim - ) - if index < (num_layers - 1): - self._layers_ordered_dict["relu_" + str(index)] = nn.ReLU() - self.in_dim = output_dim - - self.layers = nn.LayerDict(self._layers_ordered_dict) - - def forward(self, input): - for key in self.layers: - layer = self.layers[key] - output = layer(input) - input = output - return input - - -class Encoder(nn.Layer): - """Encodes node and edge features into latent features.""" - - def __init__(self, input_dim, make_mlp, latent_dim): - super(Encoder, self).__init__() - self._make_mlp = make_mlp - self._latent_dim = latent_dim - self.node_model = self._make_mlp(latent_dim, input_dim=input_dim) - self.mesh_edge_model = self._make_mlp(latent_dim, input_dim=1) - - def forward(self, graph): - node_latents = self.node_model(graph.x) - edge_latent = self.mesh_edge_model(graph.edge_attr) - - graph.x = node_latents - graph.edge_attr = edge_latent - return graph - - -class Decoder(nn.Layer): - """Decodes node features from graph. - Encodes node and edge features into latent features. - """ - - def __init__(self, make_mlp, output_dim): - super(Decoder, self).__init__() - self.model = make_mlp(output_dim, 128) - - def forward(self, node_features): - return self.model(node_features) - - -class AMGNet(nn.Layer): - """A Multi-scale Graph neural Network model - based on Encoder-Process-Decoder structure for flow field prediction. - - https://doi.org/10.1080/09540091.2022.2131737 - - Code reference: https://github.com/baoshiaijhin/amgnet - - Args: - input_keys (Tuple[str, ...]): Name of input keys, such as ("input", ). - output_keys (Tuple[str, ...]): Name of output keys, such as ("pred", ). - input_dim (int): Number of input dimension. - output_dim (int): Number of output dimension. - latent_dim (int): Number of hidden(feature) dimension. - num_layers (int): Number of layer(s). - message_passing_aggregator (Literal["sum"]): Message aggregator method in graph. - Only "sum" available now. - message_passing_steps (int): Message passing steps in graph. - speed (str): Whether use vanilla method or fast method for graph_connectivity - computation. - - Examples: - >>> import ppsci - >>> model = ppsci.arch.AMGNet( - ... ("input", ), ("pred", ), 5, 3, 64, 2, "sum", 6, "norm", - ... ) - """ - - def __init__( - self, - input_keys: Tuple[str, ...], - output_keys: Tuple[str, ...], - input_dim: int, - output_dim: int, - latent_dim: int, - num_layers: int, - message_passing_aggregator: Literal["sum"], - message_passing_steps: int, - speed: Literal["norm", "fast"], - ): - super().__init__() - self.input_keys = input_keys - self.output_keys = output_keys - self._latent_dim = latent_dim - self.speed = speed - self._output_dim = output_dim - self._num_layers = num_layers - - self.encoder = Encoder(input_dim, self._make_mlp, latent_dim=self._latent_dim) - self.processor = Processor( - make_mlp=self._make_mlp, - output_dim=self._latent_dim, - message_passing_steps=message_passing_steps, - message_passing_aggregator=message_passing_aggregator, - use_stochastic_message_passing=False, - ) - self.post_processor = self._make_mlp(self._latent_dim, 128) - self.decoder = Decoder( - make_mlp=functools.partial(self._make_mlp, layer_norm=False), - output_dim=self._output_dim, - ) - - def forward(self, x: Dict[str, "pgl.Graph"]) -> Dict[str, paddle.Tensor]: - graphs = x[self.input_keys[0]] - latent_graph = self.encoder(graphs) - x, p = self.processor(latent_graph, speed=self.speed) - node_features = self._spa_compute(x, p) - pred_field = self.decoder(node_features) - return {self.output_keys[0]: pred_field} - - def _make_mlp(self, output_dim: int, input_dim: int = 5, layer_norm: bool = True): - widths = (self._latent_dim,) * self._num_layers + (output_dim,) - network = FullyConnectedLayer(input_dim, widths) - if layer_norm: - network = nn.Sequential(network, nn.LayerNorm(normalized_shape=widths[-1])) - return network - - def _spa_compute(self, x: List["pgl.Graph"], p): - j = len(x) - 1 - node_features = x[j].x - - for k in range(1, j + 1): - pos = p[-k] - fine_nodes = x[-(k + 1)].pos - feature = _knn_interpolate(node_features, pos, fine_nodes) - node_features = x[-(k + 1)].x + feature - node_features = self.post_processor(node_features) - - return node_features diff --git a/examples/smc_reac/ppsci/arch/base.py b/examples/smc_reac/ppsci/arch/base.py deleted file mode 100644 index 5b51efec22..0000000000 --- a/examples/smc_reac/ppsci/arch/base.py +++ /dev/null @@ -1,279 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import Callable -from typing import Dict -from typing import Tuple - -import numpy as np -import paddle -from paddle import nn - -from ppsci.utils import logger - - -class Arch(nn.Layer): - """Base class for Network.""" - - input_keys: Tuple[str, ...] - output_keys: Tuple[str, ...] - - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - self._input_transform: Callable[ - [Dict[str, paddle.Tensor]], Dict[str, paddle.Tensor] - ] = None - - self._output_transform: Callable[ - [Dict[str, paddle.Tensor], Dict[str, paddle.Tensor]], - Dict[str, paddle.Tensor], - ] = None - - def forward(self, *args, **kwargs): - raise NotImplementedError("Arch.forward is not implemented") - - @property - def num_params(self) -> int: - """Return number of parameters within network. - - Returns: - int: Number of parameters. - """ - num = 0 - for name, param in self.named_parameters(): - if hasattr(param, "shape"): - num += np.prod(list(param.shape), dtype="int") - else: - logger.warning(f"{name} has no attribute 'shape'") - return num - - @property - def num_buffers(self) -> int: - """Return number of buffers within network. - - Returns: - int: Number of buffers. - """ - num = 0 - for name, buffer in self.named_buffers(): - if hasattr(buffer, "shape"): - num += np.prod(list(buffer.shape), dtype="int") - else: - logger.warning(f"{name} has no attribute 'shape'") - return num - - @staticmethod - def concat_to_tensor( - data_dict: Dict[str, paddle.Tensor], keys: Tuple[str, ...], axis=-1 - ) -> Tuple[paddle.Tensor, ...]: - """Concatenate tensors from dict in the order of given keys. - - Args: - data_dict (Dict[str, paddle.Tensor]): Dict contains tensor. - keys (Tuple[str, ...]): Keys tensor fetched from. - axis (int, optional): Axis concatenate at. Defaults to -1. - - Returns: - Tuple[paddle.Tensor, ...]: Concatenated tensor. - - Examples: - >>> import paddle - >>> import ppsci - >>> model = ppsci.arch.Arch() - >>> # fetch one tensor - >>> out = model.concat_to_tensor({'x':paddle.rand([64, 64, 1])}, ('x',)) - >>> print(out.dtype, out.shape) - paddle.float32 [64, 64, 1] - >>> # fetch more tensors - >>> out = model.concat_to_tensor( - ... {'x1':paddle.rand([64, 64, 1]), 'x2':paddle.rand([64, 64, 1])}, - ... ('x1', 'x2'), - ... axis=2) - >>> print(out.dtype, out.shape) - paddle.float32 [64, 64, 2] - - """ - if len(keys) == 1: - return data_dict[keys[0]] - data = [data_dict[key] for key in keys] - return paddle.concat(data, axis) - - @staticmethod - def split_to_dict( - data_tensor: paddle.Tensor, keys: Tuple[str, ...], axis=-1 - ) -> Dict[str, paddle.Tensor]: - """Split tensor and wrap into a dict by given keys. - - Args: - data_tensor (paddle.Tensor): Tensor to be split. - keys (Tuple[str, ...]): Keys tensor mapping to. - axis (int, optional): Axis split at. Defaults to -1. - - Returns: - Dict[str, paddle.Tensor]: Dict contains tensor. - - Examples: - >>> import paddle - >>> import ppsci - >>> model = ppsci.arch.Arch() - >>> # split one tensor - >>> out = model.split_to_dict(paddle.rand([64, 64, 1]), ('x',)) - >>> for k, v in out.items(): - ... print(f"{k} {v.dtype} {v.shape}") - x paddle.float32 [64, 64, 1] - >>> # split more tensors - >>> out = model.split_to_dict(paddle.rand([64, 64, 2]), ('x1', 'x2'), axis=2) - >>> for k, v in out.items(): - ... print(f"{k} {v.dtype} {v.shape}") - x1 paddle.float32 [64, 64, 1] - x2 paddle.float32 [64, 64, 1] - - """ - if len(keys) == 1: - return {keys[0]: data_tensor} - data = paddle.split(data_tensor, len(keys), axis=axis) - return {key: data[i] for i, key in enumerate(keys)} - - def register_input_transform( - self, - transform: Callable[[Dict[str, paddle.Tensor]], Dict[str, paddle.Tensor]], - ): - """Register input transform. - - Args: - transform (Callable[[Dict[str, paddle.Tensor]], Dict[str, paddle.Tensor]]): - Input transform of network, receive a single tensor dict and return a single tensor dict. - - Examples: - >>> import ppsci - >>> def transform_in(in_): - ... x = in_["x"] - ... # transform input - ... x_ = 2.0 * x - ... input_trans = {"2x": x_} - ... return input_trans - >>> # `MLP` inherits from `Arch` - >>> model = ppsci.arch.MLP( - ... input_keys=("2x",), - ... output_keys=("y",), - ... num_layers=5, - ... hidden_size=32) - >>> model.register_input_transform(transform_in) - >>> out = model({"x":paddle.rand([64, 64, 1])}) - >>> for k, v in out.items(): - ... print(f"{k} {v.dtype} {v.shape}") - y paddle.float32 [64, 64, 1] - - """ - self._input_transform = transform - - def register_output_transform( - self, - transform: Callable[ - [Dict[str, paddle.Tensor], Dict[str, paddle.Tensor]], - Dict[str, paddle.Tensor], - ], - ): - """Register output transform. - - Args: - transform (Callable[[Dict[str, paddle.Tensor], Dict[str, paddle.Tensor]], Dict[str, paddle.Tensor]]): - Output transform of network, receive two single tensor dict(raw input - and raw output) and return a single tensor dict(transformed output). - - Examples: - >>> import ppsci - >>> def transform_out(in_, out): - ... x = in_["x"] - ... y = out["y"] - ... u = 2.0 * x * y - ... output_trans = {"u": u} - ... return output_trans - >>> # `MLP` inherits from `Arch` - >>> model = ppsci.arch.MLP( - ... input_keys=("x",), - ... output_keys=("y",), - ... num_layers=5, - ... hidden_size=32) - >>> model.register_output_transform(transform_out) - >>> out = model({"x":paddle.rand([64, 64, 1])}) - >>> for k, v in out.items(): - ... print(f"{k} {v.dtype} {v.shape}") - u paddle.float32 [64, 64, 1] - - """ - self._output_transform = transform - - def freeze(self): - """Freeze all parameters. - - Examples: - >>> import ppsci - >>> model = ppsci.arch.Arch() - >>> # freeze all parameters and make model `eval` - >>> model.freeze() - >>> assert not model.training - >>> for p in model.parameters(): - ... assert p.stop_gradient - - """ - for param in self.parameters(): - param.stop_gradient = True - - self.eval() - - def unfreeze(self): - """Unfreeze all parameters. - - Examples: - >>> import ppsci - >>> model = ppsci.arch.Arch() - >>> # unfreeze all parameters and make model `train` - >>> model.unfreeze() - >>> assert model.training - >>> for p in model.parameters(): - ... assert not p.stop_gradient - - """ - for param in self.parameters(): - param.stop_gradient = False - - self.train() - - def __str__(self): - num_fc = 0 - num_conv = 0 - num_bn = 0 - for layer in self.sublayers(include_self=True): - if isinstance(layer, nn.Linear): - num_fc += 1 - elif isinstance(layer, (nn.Conv2D, nn.Conv3D, nn.Conv1D)): - num_conv += 1 - elif isinstance(layer, (nn.BatchNorm, nn.BatchNorm2D, nn.BatchNorm3D)): - num_bn += 1 - - return ", ".join( - [ - self.__class__.__name__, - f"input_keys = {self.input_keys}", - f"output_keys = {self.output_keys}", - f"num_fc = {num_fc}", - f"num_conv = {num_conv}", - f"num_bn = {num_bn}", - f"num_params = {self.num_params}", - f"num_buffers = {self.num_buffers}", - ] - ) diff --git a/examples/smc_reac/ppsci/arch/cfdgcn.py b/examples/smc_reac/ppsci/arch/cfdgcn.py deleted file mode 100644 index da43d3dc8f..0000000000 --- a/examples/smc_reac/ppsci/arch/cfdgcn.py +++ /dev/null @@ -1,350 +0,0 @@ -# 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. - -import os -from typing import Callable -from typing import Dict -from typing import List -from typing import Optional -from typing import Sequence -from typing import Tuple -from typing import TypeVar -from typing import Union - -import numpy as np -import paddle -from paddle import nn -from paddle.nn import functional as F - -from ppsci.data.dataset import airfoil_dataset - -try: - import pgl -except ModuleNotFoundError: - pass - -GenTensor = TypeVar("GenTensor", paddle.Tensor, np.ndarray) - -SU2_SHAPE_IDS = { - "line": 3, - "triangle": 5, - "quad": 9, -} - - -def _knn_interpolate( - features: paddle.Tensor, coarse_nodes: paddle.Tensor, fine_nodes: paddle.Tensor -) -> paddle.Tensor: - coarse_nodes_input = paddle.repeat_interleave( - coarse_nodes.unsqueeze(0), fine_nodes.shape[0], axis=0 - ) # [6684,352,2] - fine_nodes_input = paddle.repeat_interleave( - fine_nodes.unsqueeze(1), coarse_nodes.shape[0], axis=1 - ) # [6684,352,2] - dist_w = 1.0 / ( - paddle.norm(x=coarse_nodes_input - fine_nodes_input, p=2, axis=-1) + 1e-9 - ) # [6684,352] - knn_value, knn_index = paddle.topk(dist_w, k=3, largest=True) # [6684,3],[6684,3] - weight = knn_value.unsqueeze(-2) - features_input = features[knn_index] - output = paddle.bmm(weight, features_input).squeeze(-2) / paddle.sum( - knn_value, axis=-1, keepdim=True - ) - return output - - -def is_cw( - points: paddle.Tensor, triangles: paddle.Tensor, ret_val=False -) -> Union[bool, paddle.Tensor]: - tri_pts = points[triangles] - a = tri_pts[:, 0] - tri_pts[:, 1] - b = tri_pts[:, 1] - tri_pts[:, 2] - cross = b[:, 0] * a[:, 1] - b[:, 1] * a[:, 0] - - if not ret_val: - return cross > 0 - else: - return cross - - -def left_orthogonal(v: paddle.Tensor) -> paddle.Tensor: - return paddle.stack([-v[..., 1], v[..., 0]], axis=-1) - - -def signed_dist_graph( - nodes: paddle.Tensor, marker_inds, with_sign=False -) -> paddle.Tensor: - # assumes shape is convex - # approximate signed distance by distance to closest point on surface - signed_dists = paddle.zeros([nodes.shape[0]], dtype=paddle.float32) - marker_nodes = nodes[marker_inds] - if type(marker_inds) is paddle.Tensor: - marker_inds = marker_inds.tolist() - marker_inds = set(marker_inds) - - if with_sign: - marker_surfaces = marker_nodes[:-1] - marker_nodes[1:] - last_surface = marker_nodes[-1] - marker_nodes[0] - marker_surfaces = paddle.concat([marker_surfaces, last_surface.unsqueeze(0)]) - normals = left_orthogonal(marker_surfaces) / marker_surfaces.norm( - axis=1 - ).unsqueeze(1) - for i, x in enumerate(nodes): - if i not in marker_inds: - vecs = marker_nodes - x - dists = paddle.linalg.norm(vecs, axis=1) - min_dist = dists.min() - - if with_sign: - # if sign is requested, check if inside marker shape - # dot product with normals to find if inside shape - surface_dists = (vecs * normals).sum(axis=1) - if (surface_dists < 0).unique().shape[0] == 1: - # if all point in same direction it is inside - min_dist *= -1 - - signed_dists[i] = min_dist - return signed_dists - - -def quad2tri(elems: np.array) -> Tuple[List[int], Union[List[int], paddle.Tensor]]: - new_elems = [] - new_edges = [] - for e in elems: - if len(e) <= 3: - new_elems.append(e) - else: - new_elems.append([e[0], e[1], e[2]]) - new_elems.append([e[0], e[2], e[3]]) - new_edges.append(paddle.to_tensor([[e[0]], [e[2]]], dtype=paddle.int64)) - new_edges = ( - paddle.concat(new_edges, axis=1) - if new_edges - else paddle.to_tensor([], dtype=paddle.int64) - ) - return new_elems, new_edges - - -def write_graph_mesh( - output_filename: str, - points: GenTensor, - elems_list: Sequence[Sequence[Sequence[int]]], - marker_dict: Dict[str, Sequence[Sequence[int]]], - dims: int = 2, -) -> None: - def seq2str(s: Sequence[int]) -> str: - return " ".join(str(x) for x in s) - - with open(output_filename, "w") as f: - f.write(f"NDIME={dims}\n") - - num_points = points.shape[0] - f.write(f"NPOIN={num_points}\n") - for i, p in enumerate(points): - f.write(f"{seq2str(p.tolist())} {i}\n") - f.write("\n") - - num_elems = sum([len(elems) for elems in elems_list]) - f.write(f"NELEM={num_elems}\n") - for elems in elems_list: - for e in elems: - if len(e) != 3 and len(e) != 4: - raise ValueError( - f"Meshes only support triangles and quadrilaterals, " - f"passed element had {len(e)} vertices." - ) - elem_id = ( - SU2_SHAPE_IDS["triangle"] if len(e) == 3 else SU2_SHAPE_IDS["quad"] - ) - f.write(f"{elem_id} {seq2str(e)}\n") - f.write("\n") - - num_markers = len(marker_dict) - f.write(f"NMARK={num_markers}\n") - for marker_tag in marker_dict: - f.write(f"MARKER_TAG={marker_tag}\n") - marker_elems = marker_dict[marker_tag] - f.write(f"MARKER_ELEMS={len(marker_elems)}\n") - for m in marker_elems: - f.write(f'{SU2_SHAPE_IDS["line"]} {seq2str(m)}\n') - f.write("\n") - - -class CFDGCN(nn.Layer): - """Graph Neural Networks for Fluid Flow Prediction. - - [Filipe De Avila Belbute-Peres, Thomas Economon, Zico Kolter Proceedings of the 37th International Conference on Machine Learning, PMLR 119:2402-2411, 2020.](https://proceedings.mlr.press/v119/de-avila-belbute-peres20a.html) - - Code reference: https://github.com/locuslab/cfd-gcn - - Args: - input_keys (Tuple[str, ...]): Name of input keys, such as ("input", ). - output_keys (Tuple[str, ...]): Name of output keys, such as ("pred", ). - config_file (str): Name of configuration file for su2 module. - coarse_mesh (str): Path of coarse mesh file. - fine_marker_dict (Dict[str, List[List[int]]]): Dict of fine marker. - process_sim (Callable, optional): Preprocess function. Defaults to `lambda x, y: x`. - freeze_mesh (bool, optional): Whether set `stop_gradient=True` for nodes. Defaults to False. - num_convs (int, optional): Number of conv layers. Defaults to 6. - num_end_convs (int, optional): Number of end conv layers. Defaults to 3. - hidden_channel (int, optional): Number of channels of hidden layer. Defaults to 512. - out_channel (int, optional): Number of channels of output. Defaults to 3. - su2_module (Optional[Callable]): SU2Module Object. Defaults to None. - - Examples: - >>> import ppsci - >>> import su2paddle # doctest: +SKIP - >>> model = ppsci.arch.CFDGCN( - ... input_keys=("input"), - ... output_keys=("pred"), - ... config_file="/path/to/file.cfg", - ... coarse_mesh="/path/to/file.su2", - ... process_sim=None, - ... freeze_mesh=False, - ... num_convs=6, - ... num_end_convs=3, - ... hidden_channel=512, - ... out_channel=3, - ... su2_module=su2paddle.SU2Module, - ... ) # doctest: +SKIP - """ - - def __init__( - self, - input_keys: Tuple[str, ...], - output_keys: Tuple[str, ...], - config_file: str, - coarse_mesh: str, - fine_marker_dict: Dict[str, List[List[int]]], - process_sim: Callable = lambda x, y: x, - freeze_mesh: bool = False, - num_convs: int = 6, - num_end_convs: int = 3, - hidden_channel: int = 512, - out_channel: int = 3, - su2_module: Optional[Callable] = None, - ): - - super().__init__() - self.input_keys = input_keys - self.output_keys = output_keys - meshes_temp_dir = "temp_meshes" - os.makedirs(meshes_temp_dir, exist_ok=True) - self.mesh_file = os.path.join(meshes_temp_dir, f"{str(os.getpid())}_mesh.su2") - - if not coarse_mesh: - raise ValueError("Need to provide a coarse mesh for CFD-GCN.") - nodes, edges, self.elems, self.marker_dict = airfoil_dataset._get_mesh_graph( - coarse_mesh - ) - if not freeze_mesh: - self.nodes = paddle.to_tensor(nodes, stop_gradient=False) - else: - self.nodes = paddle.to_tensor(nodes, stop_gradient=True) - - self.elems, new_edges = quad2tri(sum(self.elems, [])) - self.elems = [self.elems] - - if is_cw(self.nodes, paddle.to_tensor(self.elems[0])).nonzero().shape[0] != 0: - raise ("Mesh has flipped elems.") - - self.edges = paddle.to_tensor(edges) - self.edges = paddle.concat([self.edges, new_edges], axis=1) - self.marker_inds = paddle.to_tensor(sum(self.marker_dict.values(), [])).unique() - self.fine_marker_dict = paddle.to_tensor(fine_marker_dict["airfoil"]).unique() - self.process_sim = process_sim - - self.write_mesh_file( - self.nodes, self.elems, self.marker_dict, filename=self.mesh_file - ) - self.su2 = su2_module(config_file, mesh_file=self.mesh_file) - self.sdf = None - - self.num_convs = num_end_convs - self.convs = [] - if self.num_convs > 0: - self.convs = nn.LayerList() - in_channels = out_channel + hidden_channel - for i in range(self.num_convs - 1): - self.convs.append(pgl.nn.GCNConv(in_channels, hidden_channel)) - in_channels = hidden_channel - self.convs.append(pgl.nn.GCNConv(in_channels, out_channel)) - - self.num_pre_convs = num_convs - num_end_convs - self.pre_convs = [] - if self.num_pre_convs > 0: - in_channels = 5 + 1 # one extra channel for sdf - self.pre_convs = nn.LayerList() - for i in range(self.num_pre_convs - 1): - self.pre_convs.append(pgl.nn.GCNConv(in_channels, hidden_channel)) - in_channels = hidden_channel - self.pre_convs.append(pgl.nn.GCNConv(in_channels, hidden_channel)) - - def forward(self, x: Dict[str, "pgl.Graph"]) -> Dict[str, paddle.Tensor]: - graph = x[self.input_keys[0]] - batch_size = graph.shape[0] - x_list = paddle.split(graph.x, batch_size) - fine_x_list = [] - - for idx in range(batch_size): - x = x_list[idx] - if self.sdf is None: - with paddle.no_grad(): - self.sdf = signed_dist_graph( - x[:, :2], self.fine_marker_dict - ).unsqueeze(1) - fine_x = paddle.concat([x, self.sdf], axis=1) - for conv in self.pre_convs: - fine_x = F.relu(conv(graph, fine_x)) - fine_x_list.append(fine_x) - nodes_input = self.get_nodes().tile([batch_size, 1, 1]) - - batch_y = self.su2( - nodes_input[..., 0], - nodes_input[..., 1], - graph.aoa[..., None], - graph.mach_or_reynolds[..., None], - ) - batch_y = self.process_sim(batch_y, False) - - pred_fields = [] - for idx in range(batch_size): - coarse_y = paddle.stack([y[idx].flatten() for y in batch_y], axis=1).astype( - "float32" - ) - nodes = self.get_nodes() - x = x_list[idx] - fine_y = _knn_interpolate( - features=coarse_y, coarse_nodes=nodes[:, :2], fine_nodes=x[:, :2] - ) - fine_y = paddle.concat([fine_y, fine_x_list[idx]], axis=1) - - for conv in self.convs[:-1]: - fine_y = F.relu(conv(graph, fine_y)) - fine_y = self.convs[-1](graph, fine_y) - pred_fields.append(fine_y) - pred_fields = paddle.concat(pred_fields, axis=0) - return {self.output_keys[0]: pred_fields} - - def get_nodes(self) -> paddle.Tensor: - return self.nodes - - @staticmethod - def write_mesh_file( - x: paddle.Tensor, - elems: paddle.Tensor, - marker_dict: Dict[str, Sequence[Sequence[int]]], - filename: str = "mesh.su2", - ) -> None: - write_graph_mesh(filename, x[:, :2], elems, marker_dict) diff --git a/examples/smc_reac/ppsci/arch/chip_deeponets.py b/examples/smc_reac/ppsci/arch/chip_deeponets.py deleted file mode 100644 index 30c87c5656..0000000000 --- a/examples/smc_reac/ppsci/arch/chip_deeponets.py +++ /dev/null @@ -1,214 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import Tuple -from typing import Union - -import paddle -import paddle.nn as nn - -from ppsci.arch import activation as act_mod -from ppsci.arch import base -from ppsci.arch import mlp - - -class ChipDeepONets(base.Arch): - """Multi-branch physics-informed deep operator neural network. The network consists of three branch networks: random heat source, boundary function, and boundary type, as well as a trunk network. - - Args: - branch_input_keys (Tuple[str, ...]): Name of input data for internal heat source on branch nets. - BCtype_input_keys (Tuple[str, ...]): Name of input data for boundary types on branch nets. - BC_input_keys (Tuple[str, ...]): Name of input data for boundary on branch nets. - trunk_input_keys (Tuple[str, ...]): Name of input data for trunk net. - output_keys (Tuple[str, ...]): Output name of predicted temperature. - num_loc (int): Number of sampled input data for internal heat source. - bctype_loc (int): Number of sampled input data for boundary types. - BC_num_loc (int): Number of sampled input data for boundary. - num_features (int): Number of features extracted from trunk net, same for all branch nets. - branch_num_layers (int): Number of hidden layers of internal heat source on branch nets. - BC_num_layers (int): Number of hidden layers of boundary on branch nets. - trunk_num_layers (int): Number of hidden layers of trunk net. - branch_hidden_size (Union[int, Tuple[int, ...]]): Number of hidden size of internal heat source on branch nets. - An integer for all layers, or list of integer specify each layer's size. - BC_hidden_size (Union[int, Tuple[int, ...]]): Number of hidden size of boundary on branch nets. - An integer for all layers, or list of integer specify each layer's size. - trunk_hidden_size (Union[int, Tuple[int, ...]]): Number of hidden size of trunk net. - An integer for all layers, or list of integer specify each layer's size. - branch_skip_connection (bool, optional): Whether to use skip connection for internal heat source on branch net. Defaults to False. - BC_skip_connection (bool, optional): Whether to use skip connection for boundary on branch net. Defaults to False. - trunk_skip_connection (bool, optional): Whether to use skip connection for trunk net. Defaults to False. - branch_activation (str, optional): Name of activation function for internal heat source on branch net. Defaults to "tanh". - BC_activation (str, optional): Name of activation function for boundary on branch net. Defaults to "tanh". - trunk_activation (str, optional): Name of activation function for trunk net. Defaults to "tanh". - branch_weight_norm (bool, optional): Whether to apply weight norm on parameter(s) for internal heat source on branch net. Defaults to False. - BC_weight_norm (bool, optional): Whether to apply weight norm on parameter(s) for boundary on branch net. Defaults to False. - trunk_weight_norm (bool, optional): Whether to apply weight norm on parameter(s) for trunk net. Defaults to False. - use_bias (bool, optional): Whether to add bias on predicted G(u)(y). Defaults to True. - - Examples: - >>> import ppsci - >>> model = ppsci.arch.ChipDeepONets( - ... ('u',), - ... ('bc',), - ... ('bc_data',), - ... ("x",'y'), - ... ("T",), - ... 324, - ... 1, - ... 76, - ... 400, - ... 9, - ... 9, - ... 6, - ... 256, - ... 256, - ... 128, - ... branch_activation="swish", - ... BC_activation="swish", - ... trunk_activation="swish", - ... use_bias=True, - ... ) - """ - - def __init__( - self, - branch_input_keys: Tuple[str, ...], - BCtype_input_keys: Tuple[str, ...], - BC_input_keys: Tuple[str, ...], - trunk_input_keys: Tuple[str, ...], - output_keys: Tuple[str, ...], - num_loc: int, - bctype_loc: int, - BC_num_loc: int, - num_features: int, - branch_num_layers: int, - BC_num_layers: int, - trunk_num_layers: int, - branch_hidden_size: Union[int, Tuple[int, ...]], - BC_hidden_size: Union[int, Tuple[int, ...]], - trunk_hidden_size: Union[int, Tuple[int, ...]], - branch_skip_connection: bool = False, - BC_skip_connection: bool = False, - trunk_skip_connection: bool = False, - branch_activation: str = "tanh", - BC_activation: str = "tanh", - trunk_activation: str = "tanh", - branch_weight_norm: bool = False, - BC_weight_norm: bool = False, - trunk_weight_norm: bool = False, - use_bias: bool = True, - ): - super().__init__() - self.trunk_input_keys = trunk_input_keys - self.branch_input_keys = branch_input_keys - self.BCtype_input_keys = BCtype_input_keys - self.BC_input_keys = BC_input_keys - self.input_keys = ( - self.trunk_input_keys - + self.branch_input_keys - + self.BC_input_keys - + self.BCtype_input_keys - ) - self.output_keys = output_keys - - self.branch_net = mlp.MLP( - self.branch_input_keys, - ("b",), - branch_num_layers, - branch_hidden_size, - branch_activation, - branch_skip_connection, - branch_weight_norm, - input_dim=num_loc, - output_dim=num_features, - ) - - self.BCtype_net = mlp.MLP( - self.BCtype_input_keys, - ("bctype",), - BC_num_layers, - BC_hidden_size, - BC_activation, - BC_skip_connection, - BC_weight_norm, - input_dim=bctype_loc, - output_dim=num_features, - ) - - self.BC_net = mlp.MLP( - self.BC_input_keys, - ("bc",), - BC_num_layers, - BC_hidden_size, - BC_activation, - BC_skip_connection, - BC_weight_norm, - input_dim=BC_num_loc, - output_dim=num_features, - ) - - self.trunk_net = mlp.MLP( - self.trunk_input_keys, - ("t",), - trunk_num_layers, - trunk_hidden_size, - trunk_activation, - trunk_skip_connection, - trunk_weight_norm, - input_dim=len(self.trunk_input_keys), - output_dim=num_features, - ) - self.trunk_act = act_mod.get_activation(trunk_activation) - self.bc_act = act_mod.get_activation(BC_activation) - self.branch_act = act_mod.get_activation(branch_activation) - - self.use_bias = use_bias - if use_bias: - # register bias to parameter for updating in optimizer and storage - self.b = self.create_parameter( - shape=(1,), - attr=nn.initializer.Constant(0.0), - ) - - def forward(self, x): - - if self._input_transform is not None: - x = self._input_transform(x) - - # Branch net to encode the input function - u_features = self.branch_net(x)[self.branch_net.output_keys[0]] - bc_features = self.BC_net(x)[self.BC_net.output_keys[0]] - bctype_features = self.BCtype_net(x)[self.BCtype_net.output_keys[0]] - # Trunk net to encode the domain of the output function - y_features = self.trunk_net(x)[self.trunk_net.output_keys[0]] - y_features = self.trunk_act(y_features) - # Dot product - G_u = paddle.sum( - u_features * y_features * bc_features * bctype_features, - axis=1, - keepdim=True, - ) - # Add bias - if self.use_bias: - G_u += self.b - - result_dict = { - self.output_keys[0]: G_u, - } - if self._output_transform is not None: - result_dict = self._output_transform(x, result_dict) - - return result_dict diff --git a/examples/smc_reac/ppsci/arch/crystalgraphconvnet.py b/examples/smc_reac/ppsci/arch/crystalgraphconvnet.py deleted file mode 100644 index bb82aa0b81..0000000000 --- a/examples/smc_reac/ppsci/arch/crystalgraphconvnet.py +++ /dev/null @@ -1,167 +0,0 @@ -import paddle -import paddle.nn as nn - -from ppsci.arch import base - - -class ConvLayer(nn.Layer): - def __init__(self, atom_fea_len, nbr_fea_len): - super(ConvLayer, self).__init__() - self.atom_fea_len = atom_fea_len - self.nbr_fea_len = nbr_fea_len - self.fc_full = nn.Linear( - 2 * self.atom_fea_len + self.nbr_fea_len, 2 * self.atom_fea_len - ) - self.sigmoid = nn.Sigmoid() - self.softplus1 = nn.Softplus() - self.bn1 = nn.BatchNorm1D(2 * self.atom_fea_len) - self.bn2 = nn.BatchNorm1D(self.atom_fea_len) - self.softplus2 = nn.Softplus() - - def forward(self, atom_in_fea, nbr_fea, nbr_fea_idx): - # TODO will there be problems with the index zero padding? - N, M = nbr_fea_idx.shape - atom_nbr_fea = atom_in_fea[nbr_fea_idx, :] - total_nbr_fea = paddle.concat( - [ - paddle.expand( - atom_in_fea.unsqueeze(1), shape=[N, M, self.atom_fea_len] - ), - atom_nbr_fea, - nbr_fea, - ], - axis=2, - ) - total_gated_fea = self.fc_full(total_nbr_fea) - total_gated_fea = paddle.reshape( - self.bn1(paddle.reshape(total_gated_fea, [-1, self.atom_fea_len * 2])), - [N, M, self.atom_fea_len * 2], - ) - nbr_filter, nbr_core = paddle.chunk(total_gated_fea, chunks=2, axis=2) - nbr_filter = self.sigmoid(nbr_filter) - nbr_core = self.softplus1(nbr_core) - nbr_sumed = paddle.sum(nbr_filter * nbr_core, axis=1) - nbr_sumed = self.bn2(nbr_sumed) - out = self.softplus2(atom_in_fea + nbr_sumed) - return out - - -class CrystalGraphConvNet(base.Arch): - """ - Create a crystal graph convolutional neural network for predicting total - material properties. - - Args: - orig_atom_fea_len (int): Number of atom features in the input. - nbr_fea_len (int): Number of bond features. - atom_fea_len (int): Number of hidden atom features in the convolutional layers. - n_conv (int): Number of convolutional layers. - h_fea_len (int): Number of hidden features after pooling. - n_h (int): Number of hidden layers after pooling. - - Examples: - >>> import paddle - >>> import ppsci - >>> model = ppsci.arch.CrystalGraphConvNet( - ... orig_atom_fea_len=92, - ... nbr_fea_len=41, - ... atom_fea_len=64, - ... n_conv=3, - ... h_fea_len=128, - ... n_h=1, - ... ) - >>> input_dict = { - ... "i": [ - ... paddle.rand(shape=[45, 92]), paddle.rand(shape=[45, 12, 41]), - ... paddle.randint(high=45, shape=[45, 12]), - ... [ - ... paddle.randint(high=32, shape=[32]), paddle.randint(high=8, shape=[8]), - ... paddle.randint(high=2, shape=[2]), paddle.randint(high=3, shape=[3]) - ... ] - ... ] - ... } - >>> output_dict = model(input_dict) - >>> print(output_dict["out"].shape) - [4, 1] - """ - - def __init__( - self, - orig_atom_fea_len: int, - nbr_fea_len: int, - atom_fea_len: int, - n_conv: int, - h_fea_len: int, - n_h: int, - ): - - super().__init__() - self.embedding = nn.Linear(orig_atom_fea_len, atom_fea_len) - self.convs = nn.LayerList( - [ - ConvLayer(atom_fea_len=atom_fea_len, nbr_fea_len=nbr_fea_len) - for _ in range(n_conv) - ] - ) - self.conv_to_fc = nn.Linear(atom_fea_len, h_fea_len) - self.conv_to_fc_softplus = nn.Softplus() - if n_h > 1: - self.fcs = nn.LayerList( - [nn.Linear(h_fea_len, h_fea_len) for _ in range(n_h - 1)] - ) - self.softpluses = nn.LayerList([nn.Softplus() for _ in range(n_h - 1)]) - - self.fc_out = nn.Linear(h_fea_len, 1) - - def forward(self, input) -> paddle.Tensor: - """ - Forward pass. - - N: Total number of atoms in the batch. - M: Max number of neighbors. - N0: Total number of crystals in the batch. - - Args: - input (list): List of input, which includes the following elements: - atom_fea (paddle.Tensor): Shape (N, orig_atom_fea_len). Atom features from atom type. - nbr_fea (paddle.Tensor): Shape (N, M, nbr_fea_len). Bond features of each atom's M neighbors. - nbr_fea_idx (paddle.Tensor): Shape (N, M). Indices of M neighbors of each atom. - crystal_atom_idx (list): List of paddle.Tensor of length N0. Mapping from the crystal idx to atom idx. - - Returns: - paddle.Tensor: Shape (N,). Atom hidden features after convolution. - """ - atom_fea, nbr_fea, nbr_fea_idx, crystal_atom_idx = input["i"] - atom_fea = self.embedding(atom_fea) - for conv_func in self.convs: - atom_fea = conv_func(atom_fea, nbr_fea, nbr_fea_idx) - crys_fea = self.pooling(atom_fea, crystal_atom_idx) - crys_fea = self.conv_to_fc(self.conv_to_fc_softplus(crys_fea)) - crys_fea = self.conv_to_fc_softplus(crys_fea) - if hasattr(self, "fcs") and hasattr(self, "softpluses"): - for fc, softplus in zip(self.fcs, self.softpluses): - crys_fea = softplus(fc(crys_fea)) - out = self.fc_out(crys_fea) - out_dict = {"out": out} - return out_dict - - def pooling(self, atom_fea, crystal_atom_idx): - """ - Pooling the atom features to crystal features - - N: Total number of atoms in the batch - N0: Total number of crystals in the batch - - Args: - atom_fea (paddle.Tensor): Shape (N, atom_fea_len). Atom feature vectors of the batch. - crystal_atom_idx (List[paddle.Tensor]): Length N0. Mapping from the crystal idx to atom idx - """ - assert ( - sum([len(idx_map) for idx_map in crystal_atom_idx]) - == atom_fea.data.shape[0] - ) - summed_fea = [ - paddle.mean(atom_fea[idx_map], axis=0, keepdim=True) - for idx_map in crystal_atom_idx - ] - return paddle.concat(summed_fea, axis=0) diff --git a/examples/smc_reac/ppsci/arch/cuboid_transformer.py b/examples/smc_reac/ppsci/arch/cuboid_transformer.py deleted file mode 100644 index 08a80104bc..0000000000 --- a/examples/smc_reac/ppsci/arch/cuboid_transformer.py +++ /dev/null @@ -1,958 +0,0 @@ -from typing import Sequence -from typing import Tuple -from typing import Union - -import paddle -from paddle import nn - -import ppsci.arch.cuboid_transformer_decoder as cuboid_decoder -import ppsci.arch.cuboid_transformer_encoder as cuboid_encoder -import ppsci.arch.cuboid_transformer_utils as cuboid_utils -from ppsci.arch import activation as act_mod -from ppsci.arch import base -from ppsci.arch.cuboid_transformer_encoder import NEGATIVE_SLOPE -from ppsci.utils import initializer - -"""A space-time Transformer with Cuboid Attention""" - - -class InitialEncoder(nn.Layer): - def __init__( - self, - dim, - out_dim, - downsample_scale: Union[int, Sequence[int]], - num_conv_layers: int = 2, - activation: str = "leaky", - padding_type: str = "nearest", - conv_init_mode: str = "0", - linear_init_mode: str = "0", - norm_init_mode: str = "0", - ): - super(InitialEncoder, self).__init__() - self.num_conv_layers = num_conv_layers - self.conv_init_mode = conv_init_mode - self.linear_init_mode = linear_init_mode - self.norm_init_mode = norm_init_mode - conv_block = [] - for i in range(num_conv_layers): - if i == 0: - conv_block.append( - nn.Conv2D( - kernel_size=(3, 3), - padding=(1, 1), - in_channels=dim, - out_channels=out_dim, - ) - ) - conv_block.append(nn.GroupNorm(num_groups=16, num_channels=out_dim)) - conv_block.append( - act_mod.get_activation(activation) - if activation != "leaky_relu" - else nn.LeakyReLU(NEGATIVE_SLOPE) - ) - else: - conv_block.append( - nn.Conv2D( - kernel_size=(3, 3), - padding=(1, 1), - in_channels=out_dim, - out_channels=out_dim, - ) - ) - conv_block.append(nn.GroupNorm(num_groups=16, num_channels=out_dim)) - conv_block.append( - act_mod.get_activation(activation) - if activation != "leaky_relu" - else nn.LeakyReLU(NEGATIVE_SLOPE) - ) - self.conv_block = nn.Sequential(*conv_block) - if isinstance(downsample_scale, int): - patch_merge_downsample = (1, downsample_scale, downsample_scale) - elif len(downsample_scale) == 2: - patch_merge_downsample = (1, *downsample_scale) - elif len(downsample_scale) == 3: - patch_merge_downsample = tuple(downsample_scale) - else: - raise NotImplementedError( - f"downsample_scale {downsample_scale} format not supported!" - ) - self.patch_merge = cuboid_encoder.PatchMerging3D( - dim=out_dim, - out_dim=out_dim, - padding_type=padding_type, - downsample=patch_merge_downsample, - linear_init_mode=linear_init_mode, - norm_init_mode=norm_init_mode, - ) - self.reset_parameters() - - def reset_parameters(self): - for m in self.children(): - cuboid_utils.apply_initialization( - m, - conv_mode=self.conv_init_mode, - linear_mode=self.linear_init_mode, - norm_mode=self.norm_init_mode, - ) - - def forward(self, x): - """x --> [K x Conv2D] --> PatchMerge - - Args: - x: (B, T, H, W, C) - - Returns: - out: (B, T, H_new, W_new, C_out) - """ - - B, T, H, W, C = x.shape - - if self.num_conv_layers > 0: - x = x.reshape([B * T, H, W, C]).transpose(perm=[0, 3, 1, 2]) - x = self.conv_block(x).transpose(perm=[0, 2, 3, 1]) - x = self.patch_merge(x.reshape([B, T, H, W, -1])) - else: - x = self.patch_merge(x) - return x - - -class FinalDecoder(nn.Layer): - def __init__( - self, - target_thw: Tuple[int, ...], - dim: int, - num_conv_layers: int = 2, - activation: str = "leaky", - conv_init_mode: str = "0", - linear_init_mode: str = "0", - norm_init_mode: str = "0", - ): - super(FinalDecoder, self).__init__() - self.target_thw = target_thw - self.dim = dim - self.num_conv_layers = num_conv_layers - self.conv_init_mode = conv_init_mode - self.linear_init_mode = linear_init_mode - self.norm_init_mode = norm_init_mode - conv_block = [] - for i in range(num_conv_layers): - conv_block.append( - nn.Conv2D( - kernel_size=(3, 3), - padding=(1, 1), - in_channels=dim, - out_channels=dim, - ) - ) - conv_block.append(nn.GroupNorm(num_groups=16, num_channels=dim)) - conv_block.append( - act_mod.get_activation(activation) - if activation != "leaky_relu" - else nn.LeakyReLU(NEGATIVE_SLOPE) - ) - self.conv_block = nn.Sequential(*conv_block) - self.upsample = cuboid_decoder.Upsample3DLayer( - dim=dim, - out_dim=dim, - target_size=target_thw, - kernel_size=3, - conv_init_mode=conv_init_mode, - ) - self.reset_parameters() - - def reset_parameters(self): - for m in self.children(): - cuboid_utils.apply_initialization( - m, - conv_mode=self.conv_init_mode, - linear_mode=self.linear_init_mode, - norm_mode=self.norm_init_mode, - ) - - def forward(self, x): - """x --> Upsample --> [K x Conv2D] - - Args: - x: (B, T, H, W, C) - - Returns: - out: (B, T, H_new, W_new, C) - """ - - x = self.upsample(x) - if self.num_conv_layers > 0: - B, T, H, W, C = x.shape - x = x.reshape([B * T, H, W, C]).transpose(perm=[0, 3, 1, 2]) - x = ( - self.conv_block(x) - .transpose(perm=[0, 2, 3, 1]) - .reshape([B, T, H, W, -1]) - ) - return x - - -class InitialStackPatchMergingEncoder(nn.Layer): - def __init__( - self, - num_merge: int, - in_dim: int, - out_dim_list: Tuple[int, ...], - downsample_scale_list: Tuple[float, ...], - num_conv_per_merge_list: Tuple[int, ...] = None, - activation: str = "leaky", - padding_type: str = "nearest", - conv_init_mode: str = "0", - linear_init_mode: str = "0", - norm_init_mode: str = "0", - ): - super(InitialStackPatchMergingEncoder, self).__init__() - self.conv_init_mode = conv_init_mode - self.linear_init_mode = linear_init_mode - self.norm_init_mode = norm_init_mode - self.num_merge = num_merge - self.in_dim = in_dim - self.out_dim_list = out_dim_list[:num_merge] - self.downsample_scale_list = downsample_scale_list[:num_merge] - self.num_conv_per_merge_list = num_conv_per_merge_list - self.num_group_list = [max(1, out_dim // 4) for out_dim in self.out_dim_list] - self.conv_block_list = nn.LayerList() - self.patch_merge_list = nn.LayerList() - for i in range(num_merge): - if i == 0: - in_dim = in_dim - else: - in_dim = self.out_dim_list[i - 1] - out_dim = self.out_dim_list[i] - downsample_scale = self.downsample_scale_list[i] - conv_block = [] - for j in range(self.num_conv_per_merge_list[i]): - if j == 0: - conv_in_dim = in_dim - else: - conv_in_dim = out_dim - conv_block.append( - nn.Conv2D( - kernel_size=(3, 3), - padding=(1, 1), - in_channels=conv_in_dim, - out_channels=out_dim, - ) - ) - conv_block.append( - nn.GroupNorm( - num_groups=self.num_group_list[i], num_channels=out_dim - ) - ) - conv_block.append( - act_mod.get_activation(activation) - if activation != "leaky_relu" - else nn.LeakyReLU(NEGATIVE_SLOPE) - ) - conv_block = nn.Sequential(*conv_block) - self.conv_block_list.append(conv_block) - patch_merge = cuboid_encoder.PatchMerging3D( - dim=out_dim, - out_dim=out_dim, - padding_type=padding_type, - downsample=(1, downsample_scale, downsample_scale), - linear_init_mode=linear_init_mode, - norm_init_mode=norm_init_mode, - ) - self.patch_merge_list.append(patch_merge) - self.reset_parameters() - - def reset_parameters(self): - for m in self.children(): - cuboid_utils.apply_initialization( - m, - conv_mode=self.conv_init_mode, - linear_mode=self.linear_init_mode, - norm_mode=self.norm_init_mode, - ) - - def get_out_shape_list(self, input_shape): - out_shape_list = [] - for patch_merge in self.patch_merge_list: - input_shape = patch_merge.get_out_shape(input_shape) - out_shape_list.append(input_shape) - return out_shape_list - - def forward(self, x): - """x --> [K x Conv2D] --> PatchMerge --> ... --> [K x Conv2D] --> PatchMerge - - Args: - x: (B, T, H, W, C) - - Returns: - out: (B, T, H_new, W_new, C_out) - """ - - for i, (conv_block, patch_merge) in enumerate( - zip(self.conv_block_list, self.patch_merge_list) - ): - B, T, H, W, C = x.shape - if self.num_conv_per_merge_list[i] > 0: - x = x.reshape([B * T, H, W, C]).transpose(perm=[0, 3, 1, 2]) - x = conv_block(x).transpose(perm=[0, 2, 3, 1]).reshape([B, T, H, W, -1]) - x = patch_merge(x) - return x - - -class FinalStackUpsamplingDecoder(nn.Layer): - def __init__( - self, - target_shape_list: Tuple[Tuple[int, ...]], - in_dim: int, - num_conv_per_up_list: Tuple[int, ...] = None, - activation: str = "leaky", - conv_init_mode: str = "0", - linear_init_mode: str = "0", - norm_init_mode: str = "0", - ): - super(FinalStackUpsamplingDecoder, self).__init__() - self.conv_init_mode = conv_init_mode - self.linear_init_mode = linear_init_mode - self.norm_init_mode = norm_init_mode - self.target_shape_list = target_shape_list - self.out_dim_list = [ - target_shape[-1] for target_shape in self.target_shape_list - ] - self.num_upsample = len(target_shape_list) - self.in_dim = in_dim - self.num_conv_per_up_list = num_conv_per_up_list - self.num_group_list = [max(1, out_dim // 4) for out_dim in self.out_dim_list] - self.conv_block_list = nn.LayerList() - self.upsample_list = nn.LayerList() - for i in range(self.num_upsample): - if i == 0: - in_dim = in_dim - else: - in_dim = self.out_dim_list[i - 1] - out_dim = self.out_dim_list[i] - upsample = cuboid_decoder.Upsample3DLayer( - dim=in_dim, - out_dim=in_dim, - target_size=target_shape_list[i][:-1], - kernel_size=3, - conv_init_mode=conv_init_mode, - ) - self.upsample_list.append(upsample) - conv_block = [] - for j in range(num_conv_per_up_list[i]): - if j == 0: - conv_in_dim = in_dim - else: - conv_in_dim = out_dim - conv_block.append( - nn.Conv2D( - kernel_size=(3, 3), - padding=(1, 1), - in_channels=conv_in_dim, - out_channels=out_dim, - ) - ) - conv_block.append( - nn.GroupNorm( - num_groups=self.num_group_list[i], num_channels=out_dim - ) - ) - conv_block.append( - act_mod.get_activation(activation) - if activation != "leaky_relu" - else nn.LeakyReLU(NEGATIVE_SLOPE) - ) - conv_block = nn.Sequential(*conv_block) - self.conv_block_list.append(conv_block) - self.reset_parameters() - - def reset_parameters(self): - for m in self.children(): - cuboid_utils.apply_initialization( - m, - conv_mode=self.conv_init_mode, - linear_mode=self.linear_init_mode, - norm_mode=self.norm_init_mode, - ) - - @staticmethod - def get_init_params(enc_input_shape, enc_out_shape_list, large_channel=False): - dec_target_shape_list = list(enc_out_shape_list[:-1])[::-1] + [ - tuple(enc_input_shape) - ] - if large_channel: - dec_target_shape_list_large_channel = [] - for i, enc_out_shape in enumerate(enc_out_shape_list[::-1]): - dec_target_shape_large_channel = list(dec_target_shape_list[i]) - dec_target_shape_large_channel[-1] = enc_out_shape[-1] - dec_target_shape_list_large_channel.append( - tuple(dec_target_shape_large_channel) - ) - dec_target_shape_list = dec_target_shape_list_large_channel - dec_in_dim = enc_out_shape_list[-1][-1] - return dec_target_shape_list, dec_in_dim - - def forward(self, x): - """x --> Upsample --> [K x Conv2D] --> ... --> Upsample --> [K x Conv2D] - - Args: - x: Shape (B, T, H, W, C) - - Returns: - out: Shape (B, T, H_new, W_new, C) - """ - for i, (conv_block, upsample) in enumerate( - zip(self.conv_block_list, self.upsample_list) - ): - x = upsample(x) - if self.num_conv_per_up_list[i] > 0: - B, T, H, W, C = x.shape - x = x.reshape([B * T, H, W, C]).transpose(perm=[0, 3, 1, 2]) - x = conv_block(x).transpose(perm=[0, 2, 3, 1]).reshape([B, T, H, W, -1]) - return x - - -class CuboidTransformer(base.Arch): - """Cuboid Transformer for spatiotemporal forecasting - - We adopt the Non-autoregressive encoder-decoder architecture. - The decoder takes the multi-scale memory output from the encoder. - - The initial downsampling / upsampling layers will be - Downsampling: [K x Conv2D --> PatchMerge] - Upsampling: [Nearest Interpolation-based Upsample --> K x Conv2D] - - x --> downsample (optional) ---> (+pos_embed) ---> enc --> mem_l initial_z (+pos_embed) ---> FC - | | - |------------| - | - | - y <--- upsample (optional) <--- dec <---------- - - Args: - input_keys (Tuple[str, ...]): Name of input keys, such as ("input",). - output_keys (Tuple[str, ...]): Name of output keys, such as ("output",). - input_shape (Tuple[int, ...]): The shape of the input data. - target_shape (Tuple[int, ...]): The shape of the target data. - base_units (int, optional): The base units. Defaults to 128. - block_units (int, optional): The block units. Defaults to None. - scale_alpha (float, optional): We scale up the channels based on the formula: - - round_to(base_units * max(downsample_scale) ** units_alpha, 4). Defaults to 1.0. - num_heads (int, optional): The number of heads. Defaults to 4. - attn_drop (float, optional): The attention dropout. Defaults to 0.0. - proj_drop (float, optional): The projection dropout. Defaults to 0.0. - ffn_drop (float, optional): The ffn dropout. Defaults to 0.0. - downsample (int, optional): The rate of downsample. Defaults to 2. - downsample_type (str, optional): The type of downsample. Defaults to "patch_merge". - upsample_type (str, optional): The rate of upsample. Defaults to "upsample". - upsample_kernel_size (int, optional): The kernel size of upsample. Defaults to 3. - enc_depth (list, optional): The depth of encoder. Defaults to [4, 4, 4]. - enc_attn_patterns (str, optional): The pattern of encoder attention. Defaults to None. - enc_cuboid_size (list, optional): The cuboid size of encoder. Defaults to [(4, 4, 4), (4, 4, 4)]. - enc_cuboid_strategy (list, optional): The cuboid strategy of encoder. Defaults to [("l", "l", "l"), ("d", "d", "d")]. - enc_shift_size (list, optional): The shift size of encoder. Defaults to [(0, 0, 0), (0, 0, 0)]. - enc_use_inter_ffn (bool, optional): Whether to use intermediate FFN for encoder. Defaults to True. - dec_depth (list, optional): The depth of decoder. Defaults to [2, 2]. - dec_cross_start (int, optional): The cross start of decoder. Defaults to 0. - dec_self_attn_patterns (str, optional): The partterns of decoder. Defaults to None. - dec_self_cuboid_size (list, optional): The cuboid size of decoder. Defaults to [(4, 4, 4), (4, 4, 4)]. - dec_self_cuboid_strategy (list, optional): The strategy of decoder. Defaults to [("l", "l", "l"), ("d", "d", "d")]. - dec_self_shift_size (list, optional): The shift size of decoder. Defaults to [(1, 1, 1), (0, 0, 0)]. - dec_cross_attn_patterns (_type_, optional): The cross attention patterns of decoder. Defaults to None. - dec_cross_cuboid_hw (list, optional): The cuboid_hw of decoder. Defaults to [(4, 4), (4, 4)]. - dec_cross_cuboid_strategy (list, optional): The cuboid strategy of decoder. Defaults to [("l", "l", "l"), ("d", "l", "l")]. - dec_cross_shift_hw (list, optional): The shift_hw of decoder. Defaults to [(0, 0), (0, 0)]. - dec_cross_n_temporal (list, optional): The cross_n_temporal of decoder. Defaults to [1, 2]. - dec_cross_last_n_frames (int, optional): The cross_last_n_frames of decoder. Defaults to None. - dec_use_inter_ffn (bool, optional): Whether to use intermediate FFN for decoder. Defaults to True. - dec_hierarchical_pos_embed (bool, optional): Whether to use hierarchical pos_embed for decoder. Defaults to False. - num_global_vectors (int, optional): The num of global vectors. Defaults to 4. - use_dec_self_global (bool, optional): Whether to use global vector for decoder. Defaults to True. - dec_self_update_global (bool, optional): Whether to update global vector for decoder. Defaults to True. - use_dec_cross_global (bool, optional): Whether to use cross global vector for decoder. Defaults to True. - use_global_vector_ffn (bool, optional): Whether to use global vector FFN. Defaults to True. - use_global_self_attn (bool, optional): Whether to use global attentions. Defaults to False. - separate_global_qkv (bool, optional): Whether to separate global qkv. Defaults to False. - global_dim_ratio (int, optional): The ratio of global dim. Defaults to 1. - self_pattern (str, optional): The pattern. Defaults to "axial". - cross_self_pattern (str, optional): The self cross pattern. Defaults to "axial". - cross_pattern (str, optional): The cross pattern. Defaults to "cross_1x1". - z_init_method (str, optional): How the initial input to the decoder is initialized. Defaults to "nearest_interp". - initial_downsample_type (str, optional): The downsample type of initial. Defaults to "conv". - initial_downsample_activation (str, optional): The downsample activation of initial. Defaults to "leaky". - initial_downsample_scale (int, optional): The downsample scale of initial. Defaults to 1. - initial_downsample_conv_layers (int, optional): The conv layer of downsample of initial. Defaults to 2. - final_upsample_conv_layers (int, optional): The conv layer of final upsample. Defaults to 2. - initial_downsample_stack_conv_num_layers (int, optional): The num of stack conv layer of initial downsample. Defaults to 1. - initial_downsample_stack_conv_dim_list (list, optional): The dim list of stack conv of initial downsample. Defaults to None. - initial_downsample_stack_conv_downscale_list (list, optional): The downscale list of stack conv of initial downsample. Defaults to [1]. - initial_downsample_stack_conv_num_conv_list (list, optional): The num of stack conv list of initial downsample. Defaults to [2]. - ffn_activation (str, optional): The activation of FFN. Defaults to "leaky". - gated_ffn (bool, optional): Whether to use gate FFN. Defaults to False. - norm_layer (str, optional): The type of normilize. Defaults to "layer_norm". - padding_type (str, optional): The type of padding. Defaults to "ignore". - pos_embed_type (str, optional): The type of pos embedding. Defaults to "t+hw". - checkpoint_level (bool, optional): Whether to use checkpoint. Defaults to True. - use_relative_pos (bool, optional): Whether to use relative pose. Defaults to True. - self_attn_use_final_proj (bool, optional): Whether to use final projection. Defaults to True. - dec_use_first_self_attn (bool, optional): Whether to use first self attention for decoder. Defaults to False. - attn_linear_init_mode (str, optional): The mode of attention linear init. Defaults to "0". - ffn_linear_init_mode (str, optional): The mode of FFN linear init. Defaults to "0". - conv_init_mode (str, optional): The mode of conv init. Defaults to "0". - down_up_linear_init_mode (str, optional): The mode of downsample and upsample linear init. Defaults to "0". - norm_init_mode (str, optional): The mode of normalization init. Defaults to "0". - """ - - def __init__( - self, - input_keys: Tuple[str, ...], - output_keys: Tuple[str, ...], - input_shape: Tuple[int, ...], - target_shape: Tuple[int, ...], - base_units: int = 128, - block_units: int = None, - scale_alpha: float = 1.0, - num_heads: int = 4, - attn_drop: float = 0.0, - proj_drop: float = 0.0, - ffn_drop: float = 0.0, - downsample: int = 2, - downsample_type: str = "patch_merge", - upsample_type: str = "upsample", - upsample_kernel_size: int = 3, - enc_depth: Tuple[int, ...] = [4, 4, 4], - enc_attn_patterns: str = None, - enc_cuboid_size: Tuple[Tuple[int, ...], ...] = [(4, 4, 4), (4, 4, 4)], - enc_cuboid_strategy: Tuple[Tuple[str, ...], ...] = [ - ("l", "l", "l"), - ("d", "d", "d"), - ], - enc_shift_size: Tuple[Tuple[int, ...], ...] = [(0, 0, 0), (0, 0, 0)], - enc_use_inter_ffn: bool = True, - dec_depth: Tuple[int, ...] = [2, 2], - dec_cross_start: int = 0, - dec_self_attn_patterns: str = None, - dec_self_cuboid_size: Tuple[Tuple[int, ...], ...] = [(4, 4, 4), (4, 4, 4)], - dec_self_cuboid_strategy: Tuple[Tuple[str, ...], ...] = [ - ("l", "l", "l"), - ("d", "d", "d"), - ], - dec_self_shift_size: Tuple[Tuple[int, ...], ...] = [(1, 1, 1), (0, 0, 0)], - dec_cross_attn_patterns: str = None, - dec_cross_cuboid_hw: Tuple[Tuple[int, ...], ...] = [(4, 4), (4, 4)], - dec_cross_cuboid_strategy: Tuple[Tuple[str, ...], ...] = [ - ("l", "l", "l"), - ("d", "l", "l"), - ], - dec_cross_shift_hw: Tuple[Tuple[int, ...], ...] = [(0, 0), (0, 0)], - dec_cross_n_temporal: Tuple[int, ...] = [1, 2], - dec_cross_last_n_frames: int = None, - dec_use_inter_ffn: bool = True, - dec_hierarchical_pos_embed: bool = False, - num_global_vectors: int = 4, - use_dec_self_global: bool = True, - dec_self_update_global: bool = True, - use_dec_cross_global: bool = True, - use_global_vector_ffn: bool = True, - use_global_self_attn: bool = False, - separate_global_qkv: bool = False, - global_dim_ratio: int = 1, - self_pattern: str = "axial", - cross_self_pattern: str = "axial", - cross_pattern: str = "cross_1x1", - z_init_method: str = "nearest_interp", - initial_downsample_type: str = "conv", - initial_downsample_activation: str = "leaky", - initial_downsample_scale: int = 1, - initial_downsample_conv_layers: int = 2, - final_upsample_conv_layers: int = 2, - initial_downsample_stack_conv_num_layers: int = 1, - initial_downsample_stack_conv_dim_list: Tuple[int, ...] = None, - initial_downsample_stack_conv_downscale_list: Tuple[int, ...] = [1], - initial_downsample_stack_conv_num_conv_list: Tuple[int, ...] = [2], - ffn_activation: str = "leaky", - gated_ffn: bool = False, - norm_layer: str = "layer_norm", - padding_type: str = "ignore", - pos_embed_type: str = "t+hw", - checkpoint_level: bool = True, - use_relative_pos: bool = True, - self_attn_use_final_proj: bool = True, - dec_use_first_self_attn: bool = False, - attn_linear_init_mode: str = "0", - ffn_linear_init_mode: str = "0", - conv_init_mode: str = "0", - down_up_linear_init_mode: str = "0", - norm_init_mode: str = "0", - ): - super().__init__() - self.input_keys = input_keys - self.output_keys = output_keys - self.attn_linear_init_mode = attn_linear_init_mode - self.ffn_linear_init_mode = ffn_linear_init_mode - self.conv_init_mode = conv_init_mode - self.down_up_linear_init_mode = down_up_linear_init_mode - self.norm_init_mode = norm_init_mode - assert len(enc_depth) == len(dec_depth) - self.base_units = base_units - self.num_global_vectors = num_global_vectors - - num_blocks = len(enc_depth) - if isinstance(self_pattern, str): - enc_attn_patterns = [self_pattern] * num_blocks - - if isinstance(cross_self_pattern, str): - dec_self_attn_patterns = [cross_self_pattern] * num_blocks - - if isinstance(cross_pattern, str): - dec_cross_attn_patterns = [cross_pattern] * num_blocks - - if global_dim_ratio != 1: - assert ( - separate_global_qkv is True - ), "Setting global_dim_ratio != 1 requires separate_global_qkv == True." - self.global_dim_ratio = global_dim_ratio - self.z_init_method = z_init_method - assert self.z_init_method in ["zeros", "nearest_interp", "last", "mean"] - self.input_shape = input_shape - self.target_shape = target_shape - T_in, H_in, W_in, C_in = input_shape - T_out, H_out, W_out, C_out = target_shape - assert H_in == H_out and W_in == W_out - if self.num_global_vectors > 0: - init_data = paddle.zeros( - (self.num_global_vectors, global_dim_ratio * base_units) - ) - self.init_global_vectors = paddle.create_parameter( - shape=init_data.shape, - dtype=init_data.dtype, - default_initializer=nn.initializer.Constant(0.0), - ) - - self.init_global_vectors.stop_gradient = not True - new_input_shape = self.get_initial_encoder_final_decoder( - initial_downsample_scale=initial_downsample_scale, - initial_downsample_type=initial_downsample_type, - activation=initial_downsample_activation, - initial_downsample_conv_layers=initial_downsample_conv_layers, - final_upsample_conv_layers=final_upsample_conv_layers, - padding_type=padding_type, - initial_downsample_stack_conv_num_layers=initial_downsample_stack_conv_num_layers, - initial_downsample_stack_conv_dim_list=initial_downsample_stack_conv_dim_list, - initial_downsample_stack_conv_downscale_list=initial_downsample_stack_conv_downscale_list, - initial_downsample_stack_conv_num_conv_list=initial_downsample_stack_conv_num_conv_list, - ) - T_in, H_in, W_in, _ = new_input_shape - self.encoder = cuboid_encoder.CuboidTransformerEncoder( - input_shape=(T_in, H_in, W_in, base_units), - base_units=base_units, - block_units=block_units, - scale_alpha=scale_alpha, - depth=enc_depth, - downsample=downsample, - downsample_type=downsample_type, - block_attn_patterns=enc_attn_patterns, - block_cuboid_size=enc_cuboid_size, - block_strategy=enc_cuboid_strategy, - block_shift_size=enc_shift_size, - num_heads=num_heads, - attn_drop=attn_drop, - proj_drop=proj_drop, - ffn_drop=ffn_drop, - gated_ffn=gated_ffn, - ffn_activation=ffn_activation, - norm_layer=norm_layer, - use_inter_ffn=enc_use_inter_ffn, - padding_type=padding_type, - use_global_vector=num_global_vectors > 0, - use_global_vector_ffn=use_global_vector_ffn, - use_global_self_attn=use_global_self_attn, - separate_global_qkv=separate_global_qkv, - global_dim_ratio=global_dim_ratio, - checkpoint_level=checkpoint_level, - use_relative_pos=use_relative_pos, - self_attn_use_final_proj=self_attn_use_final_proj, - attn_linear_init_mode=attn_linear_init_mode, - ffn_linear_init_mode=ffn_linear_init_mode, - conv_init_mode=conv_init_mode, - down_linear_init_mode=down_up_linear_init_mode, - norm_init_mode=norm_init_mode, - ) - self.enc_pos_embed = cuboid_decoder.PosEmbed( - embed_dim=base_units, typ=pos_embed_type, maxH=H_in, maxW=W_in, maxT=T_in - ) - mem_shapes = self.encoder.get_mem_shapes() - self.z_proj = nn.Linear( - in_features=mem_shapes[-1][-1], out_features=mem_shapes[-1][-1] - ) - self.dec_pos_embed = cuboid_decoder.PosEmbed( - embed_dim=mem_shapes[-1][-1], - typ=pos_embed_type, - maxT=T_out, - maxH=mem_shapes[-1][1], - maxW=mem_shapes[-1][2], - ) - self.decoder = cuboid_decoder.CuboidTransformerDecoder( - target_temporal_length=T_out, - mem_shapes=mem_shapes, - cross_start=dec_cross_start, - depth=dec_depth, - upsample_type=upsample_type, - block_self_attn_patterns=dec_self_attn_patterns, - block_self_cuboid_size=dec_self_cuboid_size, - block_self_shift_size=dec_self_shift_size, - block_self_cuboid_strategy=dec_self_cuboid_strategy, - block_cross_attn_patterns=dec_cross_attn_patterns, - block_cross_cuboid_hw=dec_cross_cuboid_hw, - block_cross_shift_hw=dec_cross_shift_hw, - block_cross_cuboid_strategy=dec_cross_cuboid_strategy, - block_cross_n_temporal=dec_cross_n_temporal, - cross_last_n_frames=dec_cross_last_n_frames, - num_heads=num_heads, - attn_drop=attn_drop, - proj_drop=proj_drop, - ffn_drop=ffn_drop, - upsample_kernel_size=upsample_kernel_size, - ffn_activation=ffn_activation, - gated_ffn=gated_ffn, - norm_layer=norm_layer, - use_inter_ffn=dec_use_inter_ffn, - max_temporal_relative=T_in + T_out, - padding_type=padding_type, - hierarchical_pos_embed=dec_hierarchical_pos_embed, - pos_embed_type=pos_embed_type, - use_self_global=num_global_vectors > 0 and use_dec_self_global, - self_update_global=dec_self_update_global, - use_cross_global=num_global_vectors > 0 and use_dec_cross_global, - use_global_vector_ffn=use_global_vector_ffn, - use_global_self_attn=use_global_self_attn, - separate_global_qkv=separate_global_qkv, - global_dim_ratio=global_dim_ratio, - checkpoint_level=checkpoint_level, - use_relative_pos=use_relative_pos, - self_attn_use_final_proj=self_attn_use_final_proj, - use_first_self_attn=dec_use_first_self_attn, - attn_linear_init_mode=attn_linear_init_mode, - ffn_linear_init_mode=ffn_linear_init_mode, - conv_init_mode=conv_init_mode, - up_linear_init_mode=down_up_linear_init_mode, - norm_init_mode=norm_init_mode, - ) - self.reset_parameters() - - def get_initial_encoder_final_decoder( - self, - initial_downsample_type, - activation, - initial_downsample_scale, - initial_downsample_conv_layers, - final_upsample_conv_layers, - padding_type, - initial_downsample_stack_conv_num_layers, - initial_downsample_stack_conv_dim_list, - initial_downsample_stack_conv_downscale_list, - initial_downsample_stack_conv_num_conv_list, - ): - T_in, H_in, W_in, C_in = self.input_shape - T_out, H_out, W_out, C_out = self.target_shape - self.initial_downsample_type = initial_downsample_type - if self.initial_downsample_type == "conv": - if isinstance(initial_downsample_scale, int): - initial_downsample_scale = ( - 1, - initial_downsample_scale, - initial_downsample_scale, - ) - elif len(initial_downsample_scale) == 2: - initial_downsample_scale = 1, *initial_downsample_scale - elif len(initial_downsample_scale) == 3: - initial_downsample_scale = tuple(initial_downsample_scale) - else: - raise NotImplementedError( - f"initial_downsample_scale {initial_downsample_scale} format not supported!" - ) - self.initial_encoder = InitialEncoder( - dim=C_in, - out_dim=self.base_units, - downsample_scale=initial_downsample_scale, - num_conv_layers=initial_downsample_conv_layers, - padding_type=padding_type, - activation=activation, - conv_init_mode=self.conv_init_mode, - linear_init_mode=self.down_up_linear_init_mode, - norm_init_mode=self.norm_init_mode, - ) - - self.final_decoder = FinalDecoder( - dim=self.base_units, - target_thw=(T_out, H_out, W_out), - num_conv_layers=final_upsample_conv_layers, - activation=activation, - conv_init_mode=self.conv_init_mode, - linear_init_mode=self.down_up_linear_init_mode, - norm_init_mode=self.norm_init_mode, - ) - new_input_shape = self.initial_encoder.patch_merge.get_out_shape( - self.input_shape - ) - self.dec_final_proj = nn.Linear( - in_features=self.base_units, out_features=C_out - ) - elif self.initial_downsample_type == "stack_conv": - if initial_downsample_stack_conv_dim_list is None: - initial_downsample_stack_conv_dim_list = [ - self.base_units - ] * initial_downsample_stack_conv_num_layers - self.initial_encoder = InitialStackPatchMergingEncoder( - num_merge=initial_downsample_stack_conv_num_layers, - in_dim=C_in, - out_dim_list=initial_downsample_stack_conv_dim_list, - downsample_scale_list=initial_downsample_stack_conv_downscale_list, - num_conv_per_merge_list=initial_downsample_stack_conv_num_conv_list, - padding_type=padding_type, - activation=activation, - conv_init_mode=self.conv_init_mode, - linear_init_mode=self.down_up_linear_init_mode, - norm_init_mode=self.norm_init_mode, - ) - initial_encoder_out_shape_list = self.initial_encoder.get_out_shape_list( - self.target_shape - ) - ( - dec_target_shape_list, - dec_in_dim, - ) = FinalStackUpsamplingDecoder.get_init_params( - enc_input_shape=self.target_shape, - enc_out_shape_list=initial_encoder_out_shape_list, - large_channel=True, - ) - self.final_decoder = FinalStackUpsamplingDecoder( - target_shape_list=dec_target_shape_list, - in_dim=dec_in_dim, - num_conv_per_up_list=initial_downsample_stack_conv_num_conv_list[::-1], - activation=activation, - conv_init_mode=self.conv_init_mode, - linear_init_mode=self.down_up_linear_init_mode, - norm_init_mode=self.norm_init_mode, - ) - self.dec_final_proj = nn.Linear( - in_features=dec_target_shape_list[-1][-1], out_features=C_out - ) - new_input_shape = self.initial_encoder.get_out_shape_list(self.input_shape)[ - -1 - ] - else: - raise NotImplementedError(f"{self.initial_downsample_type} is invalid.") - self.input_shape_after_initial_downsample = new_input_shape - T_in, H_in, W_in, _ = new_input_shape - return new_input_shape - - def reset_parameters(self): - if self.num_global_vectors > 0: - self.init_global_vectors = initializer.trunc_normal_( - self.init_global_vectors, std=0.02 - ) - if hasattr(self.initial_encoder, "reset_parameters"): - self.initial_encoder.reset_parameters() - else: - cuboid_utils.apply_initialization( - self.initial_encoder, - conv_mode=self.conv_init_mode, - linear_mode=self.down_up_linear_init_mode, - norm_mode=self.norm_init_mode, - ) - if hasattr(self.final_decoder, "reset_parameters"): - self.final_decoder.reset_parameters() - else: - cuboid_utils.apply_initialization( - self.final_decoder, - conv_mode=self.conv_init_mode, - linear_mode=self.down_up_linear_init_mode, - norm_mode=self.norm_init_mode, - ) - cuboid_utils.apply_initialization( - self.dec_final_proj, linear_mode=self.down_up_linear_init_mode - ) - self.encoder.reset_parameters() - self.enc_pos_embed.reset_parameters() - self.decoder.reset_parameters() - self.dec_pos_embed.reset_parameters() - cuboid_utils.apply_initialization(self.z_proj, linear_mode="0") - - def get_initial_z(self, final_mem, T_out): - B = final_mem.shape[0] - if self.z_init_method == "zeros": - z_shape = list((1, T_out)) + final_mem.shape[2:] - initial_z = paddle.zeros(shape=z_shape, dtype=final_mem.dtype) - initial_z = self.z_proj(self.dec_pos_embed(initial_z)).expand( - shape=[B, -1, -1, -1, -1] - ) - elif self.z_init_method == "nearest_interp": - initial_z = nn.functional.interpolate( - x=final_mem.transpose(perm=[0, 4, 1, 2, 3]), - size=(T_out, final_mem.shape[2], final_mem.shape[3]), - ).transpose(perm=[0, 2, 3, 4, 1]) - initial_z = self.z_proj(initial_z) - elif self.z_init_method == "last": - initial_z = paddle.broadcast_to( - x=final_mem[:, -1:, :, :, :], shape=(B, T_out) + final_mem.shape[2:] - ) - initial_z = self.z_proj(initial_z) - elif self.z_init_method == "mean": - initial_z = paddle.broadcast_to( - x=final_mem.mean(axis=1, keepdims=True), - shape=(B, T_out) + final_mem.shape[2:], - ) - initial_z = self.z_proj(initial_z) - else: - raise NotImplementedError - return initial_z - - def forward(self, x: "paddle.Tensor", verbose: bool = False) -> "paddle.Tensor": - """ - Args: - x (paddle.Tensor): Tensor with shape (B, T, H, W, C). - verbose (bool): If True, print intermediate shapes. - - Returns: - out (paddle.Tensor): The output Shape (B, T_out, H, W, C_out) - """ - - x = self.concat_to_tensor(x, self.input_keys) - flag_ndim = x.ndim - if flag_ndim == 6: - x = x.reshape([-1, *x.shape[2:]]) - B, _, _, _, _ = x.shape - - T_out = self.target_shape[0] - x = self.initial_encoder(x) - x = self.enc_pos_embed(x) - - if self.num_global_vectors > 0: - init_global_vectors = self.init_global_vectors.expand( - shape=[ - B, - self.num_global_vectors, - self.global_dim_ratio * self.base_units, - ] - ) - mem_l, mem_global_vector_l = self.encoder(x, init_global_vectors) - else: - mem_l = self.encoder(x) - - if verbose: - for i, mem in enumerate(mem_l): - print(f"mem[{i}].shape = {mem.shape}") - initial_z = self.get_initial_z(final_mem=mem_l[-1], T_out=T_out) - - if self.num_global_vectors > 0: - dec_out = self.decoder(initial_z, mem_l, mem_global_vector_l) - else: - dec_out = self.decoder(initial_z, mem_l) - - dec_out = self.final_decoder(dec_out) - - out = self.dec_final_proj(dec_out) - if flag_ndim == 6: - out = out.reshape([-1, *out.shape]) - return {key: out for key in self.output_keys} diff --git a/examples/smc_reac/ppsci/arch/cuboid_transformer_decoder.py b/examples/smc_reac/ppsci/arch/cuboid_transformer_decoder.py deleted file mode 100644 index 8cbaf5fee9..0000000000 --- a/examples/smc_reac/ppsci/arch/cuboid_transformer_decoder.py +++ /dev/null @@ -1,1245 +0,0 @@ -from functools import lru_cache -from typing import Tuple - -import numpy as np -import paddle -import paddle.nn.functional as F -from paddle import nn -from paddle.distributed import fleet - -import ppsci.arch.cuboid_transformer_encoder as cuboid_encoder -import ppsci.arch.cuboid_transformer_utils as cuboid_utils -from ppsci.utils import initializer - - -class PosEmbed(nn.Layer): - """Pose embedding - - Args: - embed_dim (int): The dimension of embedding. - maxT (int): The embedding max time. - maxH (int): The embedding max height. - maxW (int): The embedding max width. - typ (str): - The type of the positional embedding. - - t+h+w: - Embed the spatial position to embeddings - - t+hw: - Embed the spatial position to embeddings - """ - - def __init__(self, embed_dim, maxT, maxH, maxW, typ: str = "t+h+w"): - super(PosEmbed, self).__init__() - self.typ = typ - assert self.typ in ["t+h+w", "t+hw"] - self.maxT = maxT - self.maxH = maxH - self.maxW = maxW - self.embed_dim = embed_dim - if self.typ == "t+h+w": - self.T_embed = nn.Embedding(num_embeddings=maxT, embedding_dim=embed_dim) - self.H_embed = nn.Embedding(num_embeddings=maxH, embedding_dim=embed_dim) - self.W_embed = nn.Embedding(num_embeddings=maxW, embedding_dim=embed_dim) - elif self.typ == "t+hw": - self.T_embed = nn.Embedding(num_embeddings=maxT, embedding_dim=embed_dim) - self.HW_embed = nn.Embedding( - num_embeddings=maxH * maxW, embedding_dim=embed_dim - ) - else: - raise NotImplementedError(f"{self.typ} is invalid.") - self.reset_parameters() - - def reset_parameters(self): - for m in self.children(): - cuboid_utils.apply_initialization(m, embed_mode="0") - - def forward(self, x): - """ - Args: - x : Shape (B, T, H, W, C) - - Returns: - out : the x + positional embeddings - """ - - _, T, H, W, _ = x.shape - t_idx = paddle.arange(end=T) - h_idx = paddle.arange(end=H) - w_idx = paddle.arange(end=W) - if self.typ == "t+h+w": - return ( - x - + self.T_embed(t_idx).reshape([T, 1, 1, self.embed_dim]) - + self.H_embed(h_idx).reshape([1, H, 1, self.embed_dim]) - + self.W_embed(w_idx).reshape([1, 1, W, self.embed_dim]) - ) - elif self.typ == "t+hw": - spatial_idx = h_idx.unsqueeze(axis=-1) * self.maxW + w_idx - return ( - x - + self.T_embed(t_idx).reshape([T, 1, 1, self.embed_dim]) - + self.HW_embed(spatial_idx) - ) - else: - raise NotImplementedError(f"{self.typ} is invalid.") - - -@lru_cache() -def compute_cuboid_cross_attention_mask( - T_x, T_mem, H, W, n_temporal, cuboid_hw, shift_hw, strategy, padding_type, device -): - pad_t_mem = (n_temporal - T_mem % n_temporal) % n_temporal - pad_t_x = (n_temporal - T_x % n_temporal) % n_temporal - pad_h = (cuboid_hw[0] - H % cuboid_hw[0]) % cuboid_hw[0] - pad_w = (cuboid_hw[1] - W % cuboid_hw[1]) % cuboid_hw[1] - mem_cuboid_size = ((T_mem + pad_t_mem) // n_temporal,) + cuboid_hw - x_cuboid_size = ((T_x + pad_t_x) // n_temporal,) + cuboid_hw - if pad_t_mem > 0 or pad_h > 0 or pad_w > 0: - if padding_type == "ignore": - mem_mask = paddle.ones(shape=(1, T_mem, H, W, 1), dtype="bool") - mem_mask = F.pad( - mem_mask, [0, 0, 0, pad_w, 0, pad_h, pad_t_mem, 0], data_format="NDHWC" - ) - else: - mem_mask = paddle.ones( - shape=(1, T_mem + pad_t_mem, H + pad_h, W + pad_w, 1), dtype="bool" - ) - if pad_t_x > 0 or pad_h > 0 or pad_w > 0: - if padding_type == "ignore": - x_mask = paddle.ones(shape=(1, T_x, H, W, 1), dtype="bool") - x_mask = F.pad( - x_mask, [0, 0, 0, pad_w, 0, pad_h, 0, pad_t_x], data_format="NDHWC" - ) - else: - x_mask = paddle.ones( - shape=(1, T_x + pad_t_x, H + pad_h, W + pad_w, 1), dtype="bool" - ) - if any(i > 0 for i in shift_hw): - if padding_type == "ignore": - x_mask = paddle.roll( - x=x_mask, shifts=(-shift_hw[0], -shift_hw[1]), axis=(2, 3) - ) - mem_mask = paddle.roll( - x=mem_mask, shifts=(-shift_hw[0], -shift_hw[1]), axis=(2, 3) - ) - x_mask = cuboid_encoder.cuboid_reorder(x_mask, x_cuboid_size, strategy=strategy) - x_mask = x_mask.squeeze(axis=-1).squeeze(axis=0) - num_cuboids, x_cuboid_volume = x_mask.shape - mem_mask = cuboid_encoder.cuboid_reorder( - mem_mask, mem_cuboid_size, strategy=strategy - ) - mem_mask = mem_mask.squeeze(axis=-1).squeeze(axis=0) - _, mem_cuboid_volume = mem_mask.shape - shift_mask = np.zeros(shape=(1, n_temporal, H + pad_h, W + pad_w, 1)) - cnt = 0 - for h in ( - slice(-cuboid_hw[0]), - slice(-cuboid_hw[0], -shift_hw[0]), - slice(-shift_hw[0], None), - ): - for w in ( - slice(-cuboid_hw[1]), - slice(-cuboid_hw[1], -shift_hw[1]), - slice(-shift_hw[1], None), - ): - shift_mask[:, :, h, w, :] = cnt - cnt += 1 - shift_mask = paddle.to_tensor(shift_mask) - shift_mask = cuboid_encoder.cuboid_reorder( - shift_mask, (1,) + cuboid_hw, strategy=strategy - ) - shift_mask = shift_mask.squeeze(axis=-1).squeeze(axis=0) - shift_mask = shift_mask.unsqueeze(axis=1) - shift_mask.unsqueeze(axis=2) == 0 - bh_bw = cuboid_hw[0] * cuboid_hw[1] - attn_mask = ( - shift_mask.reshape((num_cuboids, 1, bh_bw, 1, bh_bw)) - * x_mask.reshape((num_cuboids, -1, bh_bw, 1, 1)) - * mem_mask.reshape([num_cuboids, 1, 1, -1, bh_bw]) - ) - attn_mask = attn_mask.reshape([num_cuboids, x_cuboid_volume, mem_cuboid_volume]) - return attn_mask - - -class CuboidCrossAttentionLayer(nn.Layer): - """Implements the cuboid cross attention. - - The idea of Cuboid Cross Attention is to extend the idea of cuboid self attention to work for the - encoder-decoder-type cross attention. - - Assume that there is a memory tensor with shape (T1, H, W, C) and another query tensor with shape (T2, H, W, C), - - Here, we decompose the query tensor and the memory tensor into the same number of cuboids and attend the cuboid in - the query tensor with the corresponding cuboid in the memory tensor. - - For the height and width axes, we reuse the grid decomposition techniques described in the cuboid self-attention. - For the temporal axis, the layer supports the "n_temporal" parameter, that controls the number of cuboids we can - get after cutting the tensors. For example, if the temporal dilation is 2, both the query and - memory will be decomposed into 2 cuboids along the temporal axis. Like in the Cuboid Self-attention, - we support "local" and "dilated" decomposition strategy. - - The complexity of the layer is O((T2 / n_t * Bh * Bw) * (T1 / n_t * Bh * Bw) * n_t (H / Bh) (W / Bw)) = O(T2 * T1 / n_t H W Bh Bw) - - Args: - dim (int): The dimension of input tensor. - num_heads (int): The number of head. - n_temporal (int, optional): The num of temporal. Defaults to 1. - cuboid_hw (tuple, optional): The height and width of cuboid. Defaults to (7, 7). - shift_hw (tuple, optional): The height and width of shift. Defaults to (0, 0). - strategy (tuple, optional): The strategy. Defaults to ("d", "l", "l"). - padding_type (str, optional): The type of padding. Defaults to "ignore". - cross_last_n_frames (int, optional): The cross_last_n_frames of decoder. Defaults to None. - qkv_bias (bool, optional): Whether to enable bias in calculating qkv attention. Defaults to False. - qk_scale (float, optional): Whether to enable scale factor when calculating the attention. Defaults to None. - attn_drop (float, optional): The attention dropout. Defaults to 0.0. - proj_drop (float, optional): The projrction dropout. Defaults to 0.0. - max_temporal_relative (int, optional): The max temporal. Defaults to 50. - norm_layer (str, optional): The normalization layer. Defaults to "layer_norm". - use_global_vector (bool, optional): Whether to use the global vector or not. Defaults to True. - separate_global_qkv (bool, optional): Whether to use different network to calc q_global, k_global, v_global. Defaults to False. - global_dim_ratio (int, optional): The dim (channels) of global vectors is `global_dim_ratio*dim`. Defaults to 1. - checkpoint_level (int, optional): Whether to enable gradient checkpointing. Defaults to 1. - use_relative_pos (bool, optional): Whether to use relative pos. Defaults to True. - attn_linear_init_mode (str, optional): The mode of attention linear initialization. Defaults to "0". - ffn_linear_init_mode (str, optional): The mode of FFN linear initialization. Defaults to "0". - norm_init_mode (str, optional): The mode of normalization initialization. Defaults to "0". - """ - - def __init__( - self, - dim: int, - num_heads: int, - n_temporal: int = 1, - cuboid_hw: Tuple[int, ...] = (7, 7), - shift_hw: Tuple[int, ...] = (0, 0), - strategy: Tuple[str, ...] = ("d", "l", "l"), - padding_type: str = "ignore", - cross_last_n_frames: int = None, - qkv_bias: bool = False, - qk_scale: float = None, - attn_drop: float = 0.0, - proj_drop: float = 0.0, - max_temporal_relative: int = 50, - norm_layer: str = "layer_norm", - use_global_vector: bool = True, - separate_global_qkv: bool = False, - global_dim_ratio: int = 1, - checkpoint_level: int = 1, - use_relative_pos: bool = True, - attn_linear_init_mode: str = "0", - ffn_linear_init_mode: str = "0", - norm_init_mode: str = "0", - ): - super(CuboidCrossAttentionLayer, self).__init__() - self.attn_linear_init_mode = attn_linear_init_mode - self.ffn_linear_init_mode = ffn_linear_init_mode - self.norm_init_mode = norm_init_mode - self.dim = dim - self.num_heads = num_heads - self.n_temporal = n_temporal - assert n_temporal > 0 - head_dim = dim // num_heads - self.scale = qk_scale or head_dim**-0.5 - shift_hw = list(shift_hw) - if strategy[1] == "d": - shift_hw[0] = 0 - if strategy[2] == "d": - shift_hw[1] = 0 - self.cuboid_hw = cuboid_hw - self.shift_hw = tuple(shift_hw) - self.strategy = strategy - self.padding_type = padding_type - self.max_temporal_relative = max_temporal_relative - self.cross_last_n_frames = cross_last_n_frames - self.use_relative_pos = use_relative_pos - self.use_global_vector = use_global_vector - self.separate_global_qkv = separate_global_qkv - if global_dim_ratio != 1 and separate_global_qkv is False: - raise ValueError( - "Setting global_dim_ratio != 1 requires separate_global_qkv == True." - ) - self.global_dim_ratio = global_dim_ratio - if self.padding_type not in ["ignore", "zeros", "nearest"]: - raise ValueError('padding_type should be ["ignore", "zeros", "nearest"]') - if use_relative_pos: - init_data = paddle.zeros( - ( - (2 * max_temporal_relative - 1) - * (2 * cuboid_hw[0] - 1) - * (2 * cuboid_hw[1] - 1), - num_heads, - ) - ) - self.relative_position_bias_table = paddle.create_parameter( - shape=init_data.shape, - dtype=init_data.dtype, - default_initializer=nn.initializer.Constant(0.0), - ) - self.relative_position_bias_table.stop_gradient = not True - self.relative_position_bias_table = initializer.trunc_normal_( - self.relative_position_bias_table, std=0.02 - ) - - coords_t = paddle.arange(end=max_temporal_relative) - coords_h = paddle.arange(end=self.cuboid_hw[0]) - coords_w = paddle.arange(end=self.cuboid_hw[1]) - coords = paddle.stack(x=paddle.meshgrid(coords_t, coords_h, coords_w)) - coords_flatten = paddle.flatten(x=coords, start_axis=1) - relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] - relative_coords = relative_coords.transpose(perm=[1, 2, 0]) - relative_coords[:, :, 0] += max_temporal_relative - 1 - relative_coords[:, :, 1] += self.cuboid_hw[0] - 1 - relative_coords[:, :, 2] += self.cuboid_hw[1] - 1 - relative_position_index = ( - relative_coords[:, :, 0] - * (2 * self.cuboid_hw[0] - 1) - * (2 * self.cuboid_hw[1] - 1) - + relative_coords[:, :, 1] * (2 * self.cuboid_hw[1] - 1) - + relative_coords[:, :, 2] - ) - self.register_buffer( - name="relative_position_index", tensor=relative_position_index - ) - self.q_proj = nn.Linear(in_features=dim, out_features=dim, bias_attr=qkv_bias) - self.kv_proj = nn.Linear( - in_features=dim, out_features=dim * 2, bias_attr=qkv_bias - ) - self.attn_drop = nn.Dropout(p=attn_drop) - self.proj = nn.Linear(in_features=dim, out_features=dim) - self.proj_drop = nn.Dropout(p=proj_drop) - if self.use_global_vector: - if self.separate_global_qkv: - self.l2g_q_net = nn.Linear( - in_features=dim, out_features=dim, bias_attr=qkv_bias - ) - self.l2g_global_kv_net = nn.Linear( - in_features=global_dim_ratio * dim, - out_features=dim * 2, - bias_attr=qkv_bias, - ) - self.norm = cuboid_utils.get_norm_layer(norm_layer, in_channels=dim) - self._checkpoint_level = checkpoint_level - self.reset_parameters() - - def reset_parameters(self): - cuboid_utils.apply_initialization( - self.q_proj, linear_mode=self.attn_linear_init_mode - ) - cuboid_utils.apply_initialization( - self.kv_proj, linear_mode=self.attn_linear_init_mode - ) - cuboid_utils.apply_initialization( - self.proj, linear_mode=self.ffn_linear_init_mode - ) - cuboid_utils.apply_initialization(self.norm, norm_mode=self.norm_init_mode) - if self.use_global_vector: - if self.separate_global_qkv: - cuboid_utils.apply_initialization( - self.l2g_q_net, linear_mode=self.attn_linear_init_mode - ) - cuboid_utils.apply_initialization( - self.l2g_global_kv_net, linear_mode=self.attn_linear_init_mode - ) - - def forward(self, x, mem, mem_global_vectors=None): - """Calculate the forward - - Along the temporal axis, we pad the mem tensor from the left and the x tensor from the right so that the - relative position encoding can be calculated correctly. For example: - - mem: 0, 1, 2, 3, 4 - x: 0, 1, 2, 3, 4, 5 - - n_temporal = 1 - mem: 0, 1, 2, 3, 4 x: 0, 1, 2, 3, 4, 5 - - n_temporal = 2 - mem: pad, 1, 3 x: 0, 2, 4 - mem: 0, 2, 4 x: 1, 3, 5 - - n_temporal = 3 - mem: pad, 2 dec: 0, 3 - mem: 0, 3 dec: 1, 4 - mem: 1, 4 dec: 2, 5 - - Args: - x (paddle.Tensor): The input of the layer. It will have shape (B, T, H, W, C) - mem (paddle.Tensor): The memory. It will have shape (B, T_mem, H, W, C) - mem_global_vectors (paddle.Tensor): The global vectors from the memory. It will have shape (B, N, C) - - Returns: - out (paddle.Tensor): Output tensor should have shape (B, T, H, W, C_out) - """ - - if self.cross_last_n_frames is not None: - cross_last_n_frames = int(min(self.cross_last_n_frames, mem.shape[1])) - mem = mem[:, -cross_last_n_frames:, ...] - if self.use_global_vector: - _, num_global, _ = mem_global_vectors.shape - x = self.norm(x) - B, T_x, H, W, C_in = x.shape - B_mem, T_mem, H_mem, W_mem, C_mem = mem.shape - assert T_x < self.max_temporal_relative and T_mem < self.max_temporal_relative - cuboid_hw = self.cuboid_hw - n_temporal = self.n_temporal - shift_hw = self.shift_hw - assert ( - B_mem == B and H == H_mem and W == W_mem and C_in == C_mem - ), f"Shape of memory and the input tensor does not match. x.shape={x.shape}, mem.shape={mem.shape}" - pad_t_mem = (n_temporal - T_mem % n_temporal) % n_temporal - pad_t_x = (n_temporal - T_x % n_temporal) % n_temporal - pad_h = (cuboid_hw[0] - H % cuboid_hw[0]) % cuboid_hw[0] - pad_w = (cuboid_hw[1] - W % cuboid_hw[1]) % cuboid_hw[1] - mem = cuboid_utils.generalize_padding( - mem, pad_t_mem, pad_h, pad_w, self.padding_type, t_pad_left=True - ) - - x = cuboid_utils.generalize_padding( - x, pad_t_x, pad_h, pad_w, self.padding_type, t_pad_left=False - ) - - if any(i > 0 for i in shift_hw): - shifted_x = paddle.roll( - x=x, shifts=(-shift_hw[0], -shift_hw[1]), axis=(2, 3) - ) - shifted_mem = paddle.roll( - x=mem, shifts=(-shift_hw[0], -shift_hw[1]), axis=(2, 3) - ) - else: - shifted_x = x - shifted_mem = mem - mem_cuboid_size = (mem.shape[1] // n_temporal,) + cuboid_hw - x_cuboid_size = (x.shape[1] // n_temporal,) + cuboid_hw - reordered_mem = cuboid_encoder.cuboid_reorder( - shifted_mem, cuboid_size=mem_cuboid_size, strategy=self.strategy - ) - reordered_x = cuboid_encoder.cuboid_reorder( - shifted_x, cuboid_size=x_cuboid_size, strategy=self.strategy - ) - _, num_cuboids_mem, mem_cuboid_volume, _ = reordered_mem.shape - _, num_cuboids, x_cuboid_volume, _ = reordered_x.shape - assert ( - num_cuboids_mem == num_cuboids - ), f"Number of cuboids do not match. num_cuboids={num_cuboids}, num_cuboids_mem={num_cuboids_mem}" - attn_mask = compute_cuboid_cross_attention_mask( - T_x, - T_mem, - H, - W, - n_temporal, - cuboid_hw, - shift_hw, - strategy=self.strategy, - padding_type=self.padding_type, - device=x.place, - ) - head_C = C_in // self.num_heads - kv = ( - self.kv_proj(reordered_mem) - .reshape([B, num_cuboids, mem_cuboid_volume, 2, self.num_heads, head_C]) - .transpose(perm=[3, 0, 4, 1, 2, 5]) - ) - k, v = kv[0], kv[1] - q = ( - self.q_proj(reordered_x) - .reshape([B, num_cuboids, x_cuboid_volume, self.num_heads, head_C]) - .transpose(perm=[0, 3, 1, 2, 4]) - ) - q = q * self.scale - perm_4 = list(range(k.ndim)) - perm_4[-2] = -1 - perm_4[-1] = -2 - attn_score = q @ k.transpose(perm=perm_4) - if self.use_relative_pos: - relative_position_bias = self.relative_position_bias_table[ - self.relative_position_index[ - :x_cuboid_volume, :mem_cuboid_volume - ].reshape([-1]) - ].reshape([x_cuboid_volume, mem_cuboid_volume, -1]) - relative_position_bias = relative_position_bias.transpose( - perm=[2, 0, 1] - ).unsqueeze(axis=1) - attn_score = attn_score + relative_position_bias - if self.use_global_vector: - if self.separate_global_qkv: - l2g_q = ( - self.l2g_q_net(reordered_x) - .reshape([B, num_cuboids, x_cuboid_volume, self.num_heads, head_C]) - .transpose(perm=[0, 3, 1, 2, 4]) - ) - l2g_q = l2g_q * self.scale - l2g_global_kv = ( - self.l2g_global_kv_net(mem_global_vectors) - .reshape([B, 1, num_global, 2, self.num_heads, head_C]) - .transpose(perm=[3, 0, 4, 1, 2, 5]) - ) - l2g_global_k, l2g_global_v = l2g_global_kv[0], l2g_global_kv[1] - else: - kv_global = ( - self.kv_proj(mem_global_vectors) - .reshape([B, 1, num_global, 2, self.num_heads, head_C]) - .transpose(perm=[3, 0, 4, 1, 2, 5]) - ) - l2g_global_k, l2g_global_v = kv_global[0], kv_global[1] - l2g_q = q - perm_5 = list(range(l2g_global_k.ndim)) - perm_5[-2] = -1 - perm_5[-1] = -2 - l2g_attn_score = l2g_q @ l2g_global_k.transpose(perm=perm_5) - attn_score_l2l_l2g = paddle.concat(x=(attn_score, l2g_attn_score), axis=-1) - if attn_mask.ndim == 5: - attn_mask_l2l_l2g = F.pad( - attn_mask, [0, num_global], "constant", 1, data_format="NDHWC" - ) - else: - attn_mask_l2l_l2g = F.pad(attn_mask, [0, num_global], "constant", 1) - v_l_g = paddle.concat( - x=( - v, - l2g_global_v.expand( - shape=[B, self.num_heads, num_cuboids, num_global, head_C] - ), - ), - axis=3, - ) - attn_score_l2l_l2g = cuboid_encoder.masked_softmax( - attn_score_l2l_l2g, mask=attn_mask_l2l_l2g - ) - attn_score_l2l_l2g = self.attn_drop(attn_score_l2l_l2g) - reordered_x = ( - (attn_score_l2l_l2g @ v_l_g) - .transpose(perm=[0, 2, 3, 1, 4]) - .reshape(B, num_cuboids, x_cuboid_volume, self.dim) - ) - else: - attn_score = cuboid_encoder.masked_softmax(attn_score, mask=attn_mask) - attn_score = self.attn_drop(attn_score) - reordered_x = ( - (attn_score @ v) - .transpose(perm=[0, 2, 3, 1, 4]) - .reshape([B, num_cuboids, x_cuboid_volume, self.dim]) - ) - reordered_x = paddle.cast(reordered_x, dtype="float32") - reordered_x = self.proj_drop(self.proj(reordered_x)) - shifted_x = cuboid_encoder.cuboid_reorder_reverse( - reordered_x, - cuboid_size=x_cuboid_size, - strategy=self.strategy, - orig_data_shape=(x.shape[1], x.shape[2], x.shape[3]), - ) - if any(i > 0 for i in shift_hw): - x = paddle.roll(x=shifted_x, shifts=(shift_hw[0], shift_hw[1]), axis=(2, 3)) - else: - x = shifted_x - x = cuboid_utils.generalize_unpadding( - x, pad_t=pad_t_x, pad_h=pad_h, pad_w=pad_w, padding_type=self.padding_type - ) - return x - - -class StackCuboidCrossAttentionBlock(nn.Layer): - """A stack of cuboid cross attention layers. - - The advantage of cuboid attention is that we can combine cuboid attention building blocks with different - hyper-parameters to mimic a broad range of space-time correlation patterns. - - - "use_inter_ffn" is True - x, mem --> attn1 -----+-------> ffn1 ---+---> attn2 --> ... --> ffn_k --> out - | ^ | ^ - | | | | - |-------------|----|-------------| - - "use_inter_ffn" is False - x, mem --> attn1 -----+------> attn2 --> ... attnk --+----> ffnk ---+---> out, mem - | ^ | ^ ^ | ^ - | | | | | | | - |-------------|----|------------|-- ----------|--|-----------| - - Args: - dim (int): The dimension of the input. - num_heads (int): The number of head. - block_cuboid_hw (list, optional): The height and width of block cuboid.Defaults to [(4, 4), (4, 4)]. - block_shift_hw (list, optional): The height and width of shift cuboid . Defaults to [(0, 0), (2, 2)]. - block_n_temporal (list, optional): The length of block temporal. Defaults to [1, 2]. - block_strategy (list, optional): The strategy of block. Defaults to [("d", "d", "d"), ("l", "l", "l")]. - padding_type (str, optional): The type of paddling. Defaults to "ignore". - cross_last_n_frames (int, optional): The num of cross_last_n_frames. Defaults to None. - qkv_bias (bool, optional): Whether to enable bias in calculating qkv attention. Defaults to False. - qk_scale (float, optional): Whether to enable scale factor when calculating the attention. Defaults to None. - attn_drop (float, optional): The attention dropout. Defaults to 0.0. - proj_drop (float, optional): The projection dropout. Defaults to 0.0. - ffn_drop (float, optional): The ratio of FFN dropout. Defaults to 0.0. - activation (str, optional): The activation. Defaults to "leaky". - gated_ffn (bool, optional): Whether to use gate FFN. Defaults to False. - norm_layer (str, optional): The normalization layer. Defaults to "layer_norm". - use_inter_ffn (bool, optional): Whether to use inter FFN. Defaults to True. - max_temporal_relative (int, optional): The max temporal. Defaults to 50. - checkpoint_level (int, optional): Whether to enable gradient checkpointing. Defaults to 1. - use_relative_pos (bool, optional): Whether to use relative pos. Defaults to True. - use_global_vector (bool, optional): Whether to use the global vector or not. Defaults to False. - separate_global_qkv (bool, optional): Whether to use different network to calc q_global, k_global, v_global. Defaults to False. - global_dim_ratio (int, optional): The dim (channels) of global vectors is `global_dim_ratio*dim`. Defaults to 1. - attn_linear_init_mode (str, optional): The mode of attention linear initialization. Defaults to "0". - ffn_linear_init_mode (str, optional): The mode of FFN linear initialization. Defaults to "0". - norm_init_mode (str, optional): The mode of normalization. Defaults to "0". - """ - - def __init__( - self, - dim: int, - num_heads: int, - block_cuboid_hw: Tuple[Tuple[int, ...], ...] = [(4, 4), (4, 4)], - block_shift_hw: Tuple[Tuple[int, ...], ...] = [(0, 0), (2, 2)], - block_n_temporal: Tuple[int, ...] = [1, 2], - block_strategy: Tuple[Tuple[str, ...], ...] = [ - ("d", "d", "d"), - ("l", "l", "l"), - ], - padding_type: str = "ignore", - cross_last_n_frames: int = None, - qkv_bias: bool = False, - qk_scale: float = None, - attn_drop: float = 0.0, - proj_drop: float = 0.0, - ffn_drop: float = 0.0, - activation: str = "leaky", - gated_ffn: bool = False, - norm_layer: str = "layer_norm", - use_inter_ffn: bool = True, - max_temporal_relative: int = 50, - checkpoint_level: int = 1, - use_relative_pos: bool = True, - use_global_vector: bool = False, - separate_global_qkv: bool = False, - global_dim_ratio: int = 1, - attn_linear_init_mode: str = "0", - ffn_linear_init_mode: str = "0", - norm_init_mode: str = "0", - ): - super(StackCuboidCrossAttentionBlock, self).__init__() - self.attn_linear_init_mode = attn_linear_init_mode - self.ffn_linear_init_mode = ffn_linear_init_mode - self.norm_init_mode = norm_init_mode - if ( - len(block_cuboid_hw[0]) <= 0 - or len(block_shift_hw) <= 0 - or len(block_strategy) <= 0 - ): - raise ValueError( - "Incorrect format.The lengths of block_cuboid_hw[0], block_shift_hw, and block_strategy must be greater than zero." - ) - if len(block_cuboid_hw) != len(block_shift_hw) and len(block_shift_hw) == len( - block_strategy - ): - raise ValueError( - "The lengths of block_cuboid_size, block_shift_size, and block_strategy must be equal." - ) - - self.num_attn = len(block_cuboid_hw) - self.checkpoint_level = checkpoint_level - self.use_inter_ffn = use_inter_ffn - self.use_global_vector = use_global_vector - if self.use_inter_ffn: - self.ffn_l = nn.LayerList( - sublayers=[ - cuboid_encoder.PositionwiseFFN( - units=dim, - hidden_size=4 * dim, - activation_dropout=ffn_drop, - dropout=ffn_drop, - gated_proj=gated_ffn, - activation=activation, - normalization=norm_layer, - pre_norm=True, - linear_init_mode=ffn_linear_init_mode, - norm_init_mode=norm_init_mode, - ) - for _ in range(self.num_attn) - ] - ) - else: - self.ffn_l = nn.LayerList( - sublayers=[ - cuboid_encoder.PositionwiseFFN( - units=dim, - hidden_size=4 * dim, - activation_dropout=ffn_drop, - dropout=ffn_drop, - gated_proj=gated_ffn, - activation=activation, - normalization=norm_layer, - pre_norm=True, - linear_init_mode=ffn_linear_init_mode, - norm_init_mode=norm_init_mode, - ) - ] - ) - self.attn_l = nn.LayerList( - sublayers=[ - CuboidCrossAttentionLayer( - dim=dim, - num_heads=num_heads, - cuboid_hw=ele_cuboid_hw, - shift_hw=ele_shift_hw, - strategy=ele_strategy, - n_temporal=ele_n_temporal, - cross_last_n_frames=cross_last_n_frames, - padding_type=padding_type, - qkv_bias=qkv_bias, - qk_scale=qk_scale, - attn_drop=attn_drop, - proj_drop=proj_drop, - norm_layer=norm_layer, - max_temporal_relative=max_temporal_relative, - use_global_vector=use_global_vector, - separate_global_qkv=separate_global_qkv, - global_dim_ratio=global_dim_ratio, - checkpoint_level=checkpoint_level, - use_relative_pos=use_relative_pos, - attn_linear_init_mode=attn_linear_init_mode, - ffn_linear_init_mode=ffn_linear_init_mode, - norm_init_mode=norm_init_mode, - ) - for ele_cuboid_hw, ele_shift_hw, ele_strategy, ele_n_temporal in zip( - block_cuboid_hw, block_shift_hw, block_strategy, block_n_temporal - ) - ] - ) - - def reset_parameters(self): - for m in self.ffn_l: - m.reset_parameters() - for m in self.attn_l: - m.reset_parameters() - - def forward(self, x, mem, mem_global_vector=None): - """ - Args: - x (paddle.Tensor): Shape (B, T_x, H, W, C) - mem (paddle.Tensor): Shape (B, T_mem, H, W, C) - mem_global_vector (paddle.Tensor): Shape (B, N_global, C) - - Returns: - out (paddle.Tensor): (B, T_x, H, W, C_out) - """ - - if self.use_inter_ffn: - for attn, ffn in zip(self.attn_l, self.ffn_l): - if self.checkpoint_level >= 2 and self.training: - x = x + fleet.utils.recompute(attn, x, mem, mem_global_vector) - else: - x = x + attn(x, mem, mem_global_vector) - if self.checkpoint_level >= 1 and self.training: - x = fleet.utils.recompute(ffn, x) - else: - x = ffn(x) - return x - else: - for attn in self.attn_l: - if self.checkpoint_level >= 2 and self.training: - x = x + fleet.utils.recompute(attn, x, mem, mem_global_vector) - else: - x = x + attn(x, mem, mem_global_vector) - if self.checkpoint_level >= 1 and self.training: - x = fleet.utils.recompute(self.ffn_l[0], x) - else: - x = self.ffn_l[0](x) - return x - - -class Upsample3DLayer(nn.Layer): - """Upsampling based on nn.UpSampling and Conv3x3. - - If the temporal dimension remains the same: - x --> interpolation-2d (nearest) --> conv3x3(dim, out_dim) - Else: - x --> interpolation-3d (nearest) --> conv3x3x3(dim, out_dim) - - Args: - dim (int): The dimension of the input tensor. - out_dim (int): The dimension of the output tensor. - target_size (Tuple[int,...]): The size of output tensor. - temporal_upsample (bool, optional): Whether the temporal axis will go through upsampling. Defaults to False. - kernel_size (int, optional): The kernel size of the Conv2D layer. Defaults to 3. - layout (str, optional): The layout of the inputs. Defaults to "THWC". - conv_init_mode (str, optional): The mode of conv initialization. Defaults to "0". - """ - - def __init__( - self, - dim: int, - out_dim: int, - target_size: Tuple[int, ...], - temporal_upsample: bool = False, - kernel_size: int = 3, - layout: str = "THWC", - conv_init_mode: str = "0", - ): - super(Upsample3DLayer, self).__init__() - self.conv_init_mode = conv_init_mode - self.target_size = target_size - self.out_dim = out_dim - self.temporal_upsample = temporal_upsample - if temporal_upsample: - self.up = nn.Upsample(size=target_size, mode="nearest") - else: - self.up = nn.Upsample(size=(target_size[1], target_size[2]), mode="nearest") - self.conv = nn.Conv2D( - in_channels=dim, - out_channels=out_dim, - kernel_size=(kernel_size, kernel_size), - padding=(kernel_size // 2, kernel_size // 2), - ) - assert layout in ["THWC", "CTHW"] - self.layout = layout - self.reset_parameters() - - def reset_parameters(self): - for m in self.children(): - cuboid_utils.apply_initialization(m, conv_mode=self.conv_init_mode) - - def forward(self, x): - """ - - Args: - x : (B, T, H, W, C) or (B, C, T, H, W) - - Returns: - out : (B, T, H_new, W_out, C_out) or (B, C, T, H_out, W_out) - """ - - if self.layout == "THWC": - B, T, H, W, C = x.shape - if self.temporal_upsample: - x = x.transpose(perm=[0, 4, 1, 2, 3]) - return self.conv(self.up(x)).transpose(perm=[0, 2, 3, 4, 1]) - else: - assert self.target_size[0] == T - x = x.reshape([B * T, H, W, C]).transpose(perm=[0, 3, 1, 2]) - x = self.up(x) - return ( - self.conv(x) - .transpose(perm=[0, 2, 3, 1]) - .reshape(list((B,) + self.target_size + (self.out_dim,))) - ) - elif self.layout == "CTHW": - B, C, T, H, W = x.shape - if self.temporal_upsample: - return self.conv(self.up(x)) - else: - assert self.output_size[0] == T - x = x.transpose(perm=[0, 2, 1, 3, 4]) - x = x.reshape([B * T, C, H, W]) - return ( - self.conv(self.up(x)) - .reshape( - [ - B, - self.target_size[0], - self.out_dim, - self.target_size[1], - self.target_size[2], - ] - ) - .transpose(perm=[0, 2, 1, 3, 4]) - ) - - -class CuboidTransformerDecoder(nn.Layer): - """Decoder of the CuboidTransformer. - - For each block, we first apply the StackCuboidSelfAttention and then apply the StackCuboidCrossAttention - - Repeat the following structure K times - - x --> StackCuboidSelfAttention --> | - |----> StackCuboidCrossAttention (If used) --> out - mem --> | - - Args: - target_temporal_length (int): The temporal length of the target. - mem_shapes (Tuple[int,...]): The mem shapes of the decoder. - cross_start (int, optional): The block to start cross attention. Defaults to 0. - depth (list, optional): The number of layers for each block. Defaults to [2, 2]. - upsample_type (str, optional): The type of upsample. Defaults to "upsample". - upsample_kernel_size (int, optional): The kernel size of upsample. Defaults to 3. - block_self_attn_patterns (str, optional): The patterns of block attention. Defaults to None. - block_self_cuboid_size (list, optional): The size of cuboid block. Defaults to [(4, 4, 4), (4, 4, 4)]. - block_self_cuboid_strategy (list, optional): The strategy of cuboid. Defaults to [("l", "l", "l"), ("d", "d", "d")]. - block_self_shift_size (list, optional): The size of shift. Defaults to [(1, 1, 1), (0, 0, 0)]. - block_cross_attn_patterns (str, optional): The patterns of cross attentions. Defaults to None. - block_cross_cuboid_hw (list, optional): The height and width of cross cuboid. Defaults to [(4, 4), (4, 4)]. - block_cross_cuboid_strategy (list, optional): The strategy of cross cuboid. Defaults to [("l", "l", "l"), ("d", "l", "l")]. - block_cross_shift_hw (list, optional): The height and width of cross shift. Defaults to [(0, 0), (0, 0)]. - block_cross_n_temporal (list, optional): The cross temporal of block. Defaults to [1, 2]. - cross_last_n_frames (int, optional): The num of cross last frames. Defaults to None. - num_heads (int, optional): The num of head. Defaults to 4. - attn_drop (float, optional): The ratio of attention dropout. Defaults to 0.0. - proj_drop (float, optional): The ratio of projection dropout. Defaults to 0.0. - ffn_drop (float, optional): The ratio of FFN dropout. Defaults to 0.0. - ffn_activation (str, optional): The activation layer of FFN. Defaults to "leaky". - gated_ffn (bool, optional): Whether to use gate FFN. Defaults to False. - norm_layer (str, optional): The normalization layer. Defaults to "layer_norm". - use_inter_ffn (bool, optional): Whether to use inter FFN. Defaults to False. - hierarchical_pos_embed (bool, optional): Whether to use hierarchical pos_embed. Defaults to False. - pos_embed_type (str, optional): The type of pos embedding. Defaults to "t+hw". - max_temporal_relative (int, optional): The max number of teemporal relative. Defaults to 50. - padding_type (str, optional): The type of padding. Defaults to "ignore". - checkpoint_level (bool, optional): Whether to enable gradient checkpointing. Defaults to True. - use_relative_pos (bool, optional): Whether to use relative pos. Defaults to True. - self_attn_use_final_proj (bool, optional): Whether to use self attention for final projection. Defaults to True. - use_first_self_attn (bool, optional): Whether to use first self attention. Defaults to False. - use_self_global (bool, optional): Whether to use self global vector. Defaults to False. - self_update_global (bool, optional): Whether to update global vector. Defaults to True. - use_cross_global (bool, optional): Whether to use cross global vector. Defaults to False. - use_global_vector_ffn (bool, optional): Whether to use FFN global vectors. Defaults to True. - use_global_self_attn (bool, optional): Whether to use global self attention. Defaults to False. - separate_global_qkv (bool, optional): Whether to use different network to calc q_global, k_global, v_global. Defaults to False. - global_dim_ratio (int, optional): The dim (channels) of global vectors is `global_dim_ratio*dim`. Defaults to 1. - attn_linear_init_mode (str, optional): The mode of attention linear initialization. Defaults to "0". - ffn_linear_init_mode (str, optional): The mode of FFN linear initialization. Defaults to "0". - conv_init_mode (str, optional): The mode of conv initialization. Defaults to "0". - up_linear_init_mode (str, optional): The mode of up linear initialization. Defaults to "0". - norm_init_mode (str, optional): The mode of normalization initialization. Defaults to "0". - """ - - def __init__( - self, - target_temporal_length: int, - mem_shapes: Tuple[int, ...], - cross_start: int = 0, - depth: Tuple[int, ...] = [2, 2], - upsample_type: str = "upsample", - upsample_kernel_size: int = 3, - block_self_attn_patterns: str = None, - block_self_cuboid_size: Tuple[Tuple[int, ...], ...] = [(4, 4, 4), (4, 4, 4)], - block_self_cuboid_strategy: Tuple[Tuple[str, ...], ...] = [ - ("l", "l", "l"), - ("d", "d", "d"), - ], - block_self_shift_size: Tuple[Tuple[int, ...], ...] = [(1, 1, 1), (0, 0, 0)], - block_cross_attn_patterns: str = None, - block_cross_cuboid_hw: Tuple[Tuple[int, ...], ...] = [(4, 4), (4, 4)], - block_cross_cuboid_strategy: Tuple[Tuple[str, ...], ...] = [ - ("l", "l", "l"), - ("d", "l", "l"), - ], - block_cross_shift_hw: Tuple[Tuple[int, ...], ...] = [(0, 0), (0, 0)], - block_cross_n_temporal: Tuple[int, ...] = [1, 2], - cross_last_n_frames: int = None, - num_heads: int = 4, - attn_drop: float = 0.0, - proj_drop: float = 0.0, - ffn_drop: float = 0.0, - ffn_activation: str = "leaky", - gated_ffn: bool = False, - norm_layer: str = "layer_norm", - use_inter_ffn: bool = False, - hierarchical_pos_embed: bool = False, - pos_embed_type: str = "t+hw", - max_temporal_relative: int = 50, - padding_type: str = "ignore", - checkpoint_level: bool = True, - use_relative_pos: bool = True, - self_attn_use_final_proj: bool = True, - use_first_self_attn: bool = False, - use_self_global: bool = False, - self_update_global: bool = True, - use_cross_global: bool = False, - use_global_vector_ffn: bool = True, - use_global_self_attn: bool = False, - separate_global_qkv: bool = False, - global_dim_ratio: int = 1, - attn_linear_init_mode: str = "0", - ffn_linear_init_mode: str = "0", - conv_init_mode: str = "0", - up_linear_init_mode: str = "0", - norm_init_mode: str = "0", - ): - super(CuboidTransformerDecoder, self).__init__() - self.attn_linear_init_mode = attn_linear_init_mode - self.ffn_linear_init_mode = ffn_linear_init_mode - self.conv_init_mode = conv_init_mode - self.up_linear_init_mode = up_linear_init_mode - self.norm_init_mode = norm_init_mode - assert len(depth) == len(mem_shapes) - self.target_temporal_length = target_temporal_length - self.num_blocks = len(mem_shapes) - self.cross_start = cross_start - self.mem_shapes = mem_shapes - self.depth = depth - self.upsample_type = upsample_type - self.hierarchical_pos_embed = hierarchical_pos_embed - self.checkpoint_level = checkpoint_level - self.use_self_global = use_self_global - self.self_update_global = self_update_global - self.use_cross_global = use_cross_global - self.use_global_vector_ffn = use_global_vector_ffn - self.use_first_self_attn = use_first_self_attn - if block_self_attn_patterns is not None: - if isinstance(block_self_attn_patterns, (tuple, list)): - assert len(block_self_attn_patterns) == self.num_blocks - else: - block_self_attn_patterns = [ - block_self_attn_patterns for _ in range(self.num_blocks) - ] - block_self_cuboid_size = [] - block_self_cuboid_strategy = [] - block_self_shift_size = [] - for idx, key in enumerate(block_self_attn_patterns): - func = cuboid_utils.CuboidSelfAttentionPatterns.get(key) - cuboid_size, strategy, shift_size = func(mem_shapes[idx]) - block_self_cuboid_size.append(cuboid_size) - block_self_cuboid_strategy.append(strategy) - block_self_shift_size.append(shift_size) - else: - if not isinstance(block_self_cuboid_size[0][0], (list, tuple)): - block_self_cuboid_size = [ - block_self_cuboid_size for _ in range(self.num_blocks) - ] - else: - assert ( - len(block_self_cuboid_size) == self.num_blocks - ), f"Incorrect input format! Received block_self_cuboid_size={block_self_cuboid_size}" - if not isinstance(block_self_cuboid_strategy[0][0], (list, tuple)): - block_self_cuboid_strategy = [ - block_self_cuboid_strategy for _ in range(self.num_blocks) - ] - else: - assert ( - len(block_self_cuboid_strategy) == self.num_blocks - ), f"Incorrect input format! Received block_self_cuboid_strategy={block_self_cuboid_strategy}" - if not isinstance(block_self_shift_size[0][0], (list, tuple)): - block_self_shift_size = [ - block_self_shift_size for _ in range(self.num_blocks) - ] - else: - assert ( - len(block_self_shift_size) == self.num_blocks - ), f"Incorrect input format! Received block_self_shift_size={block_self_shift_size}" - self_blocks = [] - for i in range(self.num_blocks): - if not self.use_first_self_attn and i == self.num_blocks - 1: - ele_depth = depth[i] - 1 - else: - ele_depth = depth[i] - stack_cuboid_blocks = [ - cuboid_encoder.StackCuboidSelfAttentionBlock( - dim=self.mem_shapes[i][-1], - num_heads=num_heads, - block_cuboid_size=block_self_cuboid_size[i], - block_strategy=block_self_cuboid_strategy[i], - block_shift_size=block_self_shift_size[i], - attn_drop=attn_drop, - proj_drop=proj_drop, - ffn_drop=ffn_drop, - activation=ffn_activation, - gated_ffn=gated_ffn, - norm_layer=norm_layer, - use_inter_ffn=use_inter_ffn, - padding_type=padding_type, - use_global_vector=use_self_global, - use_global_vector_ffn=use_global_vector_ffn, - use_global_self_attn=use_global_self_attn, - separate_global_qkv=separate_global_qkv, - global_dim_ratio=global_dim_ratio, - checkpoint_level=checkpoint_level, - use_relative_pos=use_relative_pos, - use_final_proj=self_attn_use_final_proj, - attn_linear_init_mode=attn_linear_init_mode, - ffn_linear_init_mode=ffn_linear_init_mode, - norm_init_mode=norm_init_mode, - ) - for _ in range(ele_depth) - ] - self_blocks.append(nn.LayerList(sublayers=stack_cuboid_blocks)) - self.self_blocks = nn.LayerList(sublayers=self_blocks) - if block_cross_attn_patterns is not None: - if isinstance(block_cross_attn_patterns, (tuple, list)): - assert len(block_cross_attn_patterns) == self.num_blocks - else: - block_cross_attn_patterns = [ - block_cross_attn_patterns for _ in range(self.num_blocks) - ] - block_cross_cuboid_hw = [] - block_cross_cuboid_strategy = [] - block_cross_shift_hw = [] - block_cross_n_temporal = [] - for idx, key in enumerate(block_cross_attn_patterns): - if key == "last_frame_dst": - cuboid_hw = None - shift_hw = None - strategy = None - n_temporal = None - else: - func = cuboid_utils.CuboidCrossAttentionPatterns.get(key) - cuboid_hw, shift_hw, strategy, n_temporal = func(mem_shapes[idx]) - block_cross_cuboid_hw.append(cuboid_hw) - block_cross_cuboid_strategy.append(strategy) - block_cross_shift_hw.append(shift_hw) - block_cross_n_temporal.append(n_temporal) - else: - if not isinstance(block_cross_cuboid_hw[0][0], (list, tuple)): - block_cross_cuboid_hw = [ - block_cross_cuboid_hw for _ in range(self.num_blocks) - ] - else: - assert ( - len(block_cross_cuboid_hw) == self.num_blocks - ), f"Incorrect input format! Received block_cross_cuboid_hw={block_cross_cuboid_hw}" - if not isinstance(block_cross_cuboid_strategy[0][0], (list, tuple)): - block_cross_cuboid_strategy = [ - block_cross_cuboid_strategy for _ in range(self.num_blocks) - ] - else: - assert ( - len(block_cross_cuboid_strategy) == self.num_blocks - ), f"Incorrect input format! Received block_cross_cuboid_strategy={block_cross_cuboid_strategy}" - if not isinstance(block_cross_shift_hw[0][0], (list, tuple)): - block_cross_shift_hw = [ - block_cross_shift_hw for _ in range(self.num_blocks) - ] - else: - assert ( - len(block_cross_shift_hw) == self.num_blocks - ), f"Incorrect input format! Received block_cross_shift_hw={block_cross_shift_hw}" - if not isinstance(block_cross_n_temporal[0], (list, tuple)): - block_cross_n_temporal = [ - block_cross_n_temporal for _ in range(self.num_blocks) - ] - else: - assert ( - len(block_cross_n_temporal) == self.num_blocks - ), f"Incorrect input format! Received block_cross_n_temporal={block_cross_n_temporal}" - self.cross_blocks = nn.LayerList() - for i in range(self.cross_start, self.num_blocks): - cross_block = nn.LayerList( - sublayers=[ - StackCuboidCrossAttentionBlock( - dim=self.mem_shapes[i][-1], - num_heads=num_heads, - block_cuboid_hw=block_cross_cuboid_hw[i], - block_strategy=block_cross_cuboid_strategy[i], - block_shift_hw=block_cross_shift_hw[i], - block_n_temporal=block_cross_n_temporal[i], - cross_last_n_frames=cross_last_n_frames, - attn_drop=attn_drop, - proj_drop=proj_drop, - ffn_drop=ffn_drop, - gated_ffn=gated_ffn, - norm_layer=norm_layer, - use_inter_ffn=use_inter_ffn, - activation=ffn_activation, - max_temporal_relative=max_temporal_relative, - padding_type=padding_type, - use_global_vector=use_cross_global, - separate_global_qkv=separate_global_qkv, - global_dim_ratio=global_dim_ratio, - checkpoint_level=checkpoint_level, - use_relative_pos=use_relative_pos, - attn_linear_init_mode=attn_linear_init_mode, - ffn_linear_init_mode=ffn_linear_init_mode, - norm_init_mode=norm_init_mode, - ) - for _ in range(depth[i]) - ] - ) - self.cross_blocks.append(cross_block) - if self.num_blocks > 1: - if self.upsample_type == "upsample": - self.upsample_layers = nn.LayerList( - sublayers=[ - Upsample3DLayer( - dim=self.mem_shapes[i + 1][-1], - out_dim=self.mem_shapes[i][-1], - target_size=(target_temporal_length,) - + self.mem_shapes[i][1:3], - kernel_size=upsample_kernel_size, - temporal_upsample=False, - conv_init_mode=conv_init_mode, - ) - for i in range(self.num_blocks - 1) - ] - ) - else: - raise NotImplementedError(f"{self.upsample_type} is invalid.") - if self.hierarchical_pos_embed: - self.hierarchical_pos_embed_l = nn.LayerList( - sublayers=[ - PosEmbed( - embed_dim=self.mem_shapes[i][-1], - typ=pos_embed_type, - maxT=target_temporal_length, - maxH=self.mem_shapes[i][1], - maxW=self.mem_shapes[i][2], - ) - for i in range(self.num_blocks - 1) - ] - ) - self.reset_parameters() - - def reset_parameters(self): - for ms in self.self_blocks: - for m in ms: - m.reset_parameters() - for ms in self.cross_blocks: - for m in ms: - m.reset_parameters() - if self.num_blocks > 1: - for m in self.upsample_layers: - m.reset_parameters() - if self.hierarchical_pos_embed: - for m in self.hierarchical_pos_embed_l: - m.reset_parameters() - - def forward(self, x, mem_l, mem_global_vector_l=None): - """ - Args: - x : Shape (B, T_top, H_top, W_top, C). - mem_l : A list of memory tensors. - """ - - B, T_top, H_top, W_top, C = x.shape - assert T_top == self.target_temporal_length - assert (H_top, W_top) == (self.mem_shapes[-1][1], self.mem_shapes[-1][2]) - for i in range(self.num_blocks - 1, -1, -1): - mem_global_vector = ( - None if mem_global_vector_l is None else mem_global_vector_l[i] - ) - if not self.use_first_self_attn and i == self.num_blocks - 1: - if i >= self.cross_start: - x = self.cross_blocks[i - self.cross_start][0]( - x, mem_l[i], mem_global_vector - ) - for idx in range(self.depth[i] - 1): - if self.use_self_global: - if self.self_update_global: - x, mem_global_vector = self.self_blocks[i][idx]( - x, mem_global_vector - ) - else: - x, _ = self.self_blocks[i][idx](x, mem_global_vector) - else: - x = self.self_blocks[i][idx](x) - if i >= self.cross_start: - x = self.cross_blocks[i - self.cross_start][idx + 1]( - x, mem_l[i], mem_global_vector - ) - else: - for idx in range(self.depth[i]): - if self.use_self_global: - if self.self_update_global: - x, mem_global_vector = self.self_blocks[i][idx]( - x, mem_global_vector - ) - else: - x, _ = self.self_blocks[i][idx](x, mem_global_vector) - else: - x = self.self_blocks[i][idx](x) - if i >= self.cross_start: - x = self.cross_blocks[i - self.cross_start][idx]( - x, mem_l[i], mem_global_vector - ) - if i > 0: - x = self.upsample_layers[i - 1](x) - if self.hierarchical_pos_embed: - x = self.hierarchical_pos_embed_l[i - 1](x) - return x diff --git a/examples/smc_reac/ppsci/arch/cuboid_transformer_encoder.py b/examples/smc_reac/ppsci/arch/cuboid_transformer_encoder.py deleted file mode 100644 index 79b2e6fd1d..0000000000 --- a/examples/smc_reac/ppsci/arch/cuboid_transformer_encoder.py +++ /dev/null @@ -1,1515 +0,0 @@ -from collections import OrderedDict -from functools import lru_cache -from typing import Tuple - -import numpy as np -import paddle -import paddle.nn.functional as F -from paddle import nn -from paddle.distributed import fleet - -import ppsci.arch.cuboid_transformer_utils as cuboid_utils -from ppsci.arch import activation as act_mod -from ppsci.utils import initializer - -NEGATIVE_SLOPE = 0.1 - - -class PatchMerging3D(nn.Layer): - """Patch Merging Layer - - Args: - dim (int): Number of input channels. - out_dim (int, optional): The dim of output. Defaults to None. - downsample (tuple, optional): Downsample factor. Defaults to (1, 2, 2). - norm_layer (str, optional): The normalization layer. Defaults to "layer_norm". - padding_type (str, optional): The type of padding. Defaults to "nearest". - linear_init_mode (str, optional): The mode of linear init. Defaults to "0". - norm_init_mode (str, optional): The mode of normalization init. Defaults to "0". - """ - - def __init__( - self, - dim: int, - out_dim: int = None, - downsample: Tuple[int, ...] = (1, 2, 2), - norm_layer: str = "layer_norm", - padding_type: str = "nearest", - linear_init_mode: str = "0", - norm_init_mode: str = "0", - ): - super().__init__() - self.linear_init_mode = linear_init_mode - self.norm_init_mode = norm_init_mode - self.dim = dim - if out_dim is None: - out_dim = max(downsample) * dim - self.out_dim = out_dim - self.downsample = downsample - self.padding_type = padding_type - self.reduction = nn.Linear( - in_features=downsample[0] * downsample[1] * downsample[2] * dim, - out_features=out_dim, - bias_attr=False, - ) - self.norm = cuboid_utils.get_norm_layer( - norm_layer, in_channels=downsample[0] * downsample[1] * downsample[2] * dim - ) - self.reset_parameters() - - def reset_parameters(self): - for m in self.children(): - cuboid_utils.apply_initialization( - m, linear_mode=self.linear_init_mode, norm_mode=self.norm_init_mode - ) - - def get_out_shape(self, data_shape): - T, H, W, C_in = data_shape - pad_t = (self.downsample[0] - T % self.downsample[0]) % self.downsample[0] - pad_h = (self.downsample[1] - H % self.downsample[1]) % self.downsample[1] - pad_w = (self.downsample[2] - W % self.downsample[2]) % self.downsample[2] - return ( - (T + pad_t) // self.downsample[0], - (H + pad_h) // self.downsample[1], - (W + pad_w) // self.downsample[2], - self.out_dim, - ) - - def forward(self, x): - """ - - Args: - x : (B, T, H, W, C) - - Returns: - out : Shape (B, T // downsample[0], H // downsample[1], W // downsample[2], out_dim) - """ - - B, T, H, W, C = x.shape - pad_t = (self.downsample[0] - T % self.downsample[0]) % self.downsample[0] - pad_h = (self.downsample[1] - H % self.downsample[1]) % self.downsample[1] - pad_w = (self.downsample[2] - W % self.downsample[2]) % self.downsample[2] - if pad_h or pad_h or pad_w: - T += pad_t - H += pad_h - W += pad_w - x = cuboid_utils.generalize_padding( - x, pad_t, pad_h, pad_w, padding_type=self.padding_type - ) - x = ( - x.reshape( - ( - B, - T // self.downsample[0], - self.downsample[0], - H // self.downsample[1], - self.downsample[1], - W // self.downsample[2], - self.downsample[2], - C, - ) - ) - .transpose(perm=[0, 1, 3, 5, 2, 4, 6, 7]) - .reshape( - [ - B, - T // self.downsample[0], - H // self.downsample[1], - W // self.downsample[2], - self.downsample[0] * self.downsample[1] * self.downsample[2] * C, - ] - ) - ) - x = self.norm(x) - x = self.reduction(x) - return x - - -class PositionwiseFFN(nn.Layer): - """The Position-wise FFN layer used in Transformer-like architectures - - If pre_norm is True: - norm(data) -> fc1 -> act -> act_dropout -> fc2 -> dropout -> res(+data) - Else: - data -> fc1 -> act -> act_dropout -> fc2 -> dropout -> norm(res(+data)) - Also, if we use gated projection. We will use - fc1_1 * act(fc1_2(data)) to map the data - - Args: - units (int, optional): The units. Defaults to 512. - hidden_size (int, optional): The size of hidden layer. Defaults to 2048. - activation_dropout (float, optional): The dropout of activate. Defaults to 0.0. - dropout (float, optional): The drop ratio used in DropPat. Defaults to 0.1. - gated_proj (bool, optional): Whether to use gate projection. Defaults to False. - activation (str, optional): The activate. Defaults to "relu". - normalization (str, optional): The normalization. Defaults to "layer_norm". - layer_norm_eps (float, optional): The epsilon of layer normalization. Defaults to 1e-05. - pre_norm (bool): Pre-layer normalization as proposed in the paper: - "[ACL2018] The Best of Both Worlds: Combining Recent Advances in Neural Machine Translation" This will stabilize the training of Transformers. - You may also refer to "[Arxiv2020] Understanding the Difficulty of Training Transformers". Defaults to False. - linear_init_mode (str, optional): The mode of linear initialization. Defaults to "0". - norm_init_mode (str, optional): The mode of normalization initialization. Defaults to "0". - """ - - def __init__( - self, - units: int = 512, - hidden_size: int = 2048, - activation_dropout: float = 0.0, - dropout: float = 0.1, - gated_proj: bool = False, - activation: str = "relu", - normalization: str = "layer_norm", - layer_norm_eps: float = 1e-05, - pre_norm: bool = False, - linear_init_mode: str = "0", - norm_init_mode: str = "0", - ): - super().__init__() - self.linear_init_mode = linear_init_mode - self.norm_init_mode = norm_init_mode - self._pre_norm = pre_norm - self._gated_proj = gated_proj - self._kwargs = OrderedDict( - [ - ("units", units), - ("hidden_size", hidden_size), - ("activation_dropout", activation_dropout), - ("activation", activation), - ("dropout", dropout), - ("normalization", normalization), - ("layer_norm_eps", layer_norm_eps), - ("gated_proj", gated_proj), - ("pre_norm", pre_norm), - ] - ) - self.dropout_layer = nn.Dropout(p=dropout) - self.activation_dropout_layer = nn.Dropout(p=activation_dropout) - self.ffn_1 = nn.Linear( - in_features=units, out_features=hidden_size, bias_attr=True - ) - if self._gated_proj: - self.ffn_1_gate = nn.Linear( - in_features=units, out_features=hidden_size, bias_attr=True - ) - if activation == "leaky_relu": - self.activation = nn.LeakyReLU(NEGATIVE_SLOPE) - else: - self.activation = act_mod.get_activation(activation) - self.ffn_2 = nn.Linear( - in_features=hidden_size, out_features=units, bias_attr=True - ) - self.layer_norm = cuboid_utils.get_norm_layer( - normalization=normalization, in_channels=units, epsilon=layer_norm_eps - ) - self.reset_parameters() - - def reset_parameters(self): - cuboid_utils.apply_initialization(self.ffn_1, linear_mode=self.linear_init_mode) - if self._gated_proj: - cuboid_utils.apply_initialization( - self.ffn_1_gate, linear_mode=self.linear_init_mode - ) - cuboid_utils.apply_initialization(self.ffn_2, linear_mode=self.linear_init_mode) - cuboid_utils.apply_initialization( - self.layer_norm, norm_mode=self.norm_init_mode - ) - - def forward(self, data): - """ - Args: - x : Shape (B, seq_length, C_in) - - Returns: - out : Shape (B, seq_length, C_out) - """ - - residual = data - if self._pre_norm: - data = self.layer_norm(data) - if self._gated_proj: - out = self.activation(self.ffn_1_gate(data)) * self.ffn_1(data) - else: - out = self.activation(self.ffn_1(data)) - out = self.activation_dropout_layer(out) - out = self.ffn_2(out) - out = self.dropout_layer(out) - out = out + residual - if not self._pre_norm: - out = self.layer_norm(out) - return out - - -def update_cuboid_size_shift_size(data_shape, cuboid_size, shift_size, strategy): - """Update the cuboid_size and shift_size - - Args: - data_shape (Tuple[int,...]): The shape of the data. - cuboid_size (Tuple[int,...]): Size of the cuboid. - shift_size (Tuple[int,...]): Size of the shift. - strategy (str): The strategy of attention. - - Returns: - new_cuboid_size (Tuple[int,...]): Size of the cuboid. - new_shift_size (Tuple[int,...]): Size of the shift. - """ - - new_cuboid_size = list(cuboid_size) - new_shift_size = list(shift_size) - for i in range(len(data_shape)): - if strategy[i] == "d": - new_shift_size[i] = 0 - if data_shape[i] <= cuboid_size[i]: - new_cuboid_size[i] = data_shape[i] - new_shift_size[i] = 0 - return tuple(new_cuboid_size), tuple(new_shift_size) - - -def cuboid_reorder(data, cuboid_size, strategy): - """Reorder the tensor into (B, num_cuboids, bT * bH * bW, C) - We assume that the tensor shapes are divisible to the cuboid sizes. - - Args: - data (paddle.Tensor): The input data. - cuboid_size (Tuple[int,...]): The size of the cuboid. - strategy (Tuple[int,...]): The cuboid strategy. - - Returns: - reordered_data (paddle.Tensor): Shape will be (B, num_cuboids, bT * bH * bW, C). - num_cuboids = T / bT * H / bH * W / bW - """ - - B, T, H, W, C = data.shape - num_cuboids = T // cuboid_size[0] * H // cuboid_size[1] * W // cuboid_size[2] - cuboid_volume = cuboid_size[0] * cuboid_size[1] * cuboid_size[2] - intermediate_shape = [] - nblock_axis = [] - block_axis = [] - for i, (block_size, total_size, ele_strategy) in enumerate( - zip(cuboid_size, (T, H, W), strategy) - ): - if ele_strategy == "l": - intermediate_shape.extend([total_size // block_size, block_size]) - nblock_axis.append(2 * i + 1) - block_axis.append(2 * i + 2) - elif ele_strategy == "d": - intermediate_shape.extend([block_size, total_size // block_size]) - nblock_axis.append(2 * i + 2) - block_axis.append(2 * i + 1) - else: - raise NotImplementedError(f"{ele_strategy} is invalid.") - data = data.reshape(list((B,) + tuple(intermediate_shape) + (C,))) - reordered_data = data.transpose( - perm=(0,) + tuple(nblock_axis) + tuple(block_axis) + (7,) - ) - reordered_data = reordered_data.reshape((B, num_cuboids, cuboid_volume, C)) - return reordered_data - - -@lru_cache() -def compute_cuboid_self_attention_mask( - data_shape, cuboid_size, shift_size, strategy, padding_type, device -): - """Compute the shift window attention mask - - Args: - data_shape (Tuple[int,....]): Should be (T, H, W). - cuboid_size (Tuple[int,....]): Size of the cuboid. - shift_size (Tuple[int,....]): The shift size. - strategy (str): The decomposition strategy. - padding_type (str): Type of the padding. - device (str): The device. - - Returns: - attn_mask (paddle.Tensor): Mask with shape (num_cuboid, cuboid_vol, cuboid_vol). - The padded values will always be masked. The other masks will ensure that the shifted windows - will only attend to those in the shifted windows. - """ - T, H, W = data_shape - pad_t = (cuboid_size[0] - T % cuboid_size[0]) % cuboid_size[0] - pad_h = (cuboid_size[1] - H % cuboid_size[1]) % cuboid_size[1] - pad_w = (cuboid_size[2] - W % cuboid_size[2]) % cuboid_size[2] - data_mask = None - if pad_t > 0 or pad_h > 0 or pad_w > 0: - if padding_type == "ignore": - data_mask = paddle.ones(shape=(1, T, H, W, 1), dtype="bool") - data_mask = F.pad( - data_mask, [0, 0, 0, pad_w, 0, pad_h, 0, pad_t], data_format="NDHWC" - ) - else: - data_mask = paddle.ones( - shape=(1, T + pad_t, H + pad_h, W + pad_w, 1), dtype="bool" - ) - if any(i > 0 for i in shift_size): - if padding_type == "ignore": - data_mask = paddle.roll( - x=data_mask, - shifts=(-shift_size[0], -shift_size[1], -shift_size[2]), - axis=(1, 2, 3), - ) - if padding_type == "ignore": - data_mask = cuboid_reorder(data_mask, cuboid_size, strategy=strategy) - data_mask = data_mask.squeeze(axis=-1).squeeze(axis=0) - shift_mask = np.zeros(shape=(1, T + pad_t, H + pad_h, W + pad_w, 1)) - cnt = 0 - for t in ( - slice(-cuboid_size[0]), - slice(-cuboid_size[0], -shift_size[0]), - slice(-shift_size[0], None), - ): - for h in ( - slice(-cuboid_size[1]), - slice(-cuboid_size[1], -shift_size[1]), - slice(-shift_size[1], None), - ): - for w in ( - slice(-cuboid_size[2]), - slice(-cuboid_size[2], -shift_size[2]), - slice(-shift_size[2], None), - ): - shift_mask[:, t, h, w, :] = cnt - cnt += 1 - shift_mask = paddle.to_tensor(shift_mask) - shift_mask = cuboid_reorder(shift_mask, cuboid_size, strategy=strategy) - shift_mask = shift_mask.squeeze(axis=-1).squeeze(axis=0) - attn_mask = shift_mask.unsqueeze(axis=1) - shift_mask.unsqueeze(axis=2) == 0 - if padding_type == "ignore": - attn_mask = ( - data_mask.unsqueeze(axis=1) * data_mask.unsqueeze(axis=2) * attn_mask - ) - return attn_mask - - -def masked_softmax(att_score, mask, axis: int = -1): - """Ignore the masked elements when calculating the softmax. - The mask can be broadcastable. - - Args: - att_score (paddle.Tensor): Shape (..., length, ...) - mask (paddle.Tensor): Shape (..., length, ...) - 1 --> The element is not masked - 0 --> The element is masked - axis (int): The axis to calculate the softmax. att_score.shape[axis] must be the same as mask.shape[axis] - - Returns: - att_weights (paddle.Tensor): Shape (..., length, ...). - """ - - if mask is not None: - if att_score.dtype == paddle.float16: - att_score = att_score.masked_fill(paddle.logical_not(mask), -1e4) - else: - att_score = att_score.masked_fill(paddle.logical_not(mask), -1e18) - att_weights = nn.functional.softmax(x=att_score, axis=axis) * mask - else: - att_weights = nn.functional.softmax(x=att_score, axis=axis) - return att_weights - - -def cuboid_reorder_reverse(data, cuboid_size, strategy, orig_data_shape): - """Reverse the reordered cuboid back to the original space - - Args: - data (paddle.Tensor): The input data. - cuboid_size (Tuple[int,...]): The size of cuboid. - strategy (str): The strategy of reordering. - orig_data_shape (Tuple[int,...]): The original shape of the data. - - Returns: - data (paddle.Tensor): The recovered data - """ - - B, num_cuboids, cuboid_volume, C = data.shape - T, H, W = orig_data_shape - permutation_axis = [0] - for i, (block_size, total_size, ele_strategy) in enumerate( - zip(cuboid_size, (T, H, W), strategy) - ): - if ele_strategy == "l": - permutation_axis.append(i + 1) - permutation_axis.append(i + 4) - elif ele_strategy == "d": - permutation_axis.append(i + 4) - permutation_axis.append(i + 1) - else: - raise NotImplementedError((f"{ele_strategy} is invalid.")) - permutation_axis.append(7) - data = data.reshape( - [ - B, - T // cuboid_size[0], - H // cuboid_size[1], - W // cuboid_size[2], - cuboid_size[0], - cuboid_size[1], - cuboid_size[2], - C, - ] - ) - data = data.transpose(perm=permutation_axis) - data = data.reshape((B, T, H, W, C)) - return data - - -class CuboidSelfAttentionLayer(nn.Layer): - """Implements the cuboid self attention. - - The idea of Cuboid Self Attention is to divide the input tensor (T, H, W) into several non-overlapping cuboids. - We apply self-attention inside each cuboid and all cuboid-level self attentions are executed in parallel. - - We adopt two mechanisms for decomposing the input tensor into cuboids: - - (1) local: - We group the tensors within a local window, e.g., X[t:(t+b_t), h:(h+b_h), w:(w+b_w)]. We can also apply the - shifted window strategy proposed in "[ICCV2021] Swin Transformer: Hierarchical Vision Transformer using Shifted Windows". - (2) dilated: - Inspired by the success of dilated convolution "[ICLR2016] Multi-Scale Context Aggregation by Dilated Convolutions", - we split the tensor with dilation factors that are tied to the size of the cuboid. For example, for a cuboid that has width `b_w`, - we sample the elements starting from 0 as 0, w / b_w, 2 * w / b_w, ..., (b_w - 1) * w / b_w. - - The cuboid attention can be viewed as a generalization of the attention mechanism proposed in Video Swin Transformer, https://arxiv.org/abs/2106.13230. - The computational complexity of CuboidAttention can be simply calculated as O(T H W * b_t b_h b_w). To cover multiple correlation patterns, - we are able to combine multiple CuboidAttention layers with different configurations such as cuboid size, shift size, and local / global decomposing strategy. - - In addition, it is straight-forward to extend the cuboid attention to other types of spatiotemporal data that are not described - as regular tensors. We need to define alternative approaches to partition the data into "cuboids". - - In addition, inspired by "[NeurIPS2021] Do Transformers Really Perform Badly for Graph Representation?", - "[NeurIPS2020] Big Bird: Transformers for Longer Sequences", "[EMNLP2021] Longformer: The Long-Document Transformer", we keep - $K$ global vectors to record the global status of the spatiotemporal system. These global vectors will attend to the whole tensor and - the vectors inside each individual cuboids will also attend to the global vectors so that they can peep into the global status of the system. - - Args: - dim (int): The dimension of the input tensor. - num_heads (int): The number of heads. - cuboid_size (tuple, optional): The size of cuboid. Defaults to (2, 7, 7). - shift_size (tuple, optional): The size of shift. Defaults to (0, 0, 0). - strategy (tuple, optional): The strategy. Defaults to ("l", "l", "l"). - padding_type (str, optional): The type of padding. Defaults to "ignore". - qkv_bias (bool, optional): Whether to enable bias in calculating qkv attention. Defaults to False. - qk_scale (float, optional): Whether to enable scale factor when calculating the attention. Defaults to None. - attn_drop (float, optional): The attention dropout. Defaults to 0.0. - proj_drop (float, optional): The projection dropout. Defaults to 0.0. - use_final_proj (bool, optional): Whether to use the final projection. Defaults to True. - norm_layer (str, optional): The normalization layer. Defaults to "layer_norm". - use_global_vector (bool, optional): Whether to use the global vector or not. Defaults to False. - use_global_self_attn (bool, optional): Whether to use self attention among global vectors. Defaults to False. - separate_global_qkv (bool, optional): Whether to use different network to calc q_global, k_global, v_global. Defaults to False. - global_dim_ratio (int, optional): The dim (channels) of global vectors is `global_dim_ratio*dim`. Defaults to 1. - checkpoint_level (bool, optional): Whether to enable gradient checkpointing. Defaults to True. - use_relative_pos (bool, optional): Whether to use relative pos. Defaults to True. - attn_linear_init_mode (str, optional): The mode of attention linear initialization. Defaults to "0". - ffn_linear_init_mode (str, optional): The mode of FFN linear initialization. Defaults to "0". - norm_init_mode (str, optional): The mode of normalization initialization. Defaults to "0". - """ - - def __init__( - self, - dim: int, - num_heads: int, - cuboid_size: Tuple[int, ...] = (2, 7, 7), - shift_size: Tuple[int, ...] = (0, 0, 0), - strategy: Tuple[str, ...] = ("l", "l", "l"), - padding_type: str = "ignore", - qkv_bias: bool = False, - qk_scale: float = None, - attn_drop: float = 0.0, - proj_drop: float = 0.0, - use_final_proj: bool = True, - norm_layer: str = "layer_norm", - use_global_vector: bool = False, - use_global_self_attn: bool = False, - separate_global_qkv: bool = False, - global_dim_ratio: int = 1, - checkpoint_level: bool = True, - use_relative_pos: bool = True, - attn_linear_init_mode: str = "0", - ffn_linear_init_mode: str = "0", - norm_init_mode: str = "0", - ): - super(CuboidSelfAttentionLayer, self).__init__() - self.attn_linear_init_mode = attn_linear_init_mode - self.ffn_linear_init_mode = ffn_linear_init_mode - self.norm_init_mode = norm_init_mode - assert dim % num_heads == 0 - self.num_heads = num_heads - self.dim = dim - self.cuboid_size = cuboid_size - self.shift_size = shift_size - self.strategy = strategy - self.padding_type = padding_type - self.use_final_proj = use_final_proj - self.use_relative_pos = use_relative_pos - self.use_global_vector = use_global_vector - self.use_global_self_attn = use_global_self_attn - self.separate_global_qkv = separate_global_qkv - if global_dim_ratio != 1: - assert ( - separate_global_qkv is True - ), "Setting global_dim_ratio != 1 requires separate_global_qkv == True." - self.global_dim_ratio = global_dim_ratio - assert self.padding_type in ["ignore", "zeros", "nearest"] - head_dim = dim // num_heads - self.scale = qk_scale or head_dim**-0.5 - if use_relative_pos: - init_data = paddle.zeros( - ( - (2 * cuboid_size[0] - 1) - * (2 * cuboid_size[1] - 1) - * (2 * cuboid_size[2] - 1), - num_heads, - ) - ) - self.relative_position_bias_table = paddle.create_parameter( - shape=init_data.shape, - dtype=init_data.dtype, - default_initializer=nn.initializer.Constant(0.0), - ) - self.relative_position_bias_table.stop_gradient = not True - self.relative_position_bias_table = initializer.trunc_normal_( - self.relative_position_bias_table, std=0.02 - ) - - coords_t = paddle.arange(end=self.cuboid_size[0]) - coords_h = paddle.arange(end=self.cuboid_size[1]) - coords_w = paddle.arange(end=self.cuboid_size[2]) - coords = paddle.stack(x=paddle.meshgrid(coords_t, coords_h, coords_w)) - coords_flatten = paddle.flatten(x=coords, start_axis=1) - relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] - relative_coords = relative_coords.transpose(perm=[1, 2, 0]) - relative_coords[:, :, 0] += self.cuboid_size[0] - 1 - relative_coords[:, :, 1] += self.cuboid_size[1] - 1 - relative_coords[:, :, 2] += self.cuboid_size[2] - 1 - relative_coords[:, :, 0] *= (2 * self.cuboid_size[1] - 1) * ( - 2 * self.cuboid_size[2] - 1 - ) - relative_coords[:, :, 1] *= 2 * self.cuboid_size[2] - 1 - relative_position_index = relative_coords.sum(axis=-1) - self.register_buffer( - name="relative_position_index", tensor=relative_position_index - ) - self.qkv = nn.Linear(in_features=dim, out_features=dim * 3, bias_attr=qkv_bias) - self.attn_drop = nn.Dropout(p=attn_drop) - if self.use_global_vector: - if self.separate_global_qkv: - self.l2g_q_net = nn.Linear( - in_features=dim, out_features=dim, bias_attr=qkv_bias - ) - self.l2g_global_kv_net = nn.Linear( - in_features=global_dim_ratio * dim, - out_features=dim * 2, - bias_attr=qkv_bias, - ) - self.g2l_global_q_net = nn.Linear( - in_features=global_dim_ratio * dim, - out_features=dim, - bias_attr=qkv_bias, - ) - self.g2l_k_net = nn.Linear( - in_features=dim, out_features=dim, bias_attr=qkv_bias - ) - self.g2l_v_net = nn.Linear( - in_features=dim, - out_features=global_dim_ratio * dim, - bias_attr=qkv_bias, - ) - if self.use_global_self_attn: - self.g2g_global_qkv_net = nn.Linear( - in_features=global_dim_ratio * dim, - out_features=global_dim_ratio * dim * 3, - bias_attr=qkv_bias, - ) - else: - self.global_qkv = nn.Linear( - in_features=dim, out_features=dim * 3, bias_attr=qkv_bias - ) - self.global_attn_drop = nn.Dropout(p=attn_drop) - if use_final_proj: - self.proj = nn.Linear(in_features=dim, out_features=dim) - self.proj_drop = nn.Dropout(p=proj_drop) - if self.use_global_vector: - self.global_proj = nn.Linear( - in_features=global_dim_ratio * dim, - out_features=global_dim_ratio * dim, - ) - self.norm = cuboid_utils.get_norm_layer(norm_layer, in_channels=dim) - if self.use_global_vector: - self.global_vec_norm = cuboid_utils.get_norm_layer( - norm_layer, in_channels=global_dim_ratio * dim - ) - self.checkpoint_level = checkpoint_level - self.reset_parameters() - - def reset_parameters(self): - cuboid_utils.apply_initialization( - self.qkv, linear_mode=self.attn_linear_init_mode - ) - if self.use_final_proj: - cuboid_utils.apply_initialization( - self.proj, linear_mode=self.ffn_linear_init_mode - ) - cuboid_utils.apply_initialization(self.norm, norm_mode=self.norm_init_mode) - if self.use_global_vector: - if self.separate_global_qkv: - cuboid_utils.apply_initialization( - self.l2g_q_net, linear_mode=self.attn_linear_init_mode - ) - cuboid_utils.apply_initialization( - self.l2g_global_kv_net, linear_mode=self.attn_linear_init_mode - ) - cuboid_utils.apply_initialization( - self.g2l_global_q_net, linear_mode=self.attn_linear_init_mode - ) - cuboid_utils.apply_initialization( - self.g2l_k_net, linear_mode=self.attn_linear_init_mode - ) - cuboid_utils.apply_initialization( - self.g2l_v_net, linear_mode=self.attn_linear_init_mode - ) - if self.use_global_self_attn: - cuboid_utils.apply_initialization( - self.g2g_global_qkv_net, linear_mode=self.attn_linear_init_mode - ) - else: - cuboid_utils.apply_initialization( - self.global_qkv, linear_mode=self.attn_linear_init_mode - ) - cuboid_utils.apply_initialization( - self.global_vec_norm, norm_mode=self.norm_init_mode - ) - - def forward(self, x, global_vectors=None): - x = self.norm(x) - - B, T, H, W, C_in = x.shape - assert C_in == self.dim - if self.use_global_vector: - _, num_global, _ = global_vectors.shape - global_vectors = self.global_vec_norm(global_vectors) - cuboid_size, shift_size = update_cuboid_size_shift_size( - (T, H, W), self.cuboid_size, self.shift_size, self.strategy - ) - - pad_t = (cuboid_size[0] - T % cuboid_size[0]) % cuboid_size[0] - pad_h = (cuboid_size[1] - H % cuboid_size[1]) % cuboid_size[1] - pad_w = (cuboid_size[2] - W % cuboid_size[2]) % cuboid_size[2] - x = cuboid_utils.generalize_padding(x, pad_t, pad_h, pad_w, self.padding_type) - - if any(i > 0 for i in shift_size): - shifted_x = paddle.roll( - x=x, - shifts=(-shift_size[0], -shift_size[1], -shift_size[2]), - axis=(1, 2, 3), - ) - else: - shifted_x = x - - reordered_x = cuboid_reorder( - shifted_x, cuboid_size=cuboid_size, strategy=self.strategy - ) - - _, num_cuboids, cuboid_volume, _ = reordered_x.shape - attn_mask = compute_cuboid_self_attention_mask( - (T, H, W), - cuboid_size, - shift_size=shift_size, - strategy=self.strategy, - padding_type=self.padding_type, - device=x.place, - ) - head_C = C_in // self.num_heads - qkv = ( - self.qkv(reordered_x) - .reshape([B, num_cuboids, cuboid_volume, 3, self.num_heads, head_C]) - .transpose(perm=[3, 0, 4, 1, 2, 5]) - ) - - q, k, v = qkv[0], qkv[1], qkv[2] - q = q * self.scale - perm_0 = list(range(k.ndim)) - perm_0[-2] = -1 - perm_0[-1] = -2 - attn_score = q @ k.transpose(perm=perm_0) - - if self.use_relative_pos: - relative_position_bias = self.relative_position_bias_table[ - self.relative_position_index[:cuboid_volume, :cuboid_volume].reshape( - [-1] - ) - ].reshape([cuboid_volume, cuboid_volume, -1]) - relative_position_bias = relative_position_bias.transpose( - perm=[2, 0, 1] - ).unsqueeze(axis=1) - attn_score = attn_score + relative_position_bias - - if self.use_global_vector: - global_head_C = self.global_dim_ratio * head_C - if self.separate_global_qkv: - l2g_q = ( - self.l2g_q_net(reordered_x) - .reshape([B, num_cuboids, cuboid_volume, self.num_heads, head_C]) - .transpose(perm=[0, 3, 1, 2, 4]) - ) - l2g_q = l2g_q * self.scale - l2g_global_kv = ( - self.l2g_global_kv_net(global_vectors) - .reshape([B, 1, num_global, 2, self.num_heads, head_C]) - .transpose(perm=[3, 0, 4, 1, 2, 5]) - ) - l2g_global_k, l2g_global_v = l2g_global_kv[0], l2g_global_kv[1] - g2l_global_q = ( - self.g2l_global_q_net(global_vectors) - .reshape([B, num_global, self.num_heads, head_C]) - .transpose(perm=[0, 2, 1, 3]) - ) - g2l_global_q = g2l_global_q * self.scale - g2l_k = ( - self.g2l_k_net(reordered_x) - .reshape([B, num_cuboids, cuboid_volume, self.num_heads, head_C]) - .transpose(perm=[0, 3, 1, 2, 4]) - ) - g2l_v = ( - self.g2l_v_net(reordered_x) - .reshape( - [B, num_cuboids, cuboid_volume, self.num_heads, global_head_C] - ) - .transpose(perm=[0, 3, 1, 2, 4]) - ) - if self.use_global_self_attn: - g2g_global_qkv = ( - self.g2g_global_qkv_net(global_vectors) - .reshape([B, 1, num_global, 3, self.num_heads, global_head_C]) - .transpose(perm=[3, 0, 4, 1, 2, 5]) - ) - g2g_global_q, g2g_global_k, g2g_global_v = ( - g2g_global_qkv[0], - g2g_global_qkv[1], - g2g_global_qkv[2], - ) - g2g_global_q = g2g_global_q.squeeze(axis=2) * self.scale - else: - q_global, k_global, v_global = ( - self.global_qkv(global_vectors) - .reshape([B, 1, num_global, 3, self.num_heads, head_C]) - .transpose(perm=[3, 0, 4, 1, 2, 5]) - ) - q_global = q_global.squeeze(axis=2) * self.scale - l2g_q, g2l_k, g2l_v = q, k, v - g2l_global_q, l2g_global_k, l2g_global_v = ( - q_global, - k_global, - v_global, - ) - if self.use_global_self_attn: - g2g_global_q, g2g_global_k, g2g_global_v = ( - q_global, - k_global, - v_global, - ) - - perm_1 = list(range(l2g_global_k.ndim)) - perm_1[-2] = -1 - perm_1[-1] = -2 - l2g_attn_score = l2g_q @ l2g_global_k.transpose(perm=perm_1) - attn_score_l2l_l2g = paddle.concat(x=(attn_score, l2g_attn_score), axis=-1) - - if attn_mask.ndim == 5: - attn_mask_l2l_l2g = F.pad( - attn_mask, [0, num_global], "constant", 1, data_format="NDHWC" - ) - elif attn_mask.ndim == 3: - attn_mask = attn_mask.astype("float32") - attn_mask_l2l_l2g = F.pad( - attn_mask, [0, num_global], "constant", 1, data_format="NCL" - ) - attn_mask_l2l_l2g = attn_mask_l2l_l2g.astype("bool") - else: - attn_mask_l2l_l2g = F.pad(attn_mask, [0, num_global], "constant", 1) - - v_l_g = paddle.concat( - x=( - v, - l2g_global_v.expand( - shape=[B, self.num_heads, num_cuboids, num_global, head_C] - ), - ), - axis=3, - ) - attn_score_l2l_l2g = masked_softmax( - attn_score_l2l_l2g, mask=attn_mask_l2l_l2g - ) - attn_score_l2l_l2g = self.attn_drop(attn_score_l2l_l2g) - reordered_x = ( - (attn_score_l2l_l2g @ v_l_g) - .transpose(perm=[0, 2, 3, 1, 4]) - .reshape([B, num_cuboids, cuboid_volume, self.dim]) - ) - if self.padding_type == "ignore": - g2l_attn_mask = paddle.ones(shape=(1, T, H, W, 1)) - if pad_t > 0 or pad_h > 0 or pad_w > 0: - g2l_attn_mask = F.pad( - g2l_attn_mask, - [0, 0, 0, pad_w, 0, pad_h, 0, pad_t], - data_format="NDHWC", - ) - if any(i > 0 for i in shift_size): - g2l_attn_mask = paddle.roll( - x=g2l_attn_mask, - shifts=(-shift_size[0], -shift_size[1], -shift_size[2]), - axis=(1, 2, 3), - ) - g2l_attn_mask = g2l_attn_mask.reshape((-1,)) - else: - g2l_attn_mask = None - temp = g2l_k.reshape( - [B, self.num_heads, num_cuboids * cuboid_volume, head_C] - ) - perm_2 = list(range(temp.ndim)) - perm_2[-2] = -1 - perm_2[-1] = -2 - g2l_attn_score = g2l_global_q @ temp.transpose(perm=perm_2) - if self.use_global_self_attn: - temp = g2g_global_k.squeeze(axis=2) - perm_3 = list(range(temp.ndim)) - perm_3[-2] = -1 - perm_3[-1] = -2 - g2g_attn_score = g2g_global_q @ temp.transpose(perm=perm_3) - g2all_attn_score = paddle.concat( - x=(g2l_attn_score, g2g_attn_score), axis=-1 - ) - if g2l_attn_mask is not None: - g2all_attn_mask = F.pad( - g2l_attn_mask, - [0, num_global], - "constant", - 1, - data_format="NDHWC", - ) - else: - g2all_attn_mask = None - new_v = paddle.concat( - x=( - g2l_v.reshape( - [ - B, - self.num_heads, - num_cuboids * cuboid_volume, - global_head_C, - ] - ), - g2g_global_v.reshape( - [B, self.num_heads, num_global, global_head_C] - ), - ), - axis=2, - ) - else: - g2all_attn_score = g2l_attn_score - g2all_attn_mask = g2l_attn_mask - new_v = g2l_v.reshape( - [B, self.num_heads, num_cuboids * cuboid_volume, global_head_C] - ) - g2all_attn_score = masked_softmax(g2all_attn_score, mask=g2all_attn_mask) - g2all_attn_score = self.global_attn_drop(g2all_attn_score) - new_global_vector = ( - (g2all_attn_score @ new_v) - .transpose(perm=[0, 2, 1, 3]) - .reshape([B, num_global, self.global_dim_ratio * self.dim]) - ) - else: - attn_score = masked_softmax(attn_score, mask=attn_mask) - attn_score = self.attn_drop(attn_score) - reordered_x = ( - (attn_score @ v) - .transpose(perm=[0, 2, 3, 1, 4]) - .reshape([B, num_cuboids, cuboid_volume, self.dim]) - ) - - if self.use_final_proj: - reordered_x = paddle.cast(reordered_x, dtype="float32") - reordered_x = self.proj_drop(self.proj(reordered_x)) - if self.use_global_vector: - new_global_vector = self.proj_drop(self.global_proj(new_global_vector)) - shifted_x = cuboid_reorder_reverse( - reordered_x, - cuboid_size=cuboid_size, - strategy=self.strategy, - orig_data_shape=(T + pad_t, H + pad_h, W + pad_w), - ) - if any(i > 0 for i in shift_size): - x = paddle.roll( - x=shifted_x, - shifts=(shift_size[0], shift_size[1], shift_size[2]), - axis=(1, 2, 3), - ) - else: - x = shifted_x - x = cuboid_utils.generalize_unpadding( - x, pad_t=pad_t, pad_h=pad_h, pad_w=pad_w, padding_type=self.padding_type - ) - if self.use_global_vector: - return x, new_global_vector - else: - return x - - -class StackCuboidSelfAttentionBlock(nn.Layer): - """ - - "use_inter_ffn" is True - x --> attn1 -----+-------> ffn1 ---+---> attn2 --> ... --> ffn_k --> out - | ^ | ^ - | | | | - |-------------| |-------------| - - "use_inter_ffn" is False - x --> attn1 -----+------> attn2 --> ... attnk --+----> ffnk ---+---> out - | ^ | ^ ^ | ^ - | | | | | | | - |-------------| |------------| ----------| |-----------| - If we have enabled global memory vectors, each attention will be a - - Args: - dim (int): The dimension of the input tensor. - num_heads (int): The number of heads. - block_cuboid_size (list, optional): The size of block cuboid . Defaults to [(4, 4, 4), (4, 4, 4)]. - block_shift_size (list, optional): The shift size of block. Defaults to [(0, 0, 0), (2, 2, 2)]. - block_strategy (list, optional): The strategy of block. Defaults to [("d", "d", "d"), ("l", "l", "l")]. - padding_type (str, optional): The type of padding. Defaults to "ignore". - qkv_bias (bool, optional): Whether to enable bias in calculating qkv attention. Defaults to False. - qk_scale (float, optional): Whether to enable scale factor when calculating the attention. Defaults to None. - attn_drop (float, optional): The attention dropout. Defaults to 0.0. - proj_drop (float, optional): The projection dropout. Defaults to 0.0. - use_final_proj (bool, optional): Whether to use the final projection. Defaults to True. - norm_layer (str, optional): The normalization layer. Defaults to "layer_norm". - use_global_vector (bool, optional): Whether to use the global vector or not. Defaults to False. - use_global_self_attn (bool, optional): Whether to use self attention among global vectors. Defaults to False. - separate_global_qkv (bool, optional): Whether to use different network to calc q_global, k_global, v_global. - Defaults to False. - global_dim_ratio (int, optional): The dim (channels) of global vectors is `global_dim_ratio*dim`. - Defaults to 1. - checkpoint_level (bool, optional): Whether to enable gradient checkpointing. Defaults to True. - use_relative_pos (bool, optional): Whether to use relative pos. Defaults to True. - use_relative_pos (bool, optional): Whether to use relative pos. Defaults to True. - attn_linear_init_mode (str, optional): The mode of attention linear initialization. Defaults to "0". - ffn_linear_init_mode (str, optional): The mode of FFN linear initialization. Defaults to "0". - norm_init_mode (str, optional): The mode of normalization initialization. Defaults to "0". - """ - - def __init__( - self, - dim: int, - num_heads: int, - block_cuboid_size: Tuple[Tuple[int, ...], ...] = [(4, 4, 4), (4, 4, 4)], - block_shift_size: Tuple[Tuple[int, ...], ...] = [(0, 0, 0), (2, 2, 2)], - block_strategy: Tuple[Tuple[str, ...], ...] = [ - ("d", "d", "d"), - ("l", "l", "l"), - ], - padding_type: str = "ignore", - qkv_bias: bool = False, - qk_scale: float = None, - attn_drop: float = 0.0, - proj_drop: float = 0.0, - ffn_drop: float = 0.0, - activation: str = "leaky", - gated_ffn: bool = False, - norm_layer: str = "layer_norm", - use_inter_ffn: bool = False, - use_global_vector: bool = False, - use_global_vector_ffn: bool = True, - use_global_self_attn: bool = False, - separate_global_qkv: bool = False, - global_dim_ratio: int = 1, - checkpoint_level: bool = True, - use_relative_pos: bool = True, - use_final_proj: bool = True, - attn_linear_init_mode: str = "0", - ffn_linear_init_mode: str = "0", - norm_init_mode: str = "0", - ): - super(StackCuboidSelfAttentionBlock, self).__init__() - self.attn_linear_init_mode = attn_linear_init_mode - self.ffn_linear_init_mode = ffn_linear_init_mode - self.norm_init_mode = norm_init_mode - if ( - len(block_cuboid_size[0]) <= 0 - or len(block_shift_size) <= 0 - or len(block_strategy) <= 0 - ): - raise ValueError( - "Format of the block cuboid size is not correct. block_cuboid_size={block_cuboid_size}" - ) - if len(block_cuboid_size) != len(block_shift_size) and len( - block_cuboid_size - ) != len(block_strategy): - raise ValueError( - "The lengths of block_cuboid_size, block_shift_size, and block_strategy must be equal." - ) - - self.num_attn = len(block_cuboid_size) - self.checkpoint_level = checkpoint_level - self.use_inter_ffn = use_inter_ffn - self.use_global_vector = use_global_vector - self.use_global_vector_ffn = use_global_vector_ffn - self.use_global_self_attn = use_global_self_attn - self.global_dim_ratio = global_dim_ratio - if self.use_inter_ffn: - self.ffn_l = nn.LayerList( - sublayers=[ - PositionwiseFFN( - units=dim, - hidden_size=4 * dim, - activation_dropout=ffn_drop, - dropout=ffn_drop, - gated_proj=gated_ffn, - activation=activation, - normalization=norm_layer, - pre_norm=True, - linear_init_mode=ffn_linear_init_mode, - norm_init_mode=norm_init_mode, - ) - for _ in range(self.num_attn) - ] - ) - if self.use_global_vector_ffn and self.use_global_vector: - self.global_ffn_l = nn.LayerList( - sublayers=[ - PositionwiseFFN( - units=global_dim_ratio * dim, - hidden_size=global_dim_ratio * 4 * dim, - activation_dropout=ffn_drop, - dropout=ffn_drop, - gated_proj=gated_ffn, - activation=activation, - normalization=norm_layer, - pre_norm=True, - linear_init_mode=ffn_linear_init_mode, - norm_init_mode=norm_init_mode, - ) - for _ in range(self.num_attn) - ] - ) - else: - self.ffn_l = nn.LayerList( - sublayers=[ - PositionwiseFFN( - units=dim, - hidden_size=4 * dim, - activation_dropout=ffn_drop, - dropout=ffn_drop, - gated_proj=gated_ffn, - activation=activation, - normalization=norm_layer, - pre_norm=True, - linear_init_mode=ffn_linear_init_mode, - norm_init_mode=norm_init_mode, - ) - ] - ) - if self.use_global_vector_ffn and self.use_global_vector: - self.global_ffn_l = nn.LayerList( - sublayers=[ - PositionwiseFFN( - units=global_dim_ratio * dim, - hidden_size=global_dim_ratio * 4 * dim, - activation_dropout=ffn_drop, - dropout=ffn_drop, - gated_proj=gated_ffn, - activation=activation, - normalization=norm_layer, - pre_norm=True, - linear_init_mode=ffn_linear_init_mode, - norm_init_mode=norm_init_mode, - ) - ] - ) - self.attn_l = nn.LayerList( - sublayers=[ - CuboidSelfAttentionLayer( - dim=dim, - num_heads=num_heads, - cuboid_size=ele_cuboid_size, - shift_size=ele_shift_size, - strategy=ele_strategy, - padding_type=padding_type, - qkv_bias=qkv_bias, - qk_scale=qk_scale, - attn_drop=attn_drop, - proj_drop=proj_drop, - norm_layer=norm_layer, - use_global_vector=use_global_vector, - use_global_self_attn=use_global_self_attn, - separate_global_qkv=separate_global_qkv, - global_dim_ratio=global_dim_ratio, - checkpoint_level=checkpoint_level, - use_relative_pos=use_relative_pos, - use_final_proj=use_final_proj, - attn_linear_init_mode=attn_linear_init_mode, - ffn_linear_init_mode=ffn_linear_init_mode, - norm_init_mode=norm_init_mode, - ) - for ele_cuboid_size, ele_shift_size, ele_strategy in zip( - block_cuboid_size, block_shift_size, block_strategy - ) - ] - ) - - def reset_parameters(self): - for m in self.ffn_l: - m.reset_parameters() - if self.use_global_vector_ffn and self.use_global_vector: - for m in self.global_ffn_l: - m.reset_parameters() - for m in self.attn_l: - m.reset_parameters() - - def forward(self, x, global_vectors=None): - if self.use_inter_ffn: - if self.use_global_vector: - for idx, (attn, ffn) in enumerate(zip(self.attn_l, self.ffn_l)): - if self.checkpoint_level >= 2 and self.training: - x_out, global_vectors_out = fleet.utils.recompute( - attn, x, global_vectors - ) - else: - x_out, global_vectors_out = attn(x, global_vectors) - x = x + x_out - global_vectors = global_vectors + global_vectors_out - if self.checkpoint_level >= 1 and self.training: - x = fleet.utils.recompute(ffn, x) - if self.use_global_vector_ffn: - global_vectors = fleet.utils.recompute( - self.global_ffn_l[idx], global_vectors - ) - else: - x = ffn(x) - if self.use_global_vector_ffn: - global_vectors = self.global_ffn_l[idx](global_vectors) - return x, global_vectors - else: - for idx, (attn, ffn) in enumerate(zip(self.attn_l, self.ffn_l)): - if self.checkpoint_level >= 2 and self.training: - x = x + fleet.utils.recompute(attn, x) - else: - x = x + attn(x) - if self.checkpoint_level >= 1 and self.training: - x = fleet.utils.recompute(ffn, x) - else: - x = ffn(x) - return x - elif self.use_global_vector: - for idx, attn in enumerate(self.attn_l): - if self.checkpoint_level >= 2 and self.training: - x_out, global_vectors_out = fleet.utils.recompute( - attn, x, global_vectors - ) - else: - x_out, global_vectors_out = attn(x, global_vectors) - x = x + x_out - global_vectors = global_vectors + global_vectors_out - if self.checkpoint_level >= 1 and self.training: - x = fleet.utils.recompute(self.ffn_l[0], x) - if self.use_global_vector_ffn: - global_vectors = fleet.utils.recompute( - self.global_ffn_l[0], global_vectors - ) - else: - x = self.ffn_l[0](x) - if self.use_global_vector_ffn: - global_vectors = self.global_ffn_l[0](global_vectors) - return x, global_vectors - else: - for idx, attn in enumerate(self.attn_l): - if self.checkpoint_level >= 2 and self.training: - out = fleet.utils.recompute(attn, x) - else: - out = attn(x) - x = x + out - if self.checkpoint_level >= 1 and self.training: - x = fleet.utils.recompute(self.ffn_l[0], x) - else: - x = self.ffn_l[0](x) - return x - - -class CuboidTransformerEncoder(nn.Layer): - """Encoder of the CuboidTransformer - - x --> attn_block --> patch_merge --> attn_block --> patch_merge --> ... --> out - - Args: - input_shape (Tuple[int,...]): The shape of the input. Contains T, H, W, C - base_units (int, optional): The number of units. Defaults to 128. - block_units (int, optional): The number of block units. Defaults to None. - scale_alpha (float, optional): We scale up the channels based on the formula: - - round_to(base_units * max(downsample_scale) ** units_alpha, 4). Defaults to 1.0. - depth (list, optional): The number of layers for each block. Defaults to [4, 4, 4]. - downsample (int, optional): The downsample ratio. Defaults to 2. - downsample_type (str, optional): The type of downsample. Defaults to "patch_merge". - block_attn_patterns (str, optional): Attention pattern for the cuboid attention for each block. Defaults to None. - block_cuboid_size (list, optional): A list of cuboid size parameters. Defaults to [(4, 4, 4), (4, 4, 4)]. - block_strategy (list, optional): A list of cuboid strategies. Defaults to [("l", "l", "l"), ("d", "d", "d")]. - block_shift_size (list, optional): A list of shift sizes. Defaults to [(0, 0, 0), (0, 0, 0)]. - num_heads (int, optional): The number of heads. Defaults to 4. - attn_drop (float, optional): The ratio of attention dropout. Defaults to 0.0. - proj_drop (float, optional): The ratio of projection dropout. Defaults to 0.0. - ffn_drop (float, optional): The ratio of FFN dropout. Defaults to 0.0. - ffn_activation (str, optional): The FFN activation. Defaults to "leaky". - gated_ffn (bool, optional): Whether to use gate FFN. Defaults to False. - norm_layer (str, optional): The normalization layer. Defaults to "layer_norm". - use_inter_ffn (bool, optional): Whether to use inter FFN. Defaults to True. - padding_type (str, optional): The type of padding. Defaults to "ignore". - checkpoint_level (bool, optional): Whether to enable gradient checkpointing. Defaults to True. - use_relative_pos (bool, optional): Whether to use relative pos. Defaults to True. - self_attn_use_final_proj (bool, optional): Whether to use self attention for final projection. Defaults to True. - use_global_vector (bool, optional): Whether to use the global vector or not. Defaults to False. - use_global_vector_ffn (bool, optional): Whether to use FFN global vectors. Defaults to False. - use_global_self_attn (bool, optional): Whether to use global self attention. Defaults to False. - separate_global_qkv (bool, optional): Whether to use different network to calc q_global, k_global, v_global. - Defaults to False. - global_dim_ratio (int, optional): The dim (channels) of global vectors is `global_dim_ratio*dim`. - Defaults to 1. - attn_linear_init_mode (str, optional): The mode of attention linear initialization. Defaults to "0". - ffn_linear_init_mode (str, optional): The mode of FFN linear initialization. Defaults to "0". - conv_init_mode (str, optional): The mode of conv initialization. Defaults to "0". - down_linear_init_mode (str, optional): The mode of downsample linear initialization. Defaults to "0". - norm_init_mode (str, optional): The mode of normalization. Defaults to "0". - """ - - def __init__( - self, - input_shape: Tuple[int, ...], - base_units: int = 128, - block_units: int = None, - scale_alpha: float = 1.0, - depth: Tuple[int, ...] = [4, 4, 4], - downsample: int = 2, - downsample_type: str = "patch_merge", - block_attn_patterns: str = None, - block_cuboid_size: Tuple[Tuple[int, ...], ...] = [(4, 4, 4), (4, 4, 4)], - block_strategy: Tuple[Tuple[str, ...], ...] = [ - ("l", "l", "l"), - ("d", "d", "d"), - ], - block_shift_size: Tuple[Tuple[int, ...], ...] = [(0, 0, 0), (0, 0, 0)], - num_heads: int = 4, - attn_drop: float = 0.0, - proj_drop: float = 0.0, - ffn_drop: float = 0.0, - ffn_activation: str = "leaky", - gated_ffn: bool = False, - norm_layer: str = "layer_norm", - use_inter_ffn: bool = True, - padding_type: str = "ignore", - checkpoint_level: bool = True, - use_relative_pos: bool = True, - self_attn_use_final_proj: bool = True, - use_global_vector: bool = False, - use_global_vector_ffn: bool = True, - use_global_self_attn: bool = False, - separate_global_qkv: bool = False, - global_dim_ratio: int = 1, - attn_linear_init_mode: str = "0", - ffn_linear_init_mode: str = "0", - conv_init_mode: str = "0", - down_linear_init_mode: str = "0", - norm_init_mode: str = "0", - ): - super(CuboidTransformerEncoder, self).__init__() - self.attn_linear_init_mode = attn_linear_init_mode - self.ffn_linear_init_mode = ffn_linear_init_mode - self.conv_init_mode = conv_init_mode - self.down_linear_init_mode = down_linear_init_mode - self.norm_init_mode = norm_init_mode - self.input_shape = input_shape - self.depth = depth - self.num_blocks = len(depth) - self.base_units = base_units - self.scale_alpha = scale_alpha - if not isinstance(downsample, (tuple, list)): - downsample = 1, downsample, downsample - self.downsample = downsample - self.downsample_type = downsample_type - self.num_heads = num_heads - self.use_global_vector = use_global_vector - self.checkpoint_level = checkpoint_level - if block_units is None: - block_units = [ - cuboid_utils.round_to( - base_units * int((max(downsample) ** scale_alpha) ** i), 4 - ) - for i in range(self.num_blocks) - ] - else: - assert len(block_units) == self.num_blocks and block_units[0] == base_units - self.block_units = block_units - if self.num_blocks > 1: - if downsample_type == "patch_merge": - self.down_layers = nn.LayerList( - sublayers=[ - PatchMerging3D( - dim=self.block_units[i], - downsample=downsample, - padding_type=padding_type, - out_dim=self.block_units[i + 1], - linear_init_mode=down_linear_init_mode, - norm_init_mode=norm_init_mode, - ) - for i in range(self.num_blocks - 1) - ] - ) - else: - raise NotImplementedError(f"{downsample_type} is invalid.") - if self.use_global_vector: - self.down_layer_global_proj = nn.LayerList( - sublayers=[ - nn.Linear( - in_features=global_dim_ratio * self.block_units[i], - out_features=global_dim_ratio * self.block_units[i + 1], - ) - for i in range(self.num_blocks - 1) - ] - ) - if block_attn_patterns is not None: - mem_shapes = self.get_mem_shapes() - if isinstance(block_attn_patterns, (tuple, list)): - assert len(block_attn_patterns) == self.num_blocks - else: - block_attn_patterns = [ - block_attn_patterns for _ in range(self.num_blocks) - ] - block_cuboid_size = [] - block_strategy = [] - block_shift_size = [] - for idx, key in enumerate(block_attn_patterns): - func = cuboid_utils.CuboidSelfAttentionPatterns.get(key) - cuboid_size, strategy, shift_size = func(mem_shapes[idx]) - block_cuboid_size.append(cuboid_size) - block_strategy.append(strategy) - block_shift_size.append(shift_size) - else: - if not isinstance(block_cuboid_size[0][0], (list, tuple)): - block_cuboid_size = [block_cuboid_size for _ in range(self.num_blocks)] - else: - assert ( - len(block_cuboid_size) == self.num_blocks - ), f"Incorrect input format! Received block_cuboid_size={block_cuboid_size}" - if not isinstance(block_strategy[0][0], (list, tuple)): - block_strategy = [block_strategy for _ in range(self.num_blocks)] - else: - assert ( - len(block_strategy) == self.num_blocks - ), f"Incorrect input format! Received block_strategy={block_strategy}" - if not isinstance(block_shift_size[0][0], (list, tuple)): - block_shift_size = [block_shift_size for _ in range(self.num_blocks)] - else: - assert ( - len(block_shift_size) == self.num_blocks - ), f"Incorrect input format! Received block_shift_size={block_shift_size}" - self.block_cuboid_size = block_cuboid_size - self.block_strategy = block_strategy - self.block_shift_size = block_shift_size - self.blocks = nn.LayerList( - sublayers=[ - nn.Sequential( - *[ - StackCuboidSelfAttentionBlock( - dim=self.block_units[i], - num_heads=num_heads, - block_cuboid_size=block_cuboid_size[i], - block_strategy=block_strategy[i], - block_shift_size=block_shift_size[i], - attn_drop=attn_drop, - proj_drop=proj_drop, - ffn_drop=ffn_drop, - activation=ffn_activation, - gated_ffn=gated_ffn, - norm_layer=norm_layer, - use_inter_ffn=use_inter_ffn, - padding_type=padding_type, - use_global_vector=use_global_vector, - use_global_vector_ffn=use_global_vector_ffn, - use_global_self_attn=use_global_self_attn, - separate_global_qkv=separate_global_qkv, - global_dim_ratio=global_dim_ratio, - checkpoint_level=checkpoint_level, - use_relative_pos=use_relative_pos, - use_final_proj=self_attn_use_final_proj, - attn_linear_init_mode=attn_linear_init_mode, - ffn_linear_init_mode=ffn_linear_init_mode, - norm_init_mode=norm_init_mode, - ) - for _ in range(depth[i]) - ] - ) - for i in range(self.num_blocks) - ] - ) - self.reset_parameters() - - def reset_parameters(self): - if self.num_blocks > 1: - for m in self.down_layers: - m.reset_parameters() - if self.use_global_vector: - cuboid_utils.apply_initialization( - self.down_layer_global_proj, linear_mode=self.down_linear_init_mode - ) - for ms in self.blocks: - for m in ms: - m.reset_parameters() - - def get_mem_shapes(self): - """Get the shape of the output memory based on the input shape. This can be used for constructing the decoder. - - Returns: - mem_shapes : A list of shapes of the output memory - """ - - if self.num_blocks == 1: - return [self.input_shape] - else: - mem_shapes = [self.input_shape] - curr_shape = self.input_shape - for down_layer in self.down_layers: - curr_shape = down_layer.get_out_shape(curr_shape) - mem_shapes.append(curr_shape) - return mem_shapes - - def forward(self, x, global_vectors=None): - """ - Args: - x : Shape (B, T, H, W, C) - - Returns: - out (List[paddle.Tensor,..]): A list of tensors from the bottom layer to the top layer of the encoder. For - example, it can have shape - - (B, T, H, W, C1) - - (B, T, H // 2, W // 2, 2 * C1) - - (B, T, H // 4, W // 4, 4 * C1) - ... - global_mem_out (List,Optional): The output of the global vector. - """ - - B, T, H, W, C_in = x.shape - assert (T, H, W, C_in) == self.input_shape - - if self.use_global_vector: - out = [] - global_mem_out = [] - for i in range(self.num_blocks): - for l in self.blocks[i]: - x, global_vectors = l(x, global_vectors) - out.append(x) - global_mem_out.append(global_vectors) - if self.num_blocks > 1 and i < self.num_blocks - 1: - x = self.down_layers[i](x) - global_vectors = self.down_layer_global_proj[i](global_vectors) - return out, global_mem_out - else: - out = [] - for i in range(self.num_blocks): - x = self.blocks[i](x) - out.append(x) - if self.num_blocks > 1 and i < self.num_blocks - 1: - x = self.down_layers[i](x) - return out diff --git a/examples/smc_reac/ppsci/arch/cuboid_transformer_utils.py b/examples/smc_reac/ppsci/arch/cuboid_transformer_utils.py deleted file mode 100644 index 3f7f366bc0..0000000000 --- a/examples/smc_reac/ppsci/arch/cuboid_transformer_utils.py +++ /dev/null @@ -1,347 +0,0 @@ -import functools -from typing import Tuple - -import paddle -import paddle.nn.functional as F -from paddle import nn - -from ppsci.utils import initializer - - -def round_to(dat, c): - return dat + (dat - dat % c) % c - - -class RMSNorm(nn.Layer): - """Root Mean Square Layer Normalization proposed in "[NeurIPS2019] Root Mean Square Layer Normalization" - - Args: - d (Optional[int]): The model size. - p (float, optional): The partial RMSNorm, valid value [0, 1]. Defaults to -1.0. - eps (float, optional): The epsilon value. Defaults to 1e-08. - bias (bool, optional): Whether use bias term for RMSNorm, - because RMSNorm doesn't enforce re-centering invariance.Defaults to False. - """ - - def __init__( - self, - d: Tuple[int, ...], - p: float = -1.0, - eps: float = 1e-08, - bias: bool = False, - ): - super().__init__() - self.eps = eps - self.d = d - self.p = p - self.bias = bias - init_data = paddle.ones(d) - self.scale = paddle.create_parameter( - shape=init_data.shape, - dtype=init_data.dtype, - default_initializer=nn.initializer.Constant(1.0), - ) - self.scale.stop_gradient = False - self.add_parameter(name="scale", parameter=self.scale) - if self.bias: - init_data = paddle.zeros(d) - self.offset = paddle.create_parameter( - shape=init_data.shape, - dtype=init_data.dtype, - default_initializer=nn.initializer.Constant(0.0), - ) - self.offset.stop_gradient = False - self.add_parameter(name="offset", parameter=self.offset) - - def forward(self, x): - if self.p < 0.0 or self.p > 1.0: - norm_x = x.norm(p=2, axis=-1, keepdim=True) - d_x = self.d - else: - partial_size = int(self.d * self.p) - partial_x, _ = paddle.split( - x=x, num_or_sections=[partial_size, self.d - partial_size], axis=-1 - ) - norm_x = partial_x.norm(p=2, axis=-1, keepdim=True) - d_x = partial_size - rms_x = norm_x * d_x ** (-1.0 / 2) - x_normed = x / (rms_x + self.eps) - if self.bias: - return self.scale * x_normed + self.offset - return self.scale * x_normed - - -def get_norm_layer( - normalization: str = "layer_norm", - axis: int = -1, - epsilon: float = 1e-05, - in_channels: int = 0, - **kwargs, -): - """Get the normalization layer based on the provided type - - Args: - normalization (str): The type of the layer normalization from ['layer_norm']. - axis (float): The axis to normalize the. - epsilon (float): The epsilon of the normalization layer. - in_channels (int): Input channel. - - Returns: - norm_layer (norm): The layer normalization layer. - """ - - if isinstance(normalization, str): - if normalization == "layer_norm": - assert in_channels > 0 - assert axis == -1 - norm_layer = nn.LayerNorm( - normalized_shape=in_channels, epsilon=epsilon, **kwargs - ) - elif normalization == "rms_norm": - assert axis == -1 - norm_layer = RMSNorm(d=in_channels, eps=epsilon, **kwargs) - else: - raise NotImplementedError(f"normalization={normalization} is not supported") - return norm_layer - elif normalization is None: - return nn.Identity() - else: - raise NotImplementedError("The type of normalization must be str") - - -def generalize_padding(x, pad_t, pad_h, pad_w, padding_type, t_pad_left=False): - if pad_t == 0 and pad_h == 0 and pad_w == 0: - return x - assert padding_type in ["zeros", "ignore", "nearest"] - B, T, H, W, C = x.shape - if padding_type == "nearest": - return nn.functional.interpolate( - x=x.transpose(perm=[0, 4, 1, 2, 3]), size=(T + pad_t, H + pad_h, W + pad_w) - ).transpose(perm=[0, 2, 3, 4, 1]) - elif t_pad_left: - return F.pad(x, [0, 0, 0, pad_w, 0, pad_h, pad_t, 0], data_format="NDHWC") - else: - data_pad = F.pad( - x, [0, 0, pad_t, 0, pad_h, 0, pad_w, 0, 0, 0], data_format="NDHWC" - ) - data_pad = paddle.concat( - [data_pad[:, pad_t:, ...], data_pad[:, :pad_t, ...]], axis=1 - ) - return data_pad - - -def generalize_unpadding(x, pad_t, pad_h, pad_w, padding_type): - assert padding_type in ["zeros", "ignore", "nearest"] - B, T, H, W, C = x.shape - if pad_t == 0 and pad_h == 0 and pad_w == 0: - return x - if padding_type == "nearest": - return nn.functional.interpolate( - x=x.transpose(perm=[0, 4, 1, 2, 3]), size=(T - pad_t, H - pad_h, W - pad_w) - ).transpose(perm=[0, 2, 3, 4, 1]) - else: - return x[:, : T - pad_t, : H - pad_h, : W - pad_w, :] - - -def apply_initialization( - m: nn.Layer, - linear_mode: str = "0", - conv_mode: str = "0", - norm_mode: str = "0", - embed_mode: str = "0", -): - if isinstance(m, nn.Linear): - if linear_mode in ("0",): - m.weight = initializer.kaiming_normal_(m.weight, nonlinearity="linear") - elif linear_mode in ("1",): - m.weight = initializer.kaiming_normal_( - m.weight, a=0.1, mode="fan_out", nonlinearity="leaky_relu" - ) - else: - raise NotImplementedError(f"{linear_mode} is invalid.") - if hasattr(m, "bias") and m.bias is not None: - m.bias = initializer.zeros_(m.bias) - elif isinstance( - m, - ( - nn.Conv2D, - nn.Conv3D, - nn.Conv2DTranspose, - nn.Conv3DTranspose, - ), - ): - if conv_mode in ("0",): - m.weight = initializer.kaiming_normal_( - m.weight, a=0.1, mode="fan_out", nonlinearity="leaky_relu" - ) - else: - raise NotImplementedError(f"{conv_mode} is invalid.") - if hasattr(m, "bias") and m.bias is not None: - m.bias = initializer.zeros_(m.bias) - elif isinstance(m, nn.LayerNorm): - if norm_mode in ("0",): - m.weight = initializer.zeros_(m.weight) - m.bias = initializer.zeros_(m.bias) - else: - raise NotImplementedError(f"{norm_mode} is invalid.") - elif isinstance(m, nn.GroupNorm): - if norm_mode in ("0",): - m.weight = initializer.ones_(m.weight) - m.bias = initializer.zeros_(m.bias) - else: - raise NotImplementedError(f"{norm_mode} is invalid.") - elif isinstance(m, nn.Embedding): - if embed_mode in ("0",): - m.weight.data = initializer.trunc_normal_(m.weight.data, std=0.02) - else: - raise NotImplementedError(f"{embed_mode} is invalid.") - - else: - pass - - -class CuboidSelfAttentionPatterns: - def __init__(self): - super().__init__() - self.patterns = {} - self.patterns = { - "full": self.full_attention, - "axial": self.axial, - "divided_st": self.divided_space_time, - } - for p in [1, 2, 4, 8, 10]: - for m in [1, 2, 4, 8, 16, 32]: - key = f"video_swin_{p}x{m}" - self.patterns[key] = functools.partial(self.video_swin, P=p, M=m) - - for m in [1, 2, 4, 8, 16, 32]: - key = f"spatial_lg_{m}" - self.patterns[key] = functools.partial(self.spatial_lg_v1, M=m) - - for k in [2, 4, 8]: - key = f"axial_space_dilate_{k}" - self.patterns[key] = functools.partial(self.axial_space_dilate_K, K=k) - - def get(self, pattern_name): - return self.patterns[pattern_name] - - def full_attention(self, input_shape): - T, H, W, _ = input_shape - cuboid_size = [(T, H, W)] - strategy = [("l", "l", "l")] - shift_size = [(0, 0, 0)] - return cuboid_size, strategy, shift_size - - def axial(self, input_shape): - """Axial attention proposed in https://arxiv.org/abs/1912.12180 - - Args: - input_shape (Tuple[int,...]): The shape of the input tensor, T H W. - - Returns: - cuboid_size (Tuple[int,...]): The size of cuboid. - strategy (Tuple[str,...]): The strategy of the attention. - shift_size (Tuple[int,...]): The shift size of the attention. - """ - - T, H, W, _ = input_shape - cuboid_size = [(T, 1, 1), (1, H, 1), (1, 1, W)] - strategy = [("l", "l", "l"), ("l", "l", "l"), ("l", "l", "l")] - shift_size = [(0, 0, 0), (0, 0, 0), (0, 0, 0)] - return cuboid_size, strategy, shift_size - - def divided_space_time(self, input_shape): - T, H, W, _ = input_shape - cuboid_size = [(T, 1, 1), (1, H, W)] - strategy = [("l", "l", "l"), ("l", "l", "l")] - shift_size = [(0, 0, 0), (0, 0, 0)] - return cuboid_size, strategy, shift_size - - def video_swin(self, input_shape, P=2, M=4): - """Adopt the strategy in Video SwinTransformer https://arxiv.org/pdf/2106.13230.pdf""" - T, H, W, _ = input_shape - P = min(P, T) - M = min(M, H, W) - cuboid_size = [(P, M, M), (P, M, M)] - strategy = [("l", "l", "l"), ("l", "l", "l")] - shift_size = [(0, 0, 0), (P // 2, M // 2, M // 2)] - return cuboid_size, strategy, shift_size - - def spatial_lg_v1(self, input_shape, M=4): - T, H, W, _ = input_shape - if H <= M and W <= M: - cuboid_size = [(T, 1, 1), (1, H, W)] - strategy = [("l", "l", "l"), ("l", "l", "l")] - shift_size = [(0, 0, 0), (0, 0, 0)] - else: - cuboid_size = [(T, 1, 1), (1, M, M), (1, M, M)] - strategy = [("l", "l", "l"), ("l", "l", "l"), ("d", "d", "d")] - shift_size = [(0, 0, 0), (0, 0, 0), (0, 0, 0)] - return cuboid_size, strategy, shift_size - - def axial_space_dilate_K(self, input_shape, K=2): - T, H, W, _ = input_shape - K = min(K, H, W) - cuboid_size = [ - (T, 1, 1), - (1, H // K, 1), - (1, H // K, 1), - (1, 1, W // K), - (1, 1, W // K), - ] - strategy = [ - ("l", "l", "l"), - ("d", "d", "d"), - ("l", "l", "l"), - ("d", "d", "d"), - ("l", "l", "l"), - ] - shift_size = [(0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0)] - return cuboid_size, strategy, shift_size - - -class CuboidCrossAttentionPatterns: - def __init__(self): - super().__init__() - self.patterns = {} - for k in [1, 2, 4, 8]: - key1 = f"cross_{k}x{k}" - key2 = f"cross_{k}x{k}_lg" - key3 = f"cross_{k}x{k}_heter" - self.patterns[key1] = functools.partial(self.cross_KxK, K=k) - self.patterns[key2] = functools.partial(self.cross_KxK_lg, K=k) - self.patterns[key3] = functools.partial(self.cross_KxK_heter, K=k) - - def get(self, pattern_name): - return self.patterns[pattern_name] - - def cross_KxK(self, mem_shape, K): - T_mem, H, W, _ = mem_shape - K = min(K, H, W) - cuboid_hw = [(K, K)] - shift_hw = [(0, 0)] - strategy = [("l", "l", "l")] - n_temporal = [1] - return cuboid_hw, shift_hw, strategy, n_temporal - - def cross_KxK_lg(self, mem_shape, K): - T_mem, H, W, _ = mem_shape - K = min(K, H, W) - cuboid_hw = [(K, K), (K, K)] - shift_hw = [(0, 0), (0, 0)] - strategy = [("l", "l", "l"), ("d", "d", "d")] - n_temporal = [1, 1] - return cuboid_hw, shift_hw, strategy, n_temporal - - def cross_KxK_heter(self, mem_shape, K): - T_mem, H, W, _ = mem_shape - K = min(K, H, W) - cuboid_hw = [(K, K), (K, K), (K, K)] - shift_hw = [(0, 0), (0, 0), (K // 2, K // 2)] - strategy = [("l", "l", "l"), ("d", "d", "d"), ("l", "l", "l")] - n_temporal = [1, 1, 1] - return cuboid_hw, shift_hw, strategy, n_temporal - - -CuboidSelfAttentionPatterns = CuboidSelfAttentionPatterns() -CuboidCrossAttentionPatterns = CuboidCrossAttentionPatterns() diff --git a/examples/smc_reac/ppsci/arch/cvit.py b/examples/smc_reac/ppsci/arch/cvit.py deleted file mode 100644 index d39abd3118..0000000000 --- a/examples/smc_reac/ppsci/arch/cvit.py +++ /dev/null @@ -1,1095 +0,0 @@ -# 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. - -import importlib - -try: - import einops -except ModuleNotFoundError: - pass -from typing import Callable -from typing import Optional -from typing import Sequence -from typing import Tuple - -import paddle -from paddle import nn -from paddle.nn import functional as F - -from ppsci.arch import base -from ppsci.utils import initializer - - -# Positional embedding from masked autoencoder https://arxiv.org/abs/2111.06377 -def get_1d_sincos_pos_embed_from_grid(embed_dim: int, pos: paddle.Tensor): - if embed_dim % 2 != 0: - raise ValueError(f"embedding dimension({embed_dim}) must be divisible by 2") - - omega = paddle.arange(embed_dim // 2, dtype=paddle.float32) - omega /= embed_dim / 2.0 - omega = 1.0 / 10000**omega # (D/2,) - - pos = pos.reshape([-1]) # (M,) - out = paddle.einsum("m,d->md", pos, omega) # (M, D/2), outer product - - emb_sin = paddle.sin(out) # (M, D/2) - emb_cos = paddle.cos(out) # (M, D/2) - - emb = paddle.concat([emb_sin, emb_cos], axis=1) # (M, D) - return emb - - -def get_1d_sincos_pos_embed(embed_dim: int, length: int): - return paddle.unsqueeze( - get_1d_sincos_pos_embed_from_grid( - embed_dim, paddle.arange(length, dtype=paddle.float32) - ), - 0, - ) - - -def get_2d_sincos_pos_embed(embed_dim: int, grid_size: Tuple[int, int]): - def get_2d_sincos_pos_embed_from_grid(embed_dim, grid): - if embed_dim % 2 != 0: - raise ValueError(f"embedding dimension({embed_dim}) must be divisible by 2") - - # use half of dimensions to encode grid_h - emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0]) # (H*W, D/2) - emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1]) # (H*W, D/2) - emb = paddle.concat([emb_h, emb_w], axis=1) # (H*W, D) - return emb - - grid_h = paddle.arange(grid_size[0], dtype=paddle.float32) - grid_w = paddle.arange(grid_size[1], dtype=paddle.float32) - grid = paddle.meshgrid(grid_w, grid_h, indexing="ij") # here w goes first - grid = paddle.stack(grid, axis=0) - - grid = grid.reshape([2, 1, grid_size[0], grid_size[1]]) - pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid) - - return paddle.unsqueeze(pos_embed, 0) - - -class MlpBlock(nn.Layer): - def __init__(self, in_dim: int, dim: int = 256, out_dim: int = 256): - super().__init__() - self.in_dim = in_dim - self.dim = dim - self.out_dim = out_dim - self.linear1 = nn.Linear(self.in_dim, self.dim) - self.act = nn.GELU(True) - self.linear2 = nn.Linear(self.dim, self.out_dim) - - self._init_weights() - - def forward(self, inputs): - x = self.linear1(inputs) - x = self.act(x) - x = self.linear2(x) - return x - - def _init_weights(self) -> None: - initializer.xavier_uniform_(self.linear1.weight) - initializer.constant_(self.linear1.bias, 0) - initializer.xavier_uniform_(self.linear2.weight) - initializer.constant_(self.linear2.bias, 0) - - -class SelfAttnBlock(nn.Layer): - def __init__( - self, num_heads: int, emb_dim: int, mlp_ratio: int, layer_norm_eps: float = 1e-5 - ): - super().__init__() - self.num_heads = num_heads - self.emb_dim = emb_dim - self.mlp_ratio = mlp_ratio - self.layer_norm1 = nn.LayerNorm(emb_dim, layer_norm_eps) - self.attn_layer = MultiHeadDotProductAttention( - self.emb_dim, - num_heads=self.num_heads, - qkv_features=self.emb_dim, - ) - self.layer_norm2 = nn.LayerNorm(emb_dim, layer_norm_eps) - self.mlp = MlpBlock(self.emb_dim, self.emb_dim * self.mlp_ratio, self.emb_dim) - - def forward(self, inputs): - # inputs: # [B, L/ps, self.emb_dim] - x = self.layer_norm1(inputs) - x = self.attn_layer(x, x) - x = x + inputs - y = self.layer_norm2(x) - y = self.mlp(y) - return x + y - - -class Mlp(nn.Layer): - def __init__( - self, - num_layers: int, - hidden_dim: int, - out_dim: int, - layer_norm_eps: float = 1e-5, - ): - super().__init__() - self.num_layers = num_layers - self.hidden_dim = hidden_dim - self.out_dim = out_dim - self.layer_norm_eps = layer_norm_eps - self.linears = nn.LayerList( - [ - nn.Linear( - self.hidden_dim, - self.hidden_dim, - ) - for _ in range(self.num_layers) - ] - ) - self.gelu = nn.GELU(True) - self.norms = nn.LayerList( - [ - nn.LayerNorm(self.hidden_dim, self.layer_norm_eps) - for _ in range(self.num_layers) - ] - ) - - self.linear_out = nn.Linear(self.hidden_dim, self.out_dim) - - self._init_weights() - - def forward(self, inputs): - x = inputs - for i in range(self.num_layers): - y = self.linears[i](x) - y = self.gelu(y) - x = x + y - x = self.norms[i](x) - - x = self.linear_out(x) - return x - - def _init_weights(self) -> None: - for linear in self.linears: - initializer.xavier_uniform_(linear.weight) - initializer.constant_(linear.bias, 0) - - -class PatchEmbed1D(nn.Layer): - def __init__( - self, - in_dim: int, - patch_size: Sequence[int] = (4,), - emb_dim: int = 768, - use_norm: bool = False, - layer_norm_eps: float = 1e-5, - ): - super().__init__() - self.patch_size = patch_size - self.emb_dim = emb_dim - self.use_norm = use_norm - self.layer_norm_eps = layer_norm_eps - self.conv = nn.Conv1D( - in_dim, - self.emb_dim, - self.patch_size[0], - self.patch_size[0], - data_format="NLC", - ) - self.norm = ( - nn.LayerNorm(self.emb_dim, self.layer_norm_eps) - if self.use_norm - else nn.Identity() - ) - self._init_weights() - - def forward(self, x): - x = self.conv(x) # [B, L, C] --> [B, L/ps, self.emb_dim] - if self.use_norm: - x = self.norm(x) - return x - - def _init_weights(self) -> None: - initializer.xavier_uniform_(self.conv.weight) - initializer.constant_(self.conv.bias, 0) - - -class PatchEmbed(nn.Layer): - def __init__( - self, - in_dim: int, - spatial_dims: Sequence[int], - patch_size: Tuple[int, ...] = (1, 16, 16), - emb_dim: int = 768, - use_norm: bool = False, - layer_norm_eps: float = 1e-5, - ): - super().__init__() - self.patch_size = patch_size - self.emb_dim = emb_dim - self.use_norm = use_norm - self.layer_norm_eps = layer_norm_eps - self.conv = nn.Conv3D( - in_dim, - self.emb_dim, - (self.patch_size[0], self.patch_size[1], self.patch_size[2]), - (self.patch_size[0], self.patch_size[1], self.patch_size[2]), - data_format="NDHWC", - ) - self.norm = ( - nn.LayerNorm(self.emb_dim, self.layer_norm_eps) - if self.use_norm - else nn.Identity() - ) - t, h, w = spatial_dims - self.num_patches = [ - t // self.patch_size[0], - h // self.patch_size[1], - w // self.patch_size[2], - ] - self._init_weights() - - def forward(self, x): - b, t, h, w, c = x.shape - - x = self.conv(x) # [B, L, C] --> [B, L/ps, self.emb_dim] - x = x.reshape( - [ - b, - self.num_patches[0], - self.num_patches[1] * self.num_patches[2], - self.emb_dim, - ] - ) - if self.use_norm: - x = self.norm(x) - return x - - def _init_weights(self) -> None: - initializer.xavier_uniform_(self.conv.weight) - initializer.constant_(self.conv.bias, 0) - - -class CrossAttnBlock(nn.Layer): - def __init__( - self, - num_heads: int, - emb_dim: int, - mlp_ratio: int, - layer_norm_eps: float = 1e-5, - out_features: int = None, - qkv_features: int = None, - ): - super().__init__() - self.num_heads = num_heads - self.emb_dim = emb_dim - self.mlp_ratio = mlp_ratio - self.layer_norm_eps = layer_norm_eps - self.head_dim = self.emb_dim // self.num_heads - - self.layer_norm_q = nn.LayerNorm(self.emb_dim, epsilon=self.layer_norm_eps) - self.layer_norm_kv = nn.LayerNorm(self.emb_dim, epsilon=self.layer_norm_eps) - - self.attn_layer = MultiHeadDotProductAttention( - self.emb_dim, - num_heads=num_heads, - qkv_features=qkv_features, - out_features=out_features, - ) - self.layer_norm_y = nn.LayerNorm(self.emb_dim, epsilon=self.layer_norm_eps) - self.mlp = MlpBlock(self.emb_dim, self.emb_dim * self.mlp_ratio, self.emb_dim) - - def forward(self, q_inputs, kv_inputs): - # [B, L/ps, self.dec_emb_dim] - q = self.layer_norm_q(q_inputs) - kv = self.layer_norm_kv(kv_inputs) - x = self.attn_layer(q, kv) - x = x + q_inputs - y = self.layer_norm_y(x) - y = self.mlp(y) - return x + y - - -class Encoder1D(nn.Layer): - def __init__( - self, - in_dim: int, - spatial_dims: int, - patch_size: int = (4,), - emb_dim: int = 256, - depth: int = 3, - num_heads: int = 8, - mlp_ratio: int = 1, - layer_norm_eps: float = 1e-5, - ): - super().__init__() - self.in_dim = in_dim - self.spatial_dims = spatial_dims - self.patch_size = patch_size - self.emb_dim = emb_dim - self.depth = depth - self.num_heads = num_heads - self.mlp_ratio = mlp_ratio - self.layer_norm_eps = layer_norm_eps - self.patch_embedding = PatchEmbed1D(in_dim, self.patch_size, self.emb_dim) - - self.self_attn_blocks = nn.LayerList( - [ - SelfAttnBlock( - self.num_heads, - self.emb_dim, - self.mlp_ratio, - self.layer_norm_eps, - ) - for _ in range(self.depth) - ] - ) - pos_emb = get_1d_sincos_pos_embed( - self.emb_dim, self.spatial_dims // self.patch_size[0] - ) - self.pos_emb = self.create_parameter( - pos_emb.shape, default_initializer=nn.initializer.Assign(pos_emb) - ) - - def forward(self, x): - x = self.patch_embedding(x) - x = x + self.pos_emb - - for _, block in enumerate(self.self_attn_blocks): - x = block(x) - - return x - - -class TimeAggregation(nn.Layer): - def __init__( - self, - emb_dim: int, - depth: int, - num_heads: int = 8, - num_latents: int = 64, - mlp_ratio: int = 1, - layer_norm_eps: float = 1e-5, - ): - super().__init__() - self.emb_dim = emb_dim - self.depth = depth - self.num_heads = num_heads - self.num_latents = num_latents - self.mlp_ratio = mlp_ratio - self.layer_norm_eps = layer_norm_eps - self.latents = self.create_parameter( - [self.num_latents, self.emb_dim], - default_initializer=nn.initializer.Normal(std=1e-2), - ) - self.cross_attn_blocks = nn.LayerList( - [ - CrossAttnBlock( - self.num_heads, self.emb_dim, self.mlp_ratio, self.layer_norm_eps - ) - for _ in range(self.depth) - ] - ) - - def forward(self, x): # (B, T, S, D) --> (B, T', S, D) - latents = einops.repeat( - self.latents, "t d -> b s t d", b=x.shape[0], s=x.shape[2] - ) # (B, T', S, D) - x = einops.rearrange(x, "b t s d -> b s t d") # (B, S, T, D) - - # Transformer - for i, block in enumerate(self.cross_attn_blocks): - latents = block(latents, x) - - latents = einops.rearrange(latents, "b s t d -> b t s d") # (B, T', S, D) - return latents - - -class Encoder(nn.Layer): - def __init__( - self, - in_dim: int, - spatial_dims: Sequence[int], - patch_size: int = (1, 16, 16), - emb_dim: int = 256, - depth: int = 3, - num_heads: int = 8, - mlp_ratio: int = 1, - layer_norm_eps: float = 1e-5, - ): - super().__init__() - self.in_dim = in_dim - self.spatial_dims = spatial_dims - self.patch_size = patch_size - self.emb_dim = emb_dim - self.depth = depth - self.num_heads = num_heads - self.mlp_ratio = mlp_ratio - self.layer_norm_eps = layer_norm_eps - self.patch_embedding = PatchEmbed( - in_dim, spatial_dims, self.patch_size, self.emb_dim - ) - - self.time_aggreator = TimeAggregation( - self.emb_dim, - 2, - self.num_heads, - 1, - self.mlp_ratio, - self.layer_norm_eps, - ) - self.norm = nn.LayerNorm(self.emb_dim, epsilon=self.layer_norm_eps) - - self.self_attn_blocks = nn.LayerList( - [ - SelfAttnBlock( - self.num_heads, - self.emb_dim, - self.mlp_ratio, - self.layer_norm_eps, - ) - for _ in range(self.depth) - ] - ) - t, h, w = spatial_dims - - time_emb = get_1d_sincos_pos_embed(self.emb_dim, t // self.patch_size[0]) - self.time_emb = self.create_parameter( - time_emb.shape, default_initializer=nn.initializer.Assign(time_emb) - ) - - pos_emb = get_2d_sincos_pos_embed( - self.emb_dim, (h // self.patch_size[1], w // self.patch_size[2]) - ) - self.pos_emb = self.create_parameter( - pos_emb.shape, default_initializer=nn.initializer.Assign(pos_emb) - ) - - def forward(self, x): - # patchify - x = self.patch_embedding(x) - - # add positional embedding - x = x + self.time_emb.unsqueeze(2) + self.pos_emb.unsqueeze(1) - - # aggregate along time dimension - x = self.time_aggreator(x) - x = self.norm(x) - x = einops.rearrange(x, "b t s d -> b (t s) d") - - for _, block in enumerate(self.self_attn_blocks): - x = block(x) - - return x - - -def dot_product_attention_weights( - query: paddle.Tensor, - key: paddle.Tensor, - bias: Optional[paddle.Tensor] = None, -): - """Computes dot-product attention weights given query and key. - - Used by :func:`dot_product_attention`, which is what you'll most likely use. - But if you want access to the attention weights for introspection, then - you can directly call this function and call einsum yourself. - - Args: - query: queries for calculating attention with shape of [batch..., q_length, - num_heads, qk_depth_per_head]. - key: keys for calculating attention with shape of [batch..., kv_length, - num_heads, qk_depth_per_head]. - bias: bias for the attention weights. This should be broadcastable to the - shape [batch..., num_heads, q_length, kv_length]. This can be used for - incorporating causal masks, padding masks, proximity bias, etc. - - Returns: - Output of shape [batch..., num_heads, q_length, kv_length]. - """ - dtype = query.dtype - - if paddle.in_dynamic_mode(): - assert query.ndim == key.ndim, "q, k must have same rank." - assert query.shape[:-3] == key.shape[:-3], "q, k batch dims must match." - assert query.shape[-2] == key.shape[-2], "q, k num_heads must match." - assert query.shape[-1] == key.shape[-1], "q, k depths must match." - - # calculate attention matrix - depth = query.shape[-1] - query = query / (depth**0.5) - # attn weight shape is (batch..., num_heads, q_length, kv_length) - attn_weights = paddle.einsum("...qhd,...khd->...hqk", query, key) - - # apply attention bias: masking, dropout, proximity bias, etc. - if bias is not None: - attn_weights = attn_weights + bias - - # normalize the attention weights - attn_weights = F.softmax(attn_weights).astype(dtype) - - # apply attention dropout - return attn_weights - - -def dot_product_attention( - query: paddle.Tensor, - key: paddle.Tensor, - value: paddle.Tensor, - bias: Optional[paddle.Tensor] = None, -) -> paddle.Tensor: - """Computes dot-product attention given query, key, and value. - - This is the core function for applying attention based on - https://arxiv.org/abs/1706.03762. It calculates the attention weights given - query and key and combines the values using the attention weights. - - Note: query, key, value needn't have any batch dimensions. - - Args: - query: queries for calculating attention with shape of [batch..., q_length, - num_heads, qk_depth_per_head]. - key: keys for calculating attention with shape of [batch..., kv_length, - num_heads, qk_depth_per_head]. - value: values to be used in attention with shape of [batch..., kv_length, - num_heads, v_depth_per_head]. - bias: bias for the attention weights. This should be broadcastable to the - shape [batch..., num_heads, q_length, kv_length]. This can be used for - incorporating causal masks, padding masks, proximity bias, etc. - - Returns: - paddle.Tensor: Output of shape [batch..., q_length, num_heads, v_depth_per_head]. - """ - if paddle.in_dynamic_mode(): - assert key.ndim == query.ndim == value.ndim, "q, k, v must have same rank." - assert ( - query.shape[:-3] == key.shape[:-3] == value.shape[:-3] - ), "q, k, v batch dims must match." - assert ( - query.shape[-2] == key.shape[-2] == value.shape[-2] - ), "q, k, v num_heads must match." - assert key.shape[-3] == value.shape[-3], "k, v lengths must match." - - # compute attention weights - attn_weights = dot_product_attention_weights( - query, - key, - bias, - ) - - # return weighted sum over values for each query position - return paddle.einsum("...hqk,...khd->...qhd", attn_weights, value) - - -class MultiHeadDotProductAttention(nn.Layer): - """Multi-head dot-product attention. - - Args: - in_dim: Number of input dimensions. - num_heads: Number of attention heads. Features (i.e. inputs_q.shape[-1]) - should be divisible by the number of heads. - qkv_features: dimension of the key, query, and value. - out_features: dimension of the last projection - use_bias: bool: whether pointwise QKVO dense transforms use bias. - attention_fn: dot_product_attention or compatible function. Accepts query, - key, value, and returns output of shape [bs, dim1, dim2, ..., dimN,, - num_heads, value_channels]` - normalize_qk: should QK normalization be applied (arxiv.org/abs/2302.05442). - """ - - def __init__( - self, - in_dim, - num_heads: int, - qkv_features: Optional[int] = None, - out_features: Optional[int] = None, - use_bias: bool = True, - attention_fn: Callable[..., paddle.Tensor] = dot_product_attention, - normalize_qk: bool = False, - ): - super().__init__() - self.num_heads = num_heads - self.qkv_features = qkv_features or in_dim - self.out_features = out_features or in_dim - self.use_bias = use_bias - self.attention_fn = attention_fn - self.normalize_qk = normalize_qk - assert self.qkv_features % self.num_heads == 0, ( - f"Memory dimension ({self.qkv_features}) must be divisible by number of" - f" heads ({self.num_heads})." - ) - self.head_dim = self.qkv_features // self.num_heads - - self.linear_q = nn.Linear( - in_dim, - self.qkv_features, - bias_attr=use_bias, - ) - self.linear_k = nn.Linear( - in_dim, - self.qkv_features, - bias_attr=use_bias, - ) - self.linear_v = nn.Linear( - in_dim, - self.qkv_features, - bias_attr=use_bias, - ) - self.query_ln = ( - nn.LayerNorm(self.qkv_features) if normalize_qk else nn.Identity() - ) - self.key_ln = nn.LayerNorm(self.qkv_features) if normalize_qk else nn.Identity() - self.linear_out = nn.Linear( - self.qkv_features, - self.out_features, - bias_attr=use_bias, - ) - - def forward( - self, - inputs_q: paddle.Tensor, - inputs_kv: Optional[paddle.Tensor] = None, - ): - # project inputs_q to multi-headed q/k/v - # dimensions are then [batch..., length, n_heads, n_features_per_head] - q_attn_shape = inputs_q.shape - q_attn_shape = q_attn_shape[:-1] + [self.num_heads, self.head_dim] - - kv_attn_shape = inputs_kv.shape - kv_attn_shape = kv_attn_shape[:-1] + [self.num_heads, self.head_dim] - query, key, value = ( - self.linear_q(inputs_q).reshape(q_attn_shape), - self.linear_k(inputs_kv).reshape(kv_attn_shape), - self.linear_v(inputs_kv).reshape(kv_attn_shape), - ) - - if self.normalize_qk: - # Normalizing query and key projections stabilizes training with higher - # LR. See ViT-22B paper http://arxiv.org/abs/2302.05442 for analysis. - query = self.query_ln(query) - key = self.key_ln(key) - - # apply attention - x = self.attention_fn( - query, - key, - value, - ) - # back to the original inputs dimensions - x = x.reshape(x.shape[:-2] + [x.shape[-2] * x.shape[-1]]) - out = self.linear_out(x) - return out - - -class CVit1D(base.Arch): - """ - 1D Convolutional Vision Transformer (CVit1D) class. - - [Bridging Operator Learning and Conditioned Neural Fields: A Unifying Perspective](https://arxiv.org/abs/2405.13998) - - Args: - input_keys (Sequence[str]): Keys identifying the input tensors. - output_keys (Sequence[str]): Keys identifying the output tensors. - spatial_dims (int): The spatial dimensions of the input data. - in_dim (int): The dimensionality of the input data. - coords_dim (int): The dimensionality of the positional encoding. - patch_size (Sequence[int], optional): Size of the patches. Defaults to (4,). - grid_size (Sequence[int], optional): Size of the grid. Defaults to (200,). - latent_dim (int, optional): Dimensionality of the latent space. Defaults to 256. - emb_dim (int, optional): Dimensionality of the embedding space. Defaults to 256. - depth (int, optional): Number of transformer encoder layers. Defaults to 3. - num_heads (int, optional): Number of attention heads. Defaults to 8. - dec_emb_dim (int, optional): Dimensionality of the decoder embedding space. Defaults to 256. - dec_num_heads (int, optional): Number of decoder attention heads. Defaults to 8. - dec_depth (int, optional): Number of decoder transformer layers. Defaults to 1. - num_mlp_layers (int, optional): Number of layers in the MLP. Defaults to 1. - mlp_ratio (int, optional): Ratio for determining the size of the MLP's hidden layer. Defaults to 1. - out_dim (int, optional): Dimensionality of the output data. Defaults to 1. - layer_norm_eps (float, optional): Epsilon for layer normalization. Defaults to 1e-5. - embedding_type (str, optional): Type of embedding to use ("grid" or other options). Defaults to "grid". - - Examples: - >>> import ppsci - >>> b, l, c = 2, 32, 1 - >>> l_query = 42 - >>> c_in = 1 - >>> c_out = 1 - >>> model = ppsci.arch.CVit1D( - ... input_keys=["u", "y"], - ... output_keys=["s"], - ... in_dim=c_in, - ... coords_dim=1, - ... spatial_dims=l, - ... patch_size=[4], - ... grid_size=[l], - ... latent_dim=32, - ... emb_dim=32, - ... depth=3, - ... num_heads=8, - ... dec_emb_dim=32, - ... dec_num_heads=8, - ... dec_depth=1, - ... num_mlp_layers=1, - ... mlp_ratio=1, - ... out_dim=c_out, - ... layer_norm_eps=1e-5, - ... embedding_type="grid", - ... ) - >>> x = paddle.randn([b, l, c_in]) - >>> coords = paddle.randn([l_query, 1]) - >>> out = model({"u": x, "y": coords})["s"] - >>> print(out.shape) # output shape should be [b, l_query, c_out] - [2, 42, 1] - """ - - def __init__( - self, - input_keys: Sequence[str], - output_keys: Sequence[str], - spatial_dims: int, - in_dim: int, - coords_dim: int, - patch_size: Sequence[int] = (4,), - grid_size: Sequence[int] = (200,), - latent_dim: int = 256, - emb_dim: int = 256, - depth: int = 3, - num_heads: int = 8, - dec_emb_dim: int = 256, - dec_num_heads: int = 8, - dec_depth: int = 1, - num_mlp_layers: int = 1, - mlp_ratio: int = 1, - out_dim: int = 1, - layer_norm_eps: float = 1e-5, - embedding_type: str = "grid", - ): - if not importlib.util.find_spec("einops"): - raise ModuleNotFoundError( - "Please install `einops` by running 'pip install einops'." - ) - super().__init__() - self.input_keys = input_keys - self.output_keys = output_keys - self.spatial_dims = spatial_dims - self.in_dim = in_dim - self.coords_dim = coords_dim - self.patch_size = patch_size - self.grid_size = grid_size - self.latent_dim = latent_dim - self.emb_dim = emb_dim - self.depth = depth - self.num_heads = num_heads - self.dec_emb_dim = dec_emb_dim - self.dec_num_heads = dec_num_heads - self.dec_depth = dec_depth - self.num_mlp_layers = num_mlp_layers - self.mlp_ratio = mlp_ratio - self.out_dim = out_dim - self.layer_norm_eps = layer_norm_eps - self.embedding_type = embedding_type - - if self.embedding_type == "grid": - # Create grid and latents - n_x = self.grid_size[0] - self.grid = paddle.linspace(0, 1, n_x) - self.latents = self.create_parameter( - [n_x, self.latent_dim], - default_initializer=nn.initializer.Normal(std=1e-2), - ) - self.fc = nn.Linear(self.latent_dim, self.dec_emb_dim) - self.norm = nn.LayerNorm(self.dec_emb_dim, self.layer_norm_eps) - elif self.embedding_type == "mlp": - self.mlp = MlpBlock(self.latent_dim, self.dec_emb_dim, self.dec_emb_dim) - self.norm = nn.LayerNorm(self.dec_emb_dim, self.layer_norm_eps) - - self.encoder = Encoder1D( - self.in_dim, - self.spatial_dims, - self.patch_size, - self.emb_dim, - self.depth, - self.num_heads, - self.mlp_ratio, - self.layer_norm_eps, - ) - self.enc_norm = nn.LayerNorm(self.emb_dim, self.layer_norm_eps) - self.fc1 = nn.Linear(self.emb_dim, self.dec_emb_dim) - self.cross_attn_blocks = nn.LayerList( - [ - CrossAttnBlock( - self.dec_num_heads, - self.dec_emb_dim, - self.mlp_ratio, - self.layer_norm_eps, - self.dec_emb_dim, - self.dec_emb_dim, - ) - for _ in range(self.dec_depth) - ] - ) - self.block_norm = nn.LayerNorm(self.dec_emb_dim, self.layer_norm_eps) - self.final_mlp = Mlp( - self.num_mlp_layers, - self.dec_emb_dim, - self.out_dim, - layer_norm_eps=self.layer_norm_eps, - ) - - def forward_tensor(self, x, coords): - b, h, c = x.shape - - # process query coordinates - if self.embedding_type == "grid": - d2 = (coords - self.grid.unsqueeze(0)) ** 2 - w = paddle.exp(-1e5 * d2) / paddle.exp(-1e5 * d2).sum(axis=1, keepdim=True) - coords = paddle.einsum("ic,pi->pc", self.latents, w) - coords = self.fc(coords) - coords = self.norm(coords) - elif self.embedding_type == "mlp": - coords = self.mlp(coords) - coords = self.norm(coords) - - coords = einops.repeat(coords, "n d -> b n d", b=b) - - # process input function(encoder) - x = self.encoder(x) - x = self.enc_norm(x) - x = self.fc1(x) - - # decoder - for i, block in enumerate(self.cross_attn_blocks): - coords = block(coords, x) - - # mlp - x = self.block_norm(coords) - x = self.final_mlp(x) - - return x - - def forward(self, x_dict): - if self._input_transform is not None: - x = self._input_transform(x_dict) - - x, coords = x_dict[self.input_keys[0]], x_dict[self.input_keys[1]] - if coords.ndim >= 3: - coords = coords[0] # [b, n, c] -> [n, c] - - y = self.forward_tensor(x, coords) - - y_dict = {self.output_keys[0]: y} - if self._output_transform is not None: - y_dict = self._output_transform(x_dict, y_dict) - - return y_dict - - -class CVit(base.Arch): - """ - CVit architecture. - - [Bridging Operator Learning and Conditioned Neural Fields: A Unifying Perspective](https://arxiv.org/abs/2405.13998) - - Args: - input_keys (Sequence[str]): Input keys. - output_keys (Sequence[str]): Output keys. - in_dim (int): Dimensionality of the input data. - coords_dim (int): Dimensionality of the coordinates. - spatial_dims (Sequence[int]): Spatial dimensions. - patch_size (Sequence[int], optional): Size of the patches. Defaults to (1, 16, 16). - grid_size (Sequence[int], optional): Size of the grid. Defaults to (128, 128). - latent_dim (int, optional): Dimensionality of the latent space. Defaults to 256. - emb_dim (int, optional): Dimensionality of the embedding space. Defaults to 256. - depth (int, optional): Number of transformer encoder layers. Defaults to 3. - num_heads (int, optional): Number of attention heads. Defaults to 8. - dec_emb_dim (int, optional): Dimensionality of the decoder embedding space. Defaults to 256. - dec_num_heads (int, optional): Number of decoder attention heads. Defaults to 8. - dec_depth (int, optional): Number of decoder transformer layers. Defaults to 1. - num_mlp_layers (int, optional): Number of MLP layers. Defaults to 1. - mlp_ratio (int, optional): Ratio of hidden units. Defaults to 1. - out_dim (int, optional): Dimensionality of the output. Defaults to 1. - layer_norm_eps (float, optional): Epsilon value for layer normalization. Defaults to 1e-5. - embedding_type (str, optional): Type of embedding. Defaults to "grid". - - Examples: - >>> import ppsci - >>> b, t, h, w, c_in = 2, 4, 8, 8, 3 - >>> c_out = 3 - >>> h_query, w_query = 32, 32 - >>> model = ppsci.arch.CVit( - ... input_keys=["u", "y"], - ... output_keys=["s"], - ... in_dim=c_in, - ... coords_dim=2, - ... spatial_dims=[t, h, w], - ... patch_size=(1, 4, 4), - ... grid_size=(h, w), - ... latent_dim=32, - ... emb_dim=32, - ... depth=3, - ... num_heads=8, - ... dec_emb_dim=32, - ... dec_num_heads=8, - ... dec_depth=1, - ... num_mlp_layers=1, - ... mlp_ratio=1, - ... out_dim=c_out, - ... layer_norm_eps=1e-5, - ... embedding_type="grid", - ... ) - >>> x = paddle.randn([b, t, h, w, c_in]) - >>> coords = paddle.randn([h_query * w_query, 2]) - >>> out = model({"u": x, "y": coords})["s"] - >>> print(out.shape) # output shape should be [b, h_query * w_query, c_out] - [2, 1024, 3] - """ - - def __init__( - self, - input_keys: Sequence[str], - output_keys: Sequence[str], - in_dim: int, - coords_dim: int, - spatial_dims: Sequence[int], - patch_size: Sequence[int] = (1, 16, 16), - grid_size: Sequence[int] = (128, 128), - latent_dim: int = 256, - emb_dim: int = 256, - depth: int = 3, - num_heads: int = 8, - dec_emb_dim: int = 256, - dec_num_heads: int = 8, - dec_depth: int = 1, - num_mlp_layers: int = 1, - mlp_ratio: int = 1, - out_dim: int = 1, - layer_norm_eps: float = 1e-5, - embedding_type: str = "grid", - ): - super().__init__() - self.input_keys = input_keys - self.output_keys = output_keys - self.spatial_dims = spatial_dims - self.in_dim = in_dim - self.coords_dim = coords_dim - self.patch_size = patch_size - self.grid_size = grid_size - self.latent_dim = latent_dim - self.emb_dim = emb_dim - self.depth = depth - self.num_heads = num_heads - self.dec_emb_dim = dec_emb_dim - self.dec_num_heads = dec_num_heads - self.dec_depth = dec_depth - self.num_mlp_layers = num_mlp_layers - self.mlp_ratio = mlp_ratio - self.out_dim = out_dim - self.layer_norm_eps = layer_norm_eps - self.embedding_type = embedding_type - - if self.embedding_type == "grid": - # Create grid and latents - n_x, n_y = self.grid_size[0], self.grid_size[1] - - x = paddle.linspace(0, 1, n_x) - y = paddle.linspace(0, 1, n_y) - xx, yy = paddle.meshgrid(x, y, indexing="ij") - - self.grid = paddle.hstack([xx.flatten()[:, None], yy.flatten()[:, None]]) - self.latents = self.create_parameter( - [n_x * n_y, self.latent_dim], - default_initializer=nn.initializer.Normal(std=1e-2), - ) - self.fc = nn.Linear(self.latent_dim, self.dec_emb_dim) - self.norm = nn.LayerNorm(self.dec_emb_dim, self.layer_norm_eps) - elif self.embedding_type == "mlp": - self.mlp = MlpBlock(self.latent_dim, self.dec_emb_dim, self.dec_emb_dim) - self.norm = nn.LayerNorm(self.dec_emb_dim, self.layer_norm_eps) - - self.encoder = Encoder( - self.in_dim, - self.spatial_dims, - self.patch_size, - self.emb_dim, - self.depth, - self.num_heads, - self.mlp_ratio, - self.layer_norm_eps, - ) - self.enc_norm = nn.LayerNorm(self.emb_dim, self.layer_norm_eps) - self.fc1 = nn.Linear(self.emb_dim, self.dec_emb_dim) - self.cross_attn_blocks = nn.LayerList( - [ - CrossAttnBlock( - self.dec_num_heads, - self.dec_emb_dim, - self.mlp_ratio, - self.layer_norm_eps, - self.dec_emb_dim, - self.dec_emb_dim, - ) - for _ in range(self.dec_depth) - ] - ) - self.block_norm = nn.LayerNorm(self.dec_emb_dim, self.layer_norm_eps) - self.final_mlp = Mlp( - self.num_mlp_layers, - self.dec_emb_dim, - self.out_dim, - layer_norm_eps=self.layer_norm_eps, - ) - - def forward_tensor(self, x, coords): - b, t, h, w, c = x.shape - - # process query coordinates - if self.embedding_type == "grid": - d2 = ((coords.unsqueeze(1) - self.grid.unsqueeze(0)) ** 2).sum(axis=2) - w = paddle.exp(-1e5 * d2) / paddle.exp(-1e5 * d2).sum(axis=1, keepdim=True) - coords = paddle.einsum("ic,pi->pc", self.latents, w) - coords = self.fc(coords) - coords = self.norm(coords) - elif self.embedding_type == "mlp": - coords = self.mlp(coords) - coords = self.norm(coords) - - coords = einops.repeat(coords, "n d -> b n d", b=b) - - # process input function(encoder) - x = self.encoder(x) - x = self.enc_norm(x) - x = self.fc1(x) - - # decoder - for i, block in enumerate(self.cross_attn_blocks): - coords = block(coords, x) - - # mlp - x = self.block_norm(coords) - x = self.final_mlp(x) - - return x - - def forward(self, x_dict): - if self._input_transform is not None: - x = self._input_transform(x_dict) - - x, coords = x_dict[self.input_keys[0]], x_dict[self.input_keys[1]] - if coords.ndim >= 3: - coords = coords[0] # [b, n, c] -> [n, c] - - y = self.forward_tensor(x, coords) - - y_dict = {self.output_keys[0]: y} - if self._output_transform is not None: - y_dict = self._output_transform(x_dict, y_dict) - - return y_dict diff --git a/examples/smc_reac/ppsci/arch/deeponet.py b/examples/smc_reac/ppsci/arch/deeponet.py deleted file mode 100644 index 16a8807d81..0000000000 --- a/examples/smc_reac/ppsci/arch/deeponet.py +++ /dev/null @@ -1,154 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import Tuple -from typing import Union - -import paddle -import paddle.nn as nn - -from ppsci.arch import activation as act_mod -from ppsci.arch import base -from ppsci.arch import mlp - - -class DeepONet(base.Arch): - """Deep operator network. - - [Lu et al. Learning nonlinear operators via DeepONet based on the universal approximation theorem of operators. Nat Mach Intell, 2021.](https://doi.org/10.1038/s42256-021-00302-5) - - Args: - u_key (str): Name of function data for input function u(x). - y_key (str): Name of location data for input function G(u). - G_key (str): Output name of predicted G(u)(y). - num_loc (int): Number of sampled u(x), i.e. `m` in paper. - num_features (int): Number of features extracted from u(x), same for y. - branch_num_layers (int): Number of hidden layers of branch net. - trunk_num_layers (int): Number of hidden layers of trunk net. - branch_hidden_size (Union[int, Tuple[int, ...]]): Number of hidden size of branch net. - An integer for all layers, or list of integer specify each layer's size. - trunk_hidden_size (Union[int, Tuple[int, ...]]): Number of hidden size of trunk net. - An integer for all layers, or list of integer specify each layer's size. - branch_skip_connection (bool, optional): Whether to use skip connection for branch net. Defaults to False. - trunk_skip_connection (bool, optional): Whether to use skip connection for trunk net. Defaults to False. - branch_activation (str, optional): Name of activation function. Defaults to "tanh". - trunk_activation (str, optional): Name of activation function. Defaults to "tanh". - branch_weight_norm (bool, optional): Whether to apply weight norm on parameter(s) for branch net. Defaults to False. - trunk_weight_norm (bool, optional): Whether to apply weight norm on parameter(s) for trunk net. Defaults to False. - use_bias (bool, optional): Whether to add bias on predicted G(u)(y). Defaults to True. - - Examples: - >>> import paddle - >>> import ppsci - >>> model = ppsci.arch.DeepONet( - ... "u", "y", "G", - ... 100, 40, - ... 1, 1, - ... 40, 40, - ... branch_activation="relu", trunk_activation="relu", - ... use_bias=True, - ... ) - >>> input_dict = {"u": paddle.rand([200, 100]), - ... "y": paddle.rand([200, 1])} - >>> output_dict = model(input_dict) - >>> print(output_dict["G"].shape) - [200, 1] - """ - - def __init__( - self, - u_key: str, - y_key: str, - G_key: str, - num_loc: int, - num_features: int, - branch_num_layers: int, - trunk_num_layers: int, - branch_hidden_size: Union[int, Tuple[int, ...]], - trunk_hidden_size: Union[int, Tuple[int, ...]], - branch_skip_connection: bool = False, - trunk_skip_connection: bool = False, - branch_activation: str = "tanh", - trunk_activation: str = "tanh", - branch_weight_norm: bool = False, - trunk_weight_norm: bool = False, - use_bias: bool = True, - ): - super().__init__() - self.u_key = u_key - self.y_key = y_key - self.input_keys = (u_key, y_key) - self.output_keys = (G_key,) - - self.branch_net = mlp.MLP( - (self.u_key,), - ("b",), - branch_num_layers, - branch_hidden_size, - branch_activation, - branch_skip_connection, - branch_weight_norm, - input_dim=num_loc, - output_dim=num_features, - ) - - self.trunk_net = mlp.MLP( - (self.y_key,), - ("t",), - trunk_num_layers, - trunk_hidden_size, - trunk_activation, - trunk_skip_connection, - trunk_weight_norm, - input_dim=1, - output_dim=num_features, - ) - self.trunk_act = act_mod.get_activation(trunk_activation) - - self.use_bias = use_bias - if use_bias: - # register bias to parameter for updating in optimizer and storage - self.b = self.create_parameter( - shape=(1,), - attr=nn.initializer.Constant(0.0), - ) - - def forward(self, x): - if self._input_transform is not None: - x = self._input_transform(x) - - # Branch net to encode the input function - u_features = self.branch_net(x)[self.branch_net.output_keys[0]] - - # Trunk net to encode the domain of the output function - y_features = self.trunk_net(x) - y_features = self.trunk_act(y_features[self.trunk_net.output_keys[0]]) - - # Dot product - G_u = paddle.einsum("bi,bi->b", u_features, y_features) # [batch_size, ] - G_u = paddle.reshape(G_u, [-1, 1]) # reshape [batch_size, ] to [batch_size, 1] - - # Add bias - if self.use_bias: - G_u += self.b - - result_dict = { - self.output_keys[0]: G_u, - } - if self._output_transform is not None: - result_dict = self._output_transform(x, result_dict) - - return result_dict diff --git a/examples/smc_reac/ppsci/arch/dgmr.py b/examples/smc_reac/ppsci/arch/dgmr.py deleted file mode 100644 index dd189bb2de..0000000000 --- a/examples/smc_reac/ppsci/arch/dgmr.py +++ /dev/null @@ -1,1151 +0,0 @@ -from typing import List -from typing import Tuple - -import paddle -import paddle.nn as nn - -from ppsci.arch import base - -try: - import einops -except ModuleNotFoundError: - pass - - -class DGMR(base.Arch): - """Deep Generative Model of Radar. - Nowcasting GAN is an attempt to recreate DeepMind's Skillful Nowcasting GAN from https://arxiv.org/abs/2104.00954. - but slightly modified for multiple satellite channels - - Args: - input_keys (Tuple[str, ...]): Name of input keys, such as ("input",). - output_keys (Tuple[str, ...]): Name of output keys, such as ("output",). - forecast_steps (int, optional): Number of steps to predict in the future - input_channels (int, optional): Number of input channels per image - gen_lr (float, optional): Learning rate for the generator - disc_lr (float, optional): Learning rate for the discriminators, shared for both temporal and spatial discriminator - conv_type (str, optional): Type of 2d convolution to use, see satflow/models/utils.py for options - beta1 (float, optional): Beta1 for Adam optimizer - beta2 (float, optional): Beta2 for Adam optimizer - num_samples (int, optional): Number of samples of the latent space to sample for training/validation - grid_lambda (float, optional): Lambda for the grid regularization loss - output_shape (int, optional): Shape of the output predictions, generally should be same as the input shape - generation_steps (int, optional): Number of generation steps to use in forward pass, in paper is 6 and the best is chosen for the loss - this results in huge amounts of GPU memory though, so less might work better for training. - context_channels (int, optional): Number of output channels for the lowest block of conditioning stack - latent_channels (int, optional): Number of channels that the latent space should be reshaped to, - input dimension into ConvGRU, also affects the number of channels for other linked inputs/outputs - - Examples: - >>> import ppsci - >>> import paddle - >>> model = ppsci.arch.DGMR(("input", ), ("output", )) - >>> input_dict = {"input": paddle.randn((1, 4, 1, 256, 256))} - >>> output_dict = model(input_dict) # doctest: +SKIP - >>> print(output_dict["output"].shape) # doctest: +SKIP - [1, 18, 1, 256, 256] - """ - - def __init__( - self, - input_keys: Tuple[str, ...], - output_keys: Tuple[str, ...], - forecast_steps: int = 18, - input_channels: int = 1, - output_shape: int = 256, - gen_lr: float = 5e-05, - disc_lr: float = 0.0002, - conv_type: str = "standard", - num_samples: int = 6, - grid_lambda: float = 20.0, - beta1: float = 0.0, - beta2: float = 0.999, - latent_channels: int = 768, - context_channels: int = 384, - generation_steps: int = 6, - ): - super().__init__() - self.input_keys = input_keys - self.output_keys = output_keys - self.gen_lr = gen_lr - self.disc_lr = disc_lr - self.beta1 = beta1 - self.beta2 = beta2 - self.grid_lambda = grid_lambda - self.num_samples = num_samples - self.latent_channels = latent_channels - self.context_channels = context_channels - self.input_channels = input_channels - self.generation_steps = generation_steps - self.conditioning_stack = ContextConditioningStack( - input_channels=input_channels, - conv_type=conv_type, - output_channels=self.context_channels, - ) - self.latent_stack = LatentConditioningStack( - shape=(8 * self.input_channels, output_shape // 32, output_shape // 32), - output_channels=self.latent_channels, - ) - self.sampler = Sampler( - forecast_steps=forecast_steps, - latent_channels=self.latent_channels, - context_channels=self.context_channels, - ) - self.generator = Generator( - self.conditioning_stack, self.latent_stack, self.sampler - ) - self.discriminator = Discriminator(input_channels) - self.global_iteration = 0 - self.automatic_optimization = False - - def split_to_dict( - self, data_tensors: Tuple[paddle.Tensor, ...], keys: Tuple[str, ...] - ): - return {key: data_tensors[i] for i, key in enumerate(keys)} - - def forward(self, x): - if self._input_transform is not None: - x = self._input_transform(x) - x_tensor = self.concat_to_tensor(x, self.input_keys) - y = [self.generator(x_tensor)] - y = self.split_to_dict(y, self.output_keys) - - if self._output_transform is not None: - y = self._output_transform(x, y) - return y - - -class Sampler(nn.Layer): - """ - Sampler from the Skillful Nowcasting, see https://arxiv.org/pdf/2104.00954.pdf - - The sampler takes the output from the Latent and Context conditioning stacks and - creates one stack of ConvGRU layers per future timestep. - - Args: - forecast_steps: Number of forecast steps - latent_channels: Number of input channels to the lowest ConvGRU layer - """ - - def __init__( - self, - forecast_steps: int = 18, - latent_channels: int = 768, - context_channels: int = 384, - output_channels: int = 1, - ): - super().__init__() - self.forecast_steps = forecast_steps - self.convGRU1 = ConvGRU( - input_channels=latent_channels + context_channels, - output_channels=context_channels, - kernel_size=3, - ) - self.gru_conv_1x1 = nn.utils.spectral_norm( - layer=nn.Conv2D( - in_channels=context_channels, - out_channels=latent_channels, - kernel_size=(1, 1), - ) - ) - self.g1 = GBlock( - input_channels=latent_channels, output_channels=latent_channels - ) - self.up_g1 = UpsampleGBlock( - input_channels=latent_channels, output_channels=latent_channels // 2 - ) - self.convGRU2 = ConvGRU( - input_channels=latent_channels // 2 + context_channels // 2, - output_channels=context_channels // 2, - kernel_size=3, - ) - self.gru_conv_1x1_2 = nn.utils.spectral_norm( - layer=nn.Conv2D( - in_channels=context_channels // 2, - out_channels=latent_channels // 2, - kernel_size=(1, 1), - ) - ) - self.g2 = GBlock( - input_channels=latent_channels // 2, output_channels=latent_channels // 2 - ) - self.up_g2 = UpsampleGBlock( - input_channels=latent_channels // 2, output_channels=latent_channels // 4 - ) - self.convGRU3 = ConvGRU( - input_channels=latent_channels // 4 + context_channels // 4, - output_channels=context_channels // 4, - kernel_size=3, - ) - self.gru_conv_1x1_3 = nn.utils.spectral_norm( - layer=nn.Conv2D( - in_channels=context_channels // 4, - out_channels=latent_channels // 4, - kernel_size=(1, 1), - ) - ) - self.g3 = GBlock( - input_channels=latent_channels // 4, output_channels=latent_channels // 4 - ) - self.up_g3 = UpsampleGBlock( - input_channels=latent_channels // 4, output_channels=latent_channels // 8 - ) - self.convGRU4 = ConvGRU( - input_channels=latent_channels // 8 + context_channels // 8, - output_channels=context_channels // 8, - kernel_size=3, - ) - self.gru_conv_1x1_4 = nn.utils.spectral_norm( - layer=nn.Conv2D( - in_channels=context_channels // 8, - out_channels=latent_channels // 8, - kernel_size=(1, 1), - ) - ) - self.g4 = GBlock( - input_channels=latent_channels // 8, output_channels=latent_channels // 8 - ) - self.up_g4 = UpsampleGBlock( - input_channels=latent_channels // 8, output_channels=latent_channels // 16 - ) - self.bn = nn.BatchNorm2D(num_features=latent_channels // 16) - self.relu = nn.ReLU() - self.conv_1x1 = nn.utils.spectral_norm( - layer=nn.Conv2D( - in_channels=latent_channels // 16, - out_channels=4 * output_channels, - kernel_size=(1, 1), - ) - ) - self.depth2space = nn.PixelShuffle(upscale_factor=2) - - def forward( - self, conditioning_states: List[paddle.Tensor], latent_dim: paddle.Tensor - ) -> paddle.Tensor: - """ - Perform the sampling from Skillful Nowcasting with GANs - - Args: - conditioning_states: Outputs from the `ContextConditioningStack` with the 4 input states, ordered from largest to smallest spatially - latent_dim: Output from `LatentConditioningStack` for input into the ConvGRUs - Returns: - forecast_steps-length output of images for future timesteps - - """ - init_states = conditioning_states - latent_dim = einops.repeat( - latent_dim, "b c h w -> (repeat b) c h w", repeat=init_states[0].shape[0] - ) - hidden_states = [latent_dim] * self.forecast_steps - - hidden_states = self.convGRU1(hidden_states, init_states[3]) - hidden_states = [self.gru_conv_1x1(h) for h in hidden_states] - hidden_states = [self.g1(h) for h in hidden_states] - hidden_states = [self.up_g1(h) for h in hidden_states] - hidden_states = self.convGRU2(hidden_states, init_states[2]) - hidden_states = [self.gru_conv_1x1_2(h) for h in hidden_states] - hidden_states = [self.g2(h) for h in hidden_states] - hidden_states = [self.up_g2(h) for h in hidden_states] - hidden_states = self.convGRU3(hidden_states, init_states[1]) - hidden_states = [self.gru_conv_1x1_3(h) for h in hidden_states] - hidden_states = [self.g3(h) for h in hidden_states] - hidden_states = [self.up_g3(h) for h in hidden_states] - hidden_states = self.convGRU4(hidden_states, init_states[0]) - hidden_states = [self.gru_conv_1x1_4(h) for h in hidden_states] - hidden_states = [self.g4(h) for h in hidden_states] - hidden_states = [self.up_g4(h) for h in hidden_states] - hidden_states = [nn.functional.relu(x=self.bn(h)) for h in hidden_states] - hidden_states = [self.conv_1x1(h) for h in hidden_states] - hidden_states = [self.depth2space(h) for h in hidden_states] - forecasts = paddle.stack(x=hidden_states, axis=1) - return forecasts - - -class Generator(nn.Layer): - """ - Wraps the three parts of the generator for simpler calling - - Args: - conditioning_stack: A layer representing the conditioning stack. - latent_stack: A layer representing the latent stack. - sampler: A layer representing the sampler. - """ - - def __init__( - self, - conditioning_stack: nn.Layer, - latent_stack: nn.Layer, - sampler: nn.Layer, - ): - super().__init__() - self.conditioning_stack = conditioning_stack - self.latent_stack = latent_stack - self.sampler = sampler - - def forward(self, x): - conditioning_states = self.conditioning_stack(x) - latent_dim = self.latent_stack(x) - x = self.sampler(conditioning_states, latent_dim) - return x - - -class Discriminator(nn.Layer): - def __init__( - self, - input_channels: int = 12, - num_spatial_frames: int = 8, - conv_type: str = "standard", - ): - super().__init__() - self.spatial_discriminator = SpatialDiscriminator( - input_channels=input_channels, - num_timesteps=num_spatial_frames, - conv_type=conv_type, - ) - self.temporal_discriminator = TemporalDiscriminator( - input_channels=input_channels, conv_type=conv_type - ) - - def forward(self, x: paddle.Tensor) -> paddle.Tensor: - spatial_loss = self.spatial_discriminator(x) - temporal_loss = self.temporal_discriminator(x) - return paddle.concat(x=[spatial_loss, temporal_loss], axis=1) - - -class TemporalDiscriminator(nn.Layer): - """ - Temporal Discriminator from the Skillful Nowcasting, see https://arxiv.org/pdf/2104.00954.pdf - - Args: - input_channels: Number of channels per timestep - crop_size: Size of the crop, in the paper half the width of the input images - num_layers: Number of intermediate DBlock layers to use - conv_type: Type of 2d convolutions to use, see satflow/models/utils.py for options - """ - - def __init__( - self, - input_channels: int = 12, - num_layers: int = 3, - conv_type: str = "standard", - ): - super().__init__() - self.downsample = nn.AvgPool3D( - kernel_size=(1, 2, 2), stride=(1, 2, 2), exclusive=False - ) - self.space2depth = nn.PixelUnshuffle(downscale_factor=2) - internal_chn = 48 - self.d1 = DBlock( - input_channels=4 * input_channels, - output_channels=internal_chn * input_channels, - conv_type="3d", - first_relu=False, - ) - self.d2 = DBlock( - input_channels=internal_chn * input_channels, - output_channels=2 * internal_chn * input_channels, - conv_type="3d", - ) - self.intermediate_dblocks = nn.LayerList() - for _ in range(num_layers): - internal_chn *= 2 - self.intermediate_dblocks.append( - DBlock( - input_channels=internal_chn * input_channels, - output_channels=2 * internal_chn * input_channels, - conv_type=conv_type, - ) - ) - self.d_last = DBlock( - input_channels=2 * internal_chn * input_channels, - output_channels=2 * internal_chn * input_channels, - keep_same_output=True, - conv_type=conv_type, - ) - self.fc = nn.utils.spectral_norm( - layer=nn.Linear( - in_features=2 * internal_chn * input_channels, out_features=1 - ) - ) - self.relu = nn.ReLU() - self.bn = nn.BatchNorm1D(num_features=2 * internal_chn * input_channels) - - def forward(self, x: paddle.Tensor) -> paddle.Tensor: - x = self.downsample(x) - if len(x.shape) == 4: - x = self.space2depth(x) - elif len(x.shape) == 5: - B, T = x.shape[0], x.shape[1] - x_reshaped = paddle.reshape(x, [-1] + list(x.shape[2:])) - x = self.space2depth(x_reshaped) - x = paddle.reshape(x, [B, T] + list(x.shape[1:])) - x = paddle.transpose(x=x, perm=(0, 2, 1, 3, 4)) - x = self.d1(x) - x = self.d2(x) - x = paddle.transpose(x=x, perm=(0, 2, 1, 3, 4)) - representations = [] - for idx in range(x.shape[1]): - rep = x[:, idx, :, :, :] - for d in self.intermediate_dblocks: - rep = d(rep) - rep = self.d_last(rep) - rep = paddle.sum(x=nn.functional.relu(x=rep), axis=[2, 3]) - rep = self.bn(rep) - rep = self.fc(rep) - representations.append(rep) - x = paddle.stack(x=representations, axis=1) - x = paddle.sum(x=x, keepdim=True, axis=1) - return x - - -class SpatialDiscriminator(nn.Layer): - """ - Spatial discriminator from Skillful Nowcasting, see https://arxiv.org/pdf/2104.00954.pdf - - Args: - input_channels: Number of input channels per timestep - num_timesteps: Number of timesteps to use, in the paper 8/18 timesteps were chosen - num_layers: Number of intermediate DBlock layers to use - conv_type: Type of 2d convolutions to use, see satflow/models/utils.py for options - """ - - def __init__( - self, - input_channels: int = 12, - num_timesteps: int = 8, - num_layers: int = 4, - conv_type: str = "standard", - ): - super().__init__() - self.num_timesteps = num_timesteps - self.mean_pool = nn.AvgPool2D(kernel_size=2, exclusive=False) - self.space2depth = nn.PixelUnshuffle(downscale_factor=2) - internal_chn = 24 - self.d1 = DBlock( - input_channels=4 * input_channels, - output_channels=2 * internal_chn * input_channels, - first_relu=False, - conv_type=conv_type, - ) - self.intermediate_dblocks = nn.LayerList() - for _ in range(num_layers): - internal_chn *= 2 - self.intermediate_dblocks.append( - DBlock( - input_channels=internal_chn * input_channels, - output_channels=2 * internal_chn * input_channels, - conv_type=conv_type, - ) - ) - self.d6 = DBlock( - input_channels=2 * internal_chn * input_channels, - output_channels=2 * internal_chn * input_channels, - keep_same_output=True, - conv_type=conv_type, - ) - self.fc = nn.utils.spectral_norm( - layer=nn.Linear( - in_features=2 * internal_chn * input_channels, out_features=1 - ) - ) - self.relu = nn.ReLU() - self.bn = nn.BatchNorm1D(num_features=2 * internal_chn * input_channels) - - def forward(self, x: paddle.Tensor) -> paddle.Tensor: - idxs = paddle.randint(low=0, high=x.shape[1], shape=(self.num_timesteps,)) - representations = [] - for idx in idxs: - rep = self.mean_pool(x[:, idx, :, :, :]) - if len(rep.shape) == 4: - rep = self.space2depth(rep) - elif len(rep.shape) == 5: - B, T = rep.shape[0], rep.shape[1] - rep_reshaped = paddle.reshape(rep, [-1] + list(rep.shape[2:])) - rep = self.space2depth(rep_reshaped) - rep = paddle.reshape(rep, [B, T] + list(rep.shape[1:])) - rep = self.d1(rep) - for d in self.intermediate_dblocks: - rep = d(rep) - rep = self.d6(rep) - rep = paddle.sum(x=nn.functional.relu(x=rep), axis=[2, 3]) - rep = self.bn(rep) - rep = self.fc(rep) - """ - Pseudocode from DeepMind - # Sum-pool the representations and feed to spectrally normalized lin. layer. - y = tf.reduce_sum(tf.nn.relu(y), axis=[1, 2]) - y = layers.BatchNorm(calc_sigma=False)(y) - output_layer = layers.Linear(output_size=1) - output = output_layer(y) - - # Take the sum across the t samples. Note: we apply the ReLU to - # (1 - score_real) and (1 + score_generated) in the loss. - output = tf.reshape(output, [b, n, 1]) - output = tf.reduce_sum(output, keepdims=True, axis=1) - return output - """ - representations.append(rep) - x = paddle.stack(x=representations, axis=1) - x = paddle.sum(x=x, keepdim=True, axis=1) - return x - - -class GBlock(nn.Layer): - """Residual generator block without upsampling. G Block from Skillful Nowcasting, see https://arxiv.org/pdf/2104.00954.pdf - - Args: - input_channels: Number of input channels - output_channels: Number of output channels - conv_type: Type of convolution desired, see satflow/models/utils.py for options - """ - - def __init__( - self, - input_channels: int = 12, - output_channels: int = 12, - conv_type: str = "standard", - spectral_normalized_eps=0.0001, - ): - super().__init__() - self.output_channels = output_channels - self.bn1 = nn.BatchNorm2D(num_features=input_channels) - self.bn2 = nn.BatchNorm2D(num_features=input_channels) - self.relu = nn.ReLU() - conv2d = get_conv_layer(conv_type) - self.conv_1x1 = nn.utils.spectral_norm( - layer=conv2d( - in_channels=input_channels, out_channels=output_channels, kernel_size=1 - ), - eps=spectral_normalized_eps, - ) - self.first_conv_3x3 = nn.utils.spectral_norm( - layer=conv2d( - in_channels=input_channels, - out_channels=input_channels, - kernel_size=3, - padding=1, - ), - eps=spectral_normalized_eps, - ) - self.last_conv_3x3 = nn.utils.spectral_norm( - layer=conv2d( - in_channels=input_channels, - out_channels=output_channels, - kernel_size=3, - padding=1, - ), - eps=spectral_normalized_eps, - ) - - def forward(self, x: paddle.Tensor) -> paddle.Tensor: - if x.shape[1] != self.output_channels: - sc = self.conv_1x1(x) - else: - sc = x - x2 = self.bn1(x) - x2 = self.relu(x2) - x2 = self.first_conv_3x3(x2) - x2 = self.bn2(x2) - x2 = self.relu(x2) - x2 = self.last_conv_3x3(x2) - x = x2 + sc - return x - - -class UpsampleGBlock(nn.Layer): - """Residual generator block with upsampling - G Block from Skillful Nowcasting, see https://arxiv.org/pdf/2104.00954.pdf - - Args: - input_channels: Number of input channels - output_channels: Number of output channels - conv_type: Type of convolution desired, see satflow/models/utils.py for options - """ - - def __init__( - self, - input_channels: int = 12, - output_channels: int = 12, - conv_type: str = "standard", - spectral_normalized_eps=0.0001, - ): - super().__init__() - self.output_channels = output_channels - self.bn1 = nn.BatchNorm2D(num_features=input_channels) - self.bn2 = nn.BatchNorm2D(num_features=input_channels) - self.relu = nn.ReLU() - conv2d = get_conv_layer(conv_type) - self.conv_1x1 = nn.utils.spectral_norm( - layer=conv2d( - in_channels=input_channels, out_channels=output_channels, kernel_size=1 - ), - eps=spectral_normalized_eps, - ) - self.upsample = nn.Upsample(scale_factor=2, mode="nearest") - self.first_conv_3x3 = nn.utils.spectral_norm( - layer=conv2d( - in_channels=input_channels, - out_channels=input_channels, - kernel_size=3, - padding=1, - ), - eps=spectral_normalized_eps, - ) - self.last_conv_3x3 = nn.utils.spectral_norm( - layer=conv2d( - in_channels=input_channels, - out_channels=output_channels, - kernel_size=3, - padding=1, - ), - eps=spectral_normalized_eps, - ) - - def forward(self, x: paddle.Tensor) -> paddle.Tensor: - sc = self.upsample(x) - sc = self.conv_1x1(sc) - x2 = self.bn1(x) - x2 = self.relu(x2) - x2 = self.upsample(x2) - x2 = self.first_conv_3x3(x2) - x2 = self.bn2(x2) - x2 = self.relu(x2) - x2 = self.last_conv_3x3(x2) - x = x2 + sc - return x - - -class DBlock(nn.Layer): - """ - D and 3D Block from Skillful Nowcasting, see https://arxiv.org/pdf/2104.00954.pdf - - Args: - input_channels: Number of input channels - output_channels: Number of output channels - conv_type: Convolution type, see satflow/models/utils.py for options - first_relu: Whether to have an ReLU before the first 3x3 convolution - keep_same_output: Whether the output should have the same spatial dimensions as input, if False, downscales by 2 - """ - - def __init__( - self, - input_channels: int = 12, - output_channels: int = 12, - conv_type: str = "standard", - first_relu: bool = True, - keep_same_output: bool = False, - ): - super().__init__() - self.input_channels = input_channels - self.output_channels = output_channels - self.first_relu = first_relu - self.keep_same_output = keep_same_output - self.conv_type = conv_type - conv2d = get_conv_layer(conv_type) - if conv_type == "3d": - self.pooling = nn.AvgPool3D(kernel_size=2, stride=2, exclusive=False) - else: - self.pooling = nn.AvgPool2D(kernel_size=2, stride=2, exclusive=False) - self.conv_1x1 = nn.utils.spectral_norm( - layer=conv2d( - in_channels=input_channels, out_channels=output_channels, kernel_size=1 - ) - ) - self.first_conv_3x3 = nn.utils.spectral_norm( - layer=conv2d( - in_channels=input_channels, - out_channels=output_channels, - kernel_size=3, - padding=1, - ) - ) - self.last_conv_3x3 = nn.utils.spectral_norm( - layer=conv2d( - in_channels=output_channels, - out_channels=output_channels, - kernel_size=3, - padding=1, - stride=1, - ) - ) - self.relu = nn.ReLU() - - def forward(self, x: paddle.Tensor) -> paddle.Tensor: - if self.input_channels != self.output_channels: - x1 = self.conv_1x1(x) - if not self.keep_same_output: - x1 = self.pooling(x1) - else: - x1 = x - if self.first_relu: - x = self.relu(x) - x = self.first_conv_3x3(x) - x = self.relu(x) - x = self.last_conv_3x3(x) - if not self.keep_same_output: - x = self.pooling(x) - x = x1 + x - return x - - -class LBlock(nn.Layer): - """Residual block for the Latent Stack. - L-Block for increasing the number of channels in the input - from Skillful Nowcasting, see https://arxiv.org/pdf/2104.00954.pdf - - Args: - input_channels: Number of input channels - output_channels: Number of output channels - conv_type: Which type of convolution desired, see satflow/models/utils.py for options - """ - - def __init__( - self, - input_channels: int = 12, - output_channels: int = 12, - kernel_size: int = 3, - conv_type: str = "standard", - ): - super().__init__() - self.input_channels = input_channels - self.output_channels = output_channels - conv2d = get_conv_layer(conv_type) - self.conv_1x1 = conv2d( - in_channels=input_channels, - out_channels=output_channels - input_channels, - kernel_size=1, - ) - self.first_conv_3x3 = conv2d( - input_channels, - out_channels=output_channels, - kernel_size=kernel_size, - padding=1, - stride=1, - ) - self.relu = nn.ReLU() - self.last_conv_3x3 = conv2d( - in_channels=output_channels, - out_channels=output_channels, - kernel_size=kernel_size, - padding=1, - stride=1, - ) - - def forward(self, x) -> paddle.Tensor: - if self.input_channels < self.output_channels: - sc = self.conv_1x1(x) - sc = paddle.concat(x=[x, sc], axis=1) - else: - sc = x - x2 = self.relu(x) - x2 = self.first_conv_3x3(x2) - x2 = self.relu(x2) - x2 = self.last_conv_3x3(x2) - return x2 + sc - - -class ContextConditioningStack(nn.Layer): - """ - Conditioning Stack using the context images from Skillful Nowcasting, , see https://arxiv.org/pdf/2104.00954.pdf - - Args: - input_channels: Number of input channels per timestep - output_channels: Number of output channels for the lowest block - conv_type: Type of 2D convolution to use, see satflow/models/utils.py for options - """ - - def __init__( - self, - input_channels: int = 1, - output_channels: int = 768, - num_context_steps: int = 4, - conv_type: str = "standard", - ): - super().__init__() - conv2d = get_conv_layer(conv_type) - self.space2depth = nn.PixelUnshuffle(downscale_factor=2) - self.d1 = DBlock( - input_channels=4 * input_channels, - output_channels=output_channels // 4 * input_channels // num_context_steps, - conv_type=conv_type, - ) - self.d2 = DBlock( - input_channels=output_channels // 4 * input_channels // num_context_steps, - output_channels=output_channels // 2 * input_channels // num_context_steps, - conv_type=conv_type, - ) - self.d3 = DBlock( - input_channels=output_channels // 2 * input_channels // num_context_steps, - output_channels=output_channels * input_channels // num_context_steps, - conv_type=conv_type, - ) - self.d4 = DBlock( - input_channels=output_channels * input_channels // num_context_steps, - output_channels=output_channels * 2 * input_channels // num_context_steps, - conv_type=conv_type, - ) - self.conv1 = nn.utils.spectral_norm( - layer=conv2d( - in_channels=output_channels // 4 * input_channels, - out_channels=output_channels // 8 * input_channels, - kernel_size=3, - padding=1, - ) - ) - self.conv2 = nn.utils.spectral_norm( - layer=conv2d( - in_channels=output_channels // 2 * input_channels, - out_channels=output_channels // 4 * input_channels, - kernel_size=3, - padding=1, - ) - ) - self.conv3 = nn.utils.spectral_norm( - layer=conv2d( - in_channels=output_channels * input_channels, - out_channels=output_channels // 2 * input_channels, - kernel_size=3, - padding=1, - ) - ) - self.conv4 = nn.utils.spectral_norm( - layer=conv2d( - in_channels=output_channels * 2 * input_channels, - out_channels=output_channels * input_channels, - kernel_size=3, - padding=1, - ) - ) - self.relu = nn.ReLU() - - def forward( - self, x: paddle.Tensor - ) -> Tuple[paddle.Tensor, paddle.Tensor, paddle.Tensor, paddle.Tensor]: - if len(x.shape) == 4: - x = self.space2depth(x) - elif len(x.shape) == 5: - B, T = x.shape[0], x.shape[1] - x_reshaped = paddle.reshape(x, [-1] + list(x.shape[2:])) - x = self.space2depth(x_reshaped) - x = paddle.reshape(x, [B, T] + list(x.shape[1:])) - steps = x.shape[1] - scale_1 = [] - scale_2 = [] - scale_3 = [] - scale_4 = [] - for i in range(steps): - s1 = self.d1(x[:, i, :, :, :]) - s2 = self.d2(s1) - s3 = self.d3(s2) - s4 = self.d4(s3) - scale_1.append(s1) - scale_2.append(s2) - scale_3.append(s3) - scale_4.append(s4) - scale_1 = paddle.stack(x=scale_1, axis=1) - scale_2 = paddle.stack(x=scale_2, axis=1) - scale_3 = paddle.stack(x=scale_3, axis=1) - scale_4 = paddle.stack(x=scale_4, axis=1) - scale_1 = self._mixing_layer(scale_1, self.conv1) - scale_2 = self._mixing_layer(scale_2, self.conv2) - scale_3 = self._mixing_layer(scale_3, self.conv3) - scale_4 = self._mixing_layer(scale_4, self.conv4) - return scale_1, scale_2, scale_3, scale_4 - - def _mixing_layer(self, inputs, conv_block): - stacked_inputs = einops.rearrange(inputs, "b t c h w -> b (c t) h w") - return nn.functional.relu(x=conv_block(stacked_inputs)) - - -class LatentConditioningStack(nn.Layer): - """ - Latent conditioning stack from Skillful Nowcasting, see https://arxiv.org/pdf/2104.00954.pdf - - Args: - shape: Shape of the latent space, Should be (H/32,W/32,x) of the final image shape - output_channels: Number of output channels for the conditioning stack - use_attention: Whether to have a self-attention block or not - """ - - def __init__( - self, - shape: (int, int, int) = (8, 8, 8), - output_channels: int = 768, - use_attention: bool = True, - ): - super().__init__() - self.shape = shape - self.use_attention = use_attention - self.distribution = paddle.distribution.Normal( - loc=paddle.to_tensor(data=[0.0], dtype="float32"), - scale=paddle.to_tensor(data=[2.0], dtype="float32"), - ) - self.conv_3x3 = nn.utils.spectral_norm( - layer=nn.Conv2D( - in_channels=shape[0], - out_channels=shape[0], - kernel_size=(3, 3), - padding=1, - ) - ) - self.l_block1 = LBlock( - input_channels=shape[0], output_channels=output_channels // 32 - ) - self.l_block2 = LBlock( - input_channels=output_channels // 32, output_channels=output_channels // 16 - ) - self.l_block3 = LBlock( - input_channels=output_channels // 16, output_channels=output_channels // 4 - ) - if self.use_attention: - self.att_block = AttentionLayer( - input_channels=output_channels // 4, - output_channels=output_channels // 4, - ) - self.l_block4 = LBlock( - input_channels=output_channels // 4, output_channels=output_channels - ) - - def forward(self, x: paddle.Tensor) -> paddle.Tensor: - """ - Args: - x: tensor on the correct device, to move over the latent distribution - Returns: z - """ - z = self.distribution.sample(self.shape) - z = paddle.transpose(x=z, perm=(3, 0, 1, 2)).astype(dtype=x.dtype) - z = self.conv_3x3(z) - z = self.l_block1(z) - z = self.l_block2(z) - z = self.l_block3(z) - z = self.att_block(z) - z = self.l_block4(z) - return z - - -def attention_einsum(q, k, v): - """Apply the attention operator to tensors of shape [h, w, c].""" - k = einops.rearrange(k, "h w c -> (h w) c") - v = einops.rearrange(v, "h w c -> (h w) c") - beta = nn.functional.softmax(x=paddle.einsum("hwc, Lc->hwL", q, k), axis=-1) - out = paddle.einsum("hwL, Lc->hwc", beta, v) - return out - - -class AttentionLayer(nn.Layer): - """Attention Module""" - - def __init__( - self, input_channels: int, output_channels: int, ratio_kq=8, ratio_v=8 - ): - super().__init__() - self.ratio_kq = ratio_kq - self.ratio_v = ratio_v - self.output_channels = output_channels - self.input_channels = input_channels - self.query = nn.Conv2D( - in_channels=input_channels, - out_channels=self.output_channels // self.ratio_kq, - kernel_size=(1, 1), - padding="valid", - bias_attr=False, - ) - self.key = nn.Conv2D( - in_channels=input_channels, - out_channels=self.output_channels // self.ratio_kq, - kernel_size=(1, 1), - padding="valid", - bias_attr=False, - ) - self.value = nn.Conv2D( - in_channels=input_channels, - out_channels=self.output_channels // self.ratio_v, - kernel_size=(1, 1), - padding="valid", - bias_attr=False, - ) - self.last_conv = nn.Conv2D( - in_channels=self.output_channels // 8, - out_channels=self.output_channels, - kernel_size=(1, 1), - padding="valid", - bias_attr=False, - ) - gamma = paddle.create_parameter( - shape=paddle.zeros(shape=[1]).shape, - dtype=paddle.zeros(shape=[1]).numpy().dtype, - default_initializer=nn.initializer.Assign(paddle.zeros(shape=[1])), - ) - gamma.stop_gradient = not True - self.gamma = gamma - - def forward(self, x: paddle.Tensor) -> paddle.Tensor: - query = self.query(x) - key = self.key(x) - value = self.value(x) - out = [] - for b in range(x.shape[0]): - out.append(attention_einsum(query[b], key[b], value[b])) - out = paddle.stack(x=out, axis=0) - out = self.gamma * self.last_conv(out) - return out + x - - -class AddCoords(nn.Layer): - def __init__(self, with_r=False): - super().__init__() - self.with_r = with_r - - def forward(self, input_tensor): - """ - Args: - input_tensor: shape(batch, channel, x_dim, y_dim) - """ - batch_size, _, x_dim, y_dim = input_tensor.shape - xx_channel = paddle.arange(end=x_dim).tile([1, y_dim, 1]) - x = paddle.arange(end=y_dim).tile([1, x_dim, 1]) - perm_0 = list(range(x.ndim)) - perm_0[1] = 2 - perm_0[2] = 1 - yy_channel = x.transpose(perm=perm_0) - xx_channel = xx_channel.astype(dtype="float32") / (x_dim - 1) - yy_channel = yy_channel.astype(dtype="float32") / (y_dim - 1) - xx_channel = xx_channel * 2 - 1 - yy_channel = yy_channel * 2 - 1 - x = xx_channel.tile([batch_size, 1, 1, 1]) - perm_1 = list(range(x.ndim)) - perm_1[2] = 3 - perm_1[3] = 2 - xx_channel = x.transpose(perm=perm_1) - x = yy_channel.tile([batch_size, 1, 1, 1]) - perm_2 = list(range(x.ndim)) - perm_2[2] = 3 - perm_2[3] = 2 - yy_channel = x.transpose(perm=perm_2) - ret = paddle.concat( - x=[ - input_tensor, - xx_channel.astype(dtype=input_tensor.dtype), - yy_channel.astype(dtype=input_tensor.dtype), - ], - axis=1, - ) - if self.with_r: - rr = paddle.sqrt( - x=paddle.pow(x=xx_channel.astype(dtype=input_tensor.dtype) - 0.5, y=2) - + paddle.pow(x=yy_channel.astype(dtype=input_tensor.dtype) - 0.5, y=2) - ) - ret = paddle.concat(x=[ret, rr], axis=1) - return ret - - -class CoordConv(nn.Layer): - def __init__(self, in_channels, out_channels, with_r=False): - super().__init__() - self.addcoords = AddCoords(with_r=with_r) - in_size = in_channels + 2 - if with_r: - in_size += 1 - self.conv = nn.Conv2D(in_size, out_channels) - - def forward(self, x): - ret = self.addcoords(x) - ret = self.conv(ret) - return ret - - -class ConvGRUCell(nn.Layer): - """A ConvGRU implementation. - - Args: - kernel_size: kernel size of the convolutions. Default: 3. - sn_eps: constant for spectral normalization. Default: 1e-4. - """ - - def __init__( - self, input_channels: int, output_channels: int, kernel_size=3, sn_eps=0.0001 - ): - super().__init__() - self._kernel_size = kernel_size - self._sn_eps = sn_eps - self.read_gate_conv = nn.utils.spectral_norm( - layer=nn.Conv2D( - in_channels=input_channels, - out_channels=output_channels, - kernel_size=(kernel_size, kernel_size), - padding=1, - ), - eps=sn_eps, - ) - self.update_gate_conv = nn.utils.spectral_norm( - layer=nn.Conv2D( - in_channels=input_channels, - out_channels=output_channels, - kernel_size=(kernel_size, kernel_size), - padding=1, - ), - eps=sn_eps, - ) - self.output_conv = nn.utils.spectral_norm( - layer=nn.Conv2D( - in_channels=input_channels, - out_channels=output_channels, - kernel_size=(kernel_size, kernel_size), - padding=1, - ), - eps=sn_eps, - ) - - def forward(self, x, prev_state): - """ - ConvGRU forward, returning the current+new state - - Args: - x: Input tensor - prev_state: Previous state - - Returns: - New tensor plus the new state - """ - xh = paddle.concat(x=[x, prev_state], axis=1) - read_gate = nn.functional.sigmoid(x=self.read_gate_conv(xh)) - update_gate = nn.functional.sigmoid(x=self.update_gate_conv(xh)) - gated_input = paddle.concat(x=[x, read_gate * prev_state], axis=1) - c = nn.functional.relu(x=self.output_conv(gated_input)) - out = update_gate * prev_state + (1.0 - update_gate) * c - new_state = out - return out, new_state - - -class ConvGRU(nn.Layer): - """ConvGRU Cell wrapper to replace tf.static_rnn in TF implementation""" - - def __init__( - self, - input_channels: int, - output_channels: int, - kernel_size: int = 3, - sn_eps=0.0001, - ): - super().__init__() - self.cell = ConvGRUCell(input_channels, output_channels, kernel_size, sn_eps) - - def forward(self, x: paddle.Tensor, hidden_state=None) -> paddle.Tensor: - outputs = [] - for step in range(len(x)): - output, hidden_state = self.cell(x[step], hidden_state) - outputs.append(output) - outputs = paddle.stack(x=outputs, axis=0) - return outputs - - -def get_conv_layer(conv_type: str = "standard") -> nn.Layer: - if conv_type == "standard": - conv_layer = nn.Conv2D - elif conv_type == "coord": - conv_layer = CoordConv - elif conv_type == "3d": - conv_layer = nn.Conv3D - else: - raise ValueError(f"{conv_type} is not a recognized Conv method") - return conv_layer diff --git a/examples/smc_reac/ppsci/arch/embedding_koopman.py b/examples/smc_reac/ppsci/arch/embedding_koopman.py deleted file mode 100644 index 367b5cd3ca..0000000000 --- a/examples/smc_reac/ppsci/arch/embedding_koopman.py +++ /dev/null @@ -1,544 +0,0 @@ -# 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. - -""" -Code below is heavily based on [transformer-physx](https://github.com/zabaras/transformer-physx) -""" - -from __future__ import annotations - -from typing import Optional -from typing import Tuple - -import numpy as np -import paddle -from paddle import nn -from paddle.nn.initializer import Constant -from paddle.nn.initializer import Uniform - -from ppsci.arch import base - -zeros_ = Constant(value=0.0) -ones_ = Constant(value=1.0) - - -class LorenzEmbedding(base.Arch): - """Embedding Koopman model for the Lorenz ODE system. - - Args: - input_keys (Tuple[str, ...]): Input keys, such as ("states",). - output_keys (Tuple[str, ...]): Output keys, such as ("pred_states", "recover_states"). - mean (Optional[Tuple[float, ...]]): Mean of training dataset. Defaults to None. - std (Optional[Tuple[float, ...]]): Standard Deviation of training dataset. Defaults to None. - input_size (int, optional): Size of input data. Defaults to 3. - hidden_size (int, optional): Number of hidden size. Defaults to 500. - embed_size (int, optional): Number of embedding size. Defaults to 32. - drop (float, optional): Probability of dropout the units. Defaults to 0.0. - - Examples: - >>> import ppsci - >>> model = ppsci.arch.LorenzEmbedding( - ... input_keys=("x", "y"), - ... output_keys=("u", "v"), - ... input_size=3, - ... hidden_size=500, - ... embed_size=32, - ... drop=0.0, - ... mean=None, - ... std=None, - ... ) - >>> x_shape = [8, 3, 2] - >>> y_shape = [8, 3, 1] - >>> input_dict = {"x": paddle.rand(x_shape), - ... "y": paddle.rand(y_shape)} - >>> output_dict = model(input_dict) - >>> print(output_dict["u"].shape) - [8, 2, 3] - >>> print(output_dict["v"].shape) - [8, 3, 3] - """ - - def __init__( - self, - input_keys: Tuple[str, ...], - output_keys: Tuple[str, ...], - mean: Optional[Tuple[float, ...]] = None, - std: Optional[Tuple[float, ...]] = None, - input_size: int = 3, - hidden_size: int = 500, - embed_size: int = 32, - drop: float = 0.0, - ): - super().__init__() - self.input_keys = input_keys - self.output_keys = output_keys - self.input_size = input_size - self.hidden_size = hidden_size - self.embed_size = embed_size - - # build observable network - self.encoder_net = self.build_encoder(input_size, hidden_size, embed_size, drop) - # build koopman operator - self.k_diag, self.k_ut = self.build_koopman_operator(embed_size) - # build recovery network - self.decoder_net = self.build_decoder(input_size, hidden_size, embed_size) - - mean = [0.0, 0.0, 0.0] if mean is None else mean - std = [1.0, 1.0, 1.0] if std is None else std - self.register_buffer("mean", paddle.to_tensor(mean).reshape([1, 3])) - self.register_buffer("std", paddle.to_tensor(std).reshape([1, 3])) - - self.apply(self._init_weights) - - def _init_weights(self, m: nn.Layer): - if isinstance(m, nn.Linear): - k = 1 / m.weight.shape[0] - uniform = Uniform(-(k**0.5), k**0.5) - uniform(m.weight) - if m.bias is not None: - uniform(m.bias) - elif isinstance(m, nn.LayerNorm): - zeros_(m.bias) - ones_(m.weight) - - def build_encoder( - self, input_size: int, hidden_size: int, embed_size: int, drop: float = 0.0 - ): - net = nn.Sequential( - nn.Linear(input_size, hidden_size), - nn.ReLU(), - nn.Linear(hidden_size, embed_size), - nn.LayerNorm(embed_size), - nn.Dropout(drop), - ) - return net - - def build_decoder(self, input_size: int, hidden_size: int, embed_size: int): - net = nn.Sequential( - nn.Linear(embed_size, hidden_size), - nn.ReLU(), - nn.Linear(hidden_size, input_size), - ) - return net - - def build_koopman_operator(self, embed_size: int): - # Learned Koopman operator - data = paddle.linspace(1, 0, embed_size) - k_diag = paddle.create_parameter( - shape=data.shape, - dtype=paddle.get_default_dtype(), - default_initializer=nn.initializer.Assign(data), - ) - - data = 0.1 * paddle.rand([2 * embed_size - 3]) - k_ut = paddle.create_parameter( - shape=data.shape, - dtype=paddle.get_default_dtype(), - default_initializer=nn.initializer.Assign(data), - ) - return k_diag, k_ut - - def encoder(self, x: paddle.Tensor): - x = self._normalize(x) - g = self.encoder_net(x) - return g - - def decoder(self, g: paddle.Tensor): - out = self.decoder_net(g) - x = self._unnormalize(out) - return x - - def koopman_operation(self, embed_data: paddle.Tensor, k_matrix: paddle.Tensor): - # Apply Koopman operation - embed_pred_data = paddle.bmm( - k_matrix.expand( - [embed_data.shape[0], k_matrix.shape[0], k_matrix.shape[1]] - ), - embed_data.transpose([0, 2, 1]), - ).transpose([0, 2, 1]) - return embed_pred_data - - def _normalize(self, x: paddle.Tensor): - return (x - self.mean) / self.std - - def _unnormalize(self, x: paddle.Tensor): - return self.std * x + self.mean - - def get_koopman_matrix(self): - # # Koopman operator - k_ut_tensor = self.k_ut * 1 - k_ut_tensor = paddle.diag( - k_ut_tensor[0 : self.embed_size - 1], offset=1 - ) + paddle.diag(k_ut_tensor[self.embed_size - 1 :], offset=2) - k_matrix = k_ut_tensor + (-1) * k_ut_tensor.t() - k_matrix = k_matrix + paddle.diag(self.k_diag) - return k_matrix - - def forward_tensor(self, x): - k_matrix = self.get_koopman_matrix() - embed_data = self.encoder(x) - recover_data = self.decoder(embed_data) - - embed_pred_data = self.koopman_operation(embed_data, k_matrix) - pred_data = self.decoder(embed_pred_data) - - return (pred_data[:, :-1, :], recover_data, k_matrix) - - @staticmethod - def split_to_dict(data_tensors: Tuple[paddle.Tensor, ...], keys: Tuple[str, ...]): - return {key: data_tensors[i] for i, key in enumerate(keys)} - - def forward(self, x): - if self._input_transform is not None: - x = self._input_transform(x) - - x_tensor = self.concat_to_tensor(x, self.input_keys, axis=-1) - y = self.forward_tensor(x_tensor) - y = self.split_to_dict(y, self.output_keys) - - if self._output_transform is not None: - y = self._output_transform(x, y) - return y - - -class RosslerEmbedding(LorenzEmbedding): - """Embedding Koopman model for the Rossler ODE system. - - Args: - input_keys (Tuple[str, ...]): Input keys, such as ("states",). - output_keys (Tuple[str, ...]): Output keys, such as ("pred_states", "recover_states"). - mean (Optional[Tuple[float, ...]]): Mean of training dataset. Defaults to None. - std (Optional[Tuple[float, ...]]): Standard Deviation of training dataset. Defaults to None. - input_size (int, optional): Size of input data. Defaults to 3. - hidden_size (int, optional): Number of hidden size. Defaults to 500. - embed_size (int, optional): Number of embedding size. Defaults to 32. - drop (float, optional): Probability of dropout the units. Defaults to 0.0. - - Examples: - >>> import ppsci - >>> model = ppsci.arch.RosslerEmbedding( - ... input_keys=("x", "y"), - ... output_keys=("u", "v"), - ... input_size=3, - ... hidden_size=500, - ... embed_size=32, - ... drop=0.0, - ... mean=None, - ... std=None, - ... ) - >>> x_shape = [8, 3, 2] - >>> y_shape = [8, 3, 1] - >>> input_dict = {"x": paddle.rand(x_shape), - ... "y": paddle.rand(y_shape)} - >>> output_dict = model(input_dict) - >>> print(output_dict["u"].shape) - [8, 2, 3] - >>> print(output_dict["v"].shape) - [8, 3, 3] - """ - - def __init__( - self, - input_keys: Tuple[str, ...], - output_keys: Tuple[str, ...], - mean: Optional[Tuple[float, ...]] = None, - std: Optional[Tuple[float, ...]] = None, - input_size: int = 3, - hidden_size: int = 500, - embed_size: int = 32, - drop: float = 0.0, - ): - super().__init__( - input_keys, - output_keys, - mean, - std, - input_size, - hidden_size, - embed_size, - drop, - ) - - -class CylinderEmbedding(base.Arch): - """Embedding Koopman model for the Cylinder system. - - Args: - input_keys (Tuple[str, ...]): Input keys, such as ("states", "visc"). - output_keys (Tuple[str, ...]): Output keys, such as ("pred_states", "recover_states"). - mean (Optional[Tuple[float, ...]]): Mean of training dataset. Defaults to None. - std (Optional[Tuple[float, ...]]): Standard Deviation of training dataset. Defaults to None. - embed_size (int, optional): Number of embedding size. Defaults to 128. - encoder_channels (Optional[Tuple[int, ...]]): Number of channels in encoder network. Defaults to None. - decoder_channels (Optional[Tuple[int, ...]]): Number of channels in decoder network. Defaults to None. - drop (float, optional): Probability of dropout the units. Defaults to 0.0. - - Examples: - >>> import paddle - >>> import ppsci - >>> model = ppsci.arch.CylinderEmbedding(("states", "visc"), ("pred_states", "recover_states")) - >>> states_shape = [32, 10, 3, 64, 128] - >>> visc_shape = [32, 1] - >>> input_dict = {"states" : paddle.rand(states_shape), - ... "visc" : paddle.rand(visc_shape)} - >>> out_dict = model(input_dict) - >>> print(out_dict["pred_states"].shape) - [32, 9, 3, 64, 128] - >>> print(out_dict["recover_states"].shape) - [32, 10, 3, 64, 128] - """ - - def __init__( - self, - input_keys: Tuple[str, ...], - output_keys: Tuple[str, ...], - mean: Optional[Tuple[float, ...]] = None, - std: Optional[Tuple[float, ...]] = None, - embed_size: int = 128, - encoder_channels: Optional[Tuple[int, ...]] = None, - decoder_channels: Optional[Tuple[int, ...]] = None, - drop: float = 0.0, - ): - super().__init__() - self.input_keys = input_keys - self.output_keys = output_keys - self.embed_size = embed_size - - X, Y = np.meshgrid(np.linspace(-2, 14, 128), np.linspace(-4, 4, 64)) - self.mask = paddle.to_tensor(np.sqrt(X**2 + Y**2)).unsqueeze(0).unsqueeze(0) - - encoder_channels = ( - [4, 16, 32, 64, 128] if encoder_channels is None else encoder_channels - ) - decoder_channels = ( - [embed_size // 32, 128, 64, 32, 16] - if decoder_channels is None - else decoder_channels - ) - self.encoder_net = self.build_encoder(embed_size, encoder_channels, drop) - self.k_diag_net, self.k_ut_net, self.k_lt_net = self.build_koopman_operator( - embed_size - ) - self.decoder_net = self.build_decoder(decoder_channels) - - xidx = [] - yidx = [] - for i in range(1, 5): - yidx.append(np.arange(i, embed_size)) - xidx.append(np.arange(0, embed_size - i)) - self.xidx = paddle.to_tensor(np.concatenate(xidx), dtype="int64") - self.yidx = paddle.to_tensor(np.concatenate(yidx), dtype="int64") - - mean = [0.0, 0.0, 0.0, 0.0] if mean is None else mean - std = [1.0, 1.0, 1.0, 1.0] if std is None else std - self.register_buffer("mean", paddle.to_tensor(mean).reshape([1, 4, 1, 1])) - self.register_buffer("std", paddle.to_tensor(std).reshape([1, 4, 1, 1])) - - self.apply(self._init_weights) - - def _init_weights(self, m): - if isinstance(m, nn.Linear): - k = 1 / m.weight.shape[0] - uniform = Uniform(-(k**0.5), k**0.5) - uniform(m.weight) - if m.bias is not None: - uniform(m.bias) - elif isinstance(m, nn.LayerNorm): - zeros_(m.bias) - ones_(m.weight) - elif isinstance(m, nn.Conv2D): - k = 1 / (m.weight.shape[1] * m.weight.shape[2] * m.weight.shape[3]) - uniform = Uniform(-(k**0.5), k**0.5) - uniform(m.weight) - if m.bias is not None: - uniform(m.bias) - - def _build_conv_relu_list( - self, in_channels: Tuple[int, ...], out_channels: Tuple[int, ...] - ): - net_list = [ - nn.Conv2D( - in_channels, - out_channels, - kernel_size=(3, 3), - stride=2, - padding=1, - padding_mode="replicate", - ), - nn.ReLU(), - ] - return net_list - - def build_encoder( - self, embed_size: int, channels: Tuple[int, ...], drop: float = 0.0 - ): - net = [] - for i in range(1, len(channels)): - net.extend(self._build_conv_relu_list(channels[i - 1], channels[i])) - net.append( - nn.Conv2D( - channels[-1], - embed_size // 32, - kernel_size=(3, 3), - padding=1, - padding_mode="replicate", - ) - ) - net.append( - nn.LayerNorm( - (4, 4, 8), - ) - ) - net.append(nn.Dropout(drop)) - net = nn.Sequential(*net) - return net - - def _build_upsample_conv_relu( - self, in_channels: Tuple[int, ...], out_channels: Tuple[int, ...] - ): - net_list = [ - nn.Upsample(scale_factor=2, mode="bilinear", align_corners=True), - nn.Conv2D( - in_channels, - out_channels, - kernel_size=(3, 3), - stride=1, - padding=1, - padding_mode="replicate", - ), - nn.ReLU(), - ] - return net_list - - def build_decoder(self, channels: Tuple[int, ...]): - net = [] - for i in range(1, len(channels)): - net.extend(self._build_upsample_conv_relu(channels[i - 1], channels[i])) - net.append( - nn.Conv2D( - channels[-1], - 3, - kernel_size=(3, 3), - stride=1, - padding=1, - padding_mode="replicate", - ), - ) - net = nn.Sequential(*net) - return net - - def build_koopman_operator(self, embed_size: int): - # Learned Koopman operator parameters - k_diag_net = nn.Sequential( - nn.Linear(1, 50), nn.ReLU(), nn.Linear(50, embed_size) - ) - - k_ut_net = nn.Sequential( - nn.Linear(1, 50), nn.ReLU(), nn.Linear(50, 4 * embed_size - 10) - ) - k_lt_net = nn.Sequential( - nn.Linear(1, 50), nn.ReLU(), nn.Linear(50, 4 * embed_size - 10) - ) - return k_diag_net, k_ut_net, k_lt_net - - def encoder(self, x: paddle.Tensor, viscosity: paddle.Tensor): - B, T, C, H, W = x.shape - x = x.reshape((B * T, C, H, W)) - viscosity = viscosity.repeat_interleave(T, axis=1).reshape((B * T, 1)) - x = paddle.concat( - [x, viscosity.unsqueeze(-1).unsqueeze(-1) * paddle.ones_like(x[:, :1])], - axis=1, - ) - x = self._normalize(x) - g = self.encoder_net(x) - g = g.reshape([B, T, -1]) - return g - - def decoder(self, g: paddle.Tensor): - B, T, _ = g.shape - x = self.decoder_net(g.reshape([-1, self.embed_size // 32, 4, 8])) - x = self._unnormalize(x) - mask0 = ( - self.mask.repeat_interleave(x.shape[1], axis=1).repeat_interleave( - x.shape[0], axis=0 - ) - < 1 - ) - x[mask0] = 0 - _, C, H, W = x.shape - x = x.reshape([B, T, C, H, W]) - return x - - def get_koopman_matrix(self, g: paddle.Tensor, visc: paddle.Tensor): - # # Koopman operator - kMatrix = paddle.zeros([g.shape[0], self.embed_size, self.embed_size]) - kMatrix.stop_gradient = False - # Populate the off diagonal terms - kMatrixUT_data = self.k_ut_net(100 * visc) - kMatrixLT_data = self.k_lt_net(100 * visc) - - kMatrix = kMatrix.transpose([1, 2, 0]) - kMatrixUT_data_t = kMatrixUT_data.transpose([1, 0]) - kMatrixLT_data_t = kMatrixLT_data.transpose([1, 0]) - kMatrix[self.xidx, self.yidx] = kMatrixUT_data_t - kMatrix[self.yidx, self.xidx] = kMatrixLT_data_t - - # Populate the diagonal - ind = np.diag_indices(kMatrix.shape[1]) - ind = paddle.to_tensor(ind, dtype="int64") - - kMatrixDiag = self.k_diag_net(100 * visc) - kMatrixDiag_t = kMatrixDiag.transpose([1, 0]) - kMatrix[ind[0], ind[1]] = kMatrixDiag_t - return kMatrix.transpose([2, 0, 1]) - - def koopman_operation(self, embed_data: paddle.Tensor, k_matrix: paddle.Tensor): - embed_pred_data = paddle.bmm( - k_matrix, embed_data.transpose([0, 2, 1]) - ).transpose([0, 2, 1]) - return embed_pred_data - - def _normalize(self, x: paddle.Tensor): - x = (x - self.mean) / self.std - return x - - def _unnormalize(self, x: paddle.Tensor): - return self.std[:, :3] * x + self.mean[:, :3] - - def forward_tensor(self, states, visc): - # states.shape=(B, T, C, H, W) - embed_data = self.encoder(states, visc) - recover_data = self.decoder(embed_data) - - k_matrix = self.get_koopman_matrix(embed_data, visc) - embed_pred_data = self.koopman_operation(embed_data, k_matrix) - pred_data = self.decoder(embed_pred_data) - - return (pred_data[:, :-1], recover_data, k_matrix) - - @staticmethod - def split_to_dict(data_tensors: Tuple[paddle.Tensor, ...], keys: Tuple[str, ...]): - return {key: data_tensors[i] for i, key in enumerate(keys)} - - def forward(self, x): - - if self._input_transform is not None: - x = self._input_transform(x) - - y = self.forward_tensor(**x) - y = self.split_to_dict(y, self.output_keys) - - if self._output_transform is not None: - y = self._output_transform(x, y) - return y diff --git a/examples/smc_reac/ppsci/arch/epnn.py b/examples/smc_reac/ppsci/arch/epnn.py deleted file mode 100644 index 0f6a9ffed6..0000000000 --- a/examples/smc_reac/ppsci/arch/epnn.py +++ /dev/null @@ -1,126 +0,0 @@ -# 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. - -# Elasto-Plastic Neural Network (EPNN) - -# DEVELOPED AT: -# COMPUTATIONAL GEOMECHANICS LABORATORY -# DEPARTMENT OF CIVIL ENGINEERING -# UNIVERSITY OF CALGARY, AB, CANADA -# DIRECTOR: Prof. Richard Wan - -# DEVELOPED BY: -# MAHDAD EGHBALIAN - -# MIT License - -# Copyright (c) 2022 Mahdad Eghbalian - -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: - -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. - -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. - -from typing import Tuple - -import paddle.nn as nn - -from ppsci.arch import activation as act_mod -from ppsci.arch import base - - -class Epnn(base.Arch): - """Builds a feedforward network with arbitrary layers. - - Args: - input_keys (Tuple[str, ...]): Name of input keys, such as ("x", "y", "z"). - output_keys (Tuple[str, ...]): Name of output keys, such as ("u", "v", "w"). - node_sizes (Tuple[int, ...]): The tuple of node size. - activations (Tuple[str, ...]): Name of activation functions. - drop_p (float): The parameter p of nn.Dropout. - - Examples: - >>> import ppsci - >>> ann_node_sizes_state = [1, 20] - >>> model = ppsci.arch.Epnn( - ... ("x",), - ... ("y",), - ... node_sizes=ann_node_sizes_state, - ... activations=("leaky_relu",), - ... drop_p=0.0, - ... ) - """ - - def __init__( - self, - input_keys: Tuple[str, ...], - output_keys: Tuple[str, ...], - node_sizes: Tuple[int, ...], - activations: Tuple[str, ...], - drop_p: float, - ): - super().__init__() - self.active_func = [ - act_mod.get_activation(act_name) for act_name in activations - ] - self.node_sizes = node_sizes - self.drop_p = drop_p - self.layers = [] - self.layers.append( - nn.Linear(in_features=node_sizes[0], out_features=node_sizes[1]) - ) - layer_sizes = zip(node_sizes[1:-2], node_sizes[2:-1]) - self.layers.extend( - [nn.Linear(in_features=h1, out_features=h2) for h1, h2 in layer_sizes] - ) - self.layers.append( - nn.Linear( - in_features=node_sizes[-2], out_features=node_sizes[-1], bias_attr=False - ) - ) - - self.layers = nn.LayerList(self.layers) - self.dropout = nn.Dropout(p=drop_p) - self.input_keys = input_keys - self.output_keys = output_keys - - def forward(self, x): - if self._input_transform is not None: - x = self._input_transform(x) - - y = x[self.input_keys[0]] - for ilayer in range(len(self.layers)): - y = self.layers[ilayer](y) - if ilayer != len(self.layers) - 1: - y = self.active_func[ilayer + 1](y) - if ilayer != len(self.layers) - 1: - y = self.dropout(y) - y = self.split_to_dict(y, self.output_keys, axis=-1) - - if self._output_transform is not None: - y = self._output_transform(x, y) - return y diff --git a/examples/smc_reac/ppsci/arch/extformer_moe_cuboid.py b/examples/smc_reac/ppsci/arch/extformer_moe_cuboid.py deleted file mode 100644 index 7c57d23948..0000000000 --- a/examples/smc_reac/ppsci/arch/extformer_moe_cuboid.py +++ /dev/null @@ -1,996 +0,0 @@ -from typing import Sequence -from typing import Tuple -from typing import Union - -import paddle -from paddle import nn - -import ppsci.arch.extformer_moe_cuboid_decoder as cuboid_decoder -import ppsci.arch.extformer_moe_cuboid_encoder as cuboid_encoder -import ppsci.arch.extformer_moe_cuboid_utils as cuboid_utils -from ppsci.arch import activation as act_mod -from ppsci.arch import base -from ppsci.arch import extformer_moe_utils -from ppsci.arch.extformer_moe_cuboid_encoder import NEGATIVE_SLOPE -from ppsci.utils import initializer - -"""A space-time Transformer with Cuboid Attention""" - - -class InitialEncoder(nn.Layer): - def __init__( - self, - dim, - out_dim, - downsample_scale: Union[int, Sequence[int]], - num_conv_layers: int = 2, - activation: str = "leaky", - padding_type: str = "nearest", - conv_init_mode: str = "0", - linear_init_mode: str = "0", - norm_init_mode: str = "0", - moe_config: dict = None, - ): - super(InitialEncoder, self).__init__() - self.num_conv_layers = num_conv_layers - self.conv_init_mode = conv_init_mode - self.linear_init_mode = linear_init_mode - self.norm_init_mode = norm_init_mode - conv_block = [] - for i in range(num_conv_layers): - if i == 0: - conv_block.append( - nn.Conv2D( - kernel_size=(3, 3), - padding=(1, 1), - in_channels=dim, - out_channels=out_dim, - ) - ) - conv_block.append(nn.GroupNorm(num_groups=16, num_channels=out_dim)) - conv_block.append( - act_mod.get_activation(activation) - if activation != "leaky_relu" - else nn.LeakyReLU(NEGATIVE_SLOPE) - ) - else: - conv_block.append( - nn.Conv2D( - kernel_size=(3, 3), - padding=(1, 1), - in_channels=out_dim, - out_channels=out_dim, - ) - ) - conv_block.append(nn.GroupNorm(num_groups=16, num_channels=out_dim)) - conv_block.append( - act_mod.get_activation(activation) - if activation != "leaky_relu" - else nn.LeakyReLU(NEGATIVE_SLOPE) - ) - self.conv_block = nn.Sequential(*conv_block) - if isinstance(downsample_scale, int): - patch_merge_downsample = (1, downsample_scale, downsample_scale) - elif len(downsample_scale) == 2: - patch_merge_downsample = (1, *downsample_scale) - elif len(downsample_scale) == 3: - patch_merge_downsample = tuple(downsample_scale) - else: - raise NotImplementedError( - f"downsample_scale {downsample_scale} format not supported!" - ) - self.patch_merge = cuboid_encoder.PatchMerging3D( - dim=out_dim, - out_dim=out_dim, - padding_type=padding_type, - downsample=patch_merge_downsample, - linear_init_mode=linear_init_mode, - norm_init_mode=norm_init_mode, - ) - self.reset_parameters() - - def reset_parameters(self): - for m in self.children(): - cuboid_utils.apply_initialization( - m, - conv_mode=self.conv_init_mode, - linear_mode=self.linear_init_mode, - norm_mode=self.norm_init_mode, - ) - - def forward(self, x): - """x --> [K x Conv2D] --> PatchMerge - - Args: - x: (B, T, H, W, C) - - Returns: - out: (B, T, H_new, W_new, C_out) - """ - - B, T, H, W, C = x.shape - - if self.num_conv_layers > 0: - x = x.reshape([B * T, H, W, C]).transpose(perm=[0, 3, 1, 2]) - x = self.conv_block(x).transpose(perm=[0, 2, 3, 1]) - x = self.patch_merge(x.reshape([B, T, H, W, -1])) - else: - x = self.patch_merge(x) - return x - - -class FinalDecoder(nn.Layer): - def __init__( - self, - target_thw: Tuple[int, ...], - dim: int, - num_conv_layers: int = 2, - activation: str = "leaky", - conv_init_mode: str = "0", - linear_init_mode: str = "0", - norm_init_mode: str = "0", - moe_config: dict = None, - ): - super(FinalDecoder, self).__init__() - self.target_thw = target_thw - self.dim = dim - self.num_conv_layers = num_conv_layers - self.conv_init_mode = conv_init_mode - self.linear_init_mode = linear_init_mode - self.norm_init_mode = norm_init_mode - conv_block = [] - for i in range(num_conv_layers): - conv_block.append( - nn.Conv2D( - kernel_size=(3, 3), - padding=(1, 1), - in_channels=dim, - out_channels=dim, - ) - ) - conv_block.append(nn.GroupNorm(num_groups=16, num_channels=dim)) - conv_block.append( - act_mod.get_activation(activation) - if activation != "leaky_relu" - else nn.LeakyReLU(NEGATIVE_SLOPE) - ) - self.conv_block = nn.Sequential(*conv_block) - self.upsample = cuboid_decoder.Upsample3DLayer( - dim=dim, - out_dim=dim, - target_size=target_thw, - kernel_size=3, - conv_init_mode=conv_init_mode, - ) - self.reset_parameters() - - def reset_parameters(self): - for m in self.children(): - cuboid_utils.apply_initialization( - m, - conv_mode=self.conv_init_mode, - linear_mode=self.linear_init_mode, - norm_mode=self.norm_init_mode, - ) - - def forward(self, x): - """x --> Upsample --> [K x Conv2D] - - Args: - x: (B, T, H, W, C) - - Returns: - out: (B, T, H_new, W_new, C) - """ - - x = self.upsample(x) - if self.num_conv_layers > 0: - B, T, H, W, C = x.shape - x = x.reshape([B * T, H, W, C]).transpose(perm=[0, 3, 1, 2]) - x = ( - self.conv_block(x) - .transpose(perm=[0, 2, 3, 1]) - .reshape([B, T, H, W, -1]) - ) - return x - - -class InitialStackPatchMergingEncoder(nn.Layer): - def __init__( - self, - num_merge: int, - in_dim: int, - out_dim_list: Tuple[int, ...], - downsample_scale_list: Tuple[float, ...], - num_conv_per_merge_list: Tuple[int, ...] = None, - activation: str = "leaky", - padding_type: str = "nearest", - conv_init_mode: str = "0", - linear_init_mode: str = "0", - norm_init_mode: str = "0", - moe_config: dict = None, - ): - super(InitialStackPatchMergingEncoder, self).__init__() - self.conv_init_mode = conv_init_mode - self.linear_init_mode = linear_init_mode - self.norm_init_mode = norm_init_mode - self.num_merge = num_merge - self.in_dim = in_dim - self.out_dim_list = out_dim_list[:num_merge] - self.downsample_scale_list = downsample_scale_list[:num_merge] - self.num_conv_per_merge_list = num_conv_per_merge_list - self.num_group_list = [max(1, out_dim // 4) for out_dim in self.out_dim_list] - self.conv_block_list = nn.LayerList() - self.patch_merge_list = nn.LayerList() - for i in range(num_merge): - if i == 0: - in_dim = in_dim - else: - in_dim = self.out_dim_list[i - 1] - out_dim = self.out_dim_list[i] - downsample_scale = self.downsample_scale_list[i] - conv_block = [] - for j in range(self.num_conv_per_merge_list[i]): - if j == 0: - conv_in_dim = in_dim - else: - conv_in_dim = out_dim - conv_block.append( - nn.Conv2D( - kernel_size=(3, 3), - padding=(1, 1), - in_channels=conv_in_dim, - out_channels=out_dim, - ) - ) - conv_block.append( - nn.GroupNorm( - num_groups=self.num_group_list[i], num_channels=out_dim - ) - ) - conv_block.append( - act_mod.get_activation(activation) - if activation != "leaky_relu" - else nn.LeakyReLU(NEGATIVE_SLOPE) - ) - conv_block = nn.Sequential(*conv_block) - self.conv_block_list.append(conv_block) - patch_merge = cuboid_encoder.PatchMerging3D( - dim=out_dim, - out_dim=out_dim, - padding_type=padding_type, - downsample=(1, downsample_scale, downsample_scale), - linear_init_mode=linear_init_mode, - norm_init_mode=norm_init_mode, - ) - self.patch_merge_list.append(patch_merge) - self.reset_parameters() - - def reset_parameters(self): - for m in self.children(): - cuboid_utils.apply_initialization( - m, - conv_mode=self.conv_init_mode, - linear_mode=self.linear_init_mode, - norm_mode=self.norm_init_mode, - ) - - def get_out_shape_list(self, input_shape): - out_shape_list = [] - for patch_merge in self.patch_merge_list: - input_shape = patch_merge.get_out_shape(input_shape) - out_shape_list.append(input_shape) - return out_shape_list - - def forward(self, x): - """x --> [K x Conv2D] --> PatchMerge --> ... --> [K x Conv2D] --> PatchMerge - - Args: - x: (B, T, H, W, C) - - Returns: - out: (B, T, H_new, W_new, C_out) - """ - - for i, (conv_block, patch_merge) in enumerate( - zip(self.conv_block_list, self.patch_merge_list) - ): - B, T, H, W, C = x.shape - if self.num_conv_per_merge_list[i] > 0: - x = x.reshape([B * T, H, W, C]).transpose(perm=[0, 3, 1, 2]) - x = conv_block(x).transpose(perm=[0, 2, 3, 1]).reshape([B, T, H, W, -1]) - x = patch_merge(x) - return x - - -class FinalStackUpsamplingDecoder(nn.Layer): - def __init__( - self, - target_shape_list: Tuple[Tuple[int, ...]], - in_dim: int, - num_conv_per_up_list: Tuple[int, ...] = None, - activation: str = "leaky", - conv_init_mode: str = "0", - linear_init_mode: str = "0", - norm_init_mode: str = "0", - moe_config: dict = None, - ): - super(FinalStackUpsamplingDecoder, self).__init__() - self.conv_init_mode = conv_init_mode - self.linear_init_mode = linear_init_mode - self.norm_init_mode = norm_init_mode - self.target_shape_list = target_shape_list - self.out_dim_list = [ - target_shape[-1] for target_shape in self.target_shape_list - ] - self.num_upsample = len(target_shape_list) - self.in_dim = in_dim - self.num_conv_per_up_list = num_conv_per_up_list - self.num_group_list = [max(1, out_dim // 4) for out_dim in self.out_dim_list] - self.conv_block_list = nn.LayerList() - self.upsample_list = nn.LayerList() - for i in range(self.num_upsample): - if i == 0: - in_dim = in_dim - else: - in_dim = self.out_dim_list[i - 1] - out_dim = self.out_dim_list[i] - upsample = cuboid_decoder.Upsample3DLayer( - dim=in_dim, - out_dim=in_dim, - target_size=target_shape_list[i][:-1], - kernel_size=3, - conv_init_mode=conv_init_mode, - ) - self.upsample_list.append(upsample) - conv_block = [] - for j in range(num_conv_per_up_list[i]): - if j == 0: - conv_in_dim = in_dim - else: - conv_in_dim = out_dim - conv_block.append( - nn.Conv2D( - kernel_size=(3, 3), - padding=(1, 1), - in_channels=conv_in_dim, - out_channels=out_dim, - ) - ) - conv_block.append( - nn.GroupNorm( - num_groups=self.num_group_list[i], num_channels=out_dim - ) - ) - conv_block.append( - act_mod.get_activation(activation) - if activation != "leaky_relu" - else nn.LeakyReLU(NEGATIVE_SLOPE) - ) - conv_block = nn.Sequential(*conv_block) - self.conv_block_list.append(conv_block) - self.reset_parameters() - - def reset_parameters(self): - for m in self.children(): - cuboid_utils.apply_initialization( - m, - conv_mode=self.conv_init_mode, - linear_mode=self.linear_init_mode, - norm_mode=self.norm_init_mode, - ) - - @staticmethod - def get_init_params(enc_input_shape, enc_out_shape_list, large_channel=False): - dec_target_shape_list = list(enc_out_shape_list[:-1])[::-1] + [ - tuple(enc_input_shape) - ] - if large_channel: - dec_target_shape_list_large_channel = [] - for i, enc_out_shape in enumerate(enc_out_shape_list[::-1]): - dec_target_shape_large_channel = list(dec_target_shape_list[i]) - dec_target_shape_large_channel[-1] = enc_out_shape[-1] - dec_target_shape_list_large_channel.append( - tuple(dec_target_shape_large_channel) - ) - dec_target_shape_list = dec_target_shape_list_large_channel - dec_in_dim = enc_out_shape_list[-1][-1] - return dec_target_shape_list, dec_in_dim - - def forward(self, x): - """x --> Upsample --> [K x Conv2D] --> ... --> Upsample --> [K x Conv2D] - - Args: - x: Shape (B, T, H, W, C) - - Returns: - out: Shape (B, T, H_new, W_new, C) - """ - for i, (conv_block, upsample) in enumerate( - zip(self.conv_block_list, self.upsample_list) - ): - x = upsample(x) - if self.num_conv_per_up_list[i] > 0: - B, T, H, W, C = x.shape - x = x.reshape([B * T, H, W, C]).transpose(perm=[0, 3, 1, 2]) - x = conv_block(x).transpose(perm=[0, 2, 3, 1]).reshape([B, T, H, W, -1]) - return x - - -class ExtFormerMoECuboid(base.Arch): - """Cuboid Transformer for spatiotemporal forecasting - - We adopt the Non-autoregressive encoder-decoder architecture. - The decoder takes the multi-scale memory output from the encoder. - - The initial downsampling / upsampling layers will be - Downsampling: [K x Conv2D --> PatchMerge] - Upsampling: [Nearest Interpolation-based Upsample --> K x Conv2D] - - x --> downsample (optional) ---> (+pos_embed) ---> enc --> mem_l initial_z (+pos_embed) ---> FC - | | - |------------| - | - | - y <--- upsample (optional) <--- dec <---------- - - Args: - input_keys (Tuple[str, ...]): Name of input keys, such as ("input",). - output_keys (Tuple[str, ...]): Name of output keys, such as ("output",). - input_shape (Tuple[int, ...]): The shape of the input data. - target_shape (Tuple[int, ...]): The shape of the target data. - base_units (int, optional): The base units. Defaults to 128. - block_units (int, optional): The block units. Defaults to None. - scale_alpha (float, optional): We scale up the channels based on the formula: - - round_to(base_units * max(downsample_scale) ** units_alpha, 4). Defaults to 1.0. - num_heads (int, optional): The number of heads. Defaults to 4. - attn_drop (float, optional): The attention dropout. Defaults to 0.0. - proj_drop (float, optional): The projection dropout. Defaults to 0.0. - ffn_drop (float, optional): The ffn dropout. Defaults to 0.0. - downsample (int, optional): The rate of downsample. Defaults to 2. - downsample_type (str, optional): The type of downsample. Defaults to "patch_merge". - upsample_type (str, optional): The rate of upsample. Defaults to "upsample". - upsample_kernel_size (int, optional): The kernel size of upsample. Defaults to 3. - enc_depth (list, optional): The depth of encoder. Defaults to [4, 4, 4]. - enc_attn_patterns (str, optional): The pattern of encoder attention. Defaults to None. - enc_cuboid_size (list, optional): The cuboid size of encoder. Defaults to [(4, 4, 4), (4, 4, 4)]. - enc_cuboid_strategy (list, optional): The cuboid strategy of encoder. Defaults to [("l", "l", "l"), ("d", "d", "d")]. - enc_shift_size (list, optional): The shift size of encoder. Defaults to [(0, 0, 0), (0, 0, 0)]. - enc_use_inter_ffn (bool, optional): Whether to use intermediate FFN for encoder. Defaults to True. - dec_depth (list, optional): The depth of decoder. Defaults to [2, 2]. - dec_cross_start (int, optional): The cross start of decoder. Defaults to 0. - dec_self_attn_patterns (str, optional): The partterns of decoder. Defaults to None. - dec_self_cuboid_size (list, optional): The cuboid size of decoder. Defaults to [(4, 4, 4), (4, 4, 4)]. - dec_self_cuboid_strategy (list, optional): The strategy of decoder. Defaults to [("l", "l", "l"), ("d", "d", "d")]. - dec_self_shift_size (list, optional): The shift size of decoder. Defaults to [(1, 1, 1), (0, 0, 0)]. - dec_cross_attn_patterns (_type_, optional): The cross attention patterns of decoder. Defaults to None. - dec_cross_cuboid_hw (list, optional): The cuboid_hw of decoder. Defaults to [(4, 4), (4, 4)]. - dec_cross_cuboid_strategy (list, optional): The cuboid strategy of decoder. Defaults to [("l", "l", "l"), ("d", "l", "l")]. - dec_cross_shift_hw (list, optional): The shift_hw of decoder. Defaults to [(0, 0), (0, 0)]. - dec_cross_n_temporal (list, optional): The cross_n_temporal of decoder. Defaults to [1, 2]. - dec_cross_last_n_frames (int, optional): The cross_last_n_frames of decoder. Defaults to None. - dec_use_inter_ffn (bool, optional): Whether to use intermediate FFN for decoder. Defaults to True. - dec_hierarchical_pos_embed (bool, optional): Whether to use hierarchical pos_embed for decoder. Defaults to False. - num_global_vectors (int, optional): The num of global vectors. Defaults to 4. - use_dec_self_global (bool, optional): Whether to use global vector for decoder. Defaults to True. - dec_self_update_global (bool, optional): Whether to update global vector for decoder. Defaults to True. - use_dec_cross_global (bool, optional): Whether to use cross global vector for decoder. Defaults to True. - use_global_vector_ffn (bool, optional): Whether to use global vector FFN. Defaults to True. - use_global_self_attn (bool, optional): Whether to use global attentions. Defaults to False. - separate_global_qkv (bool, optional): Whether to separate global qkv. Defaults to False. - global_dim_ratio (int, optional): The ratio of global dim. Defaults to 1. - self_pattern (str, optional): The pattern. Defaults to "axial". - cross_self_pattern (str, optional): The self cross pattern. Defaults to "axial". - cross_pattern (str, optional): The cross pattern. Defaults to "cross_1x1". - z_init_method (str, optional): How the initial input to the decoder is initialized. Defaults to "nearest_interp". - initial_downsample_type (str, optional): The downsample type of initial. Defaults to "conv". - initial_downsample_activation (str, optional): The downsample activation of initial. Defaults to "leaky". - initial_downsample_scale (int, optional): The downsample scale of initial. Defaults to 1. - initial_downsample_conv_layers (int, optional): The conv layer of downsample of initial. Defaults to 2. - final_upsample_conv_layers (int, optional): The conv layer of final upsample. Defaults to 2. - initial_downsample_stack_conv_num_layers (int, optional): The num of stack conv layer of initial downsample. Defaults to 1. - initial_downsample_stack_conv_dim_list (list, optional): The dim list of stack conv of initial downsample. Defaults to None. - initial_downsample_stack_conv_downscale_list (list, optional): The downscale list of stack conv of initial downsample. Defaults to [1]. - initial_downsample_stack_conv_num_conv_list (list, optional): The num of stack conv list of initial downsample. Defaults to [2]. - ffn_activation (str, optional): The activation of FFN. Defaults to "leaky". - gated_ffn (bool, optional): Whether to use gate FFN. Defaults to False. - norm_layer (str, optional): The type of normilize. Defaults to "layer_norm". - padding_type (str, optional): The type of padding. Defaults to "ignore". - pos_embed_type (str, optional): The type of pos embedding. Defaults to "t+hw". - checkpoint_level (bool, optional): Whether to use checkpoint. Defaults to True. - use_relative_pos (bool, optional): Whether to use relative pose. Defaults to True. - self_attn_use_final_proj (bool, optional): Whether to use final projection. Defaults to True. - dec_use_first_self_attn (bool, optional): Whether to use first self attention for decoder. Defaults to False. - attn_linear_init_mode (str, optional): The mode of attention linear init. Defaults to "0". - ffn_linear_init_mode (str, optional): The mode of FFN linear init. Defaults to "0". - conv_init_mode (str, optional): The mode of conv init. Defaults to "0". - down_up_linear_init_mode (str, optional): The mode of downsample and upsample linear init. Defaults to "0". - norm_init_mode (str, optional): The mode of normalization init. Defaults to "0". - """ - - def __init__( - self, - input_keys: Tuple[str, ...], - output_keys: Tuple[str, ...], - input_shape: Tuple[int, ...], - target_shape: Tuple[int, ...], - base_units: int = 128, - block_units: int = None, - scale_alpha: float = 1.0, - num_heads: int = 4, - attn_drop: float = 0.0, - proj_drop: float = 0.0, - ffn_drop: float = 0.0, - downsample: int = 2, - downsample_type: str = "patch_merge", - upsample_type: str = "upsample", - upsample_kernel_size: int = 3, - enc_depth: Tuple[int, ...] = [4, 4, 4], - enc_attn_patterns: str = None, - enc_cuboid_size: Tuple[Tuple[int, ...], ...] = [(4, 4, 4), (4, 4, 4)], - enc_cuboid_strategy: Tuple[Tuple[str, ...], ...] = [ - ("l", "l", "l"), - ("d", "d", "d"), - ], - enc_shift_size: Tuple[Tuple[int, ...], ...] = [(0, 0, 0), (0, 0, 0)], - enc_use_inter_ffn: bool = True, - dec_depth: Tuple[int, ...] = [2, 2], - dec_cross_start: int = 0, - dec_self_attn_patterns: str = None, - dec_self_cuboid_size: Tuple[Tuple[int, ...], ...] = [(4, 4, 4), (4, 4, 4)], - dec_self_cuboid_strategy: Tuple[Tuple[str, ...], ...] = [ - ("l", "l", "l"), - ("d", "d", "d"), - ], - dec_self_shift_size: Tuple[Tuple[int, ...], ...] = [(1, 1, 1), (0, 0, 0)], - dec_cross_attn_patterns: str = None, - dec_cross_cuboid_hw: Tuple[Tuple[int, ...], ...] = [(4, 4), (4, 4)], - dec_cross_cuboid_strategy: Tuple[Tuple[str, ...], ...] = [ - ("l", "l", "l"), - ("d", "l", "l"), - ], - dec_cross_shift_hw: Tuple[Tuple[int, ...], ...] = [(0, 0), (0, 0)], - dec_cross_n_temporal: Tuple[int, ...] = [1, 2], - dec_cross_last_n_frames: int = None, - dec_use_inter_ffn: bool = True, - dec_hierarchical_pos_embed: bool = False, - num_global_vectors: int = 4, - use_dec_self_global: bool = True, - dec_self_update_global: bool = True, - use_dec_cross_global: bool = True, - use_global_vector_ffn: bool = True, - use_global_self_attn: bool = False, - separate_global_qkv: bool = False, - global_dim_ratio: int = 1, - self_pattern: str = "axial", - cross_self_pattern: str = "axial", - cross_pattern: str = "cross_1x1", - z_init_method: str = "nearest_interp", - initial_downsample_type: str = "conv", - initial_downsample_activation: str = "leaky", - initial_downsample_scale: int = 1, - initial_downsample_conv_layers: int = 2, - final_upsample_conv_layers: int = 2, - initial_downsample_stack_conv_num_layers: int = 1, - initial_downsample_stack_conv_dim_list: Tuple[int, ...] = None, - initial_downsample_stack_conv_downscale_list: Tuple[int, ...] = [1], - initial_downsample_stack_conv_num_conv_list: Tuple[int, ...] = [2], - ffn_activation: str = "leaky", - gated_ffn: bool = False, - norm_layer: str = "layer_norm", - padding_type: str = "ignore", - pos_embed_type: str = "t+hw", - checkpoint_level: bool = True, - use_relative_pos: bool = True, - self_attn_use_final_proj: bool = True, - dec_use_first_self_attn: bool = False, - attn_linear_init_mode: str = "0", - ffn_linear_init_mode: str = "0", - conv_init_mode: str = "0", - down_up_linear_init_mode: str = "0", - norm_init_mode: str = "0", - moe_config: dict = None, - rnc_config: dict = None, - ): - super().__init__() - self.input_keys = input_keys - self.output_keys = output_keys - self.attn_linear_init_mode = attn_linear_init_mode - self.ffn_linear_init_mode = ffn_linear_init_mode - self.conv_init_mode = conv_init_mode - self.down_up_linear_init_mode = down_up_linear_init_mode - self.norm_init_mode = norm_init_mode - assert len(enc_depth) == len(dec_depth) - self.base_units = base_units - self.num_global_vectors = num_global_vectors - self.moe_config = moe_config - self.rnc_config = rnc_config - self.checkpoint_level = checkpoint_level - - num_blocks = len(enc_depth) - if isinstance(self_pattern, str): - enc_attn_patterns = [self_pattern] * num_blocks - if isinstance(cross_self_pattern, str): - dec_self_attn_patterns = [cross_self_pattern] * num_blocks - if isinstance(cross_pattern, str): - dec_cross_attn_patterns = [cross_pattern] * num_blocks - if global_dim_ratio != 1: - assert ( - separate_global_qkv is True - ), "Setting global_dim_ratio != 1 requires separate_global_qkv == True." - self.global_dim_ratio = global_dim_ratio - self.z_init_method = z_init_method - assert self.z_init_method in ["zeros", "nearest_interp", "last", "mean"] - self.input_shape = input_shape - self.target_shape = target_shape - T_in, H_in, W_in, C_in = input_shape - T_out, H_out, W_out, C_out = target_shape - assert H_in == H_out and W_in == W_out - if self.num_global_vectors > 0: - init_data = paddle.zeros( - (self.num_global_vectors, global_dim_ratio * base_units) - ) - self.init_global_vectors = paddle.create_parameter( - shape=init_data.shape, - dtype=init_data.dtype, - default_initializer=nn.initializer.Constant(0.0), - ) - - self.init_global_vectors.stop_gradient = not True - new_input_shape = self.get_initial_encoder_final_decoder( - initial_downsample_scale=initial_downsample_scale, - initial_downsample_type=initial_downsample_type, - activation=initial_downsample_activation, - initial_downsample_conv_layers=initial_downsample_conv_layers, - final_upsample_conv_layers=final_upsample_conv_layers, - padding_type=padding_type, - initial_downsample_stack_conv_num_layers=initial_downsample_stack_conv_num_layers, - initial_downsample_stack_conv_dim_list=initial_downsample_stack_conv_dim_list, - initial_downsample_stack_conv_downscale_list=initial_downsample_stack_conv_downscale_list, - initial_downsample_stack_conv_num_conv_list=initial_downsample_stack_conv_num_conv_list, - ) - T_in, H_in, W_in, _ = new_input_shape - self.encoder = cuboid_encoder.CuboidTransformerEncoder( - input_shape=(T_in, H_in, W_in, base_units), - base_units=base_units, - block_units=block_units, - scale_alpha=scale_alpha, - depth=enc_depth, - downsample=downsample, - downsample_type=downsample_type, - block_attn_patterns=enc_attn_patterns, - block_cuboid_size=enc_cuboid_size, - block_strategy=enc_cuboid_strategy, - block_shift_size=enc_shift_size, - num_heads=num_heads, - attn_drop=attn_drop, - proj_drop=proj_drop, - ffn_drop=ffn_drop, - gated_ffn=gated_ffn, - ffn_activation=ffn_activation, - norm_layer=norm_layer, - use_inter_ffn=enc_use_inter_ffn, - padding_type=padding_type, - use_global_vector=num_global_vectors > 0, - use_global_vector_ffn=use_global_vector_ffn, - use_global_self_attn=use_global_self_attn, - separate_global_qkv=separate_global_qkv, - global_dim_ratio=global_dim_ratio, - checkpoint_level=checkpoint_level, - use_relative_pos=use_relative_pos, - self_attn_use_final_proj=self_attn_use_final_proj, - attn_linear_init_mode=attn_linear_init_mode, - ffn_linear_init_mode=ffn_linear_init_mode, - conv_init_mode=conv_init_mode, - down_linear_init_mode=down_up_linear_init_mode, - norm_init_mode=norm_init_mode, - moe_config=moe_config, - ) - self.enc_pos_embed = cuboid_decoder.PosEmbed( - embed_dim=base_units, typ=pos_embed_type, maxH=H_in, maxW=W_in, maxT=T_in - ) - mem_shapes = self.encoder.get_mem_shapes() - self.z_proj = nn.Linear( - in_features=mem_shapes[-1][-1], out_features=mem_shapes[-1][-1] - ) - self.dec_pos_embed = cuboid_decoder.PosEmbed( - embed_dim=mem_shapes[-1][-1], - typ=pos_embed_type, - maxT=T_out, - maxH=mem_shapes[-1][1], - maxW=mem_shapes[-1][2], - ) - self.decoder = cuboid_decoder.CuboidTransformerDecoder( - target_temporal_length=T_out, - mem_shapes=mem_shapes, - cross_start=dec_cross_start, - depth=dec_depth, - upsample_type=upsample_type, - block_self_attn_patterns=dec_self_attn_patterns, - block_self_cuboid_size=dec_self_cuboid_size, - block_self_shift_size=dec_self_shift_size, - block_self_cuboid_strategy=dec_self_cuboid_strategy, - block_cross_attn_patterns=dec_cross_attn_patterns, - block_cross_cuboid_hw=dec_cross_cuboid_hw, - block_cross_shift_hw=dec_cross_shift_hw, - block_cross_cuboid_strategy=dec_cross_cuboid_strategy, - block_cross_n_temporal=dec_cross_n_temporal, - cross_last_n_frames=dec_cross_last_n_frames, - num_heads=num_heads, - attn_drop=attn_drop, - proj_drop=proj_drop, - ffn_drop=ffn_drop, - upsample_kernel_size=upsample_kernel_size, - ffn_activation=ffn_activation, - gated_ffn=gated_ffn, - norm_layer=norm_layer, - use_inter_ffn=dec_use_inter_ffn, - max_temporal_relative=T_in + T_out, - padding_type=padding_type, - hierarchical_pos_embed=dec_hierarchical_pos_embed, - pos_embed_type=pos_embed_type, - use_self_global=num_global_vectors > 0 and use_dec_self_global, - self_update_global=dec_self_update_global, - use_cross_global=num_global_vectors > 0 and use_dec_cross_global, - use_global_vector_ffn=use_global_vector_ffn, - use_global_self_attn=use_global_self_attn, - separate_global_qkv=separate_global_qkv, - global_dim_ratio=global_dim_ratio, - checkpoint_level=checkpoint_level, - use_relative_pos=use_relative_pos, - self_attn_use_final_proj=self_attn_use_final_proj, - use_first_self_attn=dec_use_first_self_attn, - attn_linear_init_mode=attn_linear_init_mode, - ffn_linear_init_mode=ffn_linear_init_mode, - conv_init_mode=conv_init_mode, - up_linear_init_mode=down_up_linear_init_mode, - norm_init_mode=norm_init_mode, - moe_config=moe_config, - ) - - if rnc_config["use_rnc"]: - self.rnc_cri = extformer_moe_utils.RnCLoss(rnc_config) - - self.reset_parameters() - - def get_initial_encoder_final_decoder( - self, - initial_downsample_type, - activation, - initial_downsample_scale, - initial_downsample_conv_layers, - final_upsample_conv_layers, - padding_type, - initial_downsample_stack_conv_num_layers, - initial_downsample_stack_conv_dim_list, - initial_downsample_stack_conv_downscale_list, - initial_downsample_stack_conv_num_conv_list, - ): - T_in, H_in, W_in, C_in = self.input_shape - T_out, H_out, W_out, C_out = self.target_shape - self.initial_downsample_type = initial_downsample_type - if self.initial_downsample_type == "conv": - if isinstance(initial_downsample_scale, int): - initial_downsample_scale = ( - 1, - initial_downsample_scale, - initial_downsample_scale, - ) - elif len(initial_downsample_scale) == 2: - initial_downsample_scale = 1, *initial_downsample_scale - elif len(initial_downsample_scale) == 3: - initial_downsample_scale = tuple(initial_downsample_scale) - else: - raise NotImplementedError( - f"initial_downsample_scale {initial_downsample_scale} format not supported!" - ) - self.initial_encoder = InitialEncoder( - dim=C_in, - out_dim=self.base_units, - downsample_scale=initial_downsample_scale, - num_conv_layers=initial_downsample_conv_layers, - padding_type=padding_type, - activation=activation, - conv_init_mode=self.conv_init_mode, - linear_init_mode=self.down_up_linear_init_mode, - norm_init_mode=self.norm_init_mode, - ) - - self.final_decoder = FinalDecoder( - dim=self.base_units, - target_thw=(T_out, H_out, W_out), - num_conv_layers=final_upsample_conv_layers, - activation=activation, - conv_init_mode=self.conv_init_mode, - linear_init_mode=self.down_up_linear_init_mode, - norm_init_mode=self.norm_init_mode, - ) - new_input_shape = self.initial_encoder.patch_merge.get_out_shape( - self.input_shape - ) - self.dec_final_proj = nn.Linear( - in_features=self.base_units, out_features=C_out - ) - elif self.initial_downsample_type == "stack_conv": - if initial_downsample_stack_conv_dim_list is None: - initial_downsample_stack_conv_dim_list = [ - self.base_units - ] * initial_downsample_stack_conv_num_layers - self.initial_encoder = InitialStackPatchMergingEncoder( - num_merge=initial_downsample_stack_conv_num_layers, - in_dim=C_in, - out_dim_list=initial_downsample_stack_conv_dim_list, - downsample_scale_list=initial_downsample_stack_conv_downscale_list, - num_conv_per_merge_list=initial_downsample_stack_conv_num_conv_list, - padding_type=padding_type, - activation=activation, - conv_init_mode=self.conv_init_mode, - linear_init_mode=self.down_up_linear_init_mode, - norm_init_mode=self.norm_init_mode, - ) - initial_encoder_out_shape_list = self.initial_encoder.get_out_shape_list( - self.target_shape - ) - ( - dec_target_shape_list, - dec_in_dim, - ) = FinalStackUpsamplingDecoder.get_init_params( - enc_input_shape=self.target_shape, - enc_out_shape_list=initial_encoder_out_shape_list, - large_channel=True, - ) - self.final_decoder = FinalStackUpsamplingDecoder( - target_shape_list=dec_target_shape_list, - in_dim=dec_in_dim, - num_conv_per_up_list=initial_downsample_stack_conv_num_conv_list[::-1], - activation=activation, - conv_init_mode=self.conv_init_mode, - linear_init_mode=self.down_up_linear_init_mode, - norm_init_mode=self.norm_init_mode, - ) - self.dec_final_proj = nn.Linear( - in_features=dec_target_shape_list[-1][-1], out_features=C_out - ) - new_input_shape = self.initial_encoder.get_out_shape_list(self.input_shape)[ - -1 - ] - else: - raise NotImplementedError(f"{self.initial_downsample_type} is invalid.") - self.input_shape_after_initial_downsample = new_input_shape - T_in, H_in, W_in, _ = new_input_shape - return new_input_shape - - def reset_parameters(self): - if self.num_global_vectors > 0: - self.init_global_vectors = initializer.trunc_normal_( - self.init_global_vectors, std=0.02 - ) - if hasattr(self.initial_encoder, "reset_parameters"): - self.initial_encoder.reset_parameters() - else: - cuboid_utils.apply_initialization( - self.initial_encoder, - conv_mode=self.conv_init_mode, - linear_mode=self.down_up_linear_init_mode, - norm_mode=self.norm_init_mode, - ) - if hasattr(self.final_decoder, "reset_parameters"): - self.final_decoder.reset_parameters() - else: - cuboid_utils.apply_initialization( - self.final_decoder, - conv_mode=self.conv_init_mode, - linear_mode=self.down_up_linear_init_mode, - norm_mode=self.norm_init_mode, - ) - cuboid_utils.apply_initialization( - self.dec_final_proj, linear_mode=self.down_up_linear_init_mode - ) - self.encoder.reset_parameters() - self.enc_pos_embed.reset_parameters() - self.decoder.reset_parameters() - self.dec_pos_embed.reset_parameters() - cuboid_utils.apply_initialization(self.z_proj, linear_mode="0") - - def get_initial_z(self, final_mem, T_out): - B = final_mem.shape[0] - if self.z_init_method == "zeros": - z_shape = list((1, T_out)) + final_mem.shape[2:] - initial_z = paddle.zeros(shape=z_shape, dtype=final_mem.dtype) - initial_z = self.z_proj(self.dec_pos_embed(initial_z)).expand( - shape=[B, -1, -1, -1, -1] - ) - elif self.z_init_method == "nearest_interp": - initial_z = nn.functional.interpolate( - x=final_mem.transpose(perm=[0, 4, 1, 2, 3]), - size=(T_out, final_mem.shape[2], final_mem.shape[3]), - ).transpose(perm=[0, 2, 3, 4, 1]) - initial_z = self.z_proj(initial_z) - elif self.z_init_method == "last": - initial_z = paddle.broadcast_to( - x=final_mem[:, -1:, :, :, :], shape=(B, T_out) + final_mem.shape[2:] - ) - initial_z = self.z_proj(initial_z) - elif self.z_init_method == "mean": - initial_z = paddle.broadcast_to( - x=final_mem.mean(axis=1, keepdims=True), - shape=(B, T_out) + final_mem.shape[2:], - ) - initial_z = self.z_proj(initial_z) - else: - raise NotImplementedError - return initial_z - - def forward(self, x: "paddle.Tensor", verbose: bool = False) -> "paddle.Tensor": - """ - Args: - x (paddle.Tensor): Tensor with shape (B, T, H, W, C). - verbose (bool): if True, print intermediate shapes. - - Returns: - out (paddle.Tensor): The output Shape (B, T_out, H, W, C_out) - """ - - labels = x["sst_target"] - x = self.concat_to_tensor(x, self.input_keys) - flag_ndim = x.ndim - if flag_ndim == 6: - x = x.reshape([-1, *x.shape[2:]]) - B, _, _, _, _ = x.shape - - T_out = self.target_shape[0] - x = self.initial_encoder(x) - x = self.enc_pos_embed(x) - - if self.num_global_vectors > 0: - init_global_vectors = self.init_global_vectors.expand( - shape=[ - B, - self.num_global_vectors, - self.global_dim_ratio * self.base_units, - ] - ) - mem_l, mem_global_vector_l = self.encoder(x, init_global_vectors) - else: - mem_l = self.encoder(x) - - if verbose: - for i, mem in enumerate(mem_l): - print(f"mem[{i}].shape = {mem.shape}") - initial_z = self.get_initial_z(final_mem=mem_l[-1], T_out=T_out) - - if self.num_global_vectors > 0: - dec_out = self.decoder(initial_z, mem_l, mem_global_vector_l) - else: - dec_out = self.decoder(initial_z, mem_l) - - dec_out = self.final_decoder(dec_out) - out = self.dec_final_proj(dec_out) - - if flag_ndim == 6: - out = out.reshape([-1, *out.shape]) - - out_dict = {key: out for key in self.output_keys[:2]} - - # moe loss - if self.training: - aux_losses = extformer_moe_utils.aggregate_aux_losses(self) - if len(aux_losses) > 0: - aux_loss = paddle.concat(aux_losses).mean() - else: - aux_loss = None - else: - aux_loss = None - assert "aux_loss" in self.output_keys - out_dict["aux_loss"] = aux_loss - - # rnc - if self.training and self.rnc_config["use_rnc"]: - rank_loss = self.rnc_cri(dec_out, labels) - rank_loss = rank_loss.unsqueeze(0) - else: - rank_loss = None - assert "rank_loss" in self.output_keys - out_dict["rank_loss"] = rank_loss - - return out_dict diff --git a/examples/smc_reac/ppsci/arch/extformer_moe_cuboid_decoder.py b/examples/smc_reac/ppsci/arch/extformer_moe_cuboid_decoder.py deleted file mode 100644 index b16311a0e7..0000000000 --- a/examples/smc_reac/ppsci/arch/extformer_moe_cuboid_decoder.py +++ /dev/null @@ -1,1475 +0,0 @@ -from functools import lru_cache -from typing import Tuple - -import numpy as np -import paddle -import paddle.nn.functional as F -from paddle import nn -from paddle.distributed import fleet - -import ppsci.arch.extformer_moe_cuboid_encoder as cuboid_encoder -import ppsci.arch.extformer_moe_cuboid_utils as cuboid_utils -import ppsci.arch.extformer_moe_utils as moe_utils -from ppsci.utils import initializer - - -class PosEmbed(nn.Layer): - """pose embedding - - Args: - embed_dim (int): The dimension of embedding. - maxT (int): The embedding max time. - maxH (int): The embedding max height. - maxW (int): The embedding max width. - typ (str): - The type of the positional embedding. - - t+h+w: - Embed the spatial position to embeddings - - t+hw: - Embed the spatial position to embeddings - """ - - def __init__( - self, - embed_dim, - maxT, - maxH, - maxW, - typ: str = "t+h+w", - moe_config: dict = None, - ): - super(PosEmbed, self).__init__() - self.typ = typ - assert self.typ in ["t+h+w", "t+hw"] - self.maxT = maxT - self.maxH = maxH - self.maxW = maxW - self.embed_dim = embed_dim - if self.typ == "t+h+w": - self.T_embed = nn.Embedding(num_embeddings=maxT, embedding_dim=embed_dim) - self.H_embed = nn.Embedding(num_embeddings=maxH, embedding_dim=embed_dim) - self.W_embed = nn.Embedding(num_embeddings=maxW, embedding_dim=embed_dim) - elif self.typ == "t+hw": - self.T_embed = nn.Embedding(num_embeddings=maxT, embedding_dim=embed_dim) - self.HW_embed = nn.Embedding( - num_embeddings=maxH * maxW, embedding_dim=embed_dim - ) - else: - raise NotImplementedError(f"{self.typ} is invalid.") - self.reset_parameters() - - def reset_parameters(self): - for m in self.children(): - cuboid_utils.apply_initialization(m, embed_mode="0") - - def forward(self, x): - """ - Args: - x : Shape (B, T, H, W, C) - - Returns: - out : the x + positional embeddings - """ - - _, T, H, W, _ = x.shape - t_idx = paddle.arange(end=T) - h_idx = paddle.arange(end=H) - w_idx = paddle.arange(end=W) - if self.typ == "t+h+w": - return ( - x - + self.T_embed(t_idx).reshape([T, 1, 1, self.embed_dim]) - + self.H_embed(h_idx).reshape([1, H, 1, self.embed_dim]) - + self.W_embed(w_idx).reshape([1, 1, W, self.embed_dim]) - ) - elif self.typ == "t+hw": - spatial_idx = h_idx.unsqueeze(axis=-1) * self.maxW + w_idx - return ( - x - + self.T_embed(t_idx).reshape([T, 1, 1, self.embed_dim]) - + self.HW_embed(spatial_idx) - ) - else: - raise NotImplementedError(f"{self.typ} is invalid.") - - -@lru_cache() -def compute_cuboid_cross_attention_mask( - T_x, T_mem, H, W, n_temporal, cuboid_hw, shift_hw, strategy, padding_type, device -): - pad_t_mem = (n_temporal - T_mem % n_temporal) % n_temporal - pad_t_x = (n_temporal - T_x % n_temporal) % n_temporal - pad_h = (cuboid_hw[0] - H % cuboid_hw[0]) % cuboid_hw[0] - pad_w = (cuboid_hw[1] - W % cuboid_hw[1]) % cuboid_hw[1] - mem_cuboid_size = ((T_mem + pad_t_mem) // n_temporal,) + cuboid_hw - x_cuboid_size = ((T_x + pad_t_x) // n_temporal,) + cuboid_hw - if pad_t_mem > 0 or pad_h > 0 or pad_w > 0: - if padding_type == "ignore": - mem_mask = paddle.ones(shape=(1, T_mem, H, W, 1), dtype="bool") - mem_mask = F.pad( - mem_mask, [0, 0, 0, pad_w, 0, pad_h, pad_t_mem, 0], data_format="NDHWC" - ) - else: - mem_mask = paddle.ones( - shape=(1, T_mem + pad_t_mem, H + pad_h, W + pad_w, 1), dtype="bool" - ) - if pad_t_x > 0 or pad_h > 0 or pad_w > 0: - if padding_type == "ignore": - x_mask = paddle.ones(shape=(1, T_x, H, W, 1), dtype="bool") - x_mask = F.pad( - x_mask, [0, 0, 0, pad_w, 0, pad_h, 0, pad_t_x], data_format="NDHWC" - ) - else: - x_mask = paddle.ones( - shape=(1, T_x + pad_t_x, H + pad_h, W + pad_w, 1), dtype="bool" - ) - if any(i > 0 for i in shift_hw): - if padding_type == "ignore": - x_mask = paddle.roll( - x=x_mask, shifts=(-shift_hw[0], -shift_hw[1]), axis=(2, 3) - ) - mem_mask = paddle.roll( - x=mem_mask, shifts=(-shift_hw[0], -shift_hw[1]), axis=(2, 3) - ) - x_mask = cuboid_encoder.cuboid_reorder(x_mask, x_cuboid_size, strategy=strategy) - x_mask = x_mask.squeeze(axis=-1).squeeze(axis=0) - num_cuboids, x_cuboid_volume = x_mask.shape - mem_mask = cuboid_encoder.cuboid_reorder( - mem_mask, mem_cuboid_size, strategy=strategy - ) - mem_mask = mem_mask.squeeze(axis=-1).squeeze(axis=0) - _, mem_cuboid_volume = mem_mask.shape - shift_mask = np.zeros(shape=(1, n_temporal, H + pad_h, W + pad_w, 1)) - cnt = 0 - for h in ( - slice(-cuboid_hw[0]), - slice(-cuboid_hw[0], -shift_hw[0]), - slice(-shift_hw[0], None), - ): - for w in ( - slice(-cuboid_hw[1]), - slice(-cuboid_hw[1], -shift_hw[1]), - slice(-shift_hw[1], None), - ): - shift_mask[:, :, h, w, :] = cnt - cnt += 1 - shift_mask = paddle.to_tensor(shift_mask) - shift_mask = cuboid_encoder.cuboid_reorder( - shift_mask, (1,) + cuboid_hw, strategy=strategy - ) - shift_mask = shift_mask.squeeze(axis=-1).squeeze(axis=0) - shift_mask = shift_mask.unsqueeze(axis=1) - shift_mask.unsqueeze(axis=2) == 0 - bh_bw = cuboid_hw[0] * cuboid_hw[1] - attn_mask = ( - shift_mask.reshape((num_cuboids, 1, bh_bw, 1, bh_bw)) - * x_mask.reshape((num_cuboids, -1, bh_bw, 1, 1)) - * mem_mask.reshape([num_cuboids, 1, 1, -1, bh_bw]) - ) - attn_mask = attn_mask.reshape([num_cuboids, x_cuboid_volume, mem_cuboid_volume]) - return attn_mask - - -class CuboidCrossAttentionLayer(nn.Layer): - """Implements the cuboid cross attention. - - The idea of Cuboid Cross Attention is to extend the idea of cuboid self attention to work for the - encoder-decoder-type cross attention. - - Assume that there is a memory tensor with shape (T1, H, W, C) and another query tensor with shape (T2, H, W, C), - - Here, we decompose the query tensor and the memory tensor into the same number of cuboids and attend the cuboid in - the query tensor with the corresponding cuboid in the memory tensor. - - For the height and width axes, we reuse the grid decomposition techniques described in the cuboid self-attention. - For the temporal axis, the layer supports the "n_temporal" parameter, that controls the number of cuboids we can - get after cutting the tensors. For example, if the temporal dilation is 2, both the query and - memory will be decomposed into 2 cuboids along the temporal axis. Like in the Cuboid Self-attention, - we support "local" and "dilated" decomposition strategy. - - The complexity of the layer is O((T2 / n_t * Bh * Bw) * (T1 / n_t * Bh * Bw) * n_t (H / Bh) (W / Bw)) = O(T2 * T1 / n_t H W Bh Bw) - - Args: - dim (int): The dimension of input tensor. - num_heads (int): The number of head. - n_temporal (int, optional): The num of temporal. Defaults to 1. - cuboid_hw (tuple, optional): The height and width of cuboid. Defaults to (7, 7). - shift_hw (tuple, optional): The height and width of shift. Defaults to (0, 0). - strategy (tuple, optional): The strategy. Defaults to ("d", "l", "l"). - padding_type (str, optional): The type of padding. Defaults to "ignore". - cross_last_n_frames (int, optional): The cross_last_n_frames of decoder. Defaults to None. - qkv_bias (bool, optional): Whether to enable bias in calculating qkv attention. Defaults to False. - qk_scale (float, optional): Whether to enable scale factor when calculating the attention. Defaults to None. - attn_drop (float, optional): The attention dropout. Defaults to 0.0. - proj_drop (float, optional): The projrction dropout. Defaults to 0.0. - max_temporal_relative (int, optional): The max temporal. Defaults to 50. - norm_layer (str, optional): The normalization layer. Defaults to "layer_norm". - use_global_vector (bool, optional): Whether to use the global vector or not. Defaults to True. - separate_global_qkv (bool, optional): Whether to use different network to calc q_global, k_global, v_global. Defaults to False. - global_dim_ratio (int, optional): The dim (channels) of global vectors is `global_dim_ratio*dim`. Defaults to 1. - checkpoint_level (int, optional): Whether to enable gradient checkpointing. Defaults to 1. - use_relative_pos (bool, optional): Whether to use relative pos. Defaults to True. - attn_linear_init_mode (str, optional): The mode of attention linear initialization. Defaults to "0". - ffn_linear_init_mode (str, optional): The mode of FFN linear initialization. Defaults to "0". - norm_init_mode (str, optional): The mode of normalization initialization. Defaults to "0". - """ - - def __init__( - self, - dim: int, - num_heads: int, - n_temporal: int = 1, - cuboid_hw: Tuple[int, ...] = (7, 7), - shift_hw: Tuple[int, ...] = (0, 0), - strategy: Tuple[str, ...] = ("d", "l", "l"), - padding_type: str = "ignore", - cross_last_n_frames: int = None, - qkv_bias: bool = False, - qk_scale: float = None, - attn_drop: float = 0.0, - proj_drop: float = 0.0, - max_temporal_relative: int = 50, - norm_layer: str = "layer_norm", - use_global_vector: bool = True, - separate_global_qkv: bool = False, - global_dim_ratio: int = 1, - checkpoint_level: int = 1, - use_relative_pos: bool = True, - attn_linear_init_mode: str = "0", - ffn_linear_init_mode: str = "0", - norm_init_mode: str = "0", - moe_config: dict = None, - ): - super(CuboidCrossAttentionLayer, self).__init__() - self.attn_linear_init_mode = attn_linear_init_mode - self.ffn_linear_init_mode = ffn_linear_init_mode - self.norm_init_mode = norm_init_mode - self.dim = dim - self.num_heads = num_heads - self.n_temporal = n_temporal - assert n_temporal > 0 - head_dim = dim // num_heads - self.scale = qk_scale or head_dim**-0.5 - shift_hw = list(shift_hw) - if strategy[1] == "d": - shift_hw[0] = 0 - if strategy[2] == "d": - shift_hw[1] = 0 - self.cuboid_hw = cuboid_hw - self.shift_hw = tuple(shift_hw) - self.strategy = strategy - self.padding_type = padding_type - self.max_temporal_relative = max_temporal_relative - self.cross_last_n_frames = cross_last_n_frames - self.use_relative_pos = use_relative_pos - self.use_global_vector = use_global_vector - self.separate_global_qkv = separate_global_qkv - if global_dim_ratio != 1 and separate_global_qkv is False: - raise ValueError( - "Setting global_dim_ratio != 1 requires separate_global_qkv == True." - ) - self.global_dim_ratio = global_dim_ratio - if self.padding_type not in ["ignore", "zeros", "nearest"]: - raise ValueError('padding_type should be ["ignore", "zeros", "nearest"]') - if use_relative_pos: - init_data = paddle.zeros( - ( - (2 * max_temporal_relative - 1) - * (2 * cuboid_hw[0] - 1) - * (2 * cuboid_hw[1] - 1), - num_heads, - ) - ) - self.relative_position_bias_table = paddle.create_parameter( - shape=init_data.shape, - dtype=init_data.dtype, - default_initializer=nn.initializer.Constant(0.0), - ) - self.relative_position_bias_table.stop_gradient = not True - self.relative_position_bias_table = initializer.trunc_normal_( - self.relative_position_bias_table, std=0.02 - ) - - coords_t = paddle.arange(end=max_temporal_relative) - coords_h = paddle.arange(end=self.cuboid_hw[0]) - coords_w = paddle.arange(end=self.cuboid_hw[1]) - coords = paddle.stack(x=paddle.meshgrid(coords_t, coords_h, coords_w)) - coords_flatten = paddle.flatten(x=coords, start_axis=1) - relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] - relative_coords = relative_coords.transpose(perm=[1, 2, 0]) - relative_coords[:, :, 0] += max_temporal_relative - 1 - relative_coords[:, :, 1] += self.cuboid_hw[0] - 1 - relative_coords[:, :, 2] += self.cuboid_hw[1] - 1 - relative_position_index = ( - relative_coords[:, :, 0] - * (2 * self.cuboid_hw[0] - 1) - * (2 * self.cuboid_hw[1] - 1) - + relative_coords[:, :, 1] * (2 * self.cuboid_hw[1] - 1) - + relative_coords[:, :, 2] - ) - self.register_buffer( - name="relative_position_index", tensor=relative_position_index - ) - self.q_proj = nn.Linear(in_features=dim, out_features=dim, bias_attr=qkv_bias) - self.kv_proj = nn.Linear( - in_features=dim, out_features=dim * 2, bias_attr=qkv_bias - ) - self.attn_drop = nn.Dropout(p=attn_drop) - self.proj = nn.Linear(in_features=dim, out_features=dim) - self.proj_drop = nn.Dropout(p=proj_drop) - if self.use_global_vector: - if self.separate_global_qkv: - self.l2g_q_net = nn.Linear( - in_features=dim, out_features=dim, bias_attr=qkv_bias - ) - self.l2g_global_kv_net = nn.Linear( - in_features=global_dim_ratio * dim, - out_features=dim * 2, - bias_attr=qkv_bias, - ) - self.norm = cuboid_utils.get_norm_layer(norm_layer, in_channels=dim) - self._checkpoint_level = checkpoint_level - self.reset_parameters() - - def reset_parameters(self): - cuboid_utils.apply_initialization( - self.q_proj, linear_mode=self.attn_linear_init_mode - ) - cuboid_utils.apply_initialization( - self.kv_proj, linear_mode=self.attn_linear_init_mode - ) - cuboid_utils.apply_initialization( - self.proj, linear_mode=self.ffn_linear_init_mode - ) - cuboid_utils.apply_initialization(self.norm, norm_mode=self.norm_init_mode) - if self.use_global_vector: - if self.separate_global_qkv: - cuboid_utils.apply_initialization( - self.l2g_q_net, linear_mode=self.attn_linear_init_mode - ) - cuboid_utils.apply_initialization( - self.l2g_global_kv_net, linear_mode=self.attn_linear_init_mode - ) - - def forward(self, x, mem, mem_global_vectors=None): - """Calculate the forward - - Along the temporal axis, we pad the mem tensor from the left and the x tensor from the right so that the - relative position encoding can be calculated correctly. For example: - - mem: 0, 1, 2, 3, 4 - x: 0, 1, 2, 3, 4, 5 - - n_temporal = 1 - mem: 0, 1, 2, 3, 4 x: 0, 1, 2, 3, 4, 5 - - n_temporal = 2 - mem: pad, 1, 3 x: 0, 2, 4 - mem: 0, 2, 4 x: 1, 3, 5 - - n_temporal = 3 - mem: pad, 2 dec: 0, 3 - mem: 0, 3 dec: 1, 4 - mem: 1, 4 dec: 2, 5 - - Args: - x (paddle.Tensor): The input of the layer. It will have shape (B, T, H, W, C) - mem (paddle.Tensor): The memory. It will have shape (B, T_mem, H, W, C) - mem_global_vectors (paddle.Tensor): The global vectors from the memory. It will have shape (B, N, C) - - Returns: - out (paddle.Tensor): Output tensor should have shape (B, T, H, W, C_out) - """ - - if self.cross_last_n_frames is not None: - cross_last_n_frames = int(min(self.cross_last_n_frames, mem.shape[1])) - mem = mem[:, -cross_last_n_frames:, ...] - if self.use_global_vector: - _, num_global, _ = mem_global_vectors.shape - x = self.norm(x) - B, T_x, H, W, C_in = x.shape - B_mem, T_mem, H_mem, W_mem, C_mem = mem.shape - assert T_x < self.max_temporal_relative and T_mem < self.max_temporal_relative - cuboid_hw = self.cuboid_hw - n_temporal = self.n_temporal - shift_hw = self.shift_hw - assert ( - B_mem == B and H == H_mem and W == W_mem and C_in == C_mem - ), f"Shape of memory and the input tensor does not match. x.shape={x.shape}, mem.shape={mem.shape}" - pad_t_mem = (n_temporal - T_mem % n_temporal) % n_temporal - pad_t_x = (n_temporal - T_x % n_temporal) % n_temporal - pad_h = (cuboid_hw[0] - H % cuboid_hw[0]) % cuboid_hw[0] - pad_w = (cuboid_hw[1] - W % cuboid_hw[1]) % cuboid_hw[1] - mem = cuboid_utils.generalize_padding( - mem, pad_t_mem, pad_h, pad_w, self.padding_type, t_pad_left=True - ) - - x = cuboid_utils.generalize_padding( - x, pad_t_x, pad_h, pad_w, self.padding_type, t_pad_left=False - ) - - if any(i > 0 for i in shift_hw): - shifted_x = paddle.roll( - x=x, shifts=(-shift_hw[0], -shift_hw[1]), axis=(2, 3) - ) - shifted_mem = paddle.roll( - x=mem, shifts=(-shift_hw[0], -shift_hw[1]), axis=(2, 3) - ) - else: - shifted_x = x - shifted_mem = mem - mem_cuboid_size = (mem.shape[1] // n_temporal,) + cuboid_hw - x_cuboid_size = (x.shape[1] // n_temporal,) + cuboid_hw - reordered_mem = cuboid_encoder.cuboid_reorder( - shifted_mem, cuboid_size=mem_cuboid_size, strategy=self.strategy - ) - reordered_x = cuboid_encoder.cuboid_reorder( - shifted_x, cuboid_size=x_cuboid_size, strategy=self.strategy - ) - _, num_cuboids_mem, mem_cuboid_volume, _ = reordered_mem.shape - _, num_cuboids, x_cuboid_volume, _ = reordered_x.shape - assert ( - num_cuboids_mem == num_cuboids - ), f"Number of cuboids do not match. num_cuboids={num_cuboids}, num_cuboids_mem={num_cuboids_mem}" - attn_mask = compute_cuboid_cross_attention_mask( - T_x, - T_mem, - H, - W, - n_temporal, - cuboid_hw, - shift_hw, - strategy=self.strategy, - padding_type=self.padding_type, - device=x.place, - ) - head_C = C_in // self.num_heads - kv = ( - self.kv_proj(reordered_mem) - .reshape([B, num_cuboids, mem_cuboid_volume, 2, self.num_heads, head_C]) - .transpose(perm=[3, 0, 4, 1, 2, 5]) - ) - k, v = kv[0], kv[1] - q = ( - self.q_proj(reordered_x) - .reshape([B, num_cuboids, x_cuboid_volume, self.num_heads, head_C]) - .transpose(perm=[0, 3, 1, 2, 4]) - ) - q = q * self.scale - perm_4 = list(range(k.ndim)) - perm_4[-2] = -1 - perm_4[-1] = -2 - attn_score = q @ k.transpose(perm=perm_4) - if self.use_relative_pos: - relative_position_bias = self.relative_position_bias_table[ - self.relative_position_index[ - :x_cuboid_volume, :mem_cuboid_volume - ].reshape([-1]) - ].reshape([x_cuboid_volume, mem_cuboid_volume, -1]) - relative_position_bias = relative_position_bias.transpose( - perm=[2, 0, 1] - ).unsqueeze(axis=1) - attn_score = attn_score + relative_position_bias - if self.use_global_vector: - if self.separate_global_qkv: - l2g_q = ( - self.l2g_q_net(reordered_x) - .reshape([B, num_cuboids, x_cuboid_volume, self.num_heads, head_C]) - .transpose(perm=[0, 3, 1, 2, 4]) - ) - l2g_q = l2g_q * self.scale - l2g_global_kv = ( - self.l2g_global_kv_net(mem_global_vectors) - .reshape([B, 1, num_global, 2, self.num_heads, head_C]) - .transpose(perm=[3, 0, 4, 1, 2, 5]) - ) - l2g_global_k, l2g_global_v = l2g_global_kv[0], l2g_global_kv[1] - else: - kv_global = ( - self.kv_proj(mem_global_vectors) - .reshape([B, 1, num_global, 2, self.num_heads, head_C]) - .transpose(perm=[3, 0, 4, 1, 2, 5]) - ) - l2g_global_k, l2g_global_v = kv_global[0], kv_global[1] - l2g_q = q - perm_5 = list(range(l2g_global_k.ndim)) - perm_5[-2] = -1 - perm_5[-1] = -2 - l2g_attn_score = l2g_q @ l2g_global_k.transpose(perm=perm_5) - attn_score_l2l_l2g = paddle.concat(x=(attn_score, l2g_attn_score), axis=-1) - if attn_mask.ndim == 5: - attn_mask_l2l_l2g = F.pad( - attn_mask, [0, num_global], "constant", 1, data_format="NDHWC" - ) - else: - attn_mask_l2l_l2g = F.pad(attn_mask, [0, num_global], "constant", 1) - v_l_g = paddle.concat( - x=( - v, - l2g_global_v.expand( - shape=[B, self.num_heads, num_cuboids, num_global, head_C] - ), - ), - axis=3, - ) - attn_score_l2l_l2g = cuboid_encoder.masked_softmax( - attn_score_l2l_l2g, mask=attn_mask_l2l_l2g - ) - attn_score_l2l_l2g = self.attn_drop(attn_score_l2l_l2g) - reordered_x = ( - (attn_score_l2l_l2g @ v_l_g) - .transpose(perm=[0, 2, 3, 1, 4]) - .reshape(B, num_cuboids, x_cuboid_volume, self.dim) - ) - else: - attn_score = cuboid_encoder.masked_softmax(attn_score, mask=attn_mask) - attn_score = self.attn_drop(attn_score) - reordered_x = ( - (attn_score @ v) - .transpose(perm=[0, 2, 3, 1, 4]) - .reshape([B, num_cuboids, x_cuboid_volume, self.dim]) - ) - reordered_x = paddle.cast(reordered_x, dtype="float32") - reordered_x = self.proj_drop(self.proj(reordered_x)) - shifted_x = cuboid_encoder.cuboid_reorder_reverse( - reordered_x, - cuboid_size=x_cuboid_size, - strategy=self.strategy, - orig_data_shape=(x.shape[1], x.shape[2], x.shape[3]), - ) - if any(i > 0 for i in shift_hw): - x = paddle.roll(x=shifted_x, shifts=(shift_hw[0], shift_hw[1]), axis=(2, 3)) - else: - x = shifted_x - x = cuboid_utils.generalize_unpadding( - x, pad_t=pad_t_x, pad_h=pad_h, pad_w=pad_w, padding_type=self.padding_type - ) - return x - - -class StackCuboidCrossAttentionBlock(nn.Layer): - """A stack of cuboid cross attention layers. - - The advantage of cuboid attention is that we can combine cuboid attention building blocks with different - hyper-parameters to mimic a broad range of space-time correlation patterns. - - - "use_inter_ffn" is True - x, mem --> attn1 -----+-------> ffn1 ---+---> attn2 --> ... --> ffn_k --> out - | ^ | ^ - | | | | - |-------------|----|-------------| - - "use_inter_ffn" is False - x, mem --> attn1 -----+------> attn2 --> ... attnk --+----> ffnk ---+---> out, mem - | ^ | ^ ^ | ^ - | | | | | | | - |-------------|----|------------|-- ----------|--|-----------| - - Args: - dim (int): The dimension of the input. - num_heads (int): The number of head. - block_cuboid_hw (list, optional): The height and width of block cuboid.Defaults to [(4, 4), (4, 4)]. - block_shift_hw (list, optional): The height and width of shift cuboid . Defaults to [(0, 0), (2, 2)]. - block_n_temporal (list, optional): The length of block temporal. Defaults to [1, 2]. - block_strategy (list, optional): The strategy of block. Defaults to [("d", "d", "d"), ("l", "l", "l")]. - padding_type (str, optional): The type of paddling. Defaults to "ignore". - cross_last_n_frames (int, optional): The num of cross_last_n_frames. Defaults to None. - qkv_bias (bool, optional): Whether to enable bias in calculating qkv attention. Defaults to False. - qk_scale (float, optional): Whether to enable scale factor when calculating the attention. Defaults to None. - attn_drop (float, optional): The attention dropout. Defaults to 0.0. - proj_drop (float, optional): The projection dropout. Defaults to 0.0. - ffn_drop (float, optional): The ratio of FFN dropout. Defaults to 0.0. - activation (str, optional): The activation. Defaults to "leaky". - gated_ffn (bool, optional): Whether to use gate FFN. Defaults to False. - norm_layer (str, optional): The normalization layer. Defaults to "layer_norm". - use_inter_ffn (bool, optional): Whether to use inter FFN. Defaults to True. - max_temporal_relative (int, optional): The max temporal. Defaults to 50. - checkpoint_level (int, optional): Whether to enable gradient checkpointing. Defaults to 1. - use_relative_pos (bool, optional): Whether to use relative pos. Defaults to True. - use_global_vector (bool, optional): Whether to use the global vector or not. Defaults to False. - separate_global_qkv (bool, optional): Whether to use different network to calc q_global, k_global, v_global. Defaults to False. - global_dim_ratio (int, optional): The dim (channels) of global vectors is `global_dim_ratio*dim`. Defaults to 1. - attn_linear_init_mode (str, optional): The mode of attention linear initialization. Defaults to "0". - ffn_linear_init_mode (str, optional): The mode of FFN linear initialization. Defaults to "0". - norm_init_mode (str, optional): The mode of normalization. Defaults to "0". - """ - - def __init__( - self, - dim: int, - num_heads: int, - block_cuboid_hw: Tuple[Tuple[int, ...], ...] = [(4, 4), (4, 4)], - block_shift_hw: Tuple[Tuple[int, ...], ...] = [(0, 0), (2, 2)], - block_n_temporal: Tuple[int, ...] = [1, 2], - block_strategy: Tuple[Tuple[str, ...], ...] = [ - ("d", "d", "d"), - ("l", "l", "l"), - ], - padding_type: str = "ignore", - cross_last_n_frames: int = None, - qkv_bias: bool = False, - qk_scale: float = None, - attn_drop: float = 0.0, - proj_drop: float = 0.0, - ffn_drop: float = 0.0, - activation: str = "leaky", - gated_ffn: bool = False, - norm_layer: str = "layer_norm", - use_inter_ffn: bool = True, - max_temporal_relative: int = 50, - checkpoint_level: int = 1, - use_relative_pos: bool = True, - use_global_vector: bool = False, - separate_global_qkv: bool = False, - global_dim_ratio: int = 1, - attn_linear_init_mode: str = "0", - ffn_linear_init_mode: str = "0", - norm_init_mode: str = "0", - moe_config: dict = None, - expert_shape: tuple = None, - ): - super(StackCuboidCrossAttentionBlock, self).__init__() - self.attn_linear_init_mode = attn_linear_init_mode - self.ffn_linear_init_mode = ffn_linear_init_mode - self.norm_init_mode = norm_init_mode - if ( - len(block_cuboid_hw[0]) <= 0 - or len(block_shift_hw) <= 0 - or len(block_strategy) <= 0 - ): - raise ValueError( - "Incorrect format.The lengths of block_cuboid_hw[0], block_shift_hw, and block_strategy must be greater than zero." - ) - if len(block_cuboid_hw) != len(block_shift_hw) and len(block_shift_hw) == len( - block_strategy - ): - raise ValueError( - "The lengths of block_cuboid_size, block_shift_size, and block_strategy must be equal." - ) - - self.num_attn = len(block_cuboid_hw) - self.checkpoint_level = checkpoint_level - self.use_inter_ffn = use_inter_ffn - self.use_global_vector = use_global_vector - if self.use_inter_ffn: - if moe_config["use_ffn_moe"]: - self.ffn_l = nn.LayerList( - sublayers=[ - cuboid_encoder.MixtureFFN( - units=dim, - hidden_size=4 * dim, - activation_dropout=ffn_drop, - dropout=ffn_drop, - gated_proj=gated_ffn, - activation=activation, - normalization=norm_layer, - pre_norm=True, - linear_init_mode=ffn_linear_init_mode, - norm_init_mode=norm_init_mode, - expert_shape=expert_shape, - moe_config=moe_config, - ) - for _ in range(self.num_attn) - ] - ) - else: - self.ffn_l = nn.LayerList( - sublayers=[ - cuboid_encoder.PositionwiseFFN( - units=dim, - hidden_size=4 * dim, - activation_dropout=ffn_drop, - dropout=ffn_drop, - gated_proj=gated_ffn, - activation=activation, - normalization=norm_layer, - pre_norm=True, - linear_init_mode=ffn_linear_init_mode, - norm_init_mode=norm_init_mode, - expert_shape=expert_shape, - moe_config=moe_config, - ) - for _ in range(self.num_attn) - ] - ) - else: - if moe_config["use_ffn_moe"]: - self.ffn_l = nn.LayerList( - sublayers=[ - cuboid_encoder.MixtureFFN( - units=dim, - hidden_size=4 * dim, - activation_dropout=ffn_drop, - dropout=ffn_drop, - gated_proj=gated_ffn, - activation=activation, - normalization=norm_layer, - pre_norm=True, - linear_init_mode=ffn_linear_init_mode, - norm_init_mode=norm_init_mode, - expert_shape=expert_shape, - moe_config=moe_config, - ) - ] - ) - else: - self.ffn_l = nn.LayerList( - sublayers=[ - cuboid_encoder.PositionwiseFFN( - units=dim, - hidden_size=4 * dim, - activation_dropout=ffn_drop, - dropout=ffn_drop, - gated_proj=gated_ffn, - activation=activation, - normalization=norm_layer, - pre_norm=True, - linear_init_mode=ffn_linear_init_mode, - norm_init_mode=norm_init_mode, - expert_shape=expert_shape, - moe_config=moe_config, - ) - ] - ) - - if moe_config["use_attn_moe"]: - self.attn_l = nn.LayerList( - sublayers=[ - MixtureCrossAttention( - dim=dim, - num_heads=num_heads, - cuboid_hw=ele_cuboid_hw, - shift_hw=ele_shift_hw, - strategy=ele_strategy, - n_temporal=ele_n_temporal, - cross_last_n_frames=cross_last_n_frames, - padding_type=padding_type, - qkv_bias=qkv_bias, - qk_scale=qk_scale, - attn_drop=attn_drop, - proj_drop=proj_drop, - norm_layer=norm_layer, - max_temporal_relative=max_temporal_relative, - use_global_vector=use_global_vector, - separate_global_qkv=separate_global_qkv, - global_dim_ratio=global_dim_ratio, - checkpoint_level=checkpoint_level, - use_relative_pos=use_relative_pos, - attn_linear_init_mode=attn_linear_init_mode, - ffn_linear_init_mode=ffn_linear_init_mode, - norm_init_mode=norm_init_mode, - expert_shape=expert_shape, - moe_config=moe_config, - ) - for ele_cuboid_hw, ele_shift_hw, ele_strategy, ele_n_temporal in zip( - block_cuboid_hw, - block_shift_hw, - block_strategy, - block_n_temporal, - ) - ] - ) - else: - self.attn_l = nn.LayerList( - sublayers=[ - CuboidCrossAttentionLayer( - dim=dim, - num_heads=num_heads, - cuboid_hw=ele_cuboid_hw, - shift_hw=ele_shift_hw, - strategy=ele_strategy, - n_temporal=ele_n_temporal, - cross_last_n_frames=cross_last_n_frames, - padding_type=padding_type, - qkv_bias=qkv_bias, - qk_scale=qk_scale, - attn_drop=attn_drop, - proj_drop=proj_drop, - norm_layer=norm_layer, - max_temporal_relative=max_temporal_relative, - use_global_vector=use_global_vector, - separate_global_qkv=separate_global_qkv, - global_dim_ratio=global_dim_ratio, - checkpoint_level=checkpoint_level, - use_relative_pos=use_relative_pos, - attn_linear_init_mode=attn_linear_init_mode, - ffn_linear_init_mode=ffn_linear_init_mode, - norm_init_mode=norm_init_mode, - ) - for ele_cuboid_hw, ele_shift_hw, ele_strategy, ele_n_temporal in zip( - block_cuboid_hw, - block_shift_hw, - block_strategy, - block_n_temporal, - ) - ] - ) - - def reset_parameters(self): - for m in self.ffn_l: - m.reset_parameters() - for m in self.attn_l: - m.reset_parameters() - - def forward(self, x, mem, mem_global_vector=None): - """ - Args: - x (paddle.Tensor): Shape (B, T_x, H, W, C) - mem (paddle.Tensor): Shape (B, T_mem, H, W, C) - mem_global_vector (paddle.Tensor): Shape (B, N_global, C) - - Returns: - out (paddle.Tensor): (B, T_x, H, W, C_out) - """ - - if self.use_inter_ffn: - for attn, ffn in zip(self.attn_l, self.ffn_l): - if self.checkpoint_level >= 2 and self.training: - x = x + fleet.utils.recompute(attn, x, mem, mem_global_vector) - else: - x = x + attn(x, mem, mem_global_vector) - if self.checkpoint_level >= 1 and self.training: - x = fleet.utils.recompute(ffn, x) - else: - x = ffn(x) - return x - else: - for attn in self.attn_l: - if self.checkpoint_level >= 2 and self.training: - x = x + fleet.utils.recompute(attn, x, mem, mem_global_vector) - else: - x = x + attn(x, mem, mem_global_vector) - if self.checkpoint_level >= 1 and self.training: - x = fleet.utils.recompute(self.ffn_l[0], x) - else: - x = self.ffn_l[0](x) - return x - - -class Upsample3DLayer(nn.Layer): - """Upsampling based on nn.UpSampling and Conv3x3. - - If the temporal dimension remains the same: - x --> interpolation-2d (nearest) --> conv3x3(dim, out_dim) - Else: - x --> interpolation-3d (nearest) --> conv3x3x3(dim, out_dim) - - Args: - dim (int): The dimension of the input tensor. - out_dim (int): The dimension of the output tensor. - target_size (Tuple[int,...]): The size of output tensor. - temporal_upsample (bool, optional): Whether the temporal axis will go through upsampling. Defaults to False. - kernel_size (int, optional): The kernel size of the Conv2D layer. Defaults to 3. - layout (str, optional): The layout of the inputs. Defaults to "THWC". - conv_init_mode (str, optional): The mode of conv initialization. Defaults to "0". - """ - - def __init__( - self, - dim: int, - out_dim: int, - target_size: Tuple[int, ...], - temporal_upsample: bool = False, - kernel_size: int = 3, - layout: str = "THWC", - conv_init_mode: str = "0", - moe_config: dict = None, - ): - super(Upsample3DLayer, self).__init__() - self.conv_init_mode = conv_init_mode - self.target_size = target_size - self.out_dim = out_dim - self.temporal_upsample = temporal_upsample - if temporal_upsample: - self.up = nn.Upsample(size=target_size, mode="nearest") - else: - self.up = nn.Upsample(size=(target_size[1], target_size[2]), mode="nearest") - self.conv = nn.Conv2D( - in_channels=dim, - out_channels=out_dim, - kernel_size=(kernel_size, kernel_size), - padding=(kernel_size // 2, kernel_size // 2), - ) - assert layout in ["THWC", "CTHW"] - self.layout = layout - self.reset_parameters() - - def reset_parameters(self): - for m in self.children(): - cuboid_utils.apply_initialization(m, conv_mode=self.conv_init_mode) - - def forward(self, x): - """ - - Args: - x : (B, T, H, W, C) or (B, C, T, H, W) - - Returns: - out : (B, T, H_new, W_out, C_out) or (B, C, T, H_out, W_out) - """ - - if self.layout == "THWC": - B, T, H, W, C = x.shape - if self.temporal_upsample: - x = x.transpose(perm=[0, 4, 1, 2, 3]) - return self.conv(self.up(x)).transpose(perm=[0, 2, 3, 4, 1]) - else: - assert self.target_size[0] == T - x = x.reshape([B * T, H, W, C]).transpose(perm=[0, 3, 1, 2]) - x = self.up(x) - return ( - self.conv(x) - .transpose(perm=[0, 2, 3, 1]) - .reshape(list((B,) + self.target_size + (self.out_dim,))) - ) - elif self.layout == "CTHW": - B, C, T, H, W = x.shape - if self.temporal_upsample: - return self.conv(self.up(x)) - else: - assert self.output_size[0] == T - x = x.transpose(perm=[0, 2, 1, 3, 4]) - x = x.reshape([B * T, C, H, W]) - return ( - self.conv(self.up(x)) - .reshape( - [ - B, - self.target_size[0], - self.out_dim, - self.target_size[1], - self.target_size[2], - ] - ) - .transpose(perm=[0, 2, 1, 3, 4]) - ) - - -class CuboidTransformerDecoder(nn.Layer): - """Decoder of the CuboidTransformer. - - For each block, we first apply the StackCuboidSelfAttention and then apply the StackCuboidCrossAttention - - Repeat the following structure K times - - x --> StackCuboidSelfAttention --> | - |----> StackCuboidCrossAttention (If used) --> out - mem --> | - - Args: - target_temporal_length (int): The temporal length of the target. - mem_shapes (Tuple[int,...]): The mem shapes of the decoder. - cross_start (int, optional): The block to start cross attention. Defaults to 0. - depth (list, optional): The number of layers for each block. Defaults to [2, 2]. - upsample_type (str, optional): The type of upsample. Defaults to "upsample". - upsample_kernel_size (int, optional): The kernel size of upsample. Defaults to 3. - block_self_attn_patterns (str, optional): The patterns of block attention. Defaults to None. - block_self_cuboid_size (list, optional): The size of cuboid block. Defaults to [(4, 4, 4), (4, 4, 4)]. - block_self_cuboid_strategy (list, optional): The strategy of cuboid. Defaults to [("l", "l", "l"), ("d", "d", "d")]. - block_self_shift_size (list, optional): The size of shift. Defaults to [(1, 1, 1), (0, 0, 0)]. - block_cross_attn_patterns (str, optional): The patterns of cross attentions. Defaults to None. - block_cross_cuboid_hw (list, optional): The height and width of cross cuboid. Defaults to [(4, 4), (4, 4)]. - block_cross_cuboid_strategy (list, optional): The strategy of cross cuboid. Defaults to [("l", "l", "l"), ("d", "l", "l")]. - block_cross_shift_hw (list, optional): The height and width of cross shift. Defaults to [(0, 0), (0, 0)]. - block_cross_n_temporal (list, optional): The cross temporal of block. Defaults to [1, 2]. - cross_last_n_frames (int, optional): The num of cross last frames. Defaults to None. - num_heads (int, optional): The num of head. Defaults to 4. - attn_drop (float, optional): The ratio of attention dropout. Defaults to 0.0. - proj_drop (float, optional): The ratio of projection dropout. Defaults to 0.0. - ffn_drop (float, optional): The ratio of FFN dropout. Defaults to 0.0. - ffn_activation (str, optional): The activation layer of FFN. Defaults to "leaky". - gated_ffn (bool, optional): Whether to use gate FFN. Defaults to False. - norm_layer (str, optional): The normalization layer. Defaults to "layer_norm". - use_inter_ffn (bool, optional): Whether to use inter FFN. Defaults to False. - hierarchical_pos_embed (bool, optional): Whether to use hierarchical pos_embed. Defaults to False. - pos_embed_type (str, optional): The type of pos embedding. Defaults to "t+hw". - max_temporal_relative (int, optional): The max number of teemporal relative. Defaults to 50. - padding_type (str, optional): The type of padding. Defaults to "ignore". - checkpoint_level (bool, optional): Whether to enable gradient checkpointing. Defaults to True. - use_relative_pos (bool, optional): Whether to use relative pos. Defaults to True. - self_attn_use_final_proj (bool, optional): Whether to use self attention for final projection. Defaults to True. - use_first_self_attn (bool, optional): Whether to use first self attention. Defaults to False. - use_self_global (bool, optional): Whether to use self global vector. Defaults to False. - self_update_global (bool, optional): Whether to update global vector. Defaults to True. - use_cross_global (bool, optional): Whether to use cross global vector. Defaults to False. - use_global_vector_ffn (bool, optional): Whether to use FFN global vectors. Defaults to True. - use_global_self_attn (bool, optional): Whether to use global self attention. Defaults to False. - separate_global_qkv (bool, optional): Whether to use different network to calc q_global, k_global, v_global. Defaults to False. - global_dim_ratio (int, optional): The dim (channels) of global vectors is `global_dim_ratio*dim`. Defaults to 1. - attn_linear_init_mode (str, optional): The mode of attention linear initialization. Defaults to "0". - ffn_linear_init_mode (str, optional): The mode of FFN linear initialization. Defaults to "0". - conv_init_mode (str, optional): The mode of conv initialization. Defaults to "0". - up_linear_init_mode (str, optional): The mode of up linear initialization. Defaults to "0". - norm_init_mode (str, optional): The mode of normalization initialization. Defaults to "0". - """ - - def __init__( - self, - target_temporal_length: int, - mem_shapes: Tuple[int, ...], - cross_start: int = 0, - depth: Tuple[int, ...] = [2, 2], - upsample_type: str = "upsample", - upsample_kernel_size: int = 3, - block_self_attn_patterns: str = None, - block_self_cuboid_size: Tuple[Tuple[int, ...], ...] = [(4, 4, 4), (4, 4, 4)], - block_self_cuboid_strategy: Tuple[Tuple[str, ...], ...] = [ - ("l", "l", "l"), - ("d", "d", "d"), - ], - block_self_shift_size: Tuple[Tuple[int, ...], ...] = [(1, 1, 1), (0, 0, 0)], - block_cross_attn_patterns: str = None, - block_cross_cuboid_hw: Tuple[Tuple[int, ...], ...] = [(4, 4), (4, 4)], - block_cross_cuboid_strategy: Tuple[Tuple[str, ...], ...] = [ - ("l", "l", "l"), - ("d", "l", "l"), - ], - block_cross_shift_hw: Tuple[Tuple[int, ...], ...] = [(0, 0), (0, 0)], - block_cross_n_temporal: Tuple[int, ...] = [1, 2], - cross_last_n_frames: int = None, - num_heads: int = 4, - attn_drop: float = 0.0, - proj_drop: float = 0.0, - ffn_drop: float = 0.0, - ffn_activation: str = "leaky", - gated_ffn: bool = False, - norm_layer: str = "layer_norm", - use_inter_ffn: bool = False, - hierarchical_pos_embed: bool = False, - pos_embed_type: str = "t+hw", - max_temporal_relative: int = 50, - padding_type: str = "ignore", - checkpoint_level: bool = True, - use_relative_pos: bool = True, - self_attn_use_final_proj: bool = True, - use_first_self_attn: bool = False, - use_self_global: bool = False, - self_update_global: bool = True, - use_cross_global: bool = False, - use_global_vector_ffn: bool = True, - use_global_self_attn: bool = False, - separate_global_qkv: bool = False, - global_dim_ratio: int = 1, - attn_linear_init_mode: str = "0", - ffn_linear_init_mode: str = "0", - conv_init_mode: str = "0", - up_linear_init_mode: str = "0", - norm_init_mode: str = "0", - moe_config: dict = None, - ): - super(CuboidTransformerDecoder, self).__init__() - self.attn_linear_init_mode = attn_linear_init_mode - self.ffn_linear_init_mode = ffn_linear_init_mode - self.conv_init_mode = conv_init_mode - self.up_linear_init_mode = up_linear_init_mode - self.norm_init_mode = norm_init_mode - assert len(depth) == len(mem_shapes) - self.target_temporal_length = target_temporal_length - self.num_blocks = len(mem_shapes) - self.cross_start = cross_start - self.mem_shapes = mem_shapes - self.depth = depth - self.upsample_type = upsample_type - self.hierarchical_pos_embed = hierarchical_pos_embed - self.checkpoint_level = checkpoint_level - self.use_self_global = use_self_global - self.self_update_global = self_update_global - self.use_cross_global = use_cross_global - self.use_global_vector_ffn = use_global_vector_ffn - self.use_first_self_attn = use_first_self_attn - if block_self_attn_patterns is not None: - if isinstance(block_self_attn_patterns, (tuple, list)): - assert len(block_self_attn_patterns) == self.num_blocks - else: - block_self_attn_patterns = [ - block_self_attn_patterns for _ in range(self.num_blocks) - ] - block_self_cuboid_size = [] - block_self_cuboid_strategy = [] - block_self_shift_size = [] - for idx, key in enumerate(block_self_attn_patterns): - func = cuboid_utils.CuboidSelfAttentionPatterns.get(key) - cuboid_size, strategy, shift_size = func(mem_shapes[idx]) - block_self_cuboid_size.append(cuboid_size) - block_self_cuboid_strategy.append(strategy) - block_self_shift_size.append(shift_size) - else: - if not isinstance(block_self_cuboid_size[0][0], (list, tuple)): - block_self_cuboid_size = [ - block_self_cuboid_size for _ in range(self.num_blocks) - ] - else: - assert ( - len(block_self_cuboid_size) == self.num_blocks - ), f"Incorrect input format! Received block_self_cuboid_size={block_self_cuboid_size}" - if not isinstance(block_self_cuboid_strategy[0][0], (list, tuple)): - block_self_cuboid_strategy = [ - block_self_cuboid_strategy for _ in range(self.num_blocks) - ] - else: - assert ( - len(block_self_cuboid_strategy) == self.num_blocks - ), f"Incorrect input format! Received block_self_cuboid_strategy={block_self_cuboid_strategy}" - if not isinstance(block_self_shift_size[0][0], (list, tuple)): - block_self_shift_size = [ - block_self_shift_size for _ in range(self.num_blocks) - ] - else: - assert ( - len(block_self_shift_size) == self.num_blocks - ), f"Incorrect input format! Received block_self_shift_size={block_self_shift_size}" - - expert_shape_list = [ - (target_temporal_length,) + mem_shape[1:] for mem_shape in mem_shapes - ] - self_blocks = [] - for i in range(self.num_blocks): - if not self.use_first_self_attn and i == self.num_blocks - 1: - ele_depth = depth[i] - 1 - else: - ele_depth = depth[i] - stack_cuboid_blocks = [ - cuboid_encoder.StackCuboidSelfAttentionBlock( - dim=self.mem_shapes[i][-1], - num_heads=num_heads, - block_cuboid_size=block_self_cuboid_size[i], - block_strategy=block_self_cuboid_strategy[i], - block_shift_size=block_self_shift_size[i], - attn_drop=attn_drop, - proj_drop=proj_drop, - ffn_drop=ffn_drop, - activation=ffn_activation, - gated_ffn=gated_ffn, - norm_layer=norm_layer, - use_inter_ffn=use_inter_ffn, - padding_type=padding_type, - use_global_vector=use_self_global, - use_global_vector_ffn=use_global_vector_ffn, - use_global_self_attn=use_global_self_attn, - separate_global_qkv=separate_global_qkv, - global_dim_ratio=global_dim_ratio, - checkpoint_level=checkpoint_level, - use_relative_pos=use_relative_pos, - use_final_proj=self_attn_use_final_proj, - attn_linear_init_mode=attn_linear_init_mode, - ffn_linear_init_mode=ffn_linear_init_mode, - norm_init_mode=norm_init_mode, - expert_shape=expert_shape_list[i], - moe_config=moe_config, - ) - for _ in range(ele_depth) - ] - self_blocks.append(nn.LayerList(sublayers=stack_cuboid_blocks)) - self.self_blocks = nn.LayerList(sublayers=self_blocks) - - if block_cross_attn_patterns is not None: - if isinstance(block_cross_attn_patterns, (tuple, list)): - assert len(block_cross_attn_patterns) == self.num_blocks - else: - block_cross_attn_patterns = [ - block_cross_attn_patterns for _ in range(self.num_blocks) - ] - block_cross_cuboid_hw = [] - block_cross_cuboid_strategy = [] - block_cross_shift_hw = [] - block_cross_n_temporal = [] - for idx, key in enumerate(block_cross_attn_patterns): - if key == "last_frame_dst": - cuboid_hw = None - shift_hw = None - strategy = None - n_temporal = None - else: - func = cuboid_utils.CuboidCrossAttentionPatterns.get(key) - cuboid_hw, shift_hw, strategy, n_temporal = func(mem_shapes[idx]) - block_cross_cuboid_hw.append(cuboid_hw) - block_cross_cuboid_strategy.append(strategy) - block_cross_shift_hw.append(shift_hw) - block_cross_n_temporal.append(n_temporal) - else: - if not isinstance(block_cross_cuboid_hw[0][0], (list, tuple)): - block_cross_cuboid_hw = [ - block_cross_cuboid_hw for _ in range(self.num_blocks) - ] - else: - assert ( - len(block_cross_cuboid_hw) == self.num_blocks - ), f"Incorrect input format! Received block_cross_cuboid_hw={block_cross_cuboid_hw}" - if not isinstance(block_cross_cuboid_strategy[0][0], (list, tuple)): - block_cross_cuboid_strategy = [ - block_cross_cuboid_strategy for _ in range(self.num_blocks) - ] - else: - assert ( - len(block_cross_cuboid_strategy) == self.num_blocks - ), f"Incorrect input format! Received block_cross_cuboid_strategy={block_cross_cuboid_strategy}" - if not isinstance(block_cross_shift_hw[0][0], (list, tuple)): - block_cross_shift_hw = [ - block_cross_shift_hw for _ in range(self.num_blocks) - ] - else: - assert ( - len(block_cross_shift_hw) == self.num_blocks - ), f"Incorrect input format! Received block_cross_shift_hw={block_cross_shift_hw}" - if not isinstance(block_cross_n_temporal[0], (list, tuple)): - block_cross_n_temporal = [ - block_cross_n_temporal for _ in range(self.num_blocks) - ] - else: - assert ( - len(block_cross_n_temporal) == self.num_blocks - ), f"Incorrect input format! Received block_cross_n_temporal={block_cross_n_temporal}" - self.cross_blocks = nn.LayerList() - assert self.cross_start == 0 - for i in range(self.cross_start, self.num_blocks): - cross_block = nn.LayerList( - sublayers=[ - StackCuboidCrossAttentionBlock( - dim=self.mem_shapes[i][-1], - num_heads=num_heads, - block_cuboid_hw=block_cross_cuboid_hw[i], - block_strategy=block_cross_cuboid_strategy[i], - block_shift_hw=block_cross_shift_hw[i], - block_n_temporal=block_cross_n_temporal[i], - cross_last_n_frames=cross_last_n_frames, - attn_drop=attn_drop, - proj_drop=proj_drop, - ffn_drop=ffn_drop, - gated_ffn=gated_ffn, - norm_layer=norm_layer, - use_inter_ffn=use_inter_ffn, - activation=ffn_activation, - max_temporal_relative=max_temporal_relative, - padding_type=padding_type, - use_global_vector=use_cross_global, - separate_global_qkv=separate_global_qkv, - global_dim_ratio=global_dim_ratio, - checkpoint_level=checkpoint_level, - use_relative_pos=use_relative_pos, - attn_linear_init_mode=attn_linear_init_mode, - ffn_linear_init_mode=ffn_linear_init_mode, - norm_init_mode=norm_init_mode, - expert_shape=expert_shape_list[i], - moe_config=moe_config, - ) - for _ in range(depth[i]) - ] - ) - self.cross_blocks.append(cross_block) - if self.num_blocks > 1: - if self.upsample_type == "upsample": - self.upsample_layers = nn.LayerList( - sublayers=[ - Upsample3DLayer( - dim=self.mem_shapes[i + 1][-1], - out_dim=self.mem_shapes[i][-1], - target_size=(target_temporal_length,) - + self.mem_shapes[i][1:3], - kernel_size=upsample_kernel_size, - temporal_upsample=False, - conv_init_mode=conv_init_mode, - ) - for i in range(self.num_blocks - 1) - ] - ) - else: - raise NotImplementedError(f"{self.upsample_type} is invalid.") - if self.hierarchical_pos_embed: - self.hierarchical_pos_embed_l = nn.LayerList( - sublayers=[ - PosEmbed( - embed_dim=self.mem_shapes[i][-1], - typ=pos_embed_type, - maxT=target_temporal_length, - maxH=self.mem_shapes[i][1], - maxW=self.mem_shapes[i][2], - ) - for i in range(self.num_blocks - 1) - ] - ) - self.reset_parameters() - - def reset_parameters(self): - for ms in self.self_blocks: - for m in ms: - m.reset_parameters() - for ms in self.cross_blocks: - for m in ms: - m.reset_parameters() - if self.num_blocks > 1: - for m in self.upsample_layers: - m.reset_parameters() - if self.hierarchical_pos_embed: - for m in self.hierarchical_pos_embed_l: - m.reset_parameters() - - def forward(self, x, mem_l, mem_global_vector_l=None): - """ - Args: - x : Shape (B, T_top, H_top, W_top, C). - mem_l : A list of memory tensors. - """ - - B, T_top, H_top, W_top, C = x.shape - assert T_top == self.target_temporal_length - assert (H_top, W_top) == (self.mem_shapes[-1][1], self.mem_shapes[-1][2]) - for i in range(self.num_blocks - 1, -1, -1): - mem_global_vector = ( - None if mem_global_vector_l is None else mem_global_vector_l[i] - ) - if not self.use_first_self_attn and i == self.num_blocks - 1: - if i >= self.cross_start: - x = self.cross_blocks[i - self.cross_start][0]( - x, mem_l[i], mem_global_vector - ) - for idx in range(self.depth[i] - 1): - if self.use_self_global: - if self.self_update_global: - x, mem_global_vector = self.self_blocks[i][idx]( - x, mem_global_vector - ) - else: - x, _ = self.self_blocks[i][idx](x, mem_global_vector) - else: - x = self.self_blocks[i][idx](x) - if i >= self.cross_start: - x = self.cross_blocks[i - self.cross_start][idx + 1]( - x, mem_l[i], mem_global_vector - ) - else: - for idx in range(self.depth[i]): - if self.use_self_global: - if self.self_update_global: - x, mem_global_vector = self.self_blocks[i][idx]( - x, mem_global_vector - ) - else: - x, _ = self.self_blocks[i][idx](x, mem_global_vector) - else: - x = self.self_blocks[i][idx](x) - if i >= self.cross_start: - x = self.cross_blocks[i - self.cross_start][idx]( - x, mem_l[i], mem_global_vector - ) - if i > 0: - x = self.upsample_layers[i - 1](x) - if self.hierarchical_pos_embed: - x = self.hierarchical_pos_embed_l[i - 1](x) - return x - - -class MixtureCrossAttention(nn.Layer): - def __init__( - self, - dim, - num_heads, - cuboid_hw, - shift_hw, - strategy, - n_temporal, - cross_last_n_frames, - padding_type, - qkv_bias, - qk_scale, - attn_drop, - proj_drop, - norm_layer, - max_temporal_relative, - use_global_vector, - separate_global_qkv, - global_dim_ratio, - checkpoint_level, - use_relative_pos, - attn_linear_init_mode, - ffn_linear_init_mode, - norm_init_mode, - expert_shape, - moe_config, - ): - super().__init__() - - self.in_dim = dim - self.out_dim = dim - self.expert_shape = expert_shape # T, H, W, C - self.num_experts = moe_config["num_experts"] - self.out_planes = moe_config["out_planes"] - self.moe_config = moe_config - assert expert_shape is not None and moe_config["use_attn_moe"] - assert not use_global_vector - - if moe_config["gate_style"] == "linear": - self.gate = moe_utils.LinearGatingNet(moe_config, expert_shape, dim) - elif moe_config["gate_style"] == "spatial-latent": - self.gate = moe_utils.SpatialLatentGatingNet(moe_config, expert_shape, dim) - elif moe_config["gate_style"] == "cuboid-latent": - self.gate = moe_utils.CuboidLatentGatingNet(moe_config, expert_shape, dim) - elif moe_config["gate_style"] == "spatial-latent-linear": - self.gate = moe_utils.SpatialLatentLinearGatingNet( - moe_config, expert_shape, dim - ) - elif moe_config["gate_style"] == "cuboid-latent-linear": - self.gate = moe_utils.CuboidLatentLinearGatingNet( - moe_config, expert_shape, dim - ) - else: - raise NotImplementedError - - self.experts = nn.LayerList( - [ - CuboidCrossAttentionLayer( - dim=dim, - num_heads=num_heads, - cuboid_hw=cuboid_hw, - shift_hw=shift_hw, - strategy=strategy, - n_temporal=n_temporal, - cross_last_n_frames=cross_last_n_frames, - padding_type=padding_type, - qkv_bias=qkv_bias, - qk_scale=qk_scale, - attn_drop=attn_drop, - proj_drop=proj_drop, - norm_layer=norm_layer, - max_temporal_relative=max_temporal_relative, - use_global_vector=use_global_vector, - separate_global_qkv=separate_global_qkv, - global_dim_ratio=global_dim_ratio, - checkpoint_level=checkpoint_level, - use_relative_pos=use_relative_pos, - attn_linear_init_mode=attn_linear_init_mode, - ffn_linear_init_mode=ffn_linear_init_mode, - norm_init_mode=norm_init_mode, - ) - for _ in range(self.num_experts) - ] - ) - - def forward(self, x, mem, mem_global_vectors=None): - - B, T_x, H, W, C = x.shape - _, T_m, _, _, _ = mem.shape - E = self.num_experts - assert C == self.in_dim and list(self.expert_shape)[:-1] == x.shape[1:-1] - ( - dense_routing_weights, - sparse_routing_weights, - sparse_routing_inds, - self.aux_loss, - ) = self.gate( - x - ) # dense: B, T_x, H, W, E - - dispatcher = moe_utils.DenseDispatcher( - E, - sparse_routing_weights.reshape([B * T_x * H * W, -1]), - sparse_routing_inds.reshape([B * T_x * H * W, -1]), - ) - expert_outputs = paddle.stack( - [self.experts[i](x, mem, mem_global_vectors) for i in range(E)], axis=-2 - ).reshape([B * T_x * H * W, E, C]) - y = dispatcher.combine(expert_outputs).reshape([B, T_x, H, W, C]) - - return y - - def reset_parameters(self): - - for i in range(len(self.experts)): - self.experts[i].reset_parameters() diff --git a/examples/smc_reac/ppsci/arch/extformer_moe_cuboid_encoder.py b/examples/smc_reac/ppsci/arch/extformer_moe_cuboid_encoder.py deleted file mode 100644 index c26b3837a5..0000000000 --- a/examples/smc_reac/ppsci/arch/extformer_moe_cuboid_encoder.py +++ /dev/null @@ -1,1992 +0,0 @@ -from collections import OrderedDict -from functools import lru_cache -from typing import Tuple - -import numpy as np -import paddle -import paddle.nn.functional as F -from paddle import nn -from paddle.distributed import fleet - -import ppsci.arch.extformer_moe_cuboid_utils as cuboid_utils -import ppsci.arch.extformer_moe_utils as moe_utils -from ppsci.arch import activation as act_mod -from ppsci.utils import initializer - -NEGATIVE_SLOPE = 0.1 - - -class PatchMerging3D(nn.Layer): - """Patch Merging Layer - - Args: - dim (int): Number of input channels. - out_dim (int, optional): The dim of output. Defaults to None. - downsample (tuple, optional): Downsample factor. Defaults to (1, 2, 2). - norm_layer (str, optional): The normalization layer. Defaults to "layer_norm". - padding_type (str, optional): The type of padding. Defaults to "nearest". - linear_init_mode (str, optional): The mode of linear init. Defaults to "0". - norm_init_mode (str, optional): The mode of normalization init. Defaults to "0". - """ - - def __init__( - self, - dim: int, - out_dim: int = None, - downsample: Tuple[int, ...] = (1, 2, 2), - norm_layer: str = "layer_norm", - padding_type: str = "nearest", - linear_init_mode: str = "0", - norm_init_mode: str = "0", - moe_config: dict = None, - ): - super().__init__() - self.linear_init_mode = linear_init_mode - self.norm_init_mode = norm_init_mode - self.dim = dim - if out_dim is None: - out_dim = max(downsample) * dim - self.out_dim = out_dim - self.downsample = downsample - self.padding_type = padding_type - self.reduction = nn.Linear( - in_features=downsample[0] * downsample[1] * downsample[2] * dim, - out_features=out_dim, - bias_attr=False, - ) - self.norm = cuboid_utils.get_norm_layer( - norm_layer, in_channels=downsample[0] * downsample[1] * downsample[2] * dim - ) - self.reset_parameters() - - def reset_parameters(self): - for m in self.children(): - cuboid_utils.apply_initialization( - m, linear_mode=self.linear_init_mode, norm_mode=self.norm_init_mode - ) - - def get_out_shape(self, data_shape): - T, H, W, C_in = data_shape - pad_t = (self.downsample[0] - T % self.downsample[0]) % self.downsample[0] - pad_h = (self.downsample[1] - H % self.downsample[1]) % self.downsample[1] - pad_w = (self.downsample[2] - W % self.downsample[2]) % self.downsample[2] - return ( - (T + pad_t) // self.downsample[0], - (H + pad_h) // self.downsample[1], - (W + pad_w) // self.downsample[2], - self.out_dim, - ) - - def forward(self, x): - """ - - Args: - x : (B, T, H, W, C) - - Returns: - out : Shape (B, T // downsample[0], H // downsample[1], W // downsample[2], out_dim) - """ - - B, T, H, W, C = x.shape - pad_t = (self.downsample[0] - T % self.downsample[0]) % self.downsample[0] - pad_h = (self.downsample[1] - H % self.downsample[1]) % self.downsample[1] - pad_w = (self.downsample[2] - W % self.downsample[2]) % self.downsample[2] - if pad_h or pad_h or pad_w: - T += pad_t - H += pad_h - W += pad_w - x = cuboid_utils.generalize_padding( - x, pad_t, pad_h, pad_w, padding_type=self.padding_type - ) - x = ( - x.reshape( - ( - B, - T // self.downsample[0], - self.downsample[0], - H // self.downsample[1], - self.downsample[1], - W // self.downsample[2], - self.downsample[2], - C, - ) - ) - .transpose(perm=[0, 1, 3, 5, 2, 4, 6, 7]) - .reshape( - [ - B, - T // self.downsample[0], - H // self.downsample[1], - W // self.downsample[2], - self.downsample[0] * self.downsample[1] * self.downsample[2] * C, - ] - ) - ) - x = self.norm(x) - x = self.reduction(x) - return x - - -class PositionwiseFFN(nn.Layer): - """The Position-wise FFN layer used in Transformer-like architectures - - If pre_norm is True: - norm(data) -> fc1 -> act -> act_dropout -> fc2 -> dropout -> res(+data) - Else: - data -> fc1 -> act -> act_dropout -> fc2 -> dropout -> norm(res(+data)) - Also, if we use gated projection. We will use - fc1_1 * act(fc1_2(data)) to map the data - - Args: - units (int, optional): The units. Defaults to 512. - hidden_size (int, optional): The size of hidden layer. Defaults to 2048. - activation_dropout (float, optional): The dropout of activate. Defaults to 0.0. - dropout (float, optional): The drop ratio used in DropPat. Defaults to 0.1. - gated_proj (bool, optional): Whether to use gate projection. Defaults to False. - activation (str, optional): The activate. Defaults to "relu". - normalization (str, optional): The normalization. Defaults to "layer_norm". - layer_norm_eps (float, optional): The epsilon of layer normalization. Defaults to 1e-05. - pre_norm (bool): Pre-layer normalization as proposed in the paper: - "[ACL2018] The Best of Both Worlds: Combining Recent Advances in Neural Machine Translation" This will stabilize the training of Transformers. - You may also refer to "[Arxiv2020] Understanding the Difficulty of Training Transformers". Defaults to False. - linear_init_mode (str, optional): The mode of linear initialization. Defaults to "0". - norm_init_mode (str, optional): The mode of normalization initialization. Defaults to "0". - """ - - def __init__( - self, - units: int = 512, - hidden_size: int = 2048, - activation_dropout: float = 0.0, - dropout: float = 0.1, - gated_proj: bool = False, - activation: str = "relu", - normalization: str = "layer_norm", - layer_norm_eps: float = 1e-05, - pre_norm: bool = False, - linear_init_mode: str = "0", - norm_init_mode: str = "0", - moe_config: dict = None, - expert_shape: tuple = None, - ): - super().__init__() - self.linear_init_mode = linear_init_mode - self.norm_init_mode = norm_init_mode - self._pre_norm = pre_norm - self._gated_proj = gated_proj - self._kwargs = OrderedDict( - [ - ("units", units), - ("hidden_size", hidden_size), - ("activation_dropout", activation_dropout), - ("activation", activation), - ("dropout", dropout), - ("normalization", normalization), - ("layer_norm_eps", layer_norm_eps), - ("gated_proj", gated_proj), - ("pre_norm", pre_norm), - ] - ) - self.dropout_layer = nn.Dropout(p=dropout) - self.activation_dropout_layer = nn.Dropout(p=activation_dropout) - - if moe_config["use_linear_moe"]: - self.ffn_1 = MixtureLinear( - in_dim=units, - out_dim=hidden_size, - bias_attr=True, - expert_shape=expert_shape[:-1] + (hidden_size,), - moe_config=moe_config, - ) - else: - self.ffn_1 = nn.Linear( - in_features=units, out_features=hidden_size, bias_attr=True - ) - if self._gated_proj: - self.ffn_1_gate = nn.Linear( - in_features=units, out_features=hidden_size, bias_attr=True - ) - if activation == "leaky_relu": - self.activation = nn.LeakyReLU(NEGATIVE_SLOPE) - else: - self.activation = act_mod.get_activation(activation) - - if moe_config["use_linear_moe"]: - self.ffn_2 = MixtureLinear( - in_dim=hidden_size, - out_dim=units, - bias_attr=True, - expert_shape=expert_shape, - moe_config=moe_config, - ) - else: - self.ffn_2 = nn.Linear( - in_features=hidden_size, out_features=units, bias_attr=True - ) - self.layer_norm = cuboid_utils.get_norm_layer( - normalization=normalization, in_channels=units, epsilon=layer_norm_eps - ) - self.reset_parameters() - - def reset_parameters(self): - cuboid_utils.apply_initialization(self.ffn_1, linear_mode=self.linear_init_mode) - if self._gated_proj: - cuboid_utils.apply_initialization( - self.ffn_1_gate, linear_mode=self.linear_init_mode - ) - cuboid_utils.apply_initialization(self.ffn_2, linear_mode=self.linear_init_mode) - cuboid_utils.apply_initialization( - self.layer_norm, norm_mode=self.norm_init_mode - ) - - def forward(self, data): - """ - Args: - x : Shape (B, seq_length, C_in) - - Returns: - out : Shape (B, seq_length, C_out) - """ - - residual = data - if self._pre_norm: - data = self.layer_norm(data) - if self._gated_proj: - out = self.activation(self.ffn_1_gate(data)) * self.ffn_1(data) - else: - out = self.activation(self.ffn_1(data)) - out = self.activation_dropout_layer(out) - out = self.ffn_2(out) - out = self.dropout_layer(out) - out = out + residual - if not self._pre_norm: - out = self.layer_norm(out) - return out - - -def update_cuboid_size_shift_size(data_shape, cuboid_size, shift_size, strategy): - """Update the cuboid_size and shift_size - - Args: - data_shape (Tuple[int,...]): The shape of the data. - cuboid_size (Tuple[int,...]): Size of the cuboid. - shift_size (Tuple[int,...]): Size of the shift. - strategy (str): The strategy of attention. - - Returns: - new_cuboid_size (Tuple[int,...]): Size of the cuboid. - new_shift_size (Tuple[int,...]): Size of the shift. - """ - - new_cuboid_size = list(cuboid_size) - new_shift_size = list(shift_size) - for i in range(len(data_shape)): - if strategy[i] == "d": - new_shift_size[i] = 0 - if data_shape[i] <= cuboid_size[i]: - new_cuboid_size[i] = data_shape[i] - new_shift_size[i] = 0 - return tuple(new_cuboid_size), tuple(new_shift_size) - - -def cuboid_reorder(data, cuboid_size, strategy): - """Reorder the tensor into (B, num_cuboids, bT * bH * bW, C) - We assume that the tensor shapes are divisible to the cuboid sizes. - - Args: - data (paddle.Tensor): The input data. - cuboid_size (Tuple[int,...]): The size of the cuboid. - strategy (Tuple[int,...]): The cuboid strategy. - - Returns: - reordered_data (paddle.Tensor): Shape will be (B, num_cuboids, bT * bH * bW, C). - num_cuboids = T / bT * H / bH * W / bW - """ - - B, T, H, W, C = data.shape - num_cuboids = T // cuboid_size[0] * H // cuboid_size[1] * W // cuboid_size[2] - cuboid_volume = cuboid_size[0] * cuboid_size[1] * cuboid_size[2] - intermediate_shape = [] - nblock_axis = [] - block_axis = [] - for i, (block_size, total_size, ele_strategy) in enumerate( - zip(cuboid_size, (T, H, W), strategy) - ): - if ele_strategy == "l": - intermediate_shape.extend([total_size // block_size, block_size]) - nblock_axis.append(2 * i + 1) - block_axis.append(2 * i + 2) - elif ele_strategy == "d": - intermediate_shape.extend([block_size, total_size // block_size]) - nblock_axis.append(2 * i + 2) - block_axis.append(2 * i + 1) - else: - raise NotImplementedError(f"{ele_strategy} is invalid.") - data = data.reshape(list((B,) + tuple(intermediate_shape) + (C,))) - reordered_data = data.transpose( - perm=(0,) + tuple(nblock_axis) + tuple(block_axis) + (7,) - ) - reordered_data = reordered_data.reshape((B, num_cuboids, cuboid_volume, C)) - return reordered_data - - -@lru_cache() -def compute_cuboid_self_attention_mask( - data_shape, cuboid_size, shift_size, strategy, padding_type, device -): - """Compute the shift window attention mask - - Args: - data_shape (Tuple[int,....]): Should be (T, H, W). - cuboid_size (Tuple[int,....]): Size of the cuboid. - shift_size (Tuple[int,....]): The shift size. - strategy (str): The decomposition strategy. - padding_type (str): Type of the padding. - device (str): The device. - - Returns: - attn_mask (paddle.Tensor): Mask with shape (num_cuboid, cuboid_vol, cuboid_vol). - The padded values will always be masked. The other masks will ensure that the shifted windows - will only attend to those in the shifted windows. - """ - T, H, W = data_shape - pad_t = (cuboid_size[0] - T % cuboid_size[0]) % cuboid_size[0] - pad_h = (cuboid_size[1] - H % cuboid_size[1]) % cuboid_size[1] - pad_w = (cuboid_size[2] - W % cuboid_size[2]) % cuboid_size[2] - data_mask = None - if pad_t > 0 or pad_h > 0 or pad_w > 0: - if padding_type == "ignore": - data_mask = paddle.ones(shape=(1, T, H, W, 1), dtype="bool") - data_mask = F.pad( - data_mask, [0, 0, 0, pad_w, 0, pad_h, 0, pad_t], data_format="NDHWC" - ) - else: - data_mask = paddle.ones( - shape=(1, T + pad_t, H + pad_h, W + pad_w, 1), dtype="bool" - ) - if any(i > 0 for i in shift_size): - if padding_type == "ignore": - data_mask = paddle.roll( - x=data_mask, - shifts=(-shift_size[0], -shift_size[1], -shift_size[2]), - axis=(1, 2, 3), - ) - if padding_type == "ignore": - data_mask = cuboid_reorder(data_mask, cuboid_size, strategy=strategy) - data_mask = data_mask.squeeze(axis=-1).squeeze(axis=0) - shift_mask = np.zeros(shape=(1, T + pad_t, H + pad_h, W + pad_w, 1)) - cnt = 0 - for t in ( - slice(-cuboid_size[0]), - slice(-cuboid_size[0], -shift_size[0]), - slice(-shift_size[0], None), - ): - for h in ( - slice(-cuboid_size[1]), - slice(-cuboid_size[1], -shift_size[1]), - slice(-shift_size[1], None), - ): - for w in ( - slice(-cuboid_size[2]), - slice(-cuboid_size[2], -shift_size[2]), - slice(-shift_size[2], None), - ): - shift_mask[:, t, h, w, :] = cnt - cnt += 1 - shift_mask = paddle.to_tensor(shift_mask) - shift_mask = cuboid_reorder(shift_mask, cuboid_size, strategy=strategy) - shift_mask = shift_mask.squeeze(axis=-1).squeeze(axis=0) - attn_mask = shift_mask.unsqueeze(axis=1) - shift_mask.unsqueeze(axis=2) == 0 - if padding_type == "ignore": - attn_mask = ( - data_mask.unsqueeze(axis=1) * data_mask.unsqueeze(axis=2) * attn_mask - ) - return attn_mask - - -def masked_softmax(att_score, mask, axis: int = -1): - """Ignore the masked elements when calculating the softmax. - The mask can be broadcastable. - - Args: - att_score (paddle.Tensor): Shape (..., length, ...) - mask (paddle.Tensor): Shape (..., length, ...) - 1 --> The element is not masked - 0 --> The element is masked - axis (int): The axis to calculate the softmax. att_score.shape[axis] must be the same as mask.shape[axis] - - Returns: - att_weights (paddle.Tensor): Shape (..., length, ...). - """ - - if mask is not None: - if att_score.dtype == paddle.float16: - att_score = att_score.masked_fill(paddle.logical_not(mask), -1e4) - else: - att_score = att_score.masked_fill(paddle.logical_not(mask), -1e18) - att_weights = nn.functional.softmax(x=att_score, axis=axis) * mask.astype( - att_score.dtype - ) - else: - att_weights = nn.functional.softmax(x=att_score, axis=axis) - return att_weights - - -def cuboid_reorder_reverse(data, cuboid_size, strategy, orig_data_shape): - """Reverse the reordered cuboid back to the original space - - Args: - data (paddle.Tensor): The input data. - cuboid_size (Tuple[int,...]): The size of cuboid. - strategy (str): The strategy of reordering. - orig_data_shape (Tuple[int,...]): The original shape of the data. - - Returns: - data (paddle.Tensor): The recovered data - """ - - B, num_cuboids, cuboid_volume, C = data.shape - T, H, W = orig_data_shape - permutation_axis = [0] - for i, (block_size, total_size, ele_strategy) in enumerate( - zip(cuboid_size, (T, H, W), strategy) - ): - if ele_strategy == "l": - permutation_axis.append(i + 1) - permutation_axis.append(i + 4) - elif ele_strategy == "d": - permutation_axis.append(i + 4) - permutation_axis.append(i + 1) - else: - raise NotImplementedError((f"{ele_strategy} is invalid.")) - permutation_axis.append(7) - data = data.reshape( - [ - B, - T // cuboid_size[0], - H // cuboid_size[1], - W // cuboid_size[2], - cuboid_size[0], - cuboid_size[1], - cuboid_size[2], - C, - ] - ) - data = data.transpose(perm=permutation_axis) - data = data.reshape((B, T, H, W, C)) - return data - - -class CuboidSelfAttentionLayer(nn.Layer): - """Implements the cuboid self attention. - - The idea of Cuboid Self Attention is to divide the input tensor (T, H, W) into several non-overlapping cuboids. - We apply self-attention inside each cuboid and all cuboid-level self attentions are executed in parallel. - - We adopt two mechanisms for decomposing the input tensor into cuboids: - - (1) local: - We group the tensors within a local window, e.g., X[t:(t+b_t), h:(h+b_h), w:(w+b_w)]. We can also apply the - shifted window strategy proposed in "[ICCV2021] Swin Transformer: Hierarchical Vision Transformer using Shifted Windows". - (2) dilated: - Inspired by the success of dilated convolution "[ICLR2016] Multi-Scale Context Aggregation by Dilated Convolutions", - we split the tensor with dilation factors that are tied to the size of the cuboid. For example, for a cuboid that has width `b_w`, - we sample the elements starting from 0 as 0, w / b_w, 2 * w / b_w, ..., (b_w - 1) * w / b_w. - - The cuboid attention can be viewed as a generalization of the attention mechanism proposed in Video Swin Transformer, https://arxiv.org/abs/2106.13230. - The computational complexity of CuboidAttention can be simply calculated as O(T H W * b_t b_h b_w). To cover multiple correlation patterns, - we are able to combine multiple CuboidAttention layers with different configurations such as cuboid size, shift size, and local / global decomposing strategy. - - In addition, it is straight-forward to extend the cuboid attention to other types of spatiotemporal data that are not described - as regular tensors. We need to define alternative approaches to partition the data into "cuboids". - - In addition, inspired by "[NeurIPS2021] Do Transformers Really Perform Badly for Graph Representation?", - "[NeurIPS2020] Big Bird: Transformers for Longer Sequences", "[EMNLP2021] Longformer: The Long-Document Transformer", we keep - $K$ global vectors to record the global status of the spatiotemporal system. These global vectors will attend to the whole tensor and - the vectors inside each individual cuboids will also attend to the global vectors so that they can peep into the global status of the system. - - Args: - dim (int): The dimension of the input tensor. - num_heads (int): The number of heads. - cuboid_size (tuple, optional): The size of cuboid. Defaults to (2, 7, 7). - shift_size (tuple, optional): The size of shift. Defaults to (0, 0, 0). - strategy (tuple, optional): The strategy. Defaults to ("l", "l", "l"). - padding_type (str, optional): The type of padding. Defaults to "ignore". - qkv_bias (bool, optional): Whether to enable bias in calculating qkv attention. Defaults to False. - qk_scale (float, optional): Whether to enable scale factor when calculating the attention. Defaults to None. - attn_drop (float, optional): The attention dropout. Defaults to 0.0. - proj_drop (float, optional): The projection dropout. Defaults to 0.0. - use_final_proj (bool, optional): Whether to use the final projection. Defaults to True. - norm_layer (str, optional): The normalization layer. Defaults to "layer_norm". - use_global_vector (bool, optional): Whether to use the global vector or not. Defaults to False. - use_global_self_attn (bool, optional): Whether to use self attention among global vectors. Defaults to False. - separate_global_qkv (bool, optional): Whether to use different network to calc q_global, k_global, v_global. Defaults to False. - global_dim_ratio (int, optional): The dim (channels) of global vectors is `global_dim_ratio*dim`. Defaults to 1. - checkpoint_level (bool, optional): Whether to enable gradient checkpointing. Defaults to True. - use_relative_pos (bool, optional): Whether to use relative pos. Defaults to True. - attn_linear_init_mode (str, optional): The mode of attention linear initialization. Defaults to "0". - ffn_linear_init_mode (str, optional): The mode of FFN linear initialization. Defaults to "0". - norm_init_mode (str, optional): The mode of normalization initialization. Defaults to "0". - """ - - def __init__( - self, - dim: int, - num_heads: int, - cuboid_size: Tuple[int, ...] = (2, 7, 7), - shift_size: Tuple[int, ...] = (0, 0, 0), - strategy: Tuple[str, ...] = ("l", "l", "l"), - padding_type: str = "ignore", - qkv_bias: bool = False, - qk_scale: float = None, - attn_drop: float = 0.0, - proj_drop: float = 0.0, - use_final_proj: bool = True, - norm_layer: str = "layer_norm", - use_global_vector: bool = False, - use_global_self_attn: bool = False, - separate_global_qkv: bool = False, - global_dim_ratio: int = 1, - checkpoint_level: bool = True, - use_relative_pos: bool = True, - attn_linear_init_mode: str = "0", - ffn_linear_init_mode: str = "0", - norm_init_mode: str = "0", - moe_config: dict = None, - ): - super(CuboidSelfAttentionLayer, self).__init__() - self.attn_linear_init_mode = attn_linear_init_mode - self.ffn_linear_init_mode = ffn_linear_init_mode - self.norm_init_mode = norm_init_mode - assert dim % num_heads == 0 - self.num_heads = num_heads - self.dim = dim - self.cuboid_size = cuboid_size - self.shift_size = shift_size - self.strategy = strategy - self.padding_type = padding_type - self.use_final_proj = use_final_proj - self.use_relative_pos = use_relative_pos - self.use_global_vector = use_global_vector - self.use_global_self_attn = use_global_self_attn - self.separate_global_qkv = separate_global_qkv - if global_dim_ratio != 1: - assert ( - separate_global_qkv is True - ), "Setting global_dim_ratio != 1 requires separate_global_qkv == True." - self.global_dim_ratio = global_dim_ratio - assert self.padding_type in ["ignore", "zeros", "nearest"] - head_dim = dim // num_heads - self.scale = qk_scale or head_dim**-0.5 - if use_relative_pos: - init_data = paddle.zeros( - ( - (2 * cuboid_size[0] - 1) - * (2 * cuboid_size[1] - 1) - * (2 * cuboid_size[2] - 1), - num_heads, - ) - ) - self.relative_position_bias_table = paddle.create_parameter( - shape=init_data.shape, - dtype=init_data.dtype, - default_initializer=nn.initializer.Constant(0.0), - ) - self.relative_position_bias_table.stop_gradient = not True - self.relative_position_bias_table = initializer.trunc_normal_( - self.relative_position_bias_table, std=0.02 - ) - - coords_t = paddle.arange(end=self.cuboid_size[0]) - coords_h = paddle.arange(end=self.cuboid_size[1]) - coords_w = paddle.arange(end=self.cuboid_size[2]) - coords = paddle.stack(x=paddle.meshgrid(coords_t, coords_h, coords_w)) - coords_flatten = paddle.flatten(x=coords, start_axis=1) - relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] - relative_coords = relative_coords.transpose(perm=[1, 2, 0]) - relative_coords[:, :, 0] += self.cuboid_size[0] - 1 - relative_coords[:, :, 1] += self.cuboid_size[1] - 1 - relative_coords[:, :, 2] += self.cuboid_size[2] - 1 - relative_coords[:, :, 0] *= (2 * self.cuboid_size[1] - 1) * ( - 2 * self.cuboid_size[2] - 1 - ) - relative_coords[:, :, 1] *= 2 * self.cuboid_size[2] - 1 - relative_position_index = relative_coords.sum(axis=-1) - self.register_buffer( - name="relative_position_index", tensor=relative_position_index - ) - self.qkv = nn.Linear(in_features=dim, out_features=dim * 3, bias_attr=qkv_bias) - self.attn_drop = nn.Dropout(p=attn_drop) - if self.use_global_vector: - if self.separate_global_qkv: - self.l2g_q_net = nn.Linear( - in_features=dim, out_features=dim, bias_attr=qkv_bias - ) - self.l2g_global_kv_net = nn.Linear( - in_features=global_dim_ratio * dim, - out_features=dim * 2, - bias_attr=qkv_bias, - ) - self.g2l_global_q_net = nn.Linear( - in_features=global_dim_ratio * dim, - out_features=dim, - bias_attr=qkv_bias, - ) - self.g2l_k_net = nn.Linear( - in_features=dim, out_features=dim, bias_attr=qkv_bias - ) - self.g2l_v_net = nn.Linear( - in_features=dim, - out_features=global_dim_ratio * dim, - bias_attr=qkv_bias, - ) - if self.use_global_self_attn: - self.g2g_global_qkv_net = nn.Linear( - in_features=global_dim_ratio * dim, - out_features=global_dim_ratio * dim * 3, - bias_attr=qkv_bias, - ) - else: - self.global_qkv = nn.Linear( - in_features=dim, out_features=dim * 3, bias_attr=qkv_bias - ) - self.global_attn_drop = nn.Dropout(p=attn_drop) - if use_final_proj: - self.proj = nn.Linear(in_features=dim, out_features=dim) - self.proj_drop = nn.Dropout(p=proj_drop) - if self.use_global_vector: - self.global_proj = nn.Linear( - in_features=global_dim_ratio * dim, - out_features=global_dim_ratio * dim, - ) - self.norm = cuboid_utils.get_norm_layer(norm_layer, in_channels=dim) - if self.use_global_vector: - self.global_vec_norm = cuboid_utils.get_norm_layer( - norm_layer, in_channels=global_dim_ratio * dim - ) - self.checkpoint_level = checkpoint_level - self.reset_parameters() - - def reset_parameters(self): - cuboid_utils.apply_initialization( - self.qkv, linear_mode=self.attn_linear_init_mode - ) - if self.use_final_proj: - cuboid_utils.apply_initialization( - self.proj, linear_mode=self.ffn_linear_init_mode - ) - cuboid_utils.apply_initialization(self.norm, norm_mode=self.norm_init_mode) - if self.use_global_vector: - if self.separate_global_qkv: - cuboid_utils.apply_initialization( - self.l2g_q_net, linear_mode=self.attn_linear_init_mode - ) - cuboid_utils.apply_initialization( - self.l2g_global_kv_net, linear_mode=self.attn_linear_init_mode - ) - cuboid_utils.apply_initialization( - self.g2l_global_q_net, linear_mode=self.attn_linear_init_mode - ) - cuboid_utils.apply_initialization( - self.g2l_k_net, linear_mode=self.attn_linear_init_mode - ) - cuboid_utils.apply_initialization( - self.g2l_v_net, linear_mode=self.attn_linear_init_mode - ) - if self.use_global_self_attn: - cuboid_utils.apply_initialization( - self.g2g_global_qkv_net, linear_mode=self.attn_linear_init_mode - ) - else: - cuboid_utils.apply_initialization( - self.global_qkv, linear_mode=self.attn_linear_init_mode - ) - cuboid_utils.apply_initialization( - self.global_vec_norm, norm_mode=self.norm_init_mode - ) - - def forward(self, x, global_vectors=None): - x = self.norm(x) - - B, T, H, W, C_in = x.shape - assert C_in == self.dim - if self.use_global_vector: - _, num_global, _ = global_vectors.shape - global_vectors = self.global_vec_norm(global_vectors) - cuboid_size, shift_size = update_cuboid_size_shift_size( - (T, H, W), self.cuboid_size, self.shift_size, self.strategy - ) - - pad_t = (cuboid_size[0] - T % cuboid_size[0]) % cuboid_size[0] - pad_h = (cuboid_size[1] - H % cuboid_size[1]) % cuboid_size[1] - pad_w = (cuboid_size[2] - W % cuboid_size[2]) % cuboid_size[2] - x = cuboid_utils.generalize_padding(x, pad_t, pad_h, pad_w, self.padding_type) - - if any(i > 0 for i in shift_size): - shifted_x = paddle.roll( - x=x, - shifts=(-shift_size[0], -shift_size[1], -shift_size[2]), - axis=(1, 2, 3), - ) - else: - shifted_x = x - - reordered_x = cuboid_reorder( - shifted_x, cuboid_size=cuboid_size, strategy=self.strategy - ) - - _, num_cuboids, cuboid_volume, _ = reordered_x.shape - attn_mask = compute_cuboid_self_attention_mask( - (T, H, W), - cuboid_size, - shift_size=shift_size, - strategy=self.strategy, - padding_type=self.padding_type, - device=x.place, - ) - head_C = C_in // self.num_heads - qkv = ( - self.qkv(reordered_x) - .reshape([B, num_cuboids, cuboid_volume, 3, self.num_heads, head_C]) - .transpose(perm=[3, 0, 4, 1, 2, 5]) - ) - - q, k, v = qkv[0], qkv[1], qkv[2] - q = q * self.scale - perm_0 = list(range(k.ndim)) - perm_0[-2] = -1 - perm_0[-1] = -2 - attn_score = q @ k.transpose(perm=perm_0) - - if self.use_relative_pos: - relative_position_bias = self.relative_position_bias_table[ - self.relative_position_index[:cuboid_volume, :cuboid_volume].reshape( - [-1] - ) - ].reshape([cuboid_volume, cuboid_volume, -1]) - relative_position_bias = relative_position_bias.transpose( - perm=[2, 0, 1] - ).unsqueeze(axis=1) - attn_score = attn_score + relative_position_bias - - if self.use_global_vector: - global_head_C = self.global_dim_ratio * head_C - if self.separate_global_qkv: - l2g_q = ( - self.l2g_q_net(reordered_x) - .reshape([B, num_cuboids, cuboid_volume, self.num_heads, head_C]) - .transpose(perm=[0, 3, 1, 2, 4]) - ) - l2g_q = l2g_q * self.scale - l2g_global_kv = ( - self.l2g_global_kv_net(global_vectors) - .reshape([B, 1, num_global, 2, self.num_heads, head_C]) - .transpose(perm=[3, 0, 4, 1, 2, 5]) - ) - l2g_global_k, l2g_global_v = l2g_global_kv[0], l2g_global_kv[1] - g2l_global_q = ( - self.g2l_global_q_net(global_vectors) - .reshape([B, num_global, self.num_heads, head_C]) - .transpose(perm=[0, 2, 1, 3]) - ) - g2l_global_q = g2l_global_q * self.scale - g2l_k = ( - self.g2l_k_net(reordered_x) - .reshape([B, num_cuboids, cuboid_volume, self.num_heads, head_C]) - .transpose(perm=[0, 3, 1, 2, 4]) - ) - g2l_v = ( - self.g2l_v_net(reordered_x) - .reshape( - [B, num_cuboids, cuboid_volume, self.num_heads, global_head_C] - ) - .transpose(perm=[0, 3, 1, 2, 4]) - ) - if self.use_global_self_attn: - g2g_global_qkv = ( - self.g2g_global_qkv_net(global_vectors) - .reshape([B, 1, num_global, 3, self.num_heads, global_head_C]) - .transpose(perm=[3, 0, 4, 1, 2, 5]) - ) - g2g_global_q, g2g_global_k, g2g_global_v = ( - g2g_global_qkv[0], - g2g_global_qkv[1], - g2g_global_qkv[2], - ) - g2g_global_q = g2g_global_q.squeeze(axis=2) * self.scale - else: - q_global, k_global, v_global = ( - self.global_qkv(global_vectors) - .reshape([B, 1, num_global, 3, self.num_heads, head_C]) - .transpose(perm=[3, 0, 4, 1, 2, 5]) - ) - q_global = q_global.squeeze(axis=2) * self.scale - l2g_q, g2l_k, g2l_v = q, k, v - g2l_global_q, l2g_global_k, l2g_global_v = ( - q_global, - k_global, - v_global, - ) - if self.use_global_self_attn: - g2g_global_q, g2g_global_k, g2g_global_v = ( - q_global, - k_global, - v_global, - ) - - perm_1 = list(range(l2g_global_k.ndim)) - perm_1[-2] = -1 - perm_1[-1] = -2 - l2g_attn_score = l2g_q @ l2g_global_k.transpose(perm=perm_1) - attn_score_l2l_l2g = paddle.concat(x=(attn_score, l2g_attn_score), axis=-1) - - if attn_mask.ndim == 5: - attn_mask_l2l_l2g = F.pad( - attn_mask, [0, num_global], "constant", 1, data_format="NDHWC" - ) - elif attn_mask.ndim == 3: - attn_mask = attn_mask.astype("float32") - attn_mask_l2l_l2g = F.pad( - attn_mask, [0, num_global], "constant", 1, data_format="NCL" - ) - attn_mask_l2l_l2g = attn_mask_l2l_l2g.astype("bool") - else: - attn_mask_l2l_l2g = F.pad(attn_mask, [0, num_global], "constant", 1) - - v_l_g = paddle.concat( - x=( - v, - l2g_global_v.expand( - shape=[B, self.num_heads, num_cuboids, num_global, head_C] - ), - ), - axis=3, - ) - attn_score_l2l_l2g = masked_softmax( - attn_score_l2l_l2g, mask=attn_mask_l2l_l2g - ) - attn_score_l2l_l2g = self.attn_drop(attn_score_l2l_l2g) - reordered_x = ( - (attn_score_l2l_l2g @ v_l_g) - .transpose(perm=[0, 2, 3, 1, 4]) - .reshape([B, num_cuboids, cuboid_volume, self.dim]) - ) - if self.padding_type == "ignore": - g2l_attn_mask = paddle.ones(shape=(1, T, H, W, 1)) - if pad_t > 0 or pad_h > 0 or pad_w > 0: - g2l_attn_mask = F.pad( - g2l_attn_mask, - [0, 0, 0, pad_w, 0, pad_h, 0, pad_t], - data_format="NDHWC", - ) - if any(i > 0 for i in shift_size): - g2l_attn_mask = paddle.roll( - x=g2l_attn_mask, - shifts=(-shift_size[0], -shift_size[1], -shift_size[2]), - axis=(1, 2, 3), - ) - g2l_attn_mask = g2l_attn_mask.reshape((-1,)) - else: - g2l_attn_mask = None - temp = g2l_k.reshape( - [B, self.num_heads, num_cuboids * cuboid_volume, head_C] - ) - perm_2 = list(range(temp.ndim)) - perm_2[-2] = -1 - perm_2[-1] = -2 - g2l_attn_score = g2l_global_q @ temp.transpose(perm=perm_2) - if self.use_global_self_attn: - temp = g2g_global_k.squeeze(axis=2) - perm_3 = list(range(temp.ndim)) - perm_3[-2] = -1 - perm_3[-1] = -2 - g2g_attn_score = g2g_global_q @ temp.transpose(perm=perm_3) - g2all_attn_score = paddle.concat( - x=(g2l_attn_score, g2g_attn_score), axis=-1 - ) - if g2l_attn_mask is not None: - g2all_attn_mask = F.pad( - g2l_attn_mask, - [0, num_global], - "constant", - 1, - data_format="NDHWC", - ) - else: - g2all_attn_mask = None - new_v = paddle.concat( - x=( - g2l_v.reshape( - [ - B, - self.num_heads, - num_cuboids * cuboid_volume, - global_head_C, - ] - ), - g2g_global_v.reshape( - [B, self.num_heads, num_global, global_head_C] - ), - ), - axis=2, - ) - else: - g2all_attn_score = g2l_attn_score - g2all_attn_mask = g2l_attn_mask - new_v = g2l_v.reshape( - [B, self.num_heads, num_cuboids * cuboid_volume, global_head_C] - ) - g2all_attn_score = masked_softmax(g2all_attn_score, mask=g2all_attn_mask) - g2all_attn_score = self.global_attn_drop(g2all_attn_score) - new_global_vector = ( - (g2all_attn_score @ new_v) - .transpose(perm=[0, 2, 1, 3]) - .reshape([B, num_global, self.global_dim_ratio * self.dim]) - ) - else: - attn_score = masked_softmax(attn_score, mask=attn_mask) - attn_score = self.attn_drop(attn_score) - reordered_x = ( - (attn_score @ v) - .transpose(perm=[0, 2, 3, 1, 4]) - .reshape([B, num_cuboids, cuboid_volume, self.dim]) - ) - - if self.use_final_proj: - reordered_x = paddle.cast(reordered_x, dtype="float32") - reordered_x = self.proj_drop(self.proj(reordered_x)) - if self.use_global_vector: - new_global_vector = self.proj_drop(self.global_proj(new_global_vector)) - shifted_x = cuboid_reorder_reverse( - reordered_x, - cuboid_size=cuboid_size, - strategy=self.strategy, - orig_data_shape=(T + pad_t, H + pad_h, W + pad_w), - ) - if any(i > 0 for i in shift_size): - x = paddle.roll( - x=shifted_x, - shifts=(shift_size[0], shift_size[1], shift_size[2]), - axis=(1, 2, 3), - ) - else: - x = shifted_x - x = cuboid_utils.generalize_unpadding( - x, pad_t=pad_t, pad_h=pad_h, pad_w=pad_w, padding_type=self.padding_type - ) - if self.use_global_vector: - return x, new_global_vector - else: - return x - - -class StackCuboidSelfAttentionBlock(nn.Layer): - """ - - "use_inter_ffn" is True - x --> attn1 -----+-------> ffn1 ---+---> attn2 --> ... --> ffn_k --> out - | ^ | ^ - | | | | - |-------------| |-------------| - - "use_inter_ffn" is False - x --> attn1 -----+------> attn2 --> ... attnk --+----> ffnk ---+---> out - | ^ | ^ ^ | ^ - | | | | | | | - |-------------| |------------| ----------| |-----------| - If we have enabled global memory vectors, each attention will be a - - Args: - dim (int): The dimension of the input tensor. - num_heads (int): The number of heads. - block_cuboid_size (list, optional): The size of block cuboid . Defaults to [(4, 4, 4), (4, 4, 4)]. - block_shift_size (list, optional): The shift size of block. Defaults to [(0, 0, 0), (2, 2, 2)]. - block_strategy (list, optional): The strategy of block. Defaults to [("d", "d", "d"), ("l", "l", "l")]. - padding_type (str, optional): The type of padding. Defaults to "ignore". - qkv_bias (bool, optional): Whether to enable bias in calculating qkv attention. Defaults to False. - qk_scale (float, optional): Whether to enable scale factor when calculating the attention. Defaults to None. - attn_drop (float, optional): The attention dropout. Defaults to 0.0. - proj_drop (float, optional): The projection dropout. Defaults to 0.0. - use_final_proj (bool, optional): Whether to use the final projection. Defaults to True. - norm_layer (str, optional): The normalization layer. Defaults to "layer_norm". - use_global_vector (bool, optional): Whether to use the global vector or not. Defaults to False. - use_global_self_attn (bool, optional): Whether to use self attention among global vectors. Defaults to False. - separate_global_qkv (bool, optional): Whether to use different network to calc q_global, k_global, v_global. - Defaults to False. - global_dim_ratio (int, optional): The dim (channels) of global vectors is `global_dim_ratio*dim`. - Defaults to 1. - checkpoint_level (bool, optional): Whether to enable gradient checkpointing. Defaults to True. - use_relative_pos (bool, optional): Whether to use relative pos. Defaults to True. - use_relative_pos (bool, optional): Whether to use relative pos. Defaults to True. - attn_linear_init_mode (str, optional): The mode of attention linear initialization. Defaults to "0". - ffn_linear_init_mode (str, optional): The mode of FFN linear initialization. Defaults to "0". - norm_init_mode (str, optional): The mode of normalization initialization. Defaults to "0". - """ - - def __init__( - self, - dim: int, - num_heads: int, - block_cuboid_size: Tuple[Tuple[int, ...], ...] = [(4, 4, 4), (4, 4, 4)], - block_shift_size: Tuple[Tuple[int, ...], ...] = [(0, 0, 0), (2, 2, 2)], - block_strategy: Tuple[Tuple[str, ...], ...] = [ - ("d", "d", "d"), - ("l", "l", "l"), - ], - padding_type: str = "ignore", - qkv_bias: bool = False, - qk_scale: float = None, - attn_drop: float = 0.0, - proj_drop: float = 0.0, - ffn_drop: float = 0.0, - activation: str = "leaky", - gated_ffn: bool = False, - norm_layer: str = "layer_norm", - use_inter_ffn: bool = False, - use_global_vector: bool = False, - use_global_vector_ffn: bool = True, - use_global_self_attn: bool = False, - separate_global_qkv: bool = False, - global_dim_ratio: int = 1, - checkpoint_level: bool = True, - use_relative_pos: bool = True, - use_final_proj: bool = True, - attn_linear_init_mode: str = "0", - ffn_linear_init_mode: str = "0", - norm_init_mode: str = "0", - moe_config: dict = None, - expert_shape: tuple = None, - ): - super(StackCuboidSelfAttentionBlock, self).__init__() - self.attn_linear_init_mode = attn_linear_init_mode - self.ffn_linear_init_mode = ffn_linear_init_mode - self.norm_init_mode = norm_init_mode - if ( - len(block_cuboid_size[0]) <= 0 - or len(block_shift_size) <= 0 - or len(block_strategy) <= 0 - ): - raise ValueError( - "Format of the block cuboid size is not correct. block_cuboid_size={block_cuboid_size}" - ) - if len(block_cuboid_size) != len(block_shift_size) and len( - block_cuboid_size - ) != len(block_strategy): - raise ValueError( - "The lengths of block_cuboid_size, block_shift_size, and block_strategy must be equal." - ) - - self.num_attn = len(block_cuboid_size) - self.checkpoint_level = checkpoint_level - self.use_inter_ffn = use_inter_ffn - self.use_global_vector = use_global_vector - self.use_global_vector_ffn = use_global_vector_ffn - self.use_global_self_attn = use_global_self_attn - self.global_dim_ratio = global_dim_ratio - if self.use_inter_ffn: - if moe_config["use_ffn_moe"]: - self.ffn_l = nn.LayerList( - sublayers=[ - MixtureFFN( - units=dim, - hidden_size=4 * dim, - activation_dropout=ffn_drop, - dropout=ffn_drop, - gated_proj=gated_ffn, - activation=activation, - normalization=norm_layer, - pre_norm=True, - linear_init_mode=ffn_linear_init_mode, - norm_init_mode=norm_init_mode, - expert_shape=expert_shape, - moe_config=moe_config, - ) - for _ in range(self.num_attn) - ] - ) - else: - self.ffn_l = nn.LayerList( - sublayers=[ - PositionwiseFFN( - units=dim, - hidden_size=4 * dim, - activation_dropout=ffn_drop, - dropout=ffn_drop, - gated_proj=gated_ffn, - activation=activation, - normalization=norm_layer, - pre_norm=True, - linear_init_mode=ffn_linear_init_mode, - norm_init_mode=norm_init_mode, - expert_shape=expert_shape, - moe_config=moe_config, - ) - for _ in range(self.num_attn) - ] - ) - if self.use_global_vector_ffn and self.use_global_vector: - if moe_config["use_ffn_moe"]: - self.global_ffn_l = nn.LayerList( - sublayers=[ - MixtureFFN( - units=global_dim_ratio * dim, - hidden_size=global_dim_ratio * 4 * dim, - activation_dropout=ffn_drop, - dropout=ffn_drop, - gated_proj=gated_ffn, - activation=activation, - normalization=norm_layer, - pre_norm=True, - linear_init_mode=ffn_linear_init_mode, - norm_init_mode=norm_init_mode, - expert_shape=expert_shape, - moe_config=moe_config, - ) - for _ in range(self.num_attn) - ] - ) - else: - self.global_ffn_l = nn.LayerList( - sublayers=[ - PositionwiseFFN( - units=global_dim_ratio * dim, - hidden_size=global_dim_ratio * 4 * dim, - activation_dropout=ffn_drop, - dropout=ffn_drop, - gated_proj=gated_ffn, - activation=activation, - normalization=norm_layer, - pre_norm=True, - linear_init_mode=ffn_linear_init_mode, - norm_init_mode=norm_init_mode, - expert_shape=expert_shape, - moe_config=moe_config, - ) - for _ in range(self.num_attn) - ] - ) - else: - if moe_config["use_ffn_moe"]: - self.ffn_l = nn.LayerList( - sublayers=[ - MixtureFFN( - units=dim, - hidden_size=4 * dim, - activation_dropout=ffn_drop, - dropout=ffn_drop, - gated_proj=gated_ffn, - activation=activation, - normalization=norm_layer, - pre_norm=True, - linear_init_mode=ffn_linear_init_mode, - norm_init_mode=norm_init_mode, - expert_shape=expert_shape, - moe_config=moe_config, - ) - ] - ) - else: - self.ffn_l = nn.LayerList( - sublayers=[ - PositionwiseFFN( - units=dim, - hidden_size=4 * dim, - activation_dropout=ffn_drop, - dropout=ffn_drop, - gated_proj=gated_ffn, - activation=activation, - normalization=norm_layer, - pre_norm=True, - linear_init_mode=ffn_linear_init_mode, - norm_init_mode=norm_init_mode, - expert_shape=expert_shape, - moe_config=moe_config, - ) - ] - ) - if self.use_global_vector_ffn and self.use_global_vector: - if moe_config["use_ffn_moe"]: - self.global_ffn_l = nn.LayerList( - sublayers=[ - MixtureFFN( - units=global_dim_ratio * dim, - hidden_size=global_dim_ratio * 4 * dim, - activation_dropout=ffn_drop, - dropout=ffn_drop, - gated_proj=gated_ffn, - activation=activation, - normalization=norm_layer, - pre_norm=True, - linear_init_mode=ffn_linear_init_mode, - norm_init_mode=norm_init_mode, - expert_shape=expert_shape, - moe_config=moe_config, - ) - ] - ) - else: - self.global_ffn_l = nn.LayerList( - sublayers=[ - PositionwiseFFN( - units=global_dim_ratio * dim, - hidden_size=global_dim_ratio * 4 * dim, - activation_dropout=ffn_drop, - dropout=ffn_drop, - gated_proj=gated_ffn, - activation=activation, - normalization=norm_layer, - pre_norm=True, - linear_init_mode=ffn_linear_init_mode, - norm_init_mode=norm_init_mode, - expert_shape=expert_shape, - moe_config=moe_config, - ) - ] - ) - - if moe_config["use_attn_moe"]: - self.attn_l = nn.LayerList( - sublayers=[ - MixtureSelfAttention( - dim=dim, - num_heads=num_heads, - cuboid_size=ele_cuboid_size, - shift_size=ele_shift_size, - strategy=ele_strategy, - padding_type=padding_type, - qkv_bias=qkv_bias, - qk_scale=qk_scale, - attn_drop=attn_drop, - proj_drop=proj_drop, - norm_layer=norm_layer, - use_global_vector=use_global_vector, - use_global_self_attn=use_global_self_attn, - separate_global_qkv=separate_global_qkv, - global_dim_ratio=global_dim_ratio, - checkpoint_level=checkpoint_level, - use_relative_pos=use_relative_pos, - use_final_proj=use_final_proj, - attn_linear_init_mode=attn_linear_init_mode, - ffn_linear_init_mode=ffn_linear_init_mode, - norm_init_mode=norm_init_mode, - expert_shape=expert_shape, - moe_config=moe_config, - ) - for ele_cuboid_size, ele_shift_size, ele_strategy in zip( - block_cuboid_size, block_shift_size, block_strategy - ) - ] - ) - else: - self.attn_l = nn.LayerList( - sublayers=[ - CuboidSelfAttentionLayer( - dim=dim, - num_heads=num_heads, - cuboid_size=ele_cuboid_size, - shift_size=ele_shift_size, - strategy=ele_strategy, - padding_type=padding_type, - qkv_bias=qkv_bias, - qk_scale=qk_scale, - attn_drop=attn_drop, - proj_drop=proj_drop, - norm_layer=norm_layer, - use_global_vector=use_global_vector, - use_global_self_attn=use_global_self_attn, - separate_global_qkv=separate_global_qkv, - global_dim_ratio=global_dim_ratio, - checkpoint_level=checkpoint_level, - use_relative_pos=use_relative_pos, - use_final_proj=use_final_proj, - attn_linear_init_mode=attn_linear_init_mode, - ffn_linear_init_mode=ffn_linear_init_mode, - norm_init_mode=norm_init_mode, - ) - for ele_cuboid_size, ele_shift_size, ele_strategy in zip( - block_cuboid_size, block_shift_size, block_strategy - ) - ] - ) - - def reset_parameters(self): - for m in self.ffn_l: - m.reset_parameters() - if self.use_global_vector_ffn and self.use_global_vector: - for m in self.global_ffn_l: - m.reset_parameters() - for m in self.attn_l: - m.reset_parameters() - - def forward(self, x, global_vectors=None): - if self.use_inter_ffn: - if self.use_global_vector: - for idx, (attn, ffn) in enumerate(zip(self.attn_l, self.ffn_l)): - if self.checkpoint_level >= 2 and self.training: - x_out, global_vectors_out = fleet.utils.recompute( - attn, x, global_vectors - ) - else: - x_out, global_vectors_out = attn(x, global_vectors) - x = x + x_out - global_vectors = global_vectors + global_vectors_out - if self.checkpoint_level >= 1 and self.training: - x = fleet.utils.recompute(ffn, x) - if self.use_global_vector_ffn: - global_vectors = fleet.utils.recompute( - self.global_ffn_l[idx], global_vectors - ) - else: - x = ffn(x) - if self.use_global_vector_ffn: - global_vectors = self.global_ffn_l[idx](global_vectors) - return x, global_vectors - else: - for idx, (attn, ffn) in enumerate(zip(self.attn_l, self.ffn_l)): - if self.checkpoint_level >= 2 and self.training: - x = x + fleet.utils.recompute(attn, x) - else: - x = x + attn(x) - if self.checkpoint_level >= 1 and self.training: - x = fleet.utils.recompute(ffn, x) - else: - x = ffn(x) - return x - elif self.use_global_vector: - for idx, attn in enumerate(self.attn_l): - if self.checkpoint_level >= 2 and self.training: - x_out, global_vectors_out = fleet.utils.recompute( - attn, x, global_vectors - ) - else: - x_out, global_vectors_out = attn(x, global_vectors) - x = x + x_out - global_vectors = global_vectors + global_vectors_out - if self.checkpoint_level >= 1 and self.training: - x = fleet.utils.recompute(self.ffn_l[0], x) - if self.use_global_vector_ffn: - global_vectors = fleet.utils.recompute( - self.global_ffn_l[0], global_vectors - ) - else: - x = self.ffn_l[0](x) - if self.use_global_vector_ffn: - global_vectors = self.global_ffn_l[0](global_vectors) - return x, global_vectors - else: - for idx, attn in enumerate(self.attn_l): - if self.checkpoint_level >= 2 and self.training: - out = fleet.utils.recompute(attn, x) - else: - out = attn(x) - x = x + out - if self.checkpoint_level >= 1 and self.training: - x = fleet.utils.recompute(self.ffn_l[0], x) - else: - x = self.ffn_l[0](x) - return x - - -class CuboidTransformerEncoder(nn.Layer): - """Encoder of the CuboidTransformer - - x --> attn_block --> patch_merge --> attn_block --> patch_merge --> ... --> out - - Args: - input_shape (Tuple[int,...]): The shape of the input. Contains T, H, W, C - base_units (int, optional): The number of units. Defaults to 128. - block_units (int, optional): The number of block units. Defaults to None. - scale_alpha (float, optional): We scale up the channels based on the formula: - - round_to(base_units * max(downsample_scale) ** units_alpha, 4). Defaults to 1.0. - depth (list, optional): The number of layers for each block. Defaults to [4, 4, 4]. - downsample (int, optional): The downsample ratio. Defaults to 2. - downsample_type (str, optional): The type of downsample. Defaults to "patch_merge". - block_attn_patterns (str, optional): Attention pattern for the cuboid attention for each block. Defaults to None. - block_cuboid_size (list, optional): A list of cuboid size parameters. Defaults to [(4, 4, 4), (4, 4, 4)]. - block_strategy (list, optional): A list of cuboid strategies. Defaults to [("l", "l", "l"), ("d", "d", "d")]. - block_shift_size (list, optional): A list of shift sizes. Defaults to [(0, 0, 0), (0, 0, 0)]. - num_heads (int, optional): The number of heads. Defaults to 4. - attn_drop (float, optional): The ratio of attention dropout. Defaults to 0.0. - proj_drop (float, optional): The ratio of projection dropout. Defaults to 0.0. - ffn_drop (float, optional): The ratio of FFN dropout. Defaults to 0.0. - ffn_activation (str, optional): The FFN activation. Defaults to "leaky". - gated_ffn (bool, optional): Whether to use gate FFN. Defaults to False. - norm_layer (str, optional): The normalization layer. Defaults to "layer_norm". - use_inter_ffn (bool, optional): Whether to use inter FFN. Defaults to True. - padding_type (str, optional): The type of padding. Defaults to "ignore". - checkpoint_level (bool, optional): Whether to enable gradient checkpointing. Defaults to True. - use_relative_pos (bool, optional): Whether to use relative pos. Defaults to True. - self_attn_use_final_proj (bool, optional): Whether to use self attention for final projection. Defaults to True. - use_global_vector (bool, optional): Whether to use the global vector or not. Defaults to False. - use_global_vector_ffn (bool, optional): Whether to use FFN global vectors. Defaults to False. - use_global_self_attn (bool, optional): Whether to use global self attention. Defaults to False. - separate_global_qkv (bool, optional): Whether to use different network to calc q_global, k_global, v_global. - Defaults to False. - global_dim_ratio (int, optional): The dim (channels) of global vectors is `global_dim_ratio*dim`. - Defaults to 1. - attn_linear_init_mode (str, optional): The mode of attention linear initialization. Defaults to "0". - ffn_linear_init_mode (str, optional): The mode of FFN linear initialization. Defaults to "0". - conv_init_mode (str, optional): The mode of conv initialization. Defaults to "0". - down_linear_init_mode (str, optional): The mode of downsample linear initialization. Defaults to "0". - norm_init_mode (str, optional): The mode of normalization. Defaults to "0". - """ - - def __init__( - self, - input_shape: Tuple[int, ...], - base_units: int = 128, - block_units: int = None, - scale_alpha: float = 1.0, - depth: Tuple[int, ...] = [4, 4, 4], - downsample: int = 2, - downsample_type: str = "patch_merge", - block_attn_patterns: str = None, - block_cuboid_size: Tuple[Tuple[int, ...], ...] = [(4, 4, 4), (4, 4, 4)], - block_strategy: Tuple[Tuple[str, ...], ...] = [ - ("l", "l", "l"), - ("d", "d", "d"), - ], - block_shift_size: Tuple[Tuple[int, ...], ...] = [(0, 0, 0), (0, 0, 0)], - num_heads: int = 4, - attn_drop: float = 0.0, - proj_drop: float = 0.0, - ffn_drop: float = 0.0, - ffn_activation: str = "leaky", - gated_ffn: bool = False, - norm_layer: str = "layer_norm", - use_inter_ffn: bool = True, - padding_type: str = "ignore", - checkpoint_level: bool = True, - use_relative_pos: bool = True, - self_attn_use_final_proj: bool = True, - use_global_vector: bool = False, - use_global_vector_ffn: bool = True, - use_global_self_attn: bool = False, - separate_global_qkv: bool = False, - global_dim_ratio: int = 1, - attn_linear_init_mode: str = "0", - ffn_linear_init_mode: str = "0", - conv_init_mode: str = "0", - down_linear_init_mode: str = "0", - norm_init_mode: str = "0", - moe_config: dict = None, - ): - super(CuboidTransformerEncoder, self).__init__() - self.attn_linear_init_mode = attn_linear_init_mode - self.ffn_linear_init_mode = ffn_linear_init_mode - self.conv_init_mode = conv_init_mode - self.down_linear_init_mode = down_linear_init_mode - self.norm_init_mode = norm_init_mode - self.input_shape = input_shape - self.depth = depth - self.num_blocks = len(depth) - self.base_units = base_units - self.scale_alpha = scale_alpha - if not isinstance(downsample, (tuple, list)): - downsample = 1, downsample, downsample - self.downsample = downsample - self.downsample_type = downsample_type - self.num_heads = num_heads - self.use_global_vector = use_global_vector - self.checkpoint_level = checkpoint_level - - if block_units is None: - block_units = [ - cuboid_utils.round_to( - base_units * int((max(downsample) ** scale_alpha) ** i), 4 - ) - for i in range(self.num_blocks) - ] - else: - assert len(block_units) == self.num_blocks and block_units[0] == base_units - self.block_units = block_units - if self.num_blocks > 1: - if downsample_type == "patch_merge": - self.down_layers = nn.LayerList( - sublayers=[ - PatchMerging3D( - dim=self.block_units[i], - downsample=downsample, - padding_type=padding_type, - out_dim=self.block_units[i + 1], - linear_init_mode=down_linear_init_mode, - norm_init_mode=norm_init_mode, - ) - for i in range(self.num_blocks - 1) - ] - ) - else: - raise NotImplementedError(f"{downsample_type} is invalid.") - if self.use_global_vector: - self.down_layer_global_proj = nn.LayerList( - sublayers=[ - nn.Linear( - in_features=global_dim_ratio * self.block_units[i], - out_features=global_dim_ratio * self.block_units[i + 1], - ) - for i in range(self.num_blocks - 1) - ] - ) - if block_attn_patterns is not None: - mem_shapes = self.get_mem_shapes() - if isinstance(block_attn_patterns, (tuple, list)): - assert len(block_attn_patterns) == self.num_blocks - else: - block_attn_patterns = [ - block_attn_patterns for _ in range(self.num_blocks) - ] - block_cuboid_size = [] - block_strategy = [] - block_shift_size = [] - for idx, key in enumerate(block_attn_patterns): - func = cuboid_utils.CuboidSelfAttentionPatterns.get(key) - cuboid_size, strategy, shift_size = func(mem_shapes[idx]) - block_cuboid_size.append(cuboid_size) - block_strategy.append(strategy) - block_shift_size.append(shift_size) - else: - if not isinstance(block_cuboid_size[0][0], (list, tuple)): - block_cuboid_size = [block_cuboid_size for _ in range(self.num_blocks)] - else: - assert ( - len(block_cuboid_size) == self.num_blocks - ), f"Incorrect input format! Received block_cuboid_size={block_cuboid_size}" - if not isinstance(block_strategy[0][0], (list, tuple)): - block_strategy = [block_strategy for _ in range(self.num_blocks)] - else: - assert ( - len(block_strategy) == self.num_blocks - ), f"Incorrect input format! Received block_strategy={block_strategy}" - if not isinstance(block_shift_size[0][0], (list, tuple)): - block_shift_size = [block_shift_size for _ in range(self.num_blocks)] - else: - assert ( - len(block_shift_size) == self.num_blocks - ), f"Incorrect input format! Received block_shift_size={block_shift_size}" - self.block_cuboid_size = block_cuboid_size - self.block_strategy = block_strategy - self.block_shift_size = block_shift_size - - expert_shape_list = self.get_mem_shapes() - self.blocks = nn.LayerList( - sublayers=[ - nn.Sequential( - *[ - StackCuboidSelfAttentionBlock( - dim=self.block_units[i], - num_heads=num_heads, - block_cuboid_size=block_cuboid_size[i], - block_strategy=block_strategy[i], - block_shift_size=block_shift_size[i], - attn_drop=attn_drop, - proj_drop=proj_drop, - ffn_drop=ffn_drop, - activation=ffn_activation, - gated_ffn=gated_ffn, - norm_layer=norm_layer, - use_inter_ffn=use_inter_ffn, - padding_type=padding_type, - use_global_vector=use_global_vector, - use_global_vector_ffn=use_global_vector_ffn, - use_global_self_attn=use_global_self_attn, - separate_global_qkv=separate_global_qkv, - global_dim_ratio=global_dim_ratio, - checkpoint_level=checkpoint_level, - use_relative_pos=use_relative_pos, - use_final_proj=self_attn_use_final_proj, - attn_linear_init_mode=attn_linear_init_mode, - ffn_linear_init_mode=ffn_linear_init_mode, - norm_init_mode=norm_init_mode, - expert_shape=expert_shape_list[i], - moe_config=moe_config, - ) - for _ in range(depth[i]) - ] - ) - for i in range(self.num_blocks) - ] - ) - self.reset_parameters() - - def reset_parameters(self): - if self.num_blocks > 1: - for m in self.down_layers: - m.reset_parameters() - if self.use_global_vector: - cuboid_utils.apply_initialization( - self.down_layer_global_proj, linear_mode=self.down_linear_init_mode - ) - for ms in self.blocks: - for m in ms: - m.reset_parameters() - - def get_mem_shapes(self): - """Get the shape of the output memory based on the input shape. This can be used for constructing the decoder. - - Returns: - mem_shapes : A list of shapes of the output memory - """ - - if self.num_blocks == 1: - return [self.input_shape] - else: - mem_shapes = [self.input_shape] - curr_shape = self.input_shape - for down_layer in self.down_layers: - curr_shape = down_layer.get_out_shape(curr_shape) - mem_shapes.append(curr_shape) - return mem_shapes - - def forward(self, x, global_vectors=None): - """ - Args: - x : Shape (B, T, H, W, C) - - Returns: - out (List[paddle.Tensor,..]): A list of tensors from the bottom layer to the top layer of the encoder. For - example, it can have shape - - (B, T, H, W, C1) - - (B, T, H // 2, W // 2, 2 * C1) - - (B, T, H // 4, W // 4, 4 * C1) - ... - global_mem_out (List,Optional): The output of the global vector. - """ - - B, T, H, W, C_in = x.shape - assert (T, H, W, C_in) == self.input_shape - - if self.use_global_vector: - out = [] - global_mem_out = [] - for i in range(self.num_blocks): - for l in self.blocks[i]: - x, global_vectors = l(x, global_vectors) - out.append(x) - global_mem_out.append(global_vectors) - if self.num_blocks > 1 and i < self.num_blocks - 1: - x = self.down_layers[i](x) - global_vectors = self.down_layer_global_proj[i](global_vectors) - return out, global_mem_out - else: - out = [] - for i in range(self.num_blocks): - x = self.blocks[i](x) - out.append(x) - if self.num_blocks > 1 and i < self.num_blocks - 1: - x = self.down_layers[i](x) - return out - - -class MixtureLinear(nn.Layer): - def __init__(self, in_dim, out_dim, expert_shape, moe_config, bias_attr=True): - super().__init__() - - self.in_dim = in_dim - self.out_dim = out_dim - self.bias = bias_attr - self.expert_shape = expert_shape # T, H, W, C_o - self.num_experts = moe_config["num_experts"] - self.out_planes = moe_config["out_planes"] - self.moe_config = moe_config - assert expert_shape is not None and moe_config["use_linear_moe"] - - if moe_config["gate_style"] == "linear": - self.gate = moe_utils.LinearGatingNet(moe_config, expert_shape, in_dim) - elif moe_config["gate_style"] == "spatial-latent": - self.gate = moe_utils.SpatialLatentGatingNet( - moe_config, expert_shape, in_dim - ) - elif moe_config["gate_style"] == "cuboid-latent": - self.gate = moe_utils.CuboidLatentGatingNet( - moe_config, expert_shape, in_dim - ) - elif moe_config["gate_style"] == "spatial-latent-linear": - self.gate = moe_utils.SpatialLatentLinearGatingNet( - moe_config, expert_shape, in_dim - ) - elif moe_config["gate_style"] == "cuboid-latent-linear": - self.gate = moe_utils.CuboidLatentLinearGatingNet( - moe_config, expert_shape, in_dim - ) - else: - raise NotImplementedError - - self.experts = nn.LayerList( - [ - nn.Linear(in_features=in_dim, out_features=out_dim, bias_attr=bias_attr) - for _ in range(self.num_experts) - ] - ) - - def forward(self, x): - - B, T, H, W, C = x.shape - E = self.num_experts - assert C == self.in_dim and list(self.expert_shape)[:-1] == x.shape[1:-1] - ( - dense_routing_weights, - sparse_routing_weights, - sparse_routing_inds, - self.aux_loss, - ) = self.gate( - x - ) # dense: B, T, H, W, E - - if self.moe_config["dispatch_style"] == "dense": - dispatcher = moe_utils.DenseDispatcher( - E, - sparse_routing_weights.reshape([B * T * H * W, -1]), - sparse_routing_inds.reshape([B * T * H * W, -1]), - ) - expert_outputs = paddle.stack( - [self.experts[i](x.reshape([B * T * H * W, -1])) for i in range(E)], - axis=-2, - ) - y = dispatcher.combine(expert_outputs).reshape([B, T, H, W, -1]) - elif self.moe_config["dispatch_style"] == "sparse": - dispatcher = moe_utils.SparseDispatcher( - E, - sparse_routing_weights.reshape([B * T * H * W, -1]), - sparse_routing_inds.reshape([B * T * H * W, -1]), - ) - expert_inputs = dispatcher.dispatch(x.reshape([B * T * H * W, -1])) - expert_outputs = [ - self.experts[i](expert_inputs[i]) - if expert_inputs[i].shape[0] > 0 - else paddle.zeros([0, self.out_dim]) - for i in range(E) - ] - y = dispatcher.combine(expert_outputs).reshape([B, T, H, W, -1]) - else: - raise NotImplementedError - - return y - - -class MixtureFFN(nn.Layer): - def __init__( - self, - units, - hidden_size, - activation_dropout, - dropout, - gated_proj, - activation, - normalization, - pre_norm, - linear_init_mode, - norm_init_mode, - expert_shape, - moe_config, - ): - super().__init__() - - self.in_dim = units - self.out_dim = units - self.expert_shape = expert_shape # T, H, W, C_o - self.num_experts = moe_config["num_experts"] - self.out_planes = moe_config["out_planes"] - self.moe_config = moe_config - assert expert_shape is not None and moe_config["use_ffn_moe"] - - if moe_config["gate_style"] == "linear": - self.gate = moe_utils.LinearGatingNet(moe_config, expert_shape, units) - elif moe_config["gate_style"] == "spatial-latent": - self.gate = moe_utils.SpatialLatentGatingNet( - moe_config, expert_shape, units - ) - elif moe_config["gate_style"] == "cuboid-latent": - self.gate = moe_utils.CuboidLatentGatingNet(moe_config, expert_shape, units) - elif moe_config["gate_style"] == "spatial-latent-linear": - self.gate = moe_utils.SpatialLatentLinearGatingNet( - moe_config, expert_shape, units - ) - elif moe_config["gate_style"] == "cuboid-latent-linear": - self.gate = moe_utils.CuboidLatentLinearGatingNet( - moe_config, expert_shape, units - ) - else: - raise NotImplementedError - - self.experts = nn.LayerList( - [ - PositionwiseFFN( - units=units, - hidden_size=hidden_size, - activation_dropout=activation_dropout, - dropout=dropout, - gated_proj=gated_proj, - activation=activation, - normalization=normalization, - pre_norm=pre_norm, - linear_init_mode=linear_init_mode, - norm_init_mode=norm_init_mode, - moe_config=moe_config, - expert_shape=expert_shape, - ) - for _ in range(self.num_experts) - ] - ) - - def forward(self, x): - - B, T, H, W, C = x.shape - E = self.num_experts - assert C == self.in_dim and list(self.expert_shape)[:-1] == x.shape[1:-1] - ( - dense_routing_weights, - sparse_routing_weights, - sparse_routing_inds, - self.aux_loss, - ) = self.gate( - x - ) # dense: B, T, H, W, E - - if self.moe_config["dispatch_style"] == "dense": - dispatcher = moe_utils.DenseDispatcher( - E, - sparse_routing_weights.reshape([B * T * H * W, -1]), - sparse_routing_inds.reshape([B * T * H * W, -1]), - ) - expert_outputs = paddle.stack( - [self.experts[i](x.reshape([B * T * H * W, -1])) for i in range(E)], - axis=-2, - ) - y = dispatcher.combine(expert_outputs).reshape([B, T, H, W, C]) - elif self.moe_config["dispatch_style"] == "sparse": - dispatcher = moe_utils.SparseDispatcher( - E, - sparse_routing_weights.reshape([B * T * H * W, -1]), - sparse_routing_inds.reshape([B * T * H * W, -1]), - ) - expert_inputs = dispatcher.dispatch(x.reshape([B * T * H * W, -1])) - expert_outputs = [ - self.experts[i](expert_inputs[i]) - if expert_inputs[i].shape[0] > 0 - else paddle.zeros([0, self.out_dim]) - for i in range(E) - ] - y = dispatcher.combine(expert_outputs).reshape([B, T, H, W, C]) - else: - raise NotImplementedError - - return y - - def reset_parameters(self): - - for i in range(len(self.experts)): - self.experts[i].reset_parameters() - - -class MixtureSelfAttention(nn.Layer): - def __init__( - self, - dim, - num_heads, - cuboid_size, - shift_size, - strategy, - padding_type, - qkv_bias, - qk_scale, - attn_drop, - proj_drop, - norm_layer, - use_global_vector, - use_global_self_attn, - separate_global_qkv, - global_dim_ratio, - checkpoint_level, - use_relative_pos, - use_final_proj, - attn_linear_init_mode, - ffn_linear_init_mode, - norm_init_mode, - expert_shape, - moe_config, - ): - super().__init__() - - self.in_dim = dim - self.out_dim = dim - self.expert_shape = expert_shape # T, H, W, C - self.num_experts = moe_config["num_experts"] - self.out_planes = moe_config["out_planes"] - self.moe_config = moe_config - assert expert_shape is not None and moe_config["use_attn_moe"] - assert not use_global_vector - - if moe_config["gate_style"] == "linear": - self.gate = moe_utils.LinearGatingNet(moe_config, expert_shape, dim) - elif moe_config["gate_style"] == "spatial-latent": - self.gate = moe_utils.SpatialLatentGatingNet(moe_config, expert_shape, dim) - elif moe_config["gate_style"] == "cuboid-latent": - self.gate = moe_utils.CuboidLatentGatingNet(moe_config, expert_shape, dim) - elif moe_config["gate_style"] == "spatial-latent-linear": - self.gate = moe_utils.SpatialLatentLinearGatingNet( - moe_config, expert_shape, dim - ) - elif moe_config["gate_style"] == "cuboid-latent-linear": - self.gate = moe_utils.CuboidLatentLinearGatingNet( - moe_config, expert_shape, dim - ) - else: - raise NotImplementedError - - self.experts = nn.LayerList( - [ - CuboidSelfAttentionLayer( - dim=dim, - num_heads=num_heads, - cuboid_size=cuboid_size, - shift_size=shift_size, - strategy=strategy, - padding_type=padding_type, - qkv_bias=qkv_bias, - qk_scale=qk_scale, - attn_drop=attn_drop, - proj_drop=proj_drop, - norm_layer=norm_layer, - use_global_vector=use_global_vector, - use_global_self_attn=use_global_self_attn, - separate_global_qkv=separate_global_qkv, - global_dim_ratio=global_dim_ratio, - checkpoint_level=checkpoint_level, - use_relative_pos=use_relative_pos, - use_final_proj=use_final_proj, - attn_linear_init_mode=attn_linear_init_mode, - ffn_linear_init_mode=ffn_linear_init_mode, - norm_init_mode=norm_init_mode, - ) - for _ in range(self.num_experts) - ] - ) - - def forward(self, x, global_vectors=None): - - B, T, H, W, C = x.shape - E = self.num_experts - assert C == self.in_dim and list(self.expert_shape)[:-1] == x.shape[1:-1] - ( - dense_routing_weights, - sparse_routing_weights, - sparse_routing_inds, - self.aux_loss, - ) = self.gate( - x - ) # dense: B, T, H, W, E - - dispatcher = moe_utils.DenseDispatcher( - E, - sparse_routing_weights.reshape([B * T * H * W, -1]), - sparse_routing_inds.reshape([B * T * H * W, -1]), - ) - expert_outputs = paddle.stack( - [self.experts[i](x, global_vectors) for i in range(E)], axis=-2 - ).reshape([B * T * H * W, E, C]) - y = dispatcher.combine(expert_outputs).reshape([B, T, H, W, C]) - - return y - - def reset_parameters(self): - - for i in range(len(self.experts)): - self.experts[i].reset_parameters() diff --git a/examples/smc_reac/ppsci/arch/extformer_moe_cuboid_utils.py b/examples/smc_reac/ppsci/arch/extformer_moe_cuboid_utils.py deleted file mode 100644 index 20531c82d6..0000000000 --- a/examples/smc_reac/ppsci/arch/extformer_moe_cuboid_utils.py +++ /dev/null @@ -1,350 +0,0 @@ -import functools -from typing import Tuple - -import paddle -import paddle.nn.functional as F -from paddle import nn - -from ppsci.utils import initializer - - -def round_to(dat, c): - return dat + (dat - dat % c) % c - - -class RMSNorm(nn.Layer): - """Root Mean Square Layer Normalization proposed in "[NeurIPS2019] Root Mean Square Layer Normalization" - - Args: - d (Optional[int]): The model size. - p (float, optional): The partial RMSNorm, valid value [0, 1]. Defaults to -1.0. - eps (float, optional): The epsilon value. Defaults to 1e-08. - bias (bool, optional): Whether use bias term for RMSNorm, - because RMSNorm doesn't enforce re-centering invariance.Defaults to False. - """ - - def __init__( - self, - d: Tuple[int, ...], - p: float = -1.0, - eps: float = 1e-08, - bias: bool = False, - ): - super().__init__() - self.eps = eps - self.d = d - self.p = p - self.bias = bias - init_data = paddle.ones(d) - self.scale = paddle.create_parameter( - shape=init_data.shape, - dtype=init_data.dtype, - default_initializer=nn.initializer.Constant(1.0), - ) - self.scale.stop_gradient = False - self.add_parameter(name="scale", parameter=self.scale) - if self.bias: - init_data = paddle.zeros(d) - self.offset = paddle.create_parameter( - shape=init_data.shape, - dtype=init_data.dtype, - default_initializer=nn.initializer.Constant(0.0), - ) - self.offset.stop_gradient = False - self.add_parameter(name="offset", parameter=self.offset) - - def forward(self, x): - if self.p < 0.0 or self.p > 1.0: - norm_x = x.norm(p=2, axis=-1, keepdim=True) - d_x = self.d - else: - partial_size = int(self.d * self.p) - partial_x, _ = paddle.split( - x=x, num_or_sections=[partial_size, self.d - partial_size], axis=-1 - ) - norm_x = partial_x.norm(p=2, axis=-1, keepdim=True) - d_x = partial_size - rms_x = norm_x * d_x ** (-1.0 / 2) - x_normed = x / (rms_x + self.eps) - if self.bias: - return self.scale * x_normed + self.offset - return self.scale * x_normed - - -def get_norm_layer( - normalization: str = "layer_norm", - axis: int = -1, - epsilon: float = 1e-05, - in_channels: int = 0, - **kwargs, -): - """Get the normalization layer based on the provided type - - Args: - normalization (str): The type of the layer normalization from ['layer_norm']. - axis (float): The axis to normalize the. - epsilon (float): The epsilon of the normalization layer. - in_channels (int): Input channel. - - Returns: - norm_layer (norm): The layer normalization layer. - """ - - if isinstance(normalization, str): - if normalization == "layer_norm": - assert in_channels > 0 - assert axis == -1 - norm_layer = nn.LayerNorm( - normalized_shape=in_channels, epsilon=epsilon, **kwargs - ) - elif normalization == "rms_norm": - assert axis == -1 - norm_layer = RMSNorm(d=in_channels, eps=epsilon, **kwargs) - else: - raise NotImplementedError(f"normalization={normalization} is not supported") - return norm_layer - elif normalization is None: - return nn.Identity() - else: - raise NotImplementedError("The type of normalization must be str") - - -def generalize_padding(x, pad_t, pad_h, pad_w, padding_type, t_pad_left=False): - if pad_t == 0 and pad_h == 0 and pad_w == 0: - return x - assert padding_type in ["zeros", "ignore", "nearest"] - B, T, H, W, C = x.shape - if padding_type == "nearest": - return nn.functional.interpolate( - x=x.transpose(perm=[0, 4, 1, 2, 3]), size=(T + pad_t, H + pad_h, W + pad_w) - ).transpose(perm=[0, 2, 3, 4, 1]) - elif t_pad_left: - return F.pad(x, [0, 0, 0, pad_w, 0, pad_h, pad_t, 0], data_format="NDHWC") - else: - data_pad = F.pad( - x, [0, 0, pad_t, 0, pad_h, 0, pad_w, 0, 0, 0], data_format="NDHWC" - ) - data_pad = paddle.concat( - [data_pad[:, pad_t:, ...], data_pad[:, :pad_t, ...]], axis=1 - ) - return data_pad - - -def generalize_unpadding(x, pad_t, pad_h, pad_w, padding_type): - assert padding_type in ["zeros", "ignore", "nearest"] - B, T, H, W, C = x.shape - if pad_t == 0 and pad_h == 0 and pad_w == 0: - return x - if padding_type == "nearest": - return nn.functional.interpolate( - x=x.transpose(perm=[0, 4, 1, 2, 3]), size=(T - pad_t, H - pad_h, W - pad_w) - ).transpose(perm=[0, 2, 3, 4, 1]) - else: - return x[:, : T - pad_t, : H - pad_h, : W - pad_w, :] - - -def apply_initialization( - m: nn.Layer, - linear_mode: str = "0", - conv_mode: str = "0", - norm_mode: str = "0", - embed_mode: str = "0", -): - if isinstance(m, nn.Linear): - if linear_mode in ("0",): - m.weight = initializer.kaiming_normal_(m.weight, nonlinearity="linear") - elif linear_mode in ("1",): - m.weight = initializer.kaiming_normal_( - m.weight, a=0.1, mode="fan_out", nonlinearity="leaky_relu" - ) - else: - raise NotImplementedError(f"{linear_mode} is invalid.") - if hasattr(m, "bias") and m.bias is not None: - m.bias = initializer.zeros_(m.bias) - elif isinstance( - m, - ( - nn.Conv2D, - nn.Conv3D, - nn.Conv2DTranspose, - nn.Conv3DTranspose, - ), - ): - if conv_mode in ("0",): - m.weight = initializer.kaiming_normal_( - m.weight, a=0.1, mode="fan_out", nonlinearity="leaky_relu" - ) - else: - raise NotImplementedError(f"{conv_mode} is invalid.") - if hasattr(m, "bias") and m.bias is not None: - m.bias = initializer.zeros_(m.bias) - elif isinstance(m, nn.LayerNorm): - if norm_mode in ("0",): - m.weight = initializer.zeros_(m.weight) - m.bias = initializer.zeros_(m.bias) - else: - raise NotImplementedError(f"{norm_mode} is invalid.") - elif isinstance(m, nn.GroupNorm): - if norm_mode in ("0",): - m.weight = initializer.ones_(m.weight) - m.bias = initializer.zeros_(m.bias) - else: - raise NotImplementedError(f"{norm_mode} is invalid.") - elif isinstance(m, nn.Embedding): - if embed_mode in ("0",): - m.weight.data = initializer.trunc_normal_(m.weight.data, std=0.02) - else: - raise NotImplementedError(f"{embed_mode} is invalid.") - elif isinstance(m, nn.Layer) and hasattr(m, "experts"): - for lin in m.experts: - assert isinstance(lin, nn.Linear) - apply_initialization(lin, linear_mode=linear_mode) - else: - pass - - -class CuboidSelfAttentionPatterns: - def __init__(self): - super().__init__() - self.patterns = {} - self.patterns = { - "full": self.full_attention, - "axial": self.axial, - "divided_st": self.divided_space_time, - } - for p in [1, 2, 4, 8, 10]: - for m in [1, 2, 4, 8, 16, 32]: - key = f"video_swin_{p}x{m}" - self.patterns[key] = functools.partial(self.video_swin, P=p, M=m) - - for m in [1, 2, 4, 8, 16, 32]: - key = f"spatial_lg_{m}" - self.patterns[key] = functools.partial(self.spatial_lg_v1, M=m) - - for k in [2, 4, 8]: - key = f"axial_space_dilate_{k}" - self.patterns[key] = functools.partial(self.axial_space_dilate_K, K=k) - - def get(self, pattern_name): - return self.patterns[pattern_name] - - def full_attention(self, input_shape): - T, H, W, _ = input_shape - cuboid_size = [(T, H, W)] - strategy = [("l", "l", "l")] - shift_size = [(0, 0, 0)] - return cuboid_size, strategy, shift_size - - def axial(self, input_shape): - """Axial attention proposed in https://arxiv.org/abs/1912.12180 - - Args: - input_shape (Tuple[int,...]): The shape of the input tensor, T H W. - - Returns: - cuboid_size (Tuple[int,...]): The size of cuboid. - strategy (Tuple[str,...]): The strategy of the attention. - shift_size (Tuple[int,...]): The shift size of the attention. - """ - - T, H, W, _ = input_shape - cuboid_size = [(T, 1, 1), (1, H, 1), (1, 1, W)] - strategy = [("l", "l", "l"), ("l", "l", "l"), ("l", "l", "l")] - shift_size = [(0, 0, 0), (0, 0, 0), (0, 0, 0)] - return cuboid_size, strategy, shift_size - - def divided_space_time(self, input_shape): - T, H, W, _ = input_shape - cuboid_size = [(T, 1, 1), (1, H, W)] - strategy = [("l", "l", "l"), ("l", "l", "l")] - shift_size = [(0, 0, 0), (0, 0, 0)] - return cuboid_size, strategy, shift_size - - def video_swin(self, input_shape, P=2, M=4): - """Adopt the strategy in Video SwinTransformer https://arxiv.org/pdf/2106.13230.pdf""" - T, H, W, _ = input_shape - P = min(P, T) - M = min(M, H, W) - cuboid_size = [(P, M, M), (P, M, M)] - strategy = [("l", "l", "l"), ("l", "l", "l")] - shift_size = [(0, 0, 0), (P // 2, M // 2, M // 2)] - return cuboid_size, strategy, shift_size - - def spatial_lg_v1(self, input_shape, M=4): - T, H, W, _ = input_shape - if H <= M and W <= M: - cuboid_size = [(T, 1, 1), (1, H, W)] - strategy = [("l", "l", "l"), ("l", "l", "l")] - shift_size = [(0, 0, 0), (0, 0, 0)] - else: - cuboid_size = [(T, 1, 1), (1, M, M), (1, M, M)] - strategy = [("l", "l", "l"), ("l", "l", "l"), ("d", "d", "d")] - shift_size = [(0, 0, 0), (0, 0, 0), (0, 0, 0)] - return cuboid_size, strategy, shift_size - - def axial_space_dilate_K(self, input_shape, K=2): - T, H, W, _ = input_shape - K = min(K, H, W) - cuboid_size = [ - (T, 1, 1), - (1, H // K, 1), - (1, H // K, 1), - (1, 1, W // K), - (1, 1, W // K), - ] - strategy = [ - ("l", "l", "l"), - ("d", "d", "d"), - ("l", "l", "l"), - ("d", "d", "d"), - ("l", "l", "l"), - ] - shift_size = [(0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0)] - return cuboid_size, strategy, shift_size - - -class CuboidCrossAttentionPatterns: - def __init__(self): - super().__init__() - self.patterns = {} - for k in [1, 2, 4, 8]: - key1 = f"cross_{k}x{k}" - key2 = f"cross_{k}x{k}_lg" - key3 = f"cross_{k}x{k}_heter" - self.patterns[key1] = functools.partial(self.cross_KxK, K=k) - self.patterns[key2] = functools.partial(self.cross_KxK_lg, K=k) - self.patterns[key3] = functools.partial(self.cross_KxK_heter, K=k) - - def get(self, pattern_name): - return self.patterns[pattern_name] - - def cross_KxK(self, mem_shape, K): - T_mem, H, W, _ = mem_shape - K = min(K, H, W) - cuboid_hw = [(K, K)] - shift_hw = [(0, 0)] - strategy = [("l", "l", "l")] - n_temporal = [1] - return cuboid_hw, shift_hw, strategy, n_temporal - - def cross_KxK_lg(self, mem_shape, K): - T_mem, H, W, _ = mem_shape - K = min(K, H, W) - cuboid_hw = [(K, K), (K, K)] - shift_hw = [(0, 0), (0, 0)] - strategy = [("l", "l", "l"), ("d", "d", "d")] - n_temporal = [1, 1] - return cuboid_hw, shift_hw, strategy, n_temporal - - def cross_KxK_heter(self, mem_shape, K): - T_mem, H, W, _ = mem_shape - K = min(K, H, W) - cuboid_hw = [(K, K), (K, K), (K, K)] - shift_hw = [(0, 0), (0, 0), (K // 2, K // 2)] - strategy = [("l", "l", "l"), ("d", "d", "d"), ("l", "l", "l")] - n_temporal = [1, 1, 1] - return cuboid_hw, shift_hw, strategy, n_temporal - - -CuboidSelfAttentionPatterns = CuboidSelfAttentionPatterns() -CuboidCrossAttentionPatterns = CuboidCrossAttentionPatterns() diff --git a/examples/smc_reac/ppsci/arch/extformer_moe_utils.py b/examples/smc_reac/ppsci/arch/extformer_moe_utils.py deleted file mode 100644 index 3332b356c8..0000000000 --- a/examples/smc_reac/ppsci/arch/extformer_moe_utils.py +++ /dev/null @@ -1,563 +0,0 @@ -import math - -import paddle -from paddle import nn - -# MoE Gating - - -class GatingNet(nn.Layer): - def __init__(self, moe_config, input_shape, in_channels): - super().__init__() - - self.num_experts = moe_config["num_experts"] - self.out_planes = moe_config["out_planes"] - self.aux_loss_style = moe_config["aux_loss_style"] - assert self.out_planes > 1 and self.out_planes <= self.num_experts - assert len(input_shape) == 4 - self.input_shape = input_shape - - self.noise_lin = nn.Linear( - in_features=in_channels, out_features=self.num_experts, bias_attr=False - ) - self.noise_eps = 1e-2 - self.softplus = nn.Softplus() - self.softmax = nn.Softmax(axis=-1) - - self.importance_weight = moe_config["importance_weight"] - self.load_weight = moe_config["load_weight"] - - def cv_squared(self, x, eps=1e-25): - return x.var(axis=-1) / (x.mean(axis=-1) ** 2 + eps) - - def intra_cdf(self, value, loc=0.0, scale=1.0): - return 0.5 * (1 + paddle.erf((value - loc) / scale / math.sqrt(2))) - - def importance_loss_cell(self, routing_weights): - importance_loss = self.cv_squared(routing_weights.sum(axis=0)).mean() - return importance_loss - - def load_loss_cell( - self, clean_values, noisy_values, noise_stddev, noisy_top_values - ): - B, T, H, W, E = clean_values.shape - M = noisy_top_values.shape[-1] - clean_values = clean_values.transpose([1, 2, 3, 0, 4]) - noisy_values = noisy_values.transpose([1, 2, 3, 0, 4]) - noise_stddev = noise_stddev.transpose([1, 2, 3, 0, 4]) - top_values_flat = noisy_top_values.transpose([1, 2, 3, 0, 4]).reshape( - [T, H, W, B * M] - ) - - threshold_positions_if_in = paddle.arange(B) * M + self.out_planes - threshold_if_in = paddle.take_along_axis( - top_values_flat, - axis=-1, - indices=threshold_positions_if_in.unsqueeze(axis=[0, 1, 2]), - ).unsqueeze( - -1 - ) # T, H, W, B, 1 - is_in = noisy_values > threshold_if_in # T, H, W, B, E - threshold_positions_if_out = threshold_positions_if_in - 1 - threshold_if_out = paddle.take_along_axis( - top_values_flat, - axis=-1, - indices=threshold_positions_if_out.unsqueeze(axis=[0, 1, 2]), - ).unsqueeze(-1) - - prob_if_in = self.intra_cdf( - (clean_values - threshold_if_in) / noise_stddev - ) # T, H, W, B, E - prob_if_out = self.intra_cdf( - (clean_values - threshold_if_out) / noise_stddev - ) # T, H, W, B, E - prob = paddle.where(is_in, prob_if_in, prob_if_out) # T, H, W, B, E - - load_loss = self.cv_squared(prob.sum(axis=-2)).mean() - return load_loss - - def importance_loss_all(self, routing_weights): - importance_loss = self.cv_squared(routing_weights.sum(axis=0)) - return importance_loss - - def load_loss_all(self, clean_values, noisy_values, noise_stddev, noisy_top_values): - B, E = clean_values.shape - M = noisy_top_values.shape[-1] - top_values_flat = noisy_top_values.flatten() # B * M - - threshold_positions_if_in = paddle.arange(B) * M + self.out_planes # B - threshold_if_in = paddle.take_along_axis( - top_values_flat, axis=-1, indices=threshold_positions_if_in - ).unsqueeze( - -1 - ) # B, 1 - is_in = noisy_values > threshold_if_in # B, E - threshold_positions_if_out = threshold_positions_if_in - 1 # B - threshold_if_out = paddle.take_along_axis( - top_values_flat, axis=-1, indices=threshold_positions_if_out - ).unsqueeze( - -1 - ) # B, 1 - - prob_if_in = self.intra_cdf( - (clean_values - threshold_if_in) / noise_stddev - ) # B, E - prob_if_out = self.intra_cdf( - (clean_values - threshold_if_out) / noise_stddev - ) # B, E - prob = paddle.where(is_in, prob_if_in, prob_if_out) # B, E - - load_loss = self.cv_squared(prob.sum(axis=0)) - return load_loss - - def forward(self, x, t_map=None, eps=1e-25, dense_routing=False): - assert x.shape[1:-1] == list(self.input_shape)[:-1] - B, T, H, W, C = x.shape - E = self.num_experts - - raw_logits = self.gating(x, t_map) - if self.training: - noise = self.softplus(self.noise_lin(x)) + self.noise_eps - noisy_logits = raw_logits + paddle.randn(shape=raw_logits.shape) * noise - logits = noisy_logits - else: - logits = raw_logits - - assert logits.shape[-1] == self.num_experts - logits = self.softmax(logits) # [B, T, H, W, E] - top_logits, top_indices = logits.topk( - min(self.out_planes + 1, self.num_experts), axis=-1 - ) - top_k_logits = top_logits[:, :, :, :, : self.out_planes] - top_k_indices = top_indices[:, :, :, :, : self.out_planes] - top_k_gates = top_k_logits / ( - top_k_logits.sum(axis=-1, keepdim=True) + eps - ) # normalization - - if dense_routing: - # zeros = paddle.zeros_like(logits) - # zeros.stop_gradient = False - # print(zeros.shape) - # print(top_k_gates.shape, top_k_gates[0, 0, 0, 0]) - # routing_weights = paddle.put_along_axis(zeros, axis=-1, indices=top_k_indices, values=top_k_gates) - # print(routing_weights.shape, routing_weights.stop_gradient) - pass - else: - routing_weights = None - - if self.training: - if self.aux_loss_style == "cell": - # importance_loss = self.importance_loss(routing_weights) - importance_loss = self.importance_loss_cell(logits) - load_loss = self.load_loss_cell( - raw_logits, noisy_logits, noise, top_logits - ) - elif self.aux_loss_style == "all": - importance_loss = self.importance_loss_all( - logits.reshape([B * T * H * W, E]) - ) - load_loss = self.load_loss_all( - raw_logits.reshape([B * T * H * W, E]), - noisy_logits.reshape([B * T * H * W, E]), - noise.reshape([B * T * H * W, E]), - top_logits.reshape([B * T * H * W, -1]), - ) - else: - raise NotImplementedError - loss = ( - self.importance_weight * importance_loss + self.load_weight * load_loss - ) - else: - loss = None - - return routing_weights, top_k_gates, top_k_indices, loss - - -class LinearGatingNet(GatingNet): - def __init__(self, moe_config, input_shape, in_channels): - super().__init__(moe_config, input_shape, in_channels) - assert len(input_shape) == 4 - T, H, W, C = input_shape - - self.lin = nn.Linear( - in_features=in_channels, out_features=self.num_experts, bias_attr=False - ) - - def gating(self, x, t_map=None): - routing_weights = self.lin(x) # [B, T, H, W, E] - return routing_weights - - -class SpatialLatentGatingNet(GatingNet): - def __init__(self, moe_config, input_shape, in_channels): - super().__init__(moe_config, input_shape, in_channels) - assert len(input_shape) == 4 - T, H, W, C = input_shape - - gain = 1.0 - fan = self.out_planes / self.num_experts - bound = gain * math.sqrt(3.0 / fan) - self.routing_weights = paddle.create_parameter( - shape=[H, W, self.num_experts], - dtype="float32", - default_initializer=nn.initializer.Uniform(-bound, bound), - ) - - def gating(self, x, t_map=None): - # assert t_map is not None - routing_weights = self.routing_weights.unsqueeze(0).tile( - [x.shape[0], x.shape[1], 1, 1, 1] - ) # [B, T, H, W, E] - return routing_weights - - -class SpatialLatentLinearGatingNet(GatingNet): - def __init__(self, moe_config, input_shape, in_channels): - super().__init__(moe_config, input_shape, in_channels) - assert len(input_shape) == 4 - T, H, W, C = input_shape - - gain = 1.0 - fan = self.out_planes / self.num_experts - bound = gain * math.sqrt(3.0 / fan) - self.spatial_routing_weights = paddle.create_parameter( - shape=[H, W, self.num_experts], - dtype="float32", - default_initializer=nn.initializer.Uniform(-bound, bound), - ) - self.lin = nn.Linear( - in_features=in_channels, out_features=self.num_experts, bias_attr=False - ) - - self.combine_weight = paddle.create_parameter( - shape=[H, W, self.num_experts, 2], - dtype="float32", - default_initializer=nn.initializer.Uniform(-bound, bound), - ) - - def gating(self, x, t_map=None): - # assert t_map is not None - spatial_routing_weights = self.spatial_routing_weights.tile( - [x.shape[0], x.shape[1], 1, 1, 1] - ) # [B, T, H, W, E] - linear_routing_weights = self.lin(x) # [B, T, H, W, E] - routing_weights = paddle.stack( - [spatial_routing_weights, linear_routing_weights], axis=-1 - ) # [B, T, H, W, E, 2] - combine_weight = self.combine_weight.tile( - [x.shape[0], x.shape[1], 1, 1, 1, 1] - ) # [B, T, H, W, E, 2] - routing_weights = (routing_weights * combine_weight).sum(-1) # [B, T, H, W, E] - return routing_weights - - -class CuboidLatentGatingNet(GatingNet): - def __init__(self, moe_config, input_shape, in_channels): - super().__init__(moe_config, input_shape, in_channels) - assert len(input_shape) == 4 - T, H, W, C = input_shape - - gain = 1.0 - fan = self.out_planes / self.num_experts - bound = gain * math.sqrt(3.0 / fan) - self.routing_weights = paddle.create_parameter( - shape=[T, H, W, self.num_experts], - dtype="float32", - default_initializer=nn.initializer.Uniform(-bound, bound), - ) - - def gating(self, x, t_map=None): - # assert t_map is not None - routing_weights = self.routing_weights.unsqueeze(0).tile( - [x.shape[0], 1, 1, 1, 1] - ) # [B, T, H, W, E] - return routing_weights - - -class CuboidLatentLinearGatingNet(GatingNet): - def __init__(self, moe_config, input_shape, in_channels): - super().__init__(moe_config, input_shape, in_channels) - assert len(input_shape) == 4 - T, H, W, C = input_shape - - gain = 1.0 - fan = self.out_planes / self.num_experts - bound = gain * math.sqrt(3.0 / fan) - self.cuboid_routing_weights = paddle.create_parameter( - shape=[T, H, W, self.num_experts], - dtype="float32", - default_initializer=nn.initializer.Uniform(-bound, bound), - ) - - self.lin = nn.Linear( - in_features=in_channels, out_features=self.num_experts, bias_attr=False - ) - - self.combine_weight = paddle.create_parameter( - shape=[T, H, W, self.num_experts, 2], - dtype="float32", - default_initializer=nn.initializer.Uniform(-bound, bound), - ) - - def gating(self, x, t_map=None): - # assert t_map is not None - cuboid_routing_weights = self.cuboid_routing_weights.unsqueeze(0).tile( - [x.shape[0], 1, 1, 1, 1] - ) # [B, T, H, W, E] - linear_routing_weights = self.lin(x) # [B, T, H, W, E] - routing_weights = paddle.stack( - [cuboid_routing_weights, linear_routing_weights], axis=-1 - ) # [B, T, H, W, E, 2] - combine_weight = self.combine_weight.tile( - [x.shape[0], 1, 1, 1, 1, 1] - ) # [B, T, H, W, E, 2] - routing_weights = (routing_weights * combine_weight).sum(-1) # [B, T, H, W, E] - return routing_weights - - -def aggregate_aux_losses(net): - aux_losses = [] - for module in net.sublayers(): - if hasattr(module, "aux_loss"): - aux_losses.append(module.aux_loss.unsqueeze(0)) - return aux_losses - - -# MoE Routing - - -class SparseDispatcherScatter(object): - def __init__(self, num_experts, gates): - self._gates = gates - self._num_experts = num_experts - sorted_experts, index_sorted_experts = paddle.nonzero(gates).sort( - 0 - ), paddle.nonzero(gates).argsort(0) - _, self._expert_index = sorted_experts.split(1, axis=1) - self._batch_index = paddle.nonzero(gates)[index_sorted_experts[:, 1], 0] - self._part_sizes = (gates > 0).sum(0).tolist() - gates_exp = gates[self._batch_index.flatten()] - self._nonzero_gates = paddle.take_along_axis( - gates_exp, axis=1, indices=self._expert_index - ) - - def dispatch(self, inp): - inp_exp = inp[self._batch_index].squeeze(1) - return paddle.split(inp_exp, self._part_sizes, axis=0) - - def combine(self, expert_out, multiply_by_gates=True): - stitched = paddle.concat(expert_out, 0) - if multiply_by_gates: - stitched = stitched.multiply(self._nonzero_gates) - zeros = paddle.zeros([self._gates.shape[0], expert_out[-1].shape[1]]) - zeros.stop_gradient = False - # combine samples that have been processed by the same k experts - combined = zeros.index_add(0, self._batch_index, stitched.float()) - return combined - - -class SparseDispatcher(object): - def __init__(self, num_experts, top_k_gates, top_k_indices): - self.num_experts = num_experts - self.gates = top_k_gates # [B, K] - self.gate_inds = top_k_indices # [B, K] - E = num_experts - B, K = top_k_gates.shape - self.batch_index_per_expert = paddle.stack( - [ - (top_k_indices == expert_id).sum(-1).astype("bool") - for expert_id in range(E) - ], - axis=0, - ) # [E, B] - self.gates_per_expert = paddle.concat( - [top_k_gates[top_k_indices == expert_id] for expert_id in range(E)] - ) # B * K - self.batch_index_all = paddle.nonzero(self.batch_index_per_expert)[ - :, 1 - ] # B * K - self.expert_size = self.batch_index_per_expert.sum(-1) # [E] - - def dispatch(self, x): - B, C = x.shape - dispatched_x = [ - x[batch_index] for batch_index in self.batch_index_per_expert - ] # E * [B_e, C] - return dispatched_x - - def combine(self, expert_out): - # expert_out: E * [B_e, C] - assert len(expert_out) == self.num_experts - com_res = paddle.concat(expert_out, axis=0) # [B * K, C] - zeros = paddle.zeros([self.gates.shape[0], com_res.shape[1]]) - zeros.stop_gradient = False - combined_res = zeros.index_add( - axis=0, - index=self.batch_index_all, - value=com_res * self.gates_per_expert.unsqueeze(-1), - ) - return combined_res - - -class DenseDispatcher(object): - def __init__(self, num_experts, top_k_gates, top_k_indices): - self.num_experts = num_experts - self.gates = top_k_gates # [B, K] - self.gate_inds = top_k_indices # [B, K] - - def combine(self, expert_out): - # expert_out: [B, E, C] - B, E, C = expert_out.shape - assert E == self.num_experts - selected_out = paddle.take_along_axis( - expert_out, axis=1, indices=self.gate_inds.unsqueeze(-1) - ) # [B, K, C] - combined_res = (selected_out * self.gates.unsqueeze(-1)).sum(1) - return combined_res - - -# RNC - - -class LabelDifference(nn.Layer): - def __init__(self, distance_type="l1"): - super().__init__() - self.distance_type = distance_type - - def forward(self, labels): - # labels: [bs, label_dim] - # output: [bs, bs] - assert labels.ndim == 3 - if self.distance_type == "l1": - return paddle.abs(labels[:, :, None, :] - labels[:, None, :, :]).sum( - axis=-1 - ) - else: - raise ValueError(self.distance_type) - - -class FeatureSimilarity(nn.Layer): - def __init__(self, similarity_type="l2", temperature=2): - super().__init__() - self.similarity_type = similarity_type - self.t = temperature - - def forward(self, features): - # labels: [bs, feat_dim] - # output: [bs, bs] - assert features.ndim == 3 - if self.similarity_type == "l2": - logits = -(features[:, :, None, :] - features[:, None, :, :]).norm( - 2, axis=-1 - ) - logits /= self.t - logits_max = paddle.max(logits, axis=1, keepdim=True) - logits -= logits_max.detach() - return logits - elif self.similarity_type == "cosine": - cos_func = nn.CosineSimilarity(axis=-1) - logits = cos_func(features[:, :, None, :], features[:, None, :, :]) - logits /= self.t - return logits - else: - raise ValueError(self.similarity_type) - - -class RnCLoss(nn.Layer): - def __init__(self, rnc_config): - super().__init__() - - self.rank_mode = rnc_config["rank_imbalance_style"] - self.t = rnc_config["rank_imbalance_temp"] - self.label_diff_fn = LabelDifference(rnc_config["label_difference_style"]) - self.feature_sim_fn = FeatureSimilarity( - rnc_config["feature_similarity_style"], self.t - ) - self.rnc_weight = rnc_config["rank_reg_coeff"] - self.loss_cal_mode = rnc_config["loss_cal_style"] - self.softmax_cri = nn.Softmax(axis=-1) - - def cal_loss(self, features, labels): - - B = features.shape[0] - assert B > 1 - label_diffs = self.label_diff_fn(labels) - logits = self.feature_sim_fn(features) - exp_logits = logits.exp() - n = logits.shape[1] - - # remove diagonal - logits = logits.masked_select( - (1 - paddle.eye(n)).astype("bool").unsqueeze(0).tile([B, 1, 1]) - ).reshape([B, n, n - 1]) - exp_logits = exp_logits.masked_select( - (1 - paddle.eye(n)).astype("bool").unsqueeze(0).tile([B, 1, 1]) - ).reshape([B, n, n - 1]) - label_diffs = label_diffs.masked_select( - (1 - paddle.eye(n)).astype("bool").unsqueeze(0).tile([B, 1, 1]) - ).reshape([B, n, n - 1]) - - if self.loss_cal_mode == "memory-efficient": - loss = 0.0 - for k in range(n - 1): - pos_logits = logits[:, :, k] # [B, n] - pos_label_diffs = label_diffs[:, :, k] # [B, n] - neg_mask = (label_diffs >= pos_label_diffs.unsqueeze(-1)).astype( - "float32" - ) # [B, n, n - 1] - pos_log_probs = pos_logits - paddle.log( - (neg_mask * exp_logits).sum(axis=-1) - ) # [B, n] - loss += -pos_log_probs.sum() - loss /= B * n * (n - 1) - elif self.loss_cal_mode == "computation-efficient": - neg_mask = (label_diffs.unsqueeze(-2) >= label_diffs.unsqueeze(-1)).astype( - "float32" - ) # [B, n, n - 1, n - 1] - pos_log_probs = logits - paddle.log( - (neg_mask * exp_logits.unsqueeze(-2).tile([1, 1, n - 1, 1])).sum( - axis=-1 - ) - ) # [B, n, n - 1] - loss = -pos_log_probs.mean() - else: - raise NotImplementedError - - return loss - - def forward(self, features, labels): - # features: [B, T_o, H, W, C_o] - # labels: [B, T_o, H, W, C_l] - - B, T_o, H, W, C_o = features.shape - _, _, _, _, C_l = labels.shape - - loss = None - if self.rank_mode == "batch": - features = features.reshape([B, -1, C_o]).transpose([1, 0, 2]) - labels = labels.reshape([B, -1, C_l]).transpose([1, 0, 2]) - loss = self.cal_loss(features, labels) - elif self.rank_mode == "batch+T+H+W": - feat = features.transpose([0, 2, 3, 1, 4]).reshape([-1, T_o, C_o]) - label = labels.transpose([0, 2, 3, 1, 4]).reshape([-1, T_o, C_l]) - loss_T = self.cal_loss(feat, label) - - feat = features.transpose([0, 1, 3, 2, 4]).reshape([-1, H, C_o]) - label = labels.transpose([0, 1, 3, 2, 4]).reshape([-1, H, C_l]) - loss_H = self.cal_loss(feat, label) - - feat = features.reshape([-1, W, C_o]) - label = labels.reshape([-1, W, C_l]) - loss_W = self.cal_loss(feat, label) - - feat = features.transpose([1, 2, 3, 0, 4]).reshape([-1, B, C_o]) - label = labels.transpose([1, 2, 3, 0, 4]).reshape([-1, B, C_l]) - loss_batch = self.cal_loss(feat, label) - - loss = loss_T + loss_H + loss_W + loss_batch - else: - raise NotImplementedError - - loss = self.rnc_weight * loss - - return loss diff --git a/examples/smc_reac/ppsci/arch/fno_block.py b/examples/smc_reac/ppsci/arch/fno_block.py deleted file mode 100644 index df40c36a0e..0000000000 --- a/examples/smc_reac/ppsci/arch/fno_block.py +++ /dev/null @@ -1,1269 +0,0 @@ -import itertools -from typing import Dict -from typing import List -from typing import Optional -from typing import Tuple -from typing import Union - -import omegaconf -import paddle -import paddle.nn.functional as F -from paddle import nn - -from ppsci.utils import initializer -from ppsci.utils import logger - -einsum_symbols = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ" - - -class DomainPadding(nn.Layer): - """Applies domain padding scaled automatically to the input's resolution - - Args: - domain_padding (Union[float, List[float]]): Typically, between zero and one, percentage of padding to use. - padding_mode (str, optional): Whether to pad on both sides, by default - 'one-sided'.Options are 'symmetric' or 'one-sided'。 Defaults to "one-sided". - output_scaling_factor (Union[int, List[int]], optional): Scaling factor for the - output. Defaults to 1. - """ - - def __init__( - self, - domain_padding: Union[float, List[float]], - padding_mode: str = "one-sided", - output_scaling_factor: Union[int, List[int]] = 1, - ): - super().__init__() - self.domain_padding = domain_padding - self.padding_mode = padding_mode.lower() - if output_scaling_factor is None: - output_scaling_factor = 1 - self.output_scaling_factor: Union[int, List[int]] = output_scaling_factor - - # dict(f'{resolution}'=padding) such that padded = F.pad(x, indices) - self._padding = dict() - - # dict(f'{resolution}'=indices_to_unpad) such that unpadded = x[indices] - self._unpad_indices = dict() - - def forward(self, x): - self.pad(x) - - def pad(self, x): - """Take an input and pad it by the desired fraction - - The amount of padding will be automatically scaled with the resolution - """ - resolution = x.shape[2:] - - if isinstance(self.domain_padding, (float, int)): - self.domain_padding = [float(self.domain_padding)] * len(resolution) - - assert len(self.domain_padding) == len(resolution), ( - "domain_padding length must match the number of spatial/time dimensions " - "(excluding batch, ch)" - ) - - output_scaling_factor = self.output_scaling_factor - if not isinstance(self.output_scaling_factor, list): - # if unset by the user, scaling_factor will be 1 be default, - # so `output_scaling_factor` should never be None. - output_scaling_factor: List[float] = validate_scaling_factor( - self.output_scaling_factor, len(resolution), n_layers=None - ) - - try: - padding = self._padding[f"{resolution}"] - return F.pad(x, padding, mode="constant") - except KeyError: - padding = [round(p * r) for (p, r) in zip(self.domain_padding, resolution)] - - output_pad = padding - output_pad = [ - round(i * j) for (i, j) in zip(output_scaling_factor, output_pad) - ] - - # padding is being applied in reverse order - # (so we must reverse the padding list) - padding = padding[::-1] - - # the F.pad(x, padding) function pads the tensor 'x' in reverse order of the "padding" list i.e. the last axis of tensor 'x' will be padded by the amount mention at the first position of the 'padding' vector. The details about F.pad can be found here: - # https://www.paddlepaddle.org.cn/documentation/docs/zh/api/paddle/nn/functional/pad_cn.html - - if self.padding_mode == "symmetric": - # Pad both sides - unpad_list = list() - for p in output_pad: - if p == 0: - padding_end = None - padding_start = None - else: - padding_end = p - padding_start = -p - unpad_list.append(slice(padding_end, padding_start, None)) - - unpad_indices = (Ellipsis,) + tuple( - [slice(p, -p, None) for p in padding] - ) - padding = [i for p in padding for i in (p, p)] - - elif self.padding_mode == "one-sided": - # One-side padding - unpad_list = list() - for p in output_pad: - if p == 0: - padding_start = None - else: - padding_start = -p - unpad_list.append(slice(None, padding_start, None)) - unpad_indices = (Ellipsis,) + tuple(unpad_list) - padding = [i for p in padding for i in (0, p)] - else: - raise ValueError(f"Got self.padding_mode = {self.padding_mode}") - - self._padding[f"{resolution}"] = padding - - padded = F.pad(x, padding, mode="constant") - output_shape = padded.shape[2:] - output_shape = [ - round(i * j) for (i, j) in zip(output_scaling_factor, output_shape) - ] - - self._unpad_indices[f"{[i for i in output_shape]}"] = unpad_indices - - return padded - - def unpad(self, x): - """Remove the padding from padding inputs""" - unpad_indices = self._unpad_indices[f"{x.shape[2:]}"] - - return x[unpad_indices] - - -class SoftGating(nn.Layer): - """Applies soft-gating by weighting the channels of the given input - - Given an input x of size `(batch-size, channels, height, width)`, - this returns `x * w ` - where w is of shape `(1, channels, 1, 1)` - - Args: - in_features (int): The number of input features. - out_features (int, optional): Number of output features. Defaults to None. - n_dim (int, optional): Dimensionality of the input (excluding batch-size and channels). - ``n_dim=2`` corresponds to having Module2D. Defaults to 2. - bias (bool, optional): Whether to use bias. Defaults to False. - """ - - def __init__( - self, in_features, out_features: int = None, n_dim: int = 2, bias: bool = False - ): - super().__init__() - if out_features is not None and in_features != out_features: - raise ValueError( - f"Got in_features = {in_features} and out_features = {out_features}" - "but these two must be the same for soft-gating" - ) - self.in_features = in_features - self.out_features = out_features - - self.weight = self.create_parameter( - shape=(1, self.in_features, *(1,) * n_dim), - default_initializer=nn.initializer.Constant(1.0), - ) - if bias: - self.bias = self.create_parameter( - shape=(1, self.in_features, *(1,) * n_dim), - default_initializer=nn.initializer.Constant(1.0), - ) - else: - self.bias = None - - def forward(self, x): - """Applies soft-gating to a batch of activations""" - if self.bias is not None: - return self.weight * x + self.bias - else: - return self.weight * x - - -def skip_connection( - in_features, - out_features, - n_dim: int = 2, - bias: bool = False, - type: str = "soft-gating", -): - """A wrapper for several types of skip connections. - Returns an nn.Module skip connections, one of {'identity', 'linear', soft-gating'} - - Args: - in_features (int): Number of input features. - out_features (int): Number of output features. - n_dim (int, optional): Dimensionality of the input (excluding batch-size and channels). - ``n_dim=2`` corresponds to having Module2D. . Defaults to 2. - bias (bool, optional): Whether to use a bias. Defaults to False. - type (str, optional): Kind of skip connection to use,{'identity', 'linear', soft-gating'}. - Defaults to "soft-gating". - """ - - if type.lower() == "soft-gating": - return SoftGating( - in_features=in_features, out_features=out_features, bias=bias, n_dim=n_dim - ) - elif type.lower() == "linear": - return getattr(nn, f"Conv{n_dim}D")( - in_channels=in_features, - out_channels=out_features, - kernel_size=1, - bias_attr=bias, - ) - elif type.lower() == "identity": - return nn.Identity() - else: - raise ValueError( - f"Got skip-connection type = {type}, expected one of {'soft-gating', 'linear', 'identity'}." - ) - - -class AdaIN(nn.Layer): - def __init__(self, embed_dim, in_channels, mlp=None, eps=1e-5): - super().__init__() - self.in_channels = in_channels - self.embed_dim = embed_dim - self.eps = eps - - if mlp is None: - mlp = nn.Sequential( - nn.Linear(embed_dim, 512), nn.GELU(), nn.Linear(512, 2 * in_channels) - ) - self.mlp = mlp - - self.embedding = None - - def set_embedding(self, x): - self.embedding = x.reshape( - self.embed_dim, - ) - - def forward(self, x): - assert ( - self.embedding is not None - ), "AdaIN: update embedding before running forward" - - weight, bias = paddle.split( - self.mlp(self.embedding), - self.embedding.shape[0] // self.in_channels, - axis=0, - ) - - return nn.functional.group_norm(x, self.in_channels, self.eps, weight, bias) - - -class MLP(nn.Layer): - """A Multi-Layer Perceptron, with arbitrary number of layers - - Args: - in_channels (int): The number of input channels. - out_channels (int, optional): The number of output channels. Defaults to None. - hidden_channels (int, optional): The number of hidden channels. Defaults to None. - n_layers (int, optional): The number of layers. Defaults to 2. - n_dim (int, optional): The type of convolution,2D or 3D. Defaults to 2. - non_linearity (nn.functional, optional): The activation function. Defaults to F.gelu. - dropout (float, optional): The ratio of dropout. Defaults to 0.0. - """ - - def __init__( - self, - in_channels: int, - out_channels: int = None, - hidden_channels: int = None, - n_layers: int = 2, - n_dim: int = 2, - non_linearity: nn.functional = F.gelu, - dropout: float = 0.0, - **kwargs, - ): - super().__init__() - self.n_layers = n_layers - self.in_channels = in_channels - self.out_channels = in_channels if out_channels is None else out_channels - self.hidden_channels = ( - in_channels if hidden_channels is None else hidden_channels - ) - self.non_linearity = non_linearity - self.dropout = ( - nn.LayerList([nn.Dropout(dropout) for _ in range(n_layers)]) - if dropout > 0.0 - else None - ) - - Conv = getattr(nn, f"Conv{n_dim}D") - self.fcs = nn.LayerList() - for i in range(n_layers): - if i == 0 and i == (n_layers - 1): - self.fcs.append(Conv(self.in_channels, self.out_channels, 1)) - elif i == 0: - self.fcs.append(Conv(self.in_channels, self.hidden_channels, 1)) - elif i == (n_layers - 1): - self.fcs.append(Conv(self.hidden_channels, self.out_channels, 1)) - else: - self.fcs.append(Conv(self.hidden_channels, self.hidden_channels, 1)) - - def forward(self, x): - for i, fc in enumerate(self.fcs): - x = fc(x) - if i < self.n_layers - 1: - x = self.non_linearity(x) - if self.dropout is not None: - x = self.dropout[i](x) - return x - - -def _contract_dense(x, weight, separable=False): - order = len(x.shape) - x_syms = list(einsum_symbols[:order]) - - # in_channels, out_channels, x, y... - weight_syms = list(x_syms[1:]) # no batch-size - - # batch-size, out_channels, x, y... - if separable: - out_syms = [x_syms[0]] + list(weight_syms) - else: - weight_syms.insert(1, einsum_symbols[order]) # outputs - out_syms = list(weight_syms) - out_syms[0] = x_syms[0] - - eq = "".join(x_syms) + "," + "".join(weight_syms) + "->" + "".join(out_syms) - # For the darcy flow, the only einsum is abcd,becd->aecd, where x and weights are shaped [32,32,8,8] - if not isinstance(weight, paddle.Tensor): - weight = paddle.to_tensor(weight) - - return paddle.einsum(eq, x, weight) - - -def _contract_dense_trick(x, weight_real, weight_imag, separable=False): - # the same as above function, but do the complex multiplication manually to avoid the einsum bug in paddle - order = len(x.shape) - # batch-size, in_channels, x, y... - x_syms = list(einsum_symbols[:order]) - - # in_channels, out_channels, x, y... - weight_syms = list(x_syms[1:]) # no batch-size - - # batch-size, out_channels, x, y... - if separable: - out_syms = [x_syms[0]] + list(weight_syms) - else: - weight_syms.insert(1, einsum_symbols[order]) # outputs - out_syms = list(weight_syms) - out_syms[0] = x_syms[0] - - eq = "".join(x_syms) + "," + "".join(weight_syms) + "->" + "".join(out_syms) - - o1_real = paddle.einsum(eq, x.real(), weight_real) - paddle.einsum( - eq, x.imag(), weight_imag - ) - o1_imag = paddle.einsum(eq, x.imag(), weight_real) + paddle.einsum( - eq, x.real(), weight_imag - ) - x = paddle.complex(o1_real, o1_imag) - return x - - -def _contract_dense_separable(x, weight, separable=True): - if not separable: - raise ValueError("This function is only for separable=True") - return x * weight - - -def get_contract_fun( - weight, implementation: str = "reconstructed", separable: bool = False -): - """Generic ND implementation of Fourier Spectral Conv contraction. - - Args: - weight (paddle.tensor): FactorizedTensor. - implementation (str, optional): {'reconstructed', 'factorized'}. - whether to reconstruct the weight and do a forward pass (reconstructed) - or contract directly the factors of the factorized weight with the input (factorized). Defaults to "reconstructed". - separable (bool, optional): Whether to use the separable implementation of contraction. This - arg is only checked when `implementation=reconstructed`. Defaults to False. - - Returns: - function : (x, weight) -> x * weight in Fourier space. - """ - - if implementation == "reconstructed": - if separable: - return _contract_dense_separable - else: - return _contract_dense_trick - elif implementation == "factorized": - if isinstance(weight, paddle.Tensor): - return _contract_dense_trick - - else: - raise ValueError( - f'Got implementation={implementation}, expected "reconstructed" or "factorized"' - ) - - -Number = Union[float, int] - - -def validate_scaling_factor( - scaling_factor: Union[None, Number, List[Number], List[List[Number]]], - n_dim: int, - n_layers: Optional[int] = None, -) -> Union[None, List[float], List[List[float]]]: - """Validates the format and dimensionality of the scaling factor. - - Args: - scaling_factor (Union[None, Number, List[Number], List[List[Number]]]): The - scaling factor. - n_dim (int): The required number of dimensions for expanding `scaling_factor`. - n_layers (Optional[int], optional): The number of layers for the returned - nested list. If None, return a single list (rather than a list of lists) - with `factor` repeated `dim` times. Defaults to None. - """ - - if scaling_factor is None: - return None - if isinstance(scaling_factor, (float, int)): - if n_layers is None: - return [float(scaling_factor)] * n_dim - - return [[float(scaling_factor)] * n_dim] * n_layers - if ( - isinstance(scaling_factor, list) - and len(scaling_factor) > 0 - and all([isinstance(s, (float, int)) for s in scaling_factor]) - ): - - return [[float(s)] * n_dim for s in scaling_factor] - - if ( - isinstance(scaling_factor, list) - and len(scaling_factor) > 0 - and all( - [isinstance(s, (omegaconf.listconfig.ListConfig)) for s in scaling_factor] - ) - ): - s_sub_pass = True - for s in scaling_factor: - if all([isinstance(s_sub, (int, float)) for s_sub in s]): - pass - else: - s_sub_pass = False - if s_sub_pass: - return scaling_factor - - return None - - -def resample(x, res_scale, axis, output_shape=None): - """A module for generic n-dimentional interpolation (Fourier resampling). - - Args: - x (paddle.Tensor): Input activation of size (batch_size, channels, d1, ..., dN). - res_scale (optional[int,tuple]): Scaling factor along each of the dimensions in - 'axis' parameter. If res_scale is scaler, then isotropic scaling is performed. - axis (int): Axis or dimensions along which interpolation will be performed. - output_shape (optional[None ,tuple[int]]): The output shape. Defaults to None. - """ - - if isinstance(res_scale, (float, int)): - if axis is None: - axis = list(range(2, x.ndim)) - res_scale = [res_scale] * len(axis) - elif isinstance(axis, int): - axis = [axis] - res_scale = [res_scale] - else: - res_scale = [res_scale] * len(axis) - else: - assert len(res_scale) == len(axis), "length of res_scale and axis are not same" - - old_size = x.shape[-len(axis) :] - if output_shape is None: - new_size = tuple([int(round(s * r)) for (s, r) in zip(old_size, res_scale)]) - else: - new_size = output_shape - - if len(axis) == 1: - return F.interpolate(x, size=new_size[0], mode="linear", align_corners=True) - if len(axis) == 2: - return F.interpolate(x, size=new_size, mode="bicubic", align_corners=True) - - X = paddle.fft.rfftn(x.astype("float32"), norm="forward", axes=axis) - - new_fft_size = list(new_size) - new_fft_size[-1] = new_fft_size[-1] // 2 + 1 # Redundant last coefficient - new_fft_size_c = [min(i, j) for (i, j) in zip(new_fft_size, X.shape[-len(axis) :])] - out_fft = paddle.zeros( - [x.shape[0], x.shape[1], *new_fft_size], dtype=paddle.complex64 - ) - - mode_indexing = [((None, m // 2), (-m // 2, None)) for m in new_fft_size_c[:-1]] + [ - ((None, new_fft_size_c[-1]),) - ] - for i, boundaries in enumerate(itertools.product(*mode_indexing)): - - idx_tuple = [slice(None), slice(None)] + [slice(*b) for b in boundaries] - - out_fft[idx_tuple] = X[idx_tuple] - y = paddle.fft.irfftn(out_fft, s=new_size, norm="forward", axes=axis) - - return y - - -class FactorizedTensor(nn.Layer): - def __init__(self, shape, init_scale): - super().__init__() - self.shape = shape - self.init_scale = init_scale - self.real = self.create_parameter( - shape=shape, - ) - self.real = initializer.normal_(self.real, 0, init_scale) - self.imag = self.create_parameter(shape=shape) - self.imag = initializer.normal_(self.imag, 0, init_scale) - - def __repr__(self): - return f"FactorizedTensor(shape={self.shape})" - - @property - def data(self): - return paddle.complex(self.real, self.imag) - - -class FactorizedSpectralConv(nn.Layer): - """Generic N-Dimensional Fourier Neural Operator - - Args: - in_channels (int): Number of input channels. - out_channels (int): Number of output channels. - n_modes (Tuple[int, ...]): Number of modes to use for contraction in Fourier domain during training. - .. warning:: - We take care of the redundancy in the Fourier modes, therefore, for an input - of size I_1, ..., I_N, please provide modes M_K that are I_1 < M_K <= I_N - We will automatically keep the right amount of modes: specifically, for the - last mode only, if you specify M_N modes we will use M_N // 2 + 1 modes - as the real FFT is redundant along that last dimension. - - .. note:: - Provided modes should be even integers. odd numbers will be rounded to the closest even number. - This can be updated dynamically during training. - max_n_modes (int, optional): * If not None, **maximum** number of modes to keep - in Fourier Layer, along each dim - The number of modes (`n_modes`) cannot be increased beyond that. - * If None, all the n_modes are used. Defaults to None. - bias (bool, optional): Whether to use bias in the layers. Defaults to True. - n_layers (int, optional): Number of Fourier Layers. Defaults to 1. - separable (bool, optional): Whether to use separable Fourier Conv. Defaults to False. - output_scaling_factor (Optional[Union[Number, List[Number]]], optional): Scaling factor for the - output. Defaults to None. - rank (float, optional): Rank of the tensor factorization of the Fourier weights. Defaults to 0.5. - factorization (str, optional): Tensor factorization of the parameters weight to use. - * If None, a dense tensor parametrizes the Spectral convolutions - * Otherwise, the specified tensor factorization is used. Defaults to None. - implementation (str, optional): If factorization is not None, forward mode to use. - * `reconstructed` : the full weight tensor is reconstructed from the - factorization and used for the forward pass - * `factorized` : the input is directly contracted with the factors of - the decomposition. Defaults to "reconstructed". - joint_factorization (bool, optional): Whether all the Fourier Layers should be parametrized by a - single tensor. Defaults to False. - init_std (str, optional): The std to use for the init. Defaults to "auto". - fft_norm (str, optional):The normalization mode for the FFT. Defaults to "backward". - """ - - def __init__( - self, - in_channels: int, - out_channels: int, - n_modes: Tuple[int, ...], - max_n_modes: int = None, - bias: bool = True, - n_layers: int = 1, - separable: bool = False, - output_scaling_factor: Optional[Union[Number, List[Number]]] = None, - rank: float = 0.5, - factorization: str = None, - implementation: str = "reconstructed", - joint_factorization: bool = False, - init_std: str = "auto", - fft_norm: str = "backward", - ): - super().__init__() - self.in_channels = in_channels - self.out_channels = out_channels - self.joint_factorization = joint_factorization - self.n_modes = n_modes - - self.order = len(self.n_modes) - if max_n_modes is None: - max_n_modes = self.n_modes - elif isinstance(max_n_modes, int): - max_n_modes = [max_n_modes] - self.max_n_modes = max_n_modes - self.rank = rank - self.factorization = factorization - self.n_layers = n_layers - self.implementation = implementation - - self.output_scaling_factor: Union[ - None, List[List[float]] - ] = validate_scaling_factor(output_scaling_factor, self.order, n_layers) - - if init_std == "auto": - init_std = (2 / (in_channels + out_channels)) ** 0.5 - else: - init_std = init_std - self.fft_norm = fft_norm - if factorization is None: - factorization = "Dense" - if not factorization.lower().startswith("complex"): - factorization = f"Complex{factorization}" - if separable: - if in_channels != out_channels: - raise ValueError( - f"To use separable Fourier Conv, in_channels must be equal to out_channels, but got in_channels={in_channels} and out_channels={out_channels}" - ) - weight_shape = (in_channels, *max_n_modes) - else: - weight_shape = (in_channels, out_channels, *max_n_modes) - self.separable = separable - if joint_factorization: - self.weight = paddle.create_parameter( - shape=(n_layers, *weight_shape), - dtype="float32", - ) - self.weight = initializer.normal_(self.weight, 0, init_std) - else: - self.weight = nn.LayerList( - [ - FactorizedTensor(weight_shape, init_scale=init_std) - for _ in range(n_layers) - ] - ) - - self._contract = get_contract_fun( - self.weight[0].data, implementation=implementation, separable=separable - ) - if bias: - shape = (n_layers, self.out_channels) + (1,) * self.order - init_bias = init_std * paddle.randn(shape) - self.bias = paddle.create_parameter( - shape=shape, - dtype=(init_bias.dtype), - default_initializer=nn.initializer.Assign(init_bias), - ) - self.bias.stop_gradient = False - else: - self.bias = None - - @property - def n_modes(self): - return self._n_modes - - @n_modes.setter - def n_modes(self, n_modes): - if isinstance(n_modes, int): # Should happen for 1D FNO only - n_modes = [n_modes] - else: - n_modes = list(n_modes) - # The last mode has a redundacy as we use real FFT - # As a design choice we do the operation here to avoid users dealing with the +1 - n_modes[-1] = n_modes[-1] // 2 + 1 - self._n_modes = n_modes - - def transform(self, x, layer_index=0, output_shape=None): - in_shape = list(x.shape[2:]) - - if self.output_scaling_factor is not None and output_shape is None: - out_shape = tuple( - [ - round(s * r) - for (s, r) in zip(in_shape, self.output_scaling_factor[layer_index]) - ] - ) - elif output_shape is not None: - out_shape = output_shape - else: - out_shape = in_shape - if in_shape == out_shape: - return x - else: - return resample( - x, - 1.0, - list(range(2, x.ndim)), - output_shape=out_shape, - ) - - def forward( - self, - x: paddle.Tensor, - indices: int = 0, - output_shape: Optional[Tuple[int]] = None, - ): - batchsize, channels, *mode_sizes = x.shape - fft_size = list(mode_sizes) - fft_size[-1] = fft_size[-1] // 2 + 1 - fft_dims = list(range(-self.order, 0)) - - x = paddle.fft.rfftn(x=x, norm=self.fft_norm, axes=fft_dims) - - if self.order > 1: - x = paddle.fft.fftshift(x=x, axes=fft_dims[:-1]) - - out_fft = paddle.zeros( - shape=[batchsize, self.out_channels, *fft_size], dtype=paddle.complex64 - ) - - starts = [ - (max_modes - min(size, n_mode)) - for size, n_mode, max_modes in zip(fft_size, self.n_modes, self.max_n_modes) - ] - slices_w = [slice(None), slice(None)] - slices_w += [ - (slice(start // 2, -start // 2) if start else slice(start, None)) - for start in starts[:-1] - ] - slices_w += [slice(None, -starts[-1]) if starts[-1] else slice(None)] - - w_real = self.weight[indices].real[ - slices_w[0], slices_w[1], slices_w[2], slices_w[3] - ] - w_imag = self.weight[indices].imag[ - slices_w[0], slices_w[1], slices_w[2], slices_w[3] - ] - - starts = [ - (size - min(size, n_mode)) - for (size, n_mode) in zip(list(x.shape[2:]), list(w_real.shape[2:])) - ] - slices_x = [slice(None), slice(None)] # Batch_size, channels - slices_x += [ - slice(start // 2, -start // 2) if start else slice(start, None) - for start in starts[:-1] - ] - slices_x += [ - slice(None, -starts[-1]) if starts[-1] else slice(None) - ] # The last mode already has redundant half removed - idx_tuple = slices_x - if len(idx_tuple) == 4: - out_fft[ - idx_tuple[0], idx_tuple[1], idx_tuple[2], idx_tuple[3] - ] = self._contract( - x[idx_tuple[0], idx_tuple[1], idx_tuple[2], idx_tuple[3]], - w_real, - w_imag, - separable=self.separable, - ) - elif len(idx_tuple) == 3: - out_fft[idx_tuple[0], idx_tuple[1], idx_tuple[2]] = self._contract( - x[idx_tuple[0], idx_tuple[1], idx_tuple[2]], - w_real, - w_imag, - separable=self.separable, - ) - else: - raise ValueError("Not implemented") - - if self.output_scaling_factor is not None and output_shape is None: - mode_sizes = tuple( - [ - round(s * r) - for (s, r) in zip(mode_sizes, self.output_scaling_factor[indices]) - ] - ) - - if output_shape is not None: - mode_sizes = output_shape - - if self.order > 1: - out_fft = paddle.fft.fftshift(x=out_fft, axes=fft_dims[:-1]) - - x = paddle.fft.irfftn( - x=out_fft, s=mode_sizes, axes=fft_dims, norm=self.fft_norm - ) - if self.bias is not None: - x = x + self.bias[indices, ...] - return x - - -class FactorizedSpectralConv1d(FactorizedSpectralConv): - """1D Spectral Conv - - This is provided for reference only, - see :class:`FactorizedSpectralConv` for the preferred, general implementation - """ - - def forward(self, x, indices=0): - batchsize, channels, width = x.shape - - x = paddle.fft.rfft(x, norm=self.fft_norm) - - out_fft = paddle.zeros( - shape=[batchsize, self.out_channels, width // 2 + 1], dtype=paddle.complex64 - ) - - slices = ( - slice(None), # Equivalent to: [:, - slice(None), # ............... :, - slice(None, self.n_modes[0]), # :half_n_modes[0]] - ) - - w_real = self.weight[indices].real[slices[0], slices[1], slices[2]] - w_imag = self.weight[indices].imag[slices[0], slices[1], slices[2]] - - out_fft[slices[0], slices[1], slices[2]] = self._contract( - x[slices[0], slices[1], slices[2]], - w_real, - w_imag, - separable=self.separable, - ) - - if self.output_scaling_factor is not None: - width = round(width * self.output_scaling_factor[0]) - - x = paddle.fft.irfft(out_fft, n=width, norm=self.fft_norm) - - if self.bias is not None: - x = x + self.bias[indices, ...] - - return x - - -class FactorizedSpectralConv2d(FactorizedSpectralConv): - """2D Spectral Conv. - - This is provided for reference only, - see :class:`FactorizedSpectralConv` for the preferred, general implementation - """ - - def forward(self, x, indices=0): - batchsize, channels, height, width = x.shape - - x = paddle.fft.rfft2(x.float(), norm=self.fft_norm, axes=(-2, -1)) - - # The output will be of size (batch_size, self.out_channels, - # x.size(-2), x.size(-1)//2 + 1) - out_fft = paddle.zeros( - shape=[batchsize, self.out_channels, height, width // 2 + 1], - dtype=paddle.complex64, - ) - - slices0 = ( - slice(None), # Equivalent to: [:, - slice(None), # ............... :, - slice(self.n_modes[0] // 2), # :half_n_modes[0], - slice(self.n_modes[1]), # :half_n_modes[1]] - ) - slices1 = ( - slice(None), # Equivalent to: [:, - slice(None), # ...................... :, - slice(-self.n_modes[0] // 2, None), # -half_n_modes[0]:, - slice(self.n_modes[1]), # ...... :half_n_modes[1]] - ) - logger.message( - f"2D: {x[slices0].shape=}, {self._get_weight(indices)[slices0].shape=}, {self._get_weight(indices).shape=}" - ) - - w_real = self.weight[indices].real[ - slices1[0], slices1[1], slices1[2], slices1[3] - ] - w_imag = self.weight[indices].imag[ - slices1[0], slices1[1], slices1[2], slices1[3] - ] - - """Upper block (truncate high frequencies).""" - out_fft[slices0[0], slices0[1], slices0[2], slices0[3]] = self._contract( - x[slices0[0], slices0[1], slices0[2], slices0[3]], - w_real, - w_imag, - separable=self.separable, - ) - - w_real = self.weight[indices].real[ - slices0[0], slices0[1], slices0[2], slices0[3] - ] - w_imag = self.weight[indices].imag[ - slices0[0], slices0[1], slices0[2], slices0[3] - ] - - """Lower block""" - out_fft[slices1[0], slices1[1], slices1[2], slices1[3]] = self._contract( - x[slices1[0], slices1[1], slices1[2], slices1[3]], - w_real, - w_imag, - separable=self.separable, - ) - - if self.output_scaling_factor is not None: - width = round(width * self.output_scaling_factor[indices][0]) - height = round(height * self.output_scaling_factor[indices][1]) - - x = paddle.fft.irfft2( - out_fft, s=(height, width), axes=(-2, -1), norm=self.fft_norm - ) - - if self.bias is not None: - x = x + self.bias[indices, ...] - - return x - - -class FactorizedSpectralConv3d(FactorizedSpectralConv): - """3D Spectral Conv. - - This is provided for reference only, - see :class:`FactorizedSpectralConv` for the preferred, general implementation - """ - - def forward(self, x, indices=0): - batchsize, channels, height, width, depth = x.shape - - x = paddle.fft.rfftn(x.float(), norm=self.fft_norm, axes=[-3, -2, -1]) - - out_fft = paddle.zeros( - shape=[batchsize, self.out_channels, height, width, depth // 2 + 1], - dtype=paddle.complex64, - ) - - slices0 = ( - slice(None), # Equivalent to: [:, - slice(None), # ............... :, - slice(self.n_modes[0] // 2), # :half_n_modes[0], - slice(self.n_modes[1] // 2), # :half_n_modes[1], - slice(self.n_modes[2]), # :half_n_modes[2]] - ) - slices1 = ( - slice(None), # Equivalent to: [:, - slice(None), # ...................... :, - slice(self.n_modes[0] // 2), # ...... :half_n_modes[0], - slice(-self.n_modes[1] // 2, None), # -half_n_modes[1]:, - slice(self.n_modes[2]), # ...... :half_n_modes[0]] - ) - slices2 = ( - slice(None), # Equivalent to: [:, - slice(None), # ...................... :, - slice(-self.n_modes[0] // 2, None), # -half_n_modes[0]:, - slice(self.n_modes[1] // 2), # ...... :half_n_modes[1], - slice(self.n_modes[2]), # ...... :half_n_modes[2]] - ) - slices3 = ( - slice(None), # Equivalent to: [:, - slice(None), # ...................... :, - slice(-self.n_modes[0] // 2, None), # -half_n_modes[0], - slice(-self.n_modes[1] // 2, None), # -half_n_modes[1], - slice(self.n_modes[2]), # ...... :half_n_modes[2]] - ) - - w_real = self.weight[indices].real[ - slices3[0], slices3[1], slices3[2], slices3[3], slices3[4] - ] - w_imag = self.weight[indices].imag[ - slices3[0], slices3[1], slices3[2], slices3[3], slices3[4] - ] - - """Upper block -- truncate high frequencies.""" - out_fft[ - slices0[0], slices0[1], slices0[2], slices0[3], slices0[4] - ] = self._contract( - x[slices0[0], slices0[1], slices0[2], slices0[3], slices0[4]], - w_real, - w_imag, - separable=self.separable, - ) - - w_real = self.weight[indices].real[ - slices2[0], slices2[1], slices2[2], slices2[3], slices2[4] - ] - w_imag = self.weight[indices].imag[ - slices2[0], slices2[1], slices2[2], slices2[3], slices2[4] - ] - """Low-pass filter for indices 2 & 4, and high-pass filter for index 3.""" - out_fft[ - slices1[0], slices1[1], slices1[2], slices1[3], slices1[4] - ] = self._contract( - x[slices1[0], slices1[1], slices1[2], slices1[3], slices1[4]], - w_real, - w_imag, - separable=self.separable, - ) - - w_real = self.weight[indices].real[ - slices1[0], slices1[1], slices1[2], slices1[3], slices1[4] - ] - w_imag = self.weight[indices].imag[ - slices1[0], slices1[1], slices1[2], slices1[3], slices1[4] - ] - """Low-pass filter for indices 3 & 4, and high-pass filter for index 2.""" - out_fft[ - slices2[0], slices2[1], slices2[2], slices2[3], slices2[4] - ] = self._contract( - x[slices2[0], slices2[1], slices2[2], slices2[3], slices2[4]], - w_real, - w_imag, - separable=self.separable, - ) - - w_real = self.weight[indices].real[ - slices0[0], slices0[1], slices0[2], slices0[3], slices0[4] - ] - w_imag = self.weight[indices].imag[ - slices0[0], slices0[1], slices0[2], slices0[3], slices0[4] - ] - """Lower block -- low-cut filter in indices 2 & 3 - and high-cut filter in index 4.""" - out_fft[ - slices3[0], slices3[1], slices3[2], slices3[3], slices3[4] - ] = self._contract( - x[slices3[0], slices3[1], slices3[2], slices3[3], slices3[4]], - w_real, - w_imag, - separable=self.separable, - ) - - if self.output_scaling_factor is not None: - width = round(width * self.output_scaling_factor[0]) - height = round(height * self.output_scaling_factor[1]) - depth = round(depth * self.output_scaling_factor[2]) - - x = paddle.fft.irfftn( - out_fft, s=(height, width, depth), axes=[-3, -2, -1], norm=self.fft_norm - ) - - if self.bias is not None: - x = x + self.bias[indices, ...] - return x - - -class FNOBlocks(nn.Layer): - def __init__( - self, - in_channels: int, - out_channels: int, - n_modes: Tuple[int, ...], - output_scaling_factor: Optional[Union[Number, List[Number]]] = None, - n_layers: int = 1, - max_n_modes: int = None, - use_mlp: bool = False, - mlp: Optional[Dict[str, float]] = None, - non_linearity: nn.functional = F.gelu, - stabilizer: str = None, - norm: str = None, - ada_in_features: Optional[int] = None, - preactivation: bool = False, - fno_skip: str = "linear", - mlp_skip: str = "soft-gating", - separable: bool = False, - factorization: str = None, - rank: float = 1.0, - SpectralConv: FactorizedSpectralConv = FactorizedSpectralConv, - joint_factorization: bool = False, - implementation: str = "factorized", - fft_norm: str = "forward", - **kwargs, - ): - super().__init__() - if isinstance(n_modes, int): - n_modes = [n_modes] - self._n_modes = n_modes - self.n_dim = len(n_modes) - - self.max_n_modes = max_n_modes - self.in_channels = in_channels - self.out_channels = out_channels - self.n_layers = n_layers - self.joint_factorization = joint_factorization - self.non_linearity = non_linearity - self.rank = rank - self.factorization = factorization - self.fno_skip = fno_skip - self.mlp_skip = mlp_skip - self.use_mlp = use_mlp - self.fft_norm = fft_norm - self.implementation = implementation - self.separable = separable - self.preactivation = preactivation - self.ada_in_features = ada_in_features - self.stabilizer = stabilizer - self.norm = norm - - self.convs = SpectralConv( - self.in_channels, - self.out_channels, - self.n_modes, - output_scaling_factor=output_scaling_factor, - max_n_modes=max_n_modes, - rank=rank, - implementation=implementation, - separable=separable, - factorization=factorization, - joint_factorization=joint_factorization, - n_layers=n_layers, - ) - - self.fno_skips = nn.LayerList( - [ - skip_connection( - self.in_channels, - self.out_channels, - type=fno_skip, - n_dim=self.n_dim, - ) - for _ in range(n_layers) - ] - ) - - if use_mlp: - self.mlp = nn.LayerList( - [ - MLP( - in_channels=self.out_channels, - hidden_channels=int( - round(self.out_channels * mlp["expansion"]) - ), - dropout=mlp["dropout"], - n_dim=self.n_dim, - ) - for _ in range(n_layers) - ] - ) - self.mlp_skips = nn.LayerList( - [ - skip_connection( - self.in_channels, - self.out_channels, - type=mlp_skip, - n_dim=self.n_dim, - ) - for _ in range(n_layers) - ] - ) - else: - self.mlp = None - - # Each block will have 2 norms if we also use an MLP - self.n_norms = 1 if self.mlp is None else 2 - if norm is None: - self.norm = None - elif norm == "instance_norm": - self.norm = nn.LayerList( - [ - getattr(nn, f"InstanceNorm{self.n_dim}d")( - num_features=self.out_channels - ) - for _ in range(n_layers * self.n_norms) - ] - ) - elif norm == "group_norm": - self.norm = nn.LayerList( - [ - nn.GroupNorm(num_groups=1, num_channels=self.out_channels) - for _ in range(n_layers * self.n_norms) - ] - ) - elif norm == "ada_in": - self.norm = nn.LayerList( - [ - AdaIN(ada_in_features, out_channels) - for _ in range(n_layers * self.n_norms) - ] - ) - else: - raise ValueError( - f"Got {norm} but expected None or one of [instance_norm, group_norm, layer_norm]" - ) - - def forward(self, x, index=0, output_shape=None): - if self.preactivation: - return self.forward_with_preactivation(x, index, output_shape=output_shape) - else: - return self.forward_with_postactivation(x, index, output_shape=output_shape) - - def forward_with_postactivation(self, x, index=0, output_shape=None): - x_skip_fno = self.fno_skips[index](x) - x_skip_fno = self.convs.transform(x_skip_fno, index, output_shape=output_shape) - if self.mlp is not None: - x_skip_mlp = self.mlp_skips[index](x) - x_skip_mlp = self.convs.transform( - x_skip_mlp, index, output_shape=output_shape - ) - if self.stabilizer == "tanh": - x = paddle.tanh(x) - - x_fno = self.convs(x, index, output_shape=output_shape) - if self.norm is not None: - x_fno = self.norm[self.n_norms * index](x_fno) - - x = x_fno + x_skip_fno - - if (self.mlp is not None) or (index < (self.n_layers - 1)): - x = self.non_linearity(x) - - if self.mlp is not None: - x = self.mlp[index](x) + x_skip_mlp - - if self.norm is not None: - x = self.norm[self.n_norms * index + 1](x) - - if index < (self.n_layers - 1): - x = self.non_linearity(x) - - return x - - def forward_with_preactivation(self, x, index=0, output_shape=None): - # Apply non-linear activation (and norm) - # before this block's convolution/forward pass: - x = self.non_linearity(x) - - if self.norm is not None: - x = self.norm[self.n_norms * index](x) - - x_skip_fno = self.fno_skips[index](x) - x_skip_fno = self.convs.transform(x_skip_fno, index, output_shape=output_shape) - - if self.mlp is not None: - x_skip_mlp = self.mlp_skips[index](x) - x_skip_mlp = self.convs.transform( - x_skip_mlp, index, output_shape=output_shape - ) - - if self.stabilizer == "tanh": - x = paddle.tanh(x) - - x_fno = self.convs(x, index, output_shape=output_shape) - x = x_fno + x_skip_fno - - if self.mlp is not None: - if index < (self.n_layers - 1): - x = self.non_linearity(x) - - if self.norm is not None: - x = self.norm[self.n_norms * index + 1](x) - - x = self.mlp[index](x) + x_skip_mlp - - return x - - @property - def n_modes(self): - return self._n_modes - - @n_modes.setter - def n_modes(self, n_modes): - if isinstance(n_modes, int): # Should happen for 1D FNO only - n_modes = [n_modes] - else: - n_modes = list(n_modes) - # The last mode has a redundacy as we use real FFT - # As a design choice we do the operation here to avoid users dealing with the +1 - n_modes[-1] = n_modes[-1] // 2 + 1 - self._n_modes = n_modes diff --git a/examples/smc_reac/ppsci/arch/gan.py b/examples/smc_reac/ppsci/arch/gan.py deleted file mode 100644 index a673b8e440..0000000000 --- a/examples/smc_reac/ppsci/arch/gan.py +++ /dev/null @@ -1,400 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import Dict -from typing import List -from typing import Tuple - -import paddle -import paddle.nn as nn - -from ppsci.arch import activation as act_mod -from ppsci.arch import base - - -class Conv2DBlock(nn.Layer): - def __init__( - self, - in_channel, - out_channel, - kernel_size, - stride, - use_bn, - act, - mean, - std, - value, - ): - super().__init__() - weight_attr = paddle.ParamAttr( - initializer=nn.initializer.Normal(mean=mean, std=std) - ) - bias_attr = paddle.ParamAttr(initializer=nn.initializer.Constant(value=value)) - self.conv_2d = nn.Conv2D( - in_channel, - out_channel, - kernel_size, - stride, - padding="SAME", - weight_attr=weight_attr, - bias_attr=bias_attr, - ) - self.bn = nn.BatchNorm2D(out_channel) if use_bn else None - self.act = act_mod.get_activation(act) if act else None - - def forward(self, x): - y = x - y = self.conv_2d(y) - if self.bn: - y = self.bn(y) - if self.act: - y = self.act(y) - return y - - -class VariantResBlock(nn.Layer): - def __init__( - self, - in_channel, - out_channels, - kernel_sizes, - strides, - use_bns, - acts, - mean, - std, - value, - ): - super().__init__() - self.conv_2d_0 = Conv2DBlock( - in_channel=in_channel, - out_channel=out_channels[0], - kernel_size=kernel_sizes[0], - stride=strides[0], - use_bn=use_bns[0], - act=acts[0], - mean=mean, - std=std, - value=value, - ) - self.conv_2d_1 = Conv2DBlock( - in_channel=out_channels[0], - out_channel=out_channels[1], - kernel_size=kernel_sizes[1], - stride=strides[1], - use_bn=use_bns[1], - act=acts[1], - mean=mean, - std=std, - value=value, - ) - - self.conv_2d_2 = Conv2DBlock( - in_channel=in_channel, - out_channel=out_channels[2], - kernel_size=kernel_sizes[2], - stride=strides[2], - use_bn=use_bns[2], - act=acts[2], - mean=mean, - std=std, - value=value, - ) - - self.act = act_mod.get_activation("relu") - - def forward(self, x): - y = x - y = self.conv_2d_0(y) - y = self.conv_2d_1(y) - short = self.conv_2d_2(x) - y = paddle.add(y, short) - y = self.act(y) - return y - - -class FCBlock(nn.Layer): - def __init__(self, in_channel, act, mean, std, value): - super().__init__() - self.flatten = nn.Flatten() - weight_attr = paddle.ParamAttr( - initializer=nn.initializer.Normal(mean=mean, std=std) - ) - bias_attr = paddle.ParamAttr(initializer=nn.initializer.Constant(value=value)) - self.linear = nn.Linear( - in_channel, - 1, - weight_attr=weight_attr, - bias_attr=bias_attr, - ) - self.act = act_mod.get_activation(act) if act else None - - def forward(self, x): - y = x - y = self.flatten(y) - y = self.linear(y) - if self.act: - y = self.act(y) - return y - - -class Generator(base.Arch): - """Generator Net of GAN. Attention, the net using a kind of variant of ResBlock which is - unique to "tempoGAN" example but not an open source network. - - Args: - input_keys (Tuple[str, ...]): Name of input keys, such as ("input1", "input2"). - output_keys (Tuple[str, ...]): Name of output keys, such as ("output1", "output2"). - in_channel (int): Number of input channels of the first conv layer. - out_channels_tuple (Tuple[Tuple[int, ...], ...]): Number of output channels of all conv layers, - such as [[out_res0_conv0, out_res0_conv1], [out_res1_conv0, out_res1_conv1]] - kernel_sizes_tuple (Tuple[Tuple[int, ...], ...]): Number of kernel_size of all conv layers, - such as [[kernel_size_res0_conv0, kernel_size_res0_conv1], [kernel_size_res1_conv0, kernel_size_res1_conv1]] - strides_tuple (Tuple[Tuple[int, ...], ...]): Number of stride of all conv layers, - such as [[stride_res0_conv0, stride_res0_conv1], [stride_res1_conv0, stride_res1_conv1]] - use_bns_tuple (Tuple[Tuple[bool, ...], ...]): Whether to use the batch_norm layer after each conv layer. - acts_tuple (Tuple[Tuple[str, ...], ...]): Whether to use the activation layer after each conv layer. If so, witch activation to use, - such as [[act_res0_conv0, act_res0_conv1], [act_res1_conv0, act_res1_conv1]] - - Examples: - >>> import ppsci - >>> in_channel = 1 - >>> rb_channel0 = (2, 8, 8) - >>> rb_channel1 = (128, 128, 128) - >>> rb_channel2 = (32, 8, 8) - >>> rb_channel3 = (2, 1, 1) - >>> out_channels_tuple = (rb_channel0, rb_channel1, rb_channel2, rb_channel3) - >>> kernel_sizes_tuple = (((5, 5), ) * 2 + ((1, 1), ), ) * 4 - >>> strides_tuple = ((1, 1, 1), ) * 4 - >>> use_bns_tuple = ((True, True, True), ) * 3 + ((False, False, False), ) - >>> acts_tuple = (("relu", None, None), ) * 4 - >>> model = ppsci.arch.Generator(("in",), ("out",), in_channel, out_channels_tuple, kernel_sizes_tuple, strides_tuple, use_bns_tuple, acts_tuple) - >>> batch_size = 4 - >>> height = 64 - >>> width = 64 - >>> input_data = paddle.randn([batch_size, in_channel, height, width]) - >>> input_dict = {'in': input_data} - >>> output_data = model(input_dict) - >>> print(output_data['out'].shape) - [4, 1, 64, 64] - """ - - def __init__( - self, - input_keys: Tuple[str, ...], - output_keys: Tuple[str, ...], - in_channel: int, - out_channels_tuple: Tuple[Tuple[int, ...], ...], - kernel_sizes_tuple: Tuple[Tuple[int, ...], ...], - strides_tuple: Tuple[Tuple[int, ...], ...], - use_bns_tuple: Tuple[Tuple[bool, ...], ...], - acts_tuple: Tuple[Tuple[str, ...], ...], - ): - super().__init__() - self.input_keys = input_keys - self.output_keys = output_keys - self.in_channel = in_channel - self.out_channels_tuple = out_channels_tuple - self.kernel_sizes_tuple = kernel_sizes_tuple - self.strides_tuple = strides_tuple - self.use_bns_tuple = use_bns_tuple - self.acts_tuple = acts_tuple - - self.init_blocks() - - def init_blocks(self): - blocks_list = [] - for i in range(len(self.out_channels_tuple)): - in_channel = ( - self.in_channel if i == 0 else self.out_channels_tuple[i - 1][-1] - ) - blocks_list.append( - VariantResBlock( - in_channel=in_channel, - out_channels=self.out_channels_tuple[i], - kernel_sizes=self.kernel_sizes_tuple[i], - strides=self.strides_tuple[i], - use_bns=self.use_bns_tuple[i], - acts=self.acts_tuple[i], - mean=0.0, - std=0.04, - value=0.1, - ) - ) - self.blocks = nn.LayerList(blocks_list) - - def forward_tensor(self, x): - y = x - for block in self.blocks: - y = block(y) - return y - - def forward(self, x): - if self._input_transform is not None: - x = self._input_transform(x) - - y = self.concat_to_tensor(x, self.input_keys, axis=-1) - y = self.forward_tensor(y) - y = self.split_to_dict(y, self.output_keys, axis=-1) - - if self._output_transform is not None: - y = self._output_transform(x, y) - return y - - -class Discriminator(base.Arch): - """Discriminator Net of GAN. - - Args: - input_keys (Tuple[str, ...]): Name of input keys, such as ("input1", "input2"). - output_keys (Tuple[str, ...]): Name of output keys, such as ("output1", "output2"). - in_channel (int): Number of input channels of the first conv layer. - out_channels (Tuple[int, ...]): Number of output channels of all conv layers, - such as (out_conv0, out_conv1, out_conv2). - fc_channel (int): Number of input features of linear layer. Number of output features of the layer - is set to 1 in this Net to construct a fully_connected layer. - kernel_sizes (Tuple[int, ...]): Number of kernel_size of all conv layers, - such as (kernel_size_conv0, kernel_size_conv1, kernel_size_conv2). - strides (Tuple[int, ...]): Number of stride of all conv layers, - such as (stride_conv0, stride_conv1, stride_conv2). - use_bns (Tuple[bool, ...]): Whether to use the batch_norm layer after each conv layer. - acts (Tuple[str, ...]): Whether to use the activation layer after each conv layer. If so, witch activation to use, - such as (act_conv0, act_conv1, act_conv2). - - Examples: - >>> import ppsci - >>> in_channel = 2 - >>> in_channel_tempo = 3 - >>> out_channels = (32, 64, 128, 256) - >>> fc_channel = 65536 - >>> kernel_sizes = ((4, 4), (4, 4), (4, 4), (4, 4)) - >>> strides = (2, 2, 2, 1) - >>> use_bns = (False, True, True, True) - >>> acts = ("leaky_relu", "leaky_relu", "leaky_relu", "leaky_relu", None) - >>> output_keys_disc = ("out_1", "out_2", "out_3", "out_4", "out_5", "out_6", "out_7", "out_8", "out_9", "out_10") - >>> model = ppsci.arch.Discriminator(("in_1","in_2"), output_keys_disc, in_channel, out_channels, fc_channel, kernel_sizes, strides, use_bns, acts) - >>> input_data = [paddle.to_tensor(paddle.randn([1, in_channel, 128, 128])),paddle.to_tensor(paddle.randn([1, in_channel, 128, 128]))] - >>> input_dict = {"in_1": input_data[0],"in_2": input_data[1]} - >>> out_dict = model(input_dict) - >>> for k, v in out_dict.items(): - ... print(k, v.shape) - out_1 [1, 32, 64, 64] - out_2 [1, 64, 32, 32] - out_3 [1, 128, 16, 16] - out_4 [1, 256, 16, 16] - out_5 [1, 1] - out_6 [1, 32, 64, 64] - out_7 [1, 64, 32, 32] - out_8 [1, 128, 16, 16] - out_9 [1, 256, 16, 16] - out_10 [1, 1] - """ - - def __init__( - self, - input_keys: Tuple[str, ...], - output_keys: Tuple[str, ...], - in_channel: int, - out_channels: Tuple[int, ...], - fc_channel: int, - kernel_sizes: Tuple[int, ...], - strides: Tuple[int, ...], - use_bns: Tuple[bool, ...], - acts: Tuple[str, ...], - ): - super().__init__() - self.input_keys = input_keys - self.output_keys = output_keys - self.in_channel = in_channel - self.out_channels = out_channels - self.fc_channel = fc_channel - self.kernel_sizes = kernel_sizes - self.strides = strides - self.use_bns = use_bns - self.acts = acts - - self.init_layers() - - def init_layers(self): - layers_list = [] - for i in range(len(self.out_channels)): - in_channel = self.in_channel if i == 0 else self.out_channels[i - 1] - layers_list.append( - Conv2DBlock( - in_channel=in_channel, - out_channel=self.out_channels[i], - kernel_size=self.kernel_sizes[i], - stride=self.strides[i], - use_bn=self.use_bns[i], - act=self.acts[i], - mean=0.0, - std=0.04, - value=0.1, - ) - ) - - layers_list.append( - FCBlock(self.fc_channel, self.acts[4], mean=0.0, std=0.04, value=0.1) - ) - self.layers = nn.LayerList(layers_list) - - def forward_tensor(self, x): - y = x - y_list = [] - for layer in self.layers: - y = layer(y) - y_list.append(y) - return y_list # y_conv1, y_conv2, y_conv3, y_conv4, y_fc(y_out) - - def forward(self, x): - if self._input_transform is not None: - x = self._input_transform(x) - - y_list = [] - # y1_conv1, y1_conv2, y1_conv3, y1_conv4, y1_fc, y2_conv1, y2_conv2, y2_conv3, y2_conv4, y2_fc - for k in x: - y_list.extend(self.forward_tensor(x[k])) - - y = self.split_to_dict(y_list, self.output_keys) - - if self._output_transform is not None: - y = self._output_transform(x, y) - - return y - - @staticmethod - def split_to_dict( - data_list: List[paddle.Tensor], keys: Tuple[str, ...] - ) -> Dict[str, paddle.Tensor]: - """Overwrite of split_to_dict() method belongs to Class base.Arch. - - Reason for overwriting is there is no concat_to_tensor() method called in "tempoGAN" example. - That is because input in "tempoGAN" example is not in a regular format, but a format like: - { - "input1": paddle.concat([in1, in2], axis=1), - "input2": paddle.concat([in1, in3], axis=1), - } - - Args: - data_list (List[paddle.Tensor]): The data to be split. It should be a list of tensor(s), but not a paddle.Tensor. - keys (Tuple[str, ...]): Keys of outputs. - - Returns: - Dict[str, paddle.Tensor]: Dict with split data. - """ - if len(keys) == 1: - return {keys[0]: data_list[0]} - return {key: data_list[i] for i, key in enumerate(keys)} diff --git a/examples/smc_reac/ppsci/arch/geofno.py b/examples/smc_reac/ppsci/arch/geofno.py deleted file mode 100644 index 10fe8d08b5..0000000000 --- a/examples/smc_reac/ppsci/arch/geofno.py +++ /dev/null @@ -1,205 +0,0 @@ -# Copyright (c) 2024 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. - -import paddle -import paddle.nn as nn -import paddle.nn.functional as F -import paddle.nn.initializer as Initializer - - -class SpectralConv1d(nn.Layer): - """ - 1D Fourier layer. It does FFT, linear transform, and Inverse FFT. - """ - - def __init__(self, in_channels, out_channels, modes1): - super(SpectralConv1d, self).__init__() - - self.in_channels = in_channels - self.out_channels = out_channels - # Number of Fourier modes to multiply, at most floor(N/2) + 1 - self.modes1 = modes1 - self.scale = 1 / (in_channels * out_channels) - - real = paddle.rand(shape=[in_channels, out_channels, modes1]) - real.stop_gradient = False - img = paddle.rand(shape=[in_channels, out_channels, modes1]) - img.stop_gradient = False - self.weights1_real = self.create_parameter( - [in_channels, out_channels, self.modes1], - attr=Initializer.Assign(self.scale * real), - ) - self.weights1_imag = self.create_parameter( - [in_channels, out_channels, self.modes1], - attr=Initializer.Assign(self.scale * img), - ) - - def compl_mul1d(self, op1, op2_real, op2_imag): - # (batch, in_channel, x ), (in_channel, out_channel, x) -> (batch, out_channel, x) - eq = "bix,iox->box" - op1_real = op1.real() - op1_imag = op1.imag() - result_real = paddle.unsqueeze( - paddle.einsum(eq, op1_real, op2_real) - - paddle.einsum(eq, op1_imag, op2_imag), - axis=-1, - ) - result_imag = paddle.unsqueeze( - paddle.einsum(eq, op1_real, op2_imag) - + paddle.einsum(eq, op1_imag, op2_real), - axis=-1, - ) - result_complex = paddle.as_complex( - paddle.concat([result_real, result_imag], axis=-1) - ) - return result_complex - - def forward(self, x, output_size=None): - batchsize = x.shape[0] - # Compute Fourier coefficients up to factor of e^(- something constant) - x_ft = paddle.fft.rfft(x) - - # Multiply relevant Fourier modes - out_ft_real = paddle.zeros( - [batchsize, self.out_channels, x.shape[-1] // 2 + 1], dtype="float32" - ) - out_ft_img = paddle.zeros( - [batchsize, self.out_channels, x.shape[-1] // 2 + 1], dtype="float32" - ) - out_ft = paddle.complex(out_ft_real, out_ft_img) - - out_ft[:, :, : self.modes1] = self.compl_mul1d( - x_ft[:, :, : self.modes1], self.weights1_real, self.weights1_imag - ) - - # Return to physical space - if output_size is None: - x = paddle.fft.irfft(out_ft, n=x.shape[-1]) - else: - x = paddle.fft.irfft(out_ft, n=output_size) - - return x - - -class FNO1d(nn.Layer): - """The overall network. It contains 4 layers of the Fourier layer. - 1. Lift the input to the desire channel dimension by self.fc0 . - 2. 4 layers of the integral operators u' = (W + K)(u). - W defined by self.w; K defined by self.conv . - 3. Project from the channel space to the output space by self.fc1 and self.fc2 . - - Args: - input_key (Tuple[str, ...], optional): Key to get the input tensor from the dict. Defaults to ("intput",). - output_key (Tuple[str, ...], optional): Key to save the output tensor into the dict. Defaults to ("output",). - modes (int, optional, optional): Number of Fourier modes to compute, it should be the same as - that in fft part of the code below. Defaults to 64. - width (int, optional, optional): Number of channels in each Fourier layer. Defaults to 64. - padding (int, optional, optional): How many zeros to pad to the input Tensor. Defaults to 100. - input_channel (int, optional, optional): Number of channels of the input tensor. Defaults to 2. - output_np (int, optional, optional): Number of points to sample the solution. Defaults to 2001. - - Examples: - >>> model = ppsci.arch.FNO1d() - >>> input_data = paddle.randn([100, 2001, 2]) - >>> input_dict = {"input": input_data} - >>> out_dict = model(input_dict) - >>> for k, v in out_dict.items(): - ... print(k, v.shape) - output [100, 1] - """ - - def __init__( - self, - input_key=("input",), - output_key=("output",), - modes=64, - width=64, - padding=100, - input_channel=2, - output_np=2001, - ): - super().__init__() - self.input_keys = input_key - self.output_keys = output_key - - self.output_np = output_np - self.modes1 = modes - self.width = width - self.padding = padding - self.fc0 = nn.Linear(input_channel, self.width) - - self.conv0 = SpectralConv1d(self.width, self.width, self.modes1) - self.conv1 = SpectralConv1d(self.width, self.width, self.modes1) - self.conv2 = SpectralConv1d(self.width, self.width, self.modes1) - self.conv3 = SpectralConv1d(self.width, self.width, self.modes1) - self.conv4 = SpectralConv1d(self.width, self.width, self.modes1) - - self.w0 = nn.Conv1D(self.width, self.width, 1) - self.w1 = nn.Conv1D(self.width, self.width, 1) - self.w2 = nn.Conv1D(self.width, self.width, 1) - self.w3 = nn.Conv1D(self.width, self.width, 1) - - self.fc1 = nn.Linear(self.width, 128) - self.fc2 = nn.Linear(128, 1) - - def _functional_pad(self, x, pad, mode="constant", value=0.0, data_format="NCL"): - if len(x.shape) * 2 == len(pad) and mode == "constant": - pad = ( - paddle.to_tensor(pad, dtype="float32") - .reshape((-1, 2)) - .flip([0]) - .flatten() - .tolist() - ) - return F.pad(x, pad, mode, value, data_format) - - def forward(self, x): - x = x[self.input_keys[0]] - # Dict - x = self.fc0(x) - x = paddle.transpose(x, perm=[0, 2, 1]) - # pad the domain if input is non-periodic - x = self._functional_pad(x, [0, self.padding]) - - x1 = self.conv0(x) - x2 = self.w0(x) - x = x1 + x2 - x = F.gelu(x=x, approximate=False) - - x1 = self.conv1(x) - x2 = self.w1(x) - x = x1 + x2 - x = F.gelu(x, approximate=False) - - x1 = self.conv2(x) - x2 = self.w2(x) - x = x1 + x2 - x = F.gelu(x, approximate=False) - - x1 = self.conv3(x) - x2 = self.w3(x) - x = x1 + x2 - x = F.gelu(x, approximate=False) - - x = x[..., : -self.padding] - x1 = self.conv4(x, self.output_np) - x2 = F.interpolate(x, size=[self.output_np], mode="linear", align_corners=True) - x = x1 + x2 - # Change the x-dimension to (batch, channel, 2001) - x = x.transpose(perm=[0, 2, 1]) - x = self.fc1(x) - x = F.gelu(x, approximate=False) - x = self.fc2(x) - - return {self.output_keys[0]: x} diff --git a/examples/smc_reac/ppsci/arch/graphcast.py b/examples/smc_reac/ppsci/arch/graphcast.py deleted file mode 100644 index 79a1c0aeae..0000000000 --- a/examples/smc_reac/ppsci/arch/graphcast.py +++ /dev/null @@ -1,492 +0,0 @@ -# Copyright (c) 2024 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. - -from typing import TYPE_CHECKING -from typing import Dict -from typing import Tuple - -import paddle -import paddle.nn as nn - -from ppsci.arch import base - -if TYPE_CHECKING: - import ppsci.data.dataset.atmospheric_dataset as atmospheric_dataset - - -class ResidualConnection(nn.Layer): - def __init__(self, fn): - super().__init__() - self.fn = fn - - def forward(self, inputs): - return inputs + self.fn(inputs) - - -class GraphCastMLP(nn.Layer): - def __init__( - self, in_features, out_features, latent_features=None, layer_norm=True - ): - super().__init__() - - if latent_features is None: - latent_features = out_features - - self.mlp = nn.Sequential( - nn.Linear(in_features, latent_features, bias_attr=True), - nn.Silu(), - nn.Linear(latent_features, out_features, bias_attr=True), - ) - self.layer_norm = layer_norm - if layer_norm: - self.layer_norm = nn.LayerNorm(out_features) - - def forward(self, feat): - if self.layer_norm: - out = self.layer_norm(self.mlp(feat)) - else: - out = self.mlp(feat) - return out - - -class GraphCastGNN(nn.Layer): - def __init__( - self, - grid_node_num: int, - grid_node_emb_dim: int, - mesh_node_num: int, - mesh_node_emb_dim: int, - mesh_edge_emb_dim: int, - grid2mesh_edge_emb_dim: int, - mesh2grid_edge_emb_dim: int, - src_type: str = "mesh", - dst_type: str = "mesh", - ): - super().__init__() - - self.src = src_type - self.dst = dst_type - self.grid_node_num = grid_node_num - self.mesh_node_num = mesh_node_num - self.edge_in_dim = grid_node_emb_dim + mesh_node_emb_dim - - if src_type == "mesh" and dst_type == "mesh": - self.edge_in_dim += mesh_edge_emb_dim - self.edge_out_dim = mesh_edge_emb_dim - self.node_in_dim = mesh_node_emb_dim + mesh_edge_emb_dim - self.node_out_dim = mesh_node_emb_dim - elif src_type == "grid" and dst_type == "mesh": - self.edge_in_dim += grid2mesh_edge_emb_dim - self.edge_out_dim = grid2mesh_edge_emb_dim - self.node_in_dim = mesh_node_emb_dim + grid2mesh_edge_emb_dim - self.node_out_dim = mesh_node_emb_dim - elif src_type == "mesh" and dst_type == "grid": - self.edge_in_dim += mesh2grid_edge_emb_dim - self.edge_out_dim = mesh2grid_edge_emb_dim - self.node_in_dim = grid_node_emb_dim + mesh2grid_edge_emb_dim - self.node_out_dim = grid_node_emb_dim - else: - raise ValueError - - self.edge_layer = GraphCastMLP(self.edge_in_dim, self.edge_out_dim) - self.node_layer = GraphCastMLP(self.node_in_dim, self.node_out_dim) - - def forward(self, graph: "atmospheric_dataset.GraphGridMesh"): - if self.src == "mesh" and self.dst == "mesh": - edge_feats = graph.mesh_edge_feat - src_node_feats = graph.mesh_node_feat - dst_node_feats = graph.mesh_node_feat - src_idx = graph.mesh2mesh_src_index - dst_idx = graph.mesh2mesh_dst_index - dst_node_num = self.mesh_node_num - elif self.src == "grid" and self.dst == "mesh": - edge_feats = graph.grid2mesh_edge_feat - src_node_feats = graph.grid_node_feat - dst_node_feats = graph.mesh_node_feat - src_idx = graph.grid2mesh_src_index - dst_idx = graph.grid2mesh_dst_index - dst_node_num = self.mesh_node_num - elif self.src == "mesh" and self.dst == "grid": - edge_feats = graph.mesh2grid_edge_feat - src_node_feats = graph.mesh_node_feat - dst_node_feats = graph.grid_node_feat - src_idx = graph.mesh2grid_src_index - dst_idx = graph.mesh2grid_dst_index - dst_node_num = self.grid_node_num - - # update edge features - edge_feats_concat = paddle.concat( - [ - edge_feats, - paddle.gather(src_node_feats, src_idx), - paddle.gather(dst_node_feats, dst_idx), - ], - axis=-1, - ) - edge_feats_out = self.edge_layer(edge_feats_concat) - - _, batch_dim, _ = edge_feats_out.shape - - # update node features - edge_feats_scatter = paddle.zeros([dst_node_num, batch_dim, self.edge_out_dim]) - node_feats_concat = paddle.concat( - [ - dst_node_feats, - paddle.scatter( - edge_feats_scatter, dst_idx, edge_feats_out, overwrite=False - ), - ], - axis=-1, - ) - node_feats_out = self.node_layer(node_feats_concat) - - if self.src == "mesh" and self.dst == "mesh": - graph.mesh_edge_feat += edge_feats_out - graph.mesh_node_feat += node_feats_out - elif self.src == "grid" and self.dst == "mesh": - graph.grid2mesh_edge_feat += edge_feats_out - graph.mesh_node_feat += node_feats_out - elif self.src == "mesh" and self.dst == "grid": - graph.mesh2grid_edge_feat += edge_feats_out - graph.grid_node_feat += node_feats_out - - return graph - - -class GraphCastEmbedding(nn.Layer): - def __init__( - self, - grid_node_dim: int, - grid_node_emb_dim: int, - mesh_node_dim: int, - mesh_node_emb_dim: int, - mesh_edge_dim: int, - mesh_edge_emb_dim: int, - grid2mesh_edge_dim: int, - grid2mesh_edge_emb_dim: int, - mesh2grid_edge_dim: int, - mesh2grid_edge_emb_dim: int, - ): - super().__init__() - - self.grid_node_embedding = GraphCastMLP(grid_node_dim, grid_node_emb_dim) - self.mesh_node_embedding = GraphCastMLP(mesh_node_dim, mesh_node_emb_dim) - self.mesh_edge_embedding = GraphCastMLP(mesh_edge_dim, mesh_edge_emb_dim) - self.grid2mesh_edge_embedding = GraphCastMLP( - grid2mesh_edge_dim, grid2mesh_edge_emb_dim - ) - self.mesh2grid_edge_embedding = GraphCastMLP( - mesh2grid_edge_dim, mesh2grid_edge_emb_dim - ) - - def forward(self, graph: "atmospheric_dataset.GraphGridMesh"): - grid_node_emb = self.grid_node_embedding(graph.grid_node_feat) - mesh_node_emb = self.mesh_node_embedding(graph.mesh_node_feat) - mesh_edge_emb = self.mesh_edge_embedding(graph.mesh_edge_feat) - grid2mesh_edge_emb = self.grid2mesh_edge_embedding(graph.grid2mesh_edge_feat) - mesh2grid_edge_emb = self.mesh2grid_edge_embedding(graph.mesh2grid_edge_feat) - - graph.grid_node_feat = grid_node_emb - graph.mesh_node_feat = mesh_node_emb - graph.mesh_edge_feat = mesh_edge_emb - graph.grid2mesh_edge_feat = grid2mesh_edge_emb - graph.mesh2grid_edge_feat = mesh2grid_edge_emb - - return graph - - -class GraphCastGrid2Mesh(nn.Layer): - def __init__( - self, - grid_node_num: int, - grid_node_emb_dim: int, - mesh_node_num: int, - mesh_node_emb_dim: int, - mesh_edge_emb_dim: int, - grid2mesh_edge_emb_dim: int, - mesh2grid_edge_emb_dim: int, - ): - super().__init__() - self.grid2mesh_gnn = GraphCastGNN( - grid_node_num=grid_node_num, - grid_node_emb_dim=grid_node_emb_dim, - mesh_node_num=mesh_node_num, - mesh_node_emb_dim=mesh_node_emb_dim, - mesh_edge_emb_dim=mesh_edge_emb_dim, - grid2mesh_edge_emb_dim=grid2mesh_edge_emb_dim, - mesh2grid_edge_emb_dim=mesh2grid_edge_emb_dim, - src_type="grid", - dst_type="mesh", - ) - self.grid_node_layer = ResidualConnection( - GraphCastMLP(grid_node_emb_dim, grid_node_emb_dim) - ) - - def forward(self, graph: "atmospheric_dataset.GraphGridMesh"): - graph = self.grid2mesh_gnn(graph) - graph.grid_node_feat = self.grid_node_layer(graph.grid_node_feat) - return graph - - -class GraphCastMesh2Grid(nn.Layer): - def __init__( - self, - grid_node_num: int, - grid_node_emb_dim: int, - mesh_node_num: int, - mesh_node_emb_dim: int, - mesh_edge_emb_dim: int, - grid2mesh_edge_emb_dim: int, - mesh2grid_edge_emb_dim: int, - ): - super().__init__() - self.mesh2grid_gnn = GraphCastGNN( - grid_node_num=grid_node_num, - grid_node_emb_dim=grid_node_emb_dim, - mesh_node_num=mesh_node_num, - mesh_node_emb_dim=mesh_node_emb_dim, - mesh_edge_emb_dim=mesh_edge_emb_dim, - grid2mesh_edge_emb_dim=grid2mesh_edge_emb_dim, - mesh2grid_edge_emb_dim=mesh2grid_edge_emb_dim, - src_type="mesh", - dst_type="grid", - ) - self.mesh_node_layer = ResidualConnection( - GraphCastMLP(mesh_node_emb_dim, mesh_node_emb_dim) - ) - - def forward(self, graph: "atmospheric_dataset.GraphGridMesh"): - graph = self.mesh2grid_gnn(graph) - graph.mesh_node_feat = self.mesh_node_layer(graph.mesh_node_feat) - return graph - - -class GraphCastEncoder(nn.Layer): - def __init__( - self, - grid_node_num: int, - grid_node_dim: int, - grid_node_emb_dim: int, - mesh_node_num: int, - mesh_node_dim: int, - mesh_node_emb_dim: int, - mesh_edge_dim: int, - mesh_edge_emb_dim: int, - grid2mesh_edge_dim: int, - grid2mesh_edge_emb_dim: int, - mesh2grid_edge_dim: int, - mesh2grid_edge_emb_dim: int, - ): - super().__init__() - self.embedding = GraphCastEmbedding( - grid_node_dim=grid_node_dim, - grid_node_emb_dim=grid_node_emb_dim, - mesh_node_dim=mesh_node_dim, - mesh_node_emb_dim=mesh_node_emb_dim, - mesh_edge_dim=mesh_edge_dim, - mesh_edge_emb_dim=mesh_edge_emb_dim, - grid2mesh_edge_dim=grid2mesh_edge_dim, - grid2mesh_edge_emb_dim=grid2mesh_edge_emb_dim, - mesh2grid_edge_dim=mesh2grid_edge_dim, - mesh2grid_edge_emb_dim=mesh2grid_edge_emb_dim, - ) - self.grid2mesh_gnn = GraphCastGrid2Mesh( - grid_node_num=grid_node_num, - grid_node_emb_dim=grid_node_emb_dim, - mesh_node_num=mesh_node_num, - mesh_node_emb_dim=mesh_node_emb_dim, - mesh_edge_emb_dim=mesh_edge_emb_dim, - grid2mesh_edge_emb_dim=grid2mesh_edge_emb_dim, - mesh2grid_edge_emb_dim=mesh2grid_edge_emb_dim, - ) - - def forward(self, graph: "atmospheric_dataset.GraphGridMesh"): - graph = self.embedding(graph) - graph = self.grid2mesh_gnn(graph) - return graph - - -class GraphCastDecoder(nn.Layer): - def __init__( - self, - grid_node_num: int, - grid_node_emb_dim: int, - mesh_node_num: int, - mesh_node_emb_dim: int, - mesh_edge_emb_dim: int, - grid2mesh_edge_emb_dim: int, - mesh2grid_edge_emb_dim: int, - node_output_dim: int, - ): - super().__init__() - self.mesh2grid_gnn = GraphCastMesh2Grid( - grid_node_num=grid_node_num, - grid_node_emb_dim=grid_node_emb_dim, - mesh_node_num=mesh_node_num, - mesh_node_emb_dim=mesh_node_emb_dim, - mesh_edge_emb_dim=mesh_edge_emb_dim, - grid2mesh_edge_emb_dim=grid2mesh_edge_emb_dim, - mesh2grid_edge_emb_dim=mesh2grid_edge_emb_dim, - ) - self.grid_node_layer = GraphCastMLP( - grid_node_emb_dim, - node_output_dim, - latent_features=grid_node_emb_dim, - layer_norm=False, - ) - - def forward(self, graph: "atmospheric_dataset.GraphGridMesh"): - graph = self.mesh2grid_gnn(graph) - graph.grid_node_feat = self.grid_node_layer(graph.grid_node_feat) - return graph - - -class GraphCastProcessor(nn.Layer): - def __init__( - self, - grid_node_num: int, - grid_node_emb_dim: int, - mesh_node_num: int, - mesh_node_emb_dim: int, - mesh_edge_emb_dim: int, - grid2mesh_edge_emb_dim: int, - mesh2grid_edge_emb_dim: int, - gnn_msg_steps: int, - ): - super().__init__() - - self.processor = nn.Sequential() - for idx in range(gnn_msg_steps): - self.processor.add_sublayer( - f"{idx}", - GraphCastGNN( - grid_node_num=grid_node_num, - grid_node_emb_dim=grid_node_emb_dim, - mesh_node_num=mesh_node_num, - mesh_node_emb_dim=mesh_node_emb_dim, - mesh_edge_emb_dim=mesh_edge_emb_dim, - grid2mesh_edge_emb_dim=grid2mesh_edge_emb_dim, - mesh2grid_edge_emb_dim=mesh2grid_edge_emb_dim, - src_type="mesh", - dst_type="mesh", - ), - ) - - def forward(self, graph: "atmospheric_dataset.GraphGridMesh"): - graph = self.processor(graph) - return graph - - -class GraphCastNet(base.Arch): - """GraphCast Network - - Args: - input_keys (Tuple[str, ...]): Name of input keys. - output_keys (Tuple[str, ...]): Name of output keys. - grid_node_num (int): Number of grid nodes. - grid_node_dim (int): Dimension of grid nodes. - grid_node_emb_dim (int): Dimension of emdding grid nodes. - mesh_node_num (int): Number of mesh nodes. - mesh_node_dim (int): Dimension of mesh nodes. - mesh_node_emb_dim (int): Dimension of emdding mesh nodes. - mesh_edge_dim (int): Dimension of mesh edges. - mesh_edge_emb_dim (int): Dimension of emdding mesh edges. - grid2mesh_edge_dim (int): Dimension of mesh edges in Grid2Mesh GNN. - grid2mesh_edge_emb_dim (int): Dimension of emdding mesh edges in Grid2Mesh GNN. - mesh2grid_edge_dim (int): Dimension of mesh edges in Mesh2Grid GNN. - mesh2grid_edge_emb_dim (int): Dimension of emdding mesh edges in Mesh2Grid GNN. - gnn_msg_steps (int): Step of gnn messages. - node_output_dim (int): Dimension of output nodes. - """ - - def __init__( - self, - input_keys: Tuple[str, ...], - output_keys: Tuple[str, ...], - grid_node_num: int, - grid_node_dim: int, - grid_node_emb_dim: int, - mesh_node_num: int, - mesh_node_dim: int, - mesh_node_emb_dim: int, - mesh_edge_dim: int, - mesh_edge_emb_dim: int, - grid2mesh_edge_dim: int, - grid2mesh_edge_emb_dim: int, - mesh2grid_edge_dim: int, - mesh2grid_edge_emb_dim: int, - gnn_msg_steps: int, - node_output_dim: int, - ): - super().__init__() - self.input_keys = input_keys - self.output_keys = output_keys - self.graphcast = nn.Sequential( - ( - "encoder", - GraphCastEncoder( - grid_node_num=grid_node_num, - grid_node_dim=grid_node_dim, - grid_node_emb_dim=grid_node_emb_dim, - mesh_node_num=mesh_node_num, - mesh_node_dim=mesh_node_dim, - mesh_node_emb_dim=mesh_node_emb_dim, - mesh_edge_dim=mesh_edge_dim, - mesh_edge_emb_dim=mesh_edge_emb_dim, - grid2mesh_edge_dim=grid2mesh_edge_dim, - grid2mesh_edge_emb_dim=grid2mesh_edge_emb_dim, - mesh2grid_edge_dim=mesh2grid_edge_dim, - mesh2grid_edge_emb_dim=mesh2grid_edge_emb_dim, - ), - ), - ( - "processor", - GraphCastProcessor( - grid_node_num=grid_node_num, - grid_node_emb_dim=grid_node_emb_dim, - mesh_node_num=mesh_node_num, - mesh_node_emb_dim=mesh_node_emb_dim, - mesh_edge_emb_dim=mesh_edge_emb_dim, - grid2mesh_edge_emb_dim=grid2mesh_edge_emb_dim, - mesh2grid_edge_emb_dim=mesh2grid_edge_emb_dim, - gnn_msg_steps=gnn_msg_steps, - ), - ), - ( - "decoder", - GraphCastDecoder( - grid_node_num=grid_node_num, - grid_node_emb_dim=grid_node_emb_dim, - mesh_node_num=mesh_node_num, - mesh_node_emb_dim=mesh_node_emb_dim, - mesh_edge_emb_dim=mesh_edge_emb_dim, - grid2mesh_edge_emb_dim=grid2mesh_edge_emb_dim, - mesh2grid_edge_emb_dim=mesh2grid_edge_emb_dim, - node_output_dim=node_output_dim, - ), - ), - ) - - def forward( - self, x: Dict[str, "atmospheric_dataset.GraphGridMesh"] - ) -> Dict[str, paddle.Tensor]: - if self._input_transform is not None: - x = self._input_transform(x) - - graph = x[self.input_keys[0]] - y = self.graphcast(graph) - - if self._output_transform is not None: - y = self._output_transform(x, y) - return {self.output_keys[0]: y} diff --git a/examples/smc_reac/ppsci/arch/he_deeponets.py b/examples/smc_reac/ppsci/arch/he_deeponets.py deleted file mode 100644 index 811da0d1b1..0000000000 --- a/examples/smc_reac/ppsci/arch/he_deeponets.py +++ /dev/null @@ -1,197 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import Tuple -from typing import Union - -import paddle -import paddle.nn as nn - -from ppsci.arch import activation as act_mod -from ppsci.arch import base -from ppsci.arch import mlp - - -class HEDeepONets(base.Arch): - """Physical information deep operator networks. - - Args: - heat_input_keys (Tuple[str, ...]): Name of input data for heat boundary. - cold_input_keys (Tuple[str, ...]): Name of input data for cold boundary. - trunk_input_keys (Tuple[str, ...]): Name of input data for trunk net. - output_keys (Tuple[str, ...]): Output name of predicted temperature. - heat_num_loc (int): Number of sampled input data for heat boundary. - cold_num_loc (int): Number of sampled input data for cold boundary. - num_features (int): Number of features extracted from heat boundary, same for cold boundary and trunk net. - branch_num_layers (int): Number of hidden layers of branch net. - trunk_num_layers (int): Number of hidden layers of trunk net. - branch_hidden_size (Union[int, Tuple[int, ...]]): Number of hidden size of branch net. - An integer for all layers, or list of integer specify each layer's size. - trunk_hidden_size (Union[int, Tuple[int, ...]]): Number of hidden size of trunk net. - An integer for all layers, or list of integer specify each layer's size. - branch_skip_connection (bool, optional): Whether to use skip connection for branch net. Defaults to False. - trunk_skip_connection (bool, optional): Whether to use skip connection for trunk net. Defaults to False. - branch_activation (str, optional): Name of activation function for branch net. Defaults to "tanh". - trunk_activation (str, optional): Name of activation function for trunk net. Defaults to "tanh". - branch_weight_norm (bool, optional): Whether to apply weight norm on parameter(s) for branch net. Defaults to False. - trunk_weight_norm (bool, optional): Whether to apply weight norm on parameter(s) for trunk net. Defaults to False. - use_bias (bool, optional): Whether to add bias on predicted G(u)(y). Defaults to True. - - Examples: - >>> import ppsci - >>> model = ppsci.arch.HEDeepONets( - ... ('qm_h',), - ... ('qm_c',), - ... ("x",'t'), - ... ("T_h",'T_c','T_w'), - ... 1, - ... 1, - ... 100, - ... 9, - ... 6, - ... 256, - ... 128, - ... branch_activation="swish", - ... trunk_activation="swish", - ... use_bias=True, - ... ) - """ - - def __init__( - self, - heat_input_keys: Tuple[str, ...], - cold_input_keys: Tuple[str, ...], - trunk_input_keys: Tuple[str, ...], - output_keys: Tuple[str, ...], - heat_num_loc: int, - cold_num_loc: int, - num_features: int, - branch_num_layers: int, - trunk_num_layers: int, - branch_hidden_size: Union[int, Tuple[int, ...]], - trunk_hidden_size: Union[int, Tuple[int, ...]], - branch_skip_connection: bool = False, - trunk_skip_connection: bool = False, - branch_activation: str = "tanh", - trunk_activation: str = "tanh", - branch_weight_norm: bool = False, - trunk_weight_norm: bool = False, - use_bias: bool = True, - ): - super().__init__() - self.trunk_input_keys = trunk_input_keys - self.heat_input_keys = heat_input_keys - self.cold_input_keys = cold_input_keys - self.input_keys = ( - self.trunk_input_keys + self.heat_input_keys + self.cold_input_keys - ) - self.output_keys = output_keys - self.num_features = num_features - - self.heat_net = mlp.MLP( - self.heat_input_keys, - ("h",), - branch_num_layers, - branch_hidden_size, - branch_activation, - branch_skip_connection, - branch_weight_norm, - input_dim=heat_num_loc, - output_dim=num_features * len(self.output_keys), - ) - - self.cold_net = mlp.MLP( - self.cold_input_keys, - ("c",), - branch_num_layers, - branch_hidden_size, - branch_activation, - branch_skip_connection, - branch_weight_norm, - input_dim=cold_num_loc, - output_dim=num_features * len(self.output_keys), - ) - - self.trunk_net = mlp.MLP( - self.trunk_input_keys, - ("t",), - trunk_num_layers, - trunk_hidden_size, - trunk_activation, - trunk_skip_connection, - trunk_weight_norm, - input_dim=len(self.trunk_input_keys), - output_dim=num_features * len(self.output_keys), - ) - self.trunk_act = act_mod.get_activation(trunk_activation) - self.heat_act = act_mod.get_activation(branch_activation) - self.cold_act = act_mod.get_activation(branch_activation) - - self.use_bias = use_bias - if use_bias: - # register bias to parameter for updating in optimizer and storage - self.b = self.create_parameter( - shape=(len(self.output_keys),), - attr=nn.initializer.Constant(0.0), - ) - - def forward(self, x): - if self._input_transform is not None: - x = self._input_transform(x) - - # Branch net to encode the input function - heat_features = self.heat_net(x)[self.heat_net.output_keys[0]] - cold_features = self.cold_net(x)[self.cold_net.output_keys[0]] - # Trunk net to encode the domain of the output function - y_features = self.trunk_net(x)[self.trunk_net.output_keys[0]] - y_features = self.trunk_act(y_features) - # Dot product - G_u_h = paddle.sum( - heat_features[:, : self.num_features] - * y_features[:, : self.num_features] - * cold_features[:, : self.num_features], - axis=1, - keepdim=True, - ) - G_u_c = paddle.sum( - heat_features[:, self.num_features : 2 * self.num_features] - * y_features[:, self.num_features : 2 * self.num_features] - * cold_features[:, self.num_features : 2 * self.num_features], - axis=1, - keepdim=True, - ) - G_u_w = paddle.sum( - heat_features[:, 2 * self.num_features :] - * y_features[:, 2 * self.num_features :] - * cold_features[:, 2 * self.num_features :], - axis=1, - keepdim=True, - ) - # Add bias - if self.use_bias: - G_u_h += self.b[0] - G_u_c += self.b[1] - G_u_w += self.b[2] - - result_dict = { - self.output_keys[0]: G_u_h, - self.output_keys[1]: G_u_c, - self.output_keys[2]: G_u_w, - } - if self._output_transform is not None: - result_dict = self._output_transform(x, result_dict) - - return result_dict diff --git a/examples/smc_reac/ppsci/arch/ifm_mlp.py b/examples/smc_reac/ppsci/arch/ifm_mlp.py deleted file mode 100644 index 4685b34be6..0000000000 --- a/examples/smc_reac/ppsci/arch/ifm_mlp.py +++ /dev/null @@ -1,540 +0,0 @@ -# Copyright (c) 2024 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. - -from typing import List -from typing import Tuple - -import numpy as np -import paddle -import paddle.nn as nn -import paddle.nn.functional as F - -import ppsci -from ppsci.arch import base - - -def init_parameter_uniform( - parameter: paddle.base.framework.EagerParamBase, n: int -) -> None: - ppsci.utils.initializer.uniform_(parameter, -1 / np.sqrt(n), 1 / np.sqrt(n)) - - -# inputs, hidden_units, outputs, d_out, sigma, dp_ratio, first_omega_0, hidden_omega_0, reg -class IFMMLP(base.Arch): - """Understanding the limitations of deep models for molecular property prediction: Insights and solutions. - [Xia, Jun, et al. Advances in Neural Information Processing Systems 36 (2023): 64774-64792.]https://openreview.net/forum?id=NLFqlDeuzt) - - Code reference: https://github.com/junxia97/IFM - - Args: - input_keys (Tuple[str, ...]): Name of input keys, such as ("input", ). - output_keys (Tuple[str, ...]): Name of output keys, such as ("pred", ). - hidden_units (List[int]): Units num in hidden layers. - embed_name (str): Embed name used in arch, such as "IMF", "None". - inputs (int): Input dim. - outputs (int): Output dim. - d_out (int): Embedding output dim for some architecture. - sigma (float): Hyper parameter for some architecture. - dp_ratio (float): Dropout ratio. - reg (bool): Regularization flag. - first_omega_0 (float): Frequency factor used in first layer. - hidden_omega_0 (float): Frequency factor used in hidden layer. - """ - - def __init__( - self, - input_keys: Tuple[str, ...], - output_keys: Tuple[str, ...], - hidden_units: List[int], - embed_name: str, - inputs: int, - outputs: int, - d_out: int, - sigma: float, - dp_ratio: float, - reg: bool, - first_omega_0: float, - hidden_omega_0: float, - ): - super().__init__() - self.input_keys = input_keys - self.output_keys = output_keys - - # initialization - if embed_name == "None": - my_model = MyDNN( - inputs=inputs, - hidden_units=hidden_units, - dp_ratio=dp_ratio, - outputs=outputs, - reg=reg, - ) - elif embed_name == "LE": - my_model = LE_DNN( - inputs=inputs, - hidden_units=hidden_units, - d_out=d_out + 1, - dp_ratio=dp_ratio, - outputs=outputs, - reg=reg, - ) - elif embed_name == "LSIM": - my_model = LSIM_DNN( - inputs=inputs, - hidden_units=hidden_units, - d_out=d_out + 1, - sigma=sigma, - dp_ratio=dp_ratio, - outputs=outputs, - reg=reg, - ) - elif embed_name == "IFM": - my_model = IFM_DNN( - inputs=inputs, - hidden_units=hidden_units, - outputs=outputs, - dp_ratio=dp_ratio, - first_omega_0=first_omega_0, - hidden_omega_0=hidden_omega_0, - reg=reg, - ) - elif embed_name == "GM": - my_model = GM_DNN( - inputs=inputs, - hidden_units=hidden_units, - d_out=d_out + 1, - sigma=sigma + 1, - dp_ratio=dp_ratio, - outputs=outputs, - reg=reg, - ) - elif embed_name == "SIM": - my_model = SIM_DNN( - inputs=inputs, - hidden_units=hidden_units, - d_out=d_out + 1, - sigma=sigma + 1, - dp_ratio=dp_ratio, - outputs=outputs, - reg=reg, - ) - else: - raise ValueError("Invalid Embedding Name") - - self.model = my_model - - def forward(self, x): - Xs = x[self.input_keys[0]] - ret = self.model(Xs) - return {self.output_keys[0]: ret} - - -class MyDNN(nn.Layer): - """ - Args: - inputs (int): Input dim. - hidden_units (List[int]): Units num in hidden layers. - outputs (int): Output dim. - dp_ratio (float): Dropout ratio. - reg (bool): Regularization flag. - """ - - def __init__(self, inputs, hidden_units, outputs, dp_ratio, reg): - super(MyDNN, self).__init__() - # parameters - self.reg = reg - - # layers - self.hidden1 = nn.Linear(inputs, hidden_units[0]) - self.dropout1 = nn.Dropout(dp_ratio) - - self.hidden2 = nn.Linear(hidden_units[0], hidden_units[1]) - self.dropout2 = nn.Dropout(dp_ratio) - - self.hidden3 = nn.Linear(hidden_units[1], hidden_units[2]) - self.dropout3 = nn.Dropout(dp_ratio) - - if reg: - self.output = nn.Linear(hidden_units[2], 1) - else: - self.output = nn.Linear(hidden_units[2], outputs) - - def forward(self, x): - x = self.hidden1(x) - x = F.relu(self.dropout1(x)) - - x = self.hidden2(x) - x = F.relu(self.dropout2(x)) - - x = self.hidden3(x) - x = F.relu(self.dropout3(x)) - - return self.output(x) - - -class LE(nn.Layer): - def __init__(self, n_tokens: int, d_out: int): - super().__init__() - self.weight = self.create_parameter([n_tokens, 1, d_out]) - self.bias = self.create_parameter([n_tokens, d_out]) - self.reset_parameters() - - def reset_parameters(self) -> None: - d_out = self.weight.shape[-1] - init_parameter_uniform(self.weight, d_out) - init_parameter_uniform(self.bias, d_out) - - def forward(self, x: paddle.Tensor) -> paddle.Tensor: - """ - x: (n_batch, n_features, d_in) - returns: (n_batch, n_features, d_out) - """ - x = x.unsqueeze(-1) - x = (x.unsqueeze(-2) @ self.weight[None]).squeeze(-2) - x = x + self.bias[None] - return x - - -class PLE(nn.Layer): - def __init__(self, n_num_features: int, d_out: int, sigma: float): - super().__init__() - self.d_out = d_out - self.sigma = sigma - self.coefficients = self.create_parameter([n_num_features, d_out]) - self.reset_parameters() - - def reset_parameters(self) -> None: - ppsci.utils.initializer.normal_(self.coefficients, 0.0, self.sigma) - - def forward(self, x: paddle.Tensor) -> paddle.Tensor: - x = 2 * np.pi * self.coefficients[None] * x[..., None] - return paddle.concat([paddle.cos(x), paddle.sin(x)], -1) - - -class LE_DNN(nn.Layer): - def __init__(self, inputs, hidden_units, outputs, d_out, dp_ratio, reg): - super(LE_DNN, self).__init__() - # parameters - self.reg = reg - # layers - self.hidden1 = nn.Linear(inputs * d_out, hidden_units[0]) - self.dropout1 = nn.Dropout(dp_ratio) - - self.hidden2 = nn.Linear(hidden_units[0], hidden_units[1]) - self.dropout2 = nn.Dropout(dp_ratio) - - self.hidden3 = nn.Linear(hidden_units[1], hidden_units[2]) - self.dropout3 = nn.Dropout(dp_ratio) - - if reg: - self.output = nn.Linear(hidden_units[2], 1) - else: - self.output = nn.Linear(hidden_units[2], outputs) - self.embedding = LE(inputs, d_out) - - def forward(self, x): - x = self.embedding(x).view([x.size(0), -1]) - x = self.hidden1(x) - x = F.relu(self.dropout1(x)) - - x = self.hidden2(x) - x = F.relu(self.dropout2(x)) - - x = self.hidden3(x) - x = F.relu(self.dropout3(x)) - - return self.output(x) - - -class LSIM_DNN(nn.Layer): - def __init__(self, inputs, hidden_units, outputs, d_out, sigma, dp_ratio, reg): - super(LSIM_DNN, self).__init__() - # parameters - self.reg = reg - # layers - self.hidden1 = nn.Linear(inputs, hidden_units[0]) - self.dropout1 = nn.Dropout(dp_ratio) - - self.hidden2 = nn.Linear(hidden_units[0], hidden_units[1]) - self.dropout2 = nn.Dropout(dp_ratio) - - self.hidden3 = nn.Linear(hidden_units[1], hidden_units[2]) - self.dropout3 = nn.Dropout(dp_ratio) - - if reg: - self.output = nn.Linear(hidden_units[2], 1) - else: - self.output = nn.Linear(hidden_units[2], outputs) - self.embedding = PLE(inputs, d_out, sigma) - self.linear = nn.Linear(d_out * 2, inputs) - - def forward(self, x): - x = self.embedding(x).sum(1) - x = F.relu(self.linear(x)) - x = self.hidden1(x) - x = F.relu(self.dropout1(x)) - - x = self.hidden2(x) - x = F.relu(self.dropout2(x)) - - x = self.hidden3(x) - x = F.relu(self.dropout3(x)) - - return self.output(x) - - -class GaussianEncoding(nn.Layer): - def __init__(self, n_num_features: int, d_out: int, sigma: float) -> None: - super().__init__() - self.d_out = d_out - self.sigma = sigma - self.n_num_features = n_num_features - self.size = (d_out, n_num_features) - self.B = paddle.randn(self.size) * sigma - - def forward(self, x: paddle.Tensor) -> paddle.Tensor: - """ - x: (n_batch, n_features) - returns: (n_batch, n_features * 2 * d_out) - """ - self.B = self.B.to(x.place) - xp = 2 * np.pi * x @ self.B.T - return paddle.concat((paddle.cos(xp), paddle.sin(xp)), axis=-1) - - -class GM_DNN(nn.Layer): - """ - Args: - inputs (int): Input dim. - hidden_units (List[int]): Units num in hidden layers. - outputs (int): Output dim. - d_out (int): Embedding output dim for some architecture. - sigma (float): Hyper parameter for some architecture. - dp_ratio (float): Dropout ratio. - reg (bool): Regularization flag. - """ - - def __init__(self, inputs, hidden_units, outputs, d_out, sigma, dp_ratio, reg): - super(GM_DNN, self).__init__() - # parameters - self.reg = reg - self.d_out = d_out - self.sigma = sigma - # layers - self.hidden1 = nn.Linear(d_out * 2, hidden_units[0]) - self.dropout1 = nn.Dropout(dp_ratio) - - self.hidden2 = nn.Linear(hidden_units[0], hidden_units[1]) - self.dropout2 = nn.Dropout(dp_ratio) - - self.hidden3 = nn.Linear(hidden_units[1], hidden_units[2]) - self.dropout3 = nn.Dropout(dp_ratio) - - if reg: - self.output = nn.Linear(hidden_units[2], 1) - else: - self.output = nn.Linear(hidden_units[2], outputs) - - self.embedding = GaussianEncoding(inputs, d_out, sigma) - - def forward(self, x): - x = self.embedding(x) - x = self.hidden1(x) - x = F.relu(self.dropout1(x)) - - x = self.hidden2(x) - x = F.relu(self.dropout2(x)) - - x = self.hidden3(x) - x = F.relu(self.dropout3(x)) - - return self.output(x) - - -class SineLayer(nn.Layer): - # If is_first=True, omega_0 is a frequency factor which simply multiplies the activations before the - # nonlinearity. Different signals may require different omega_0 in the first layer - this is a - # hyperparameter. - - # If is_first=False, then the weights will be divided by omega_0 so as to keep the magnitude of - # activations constant, but boost gradients to the weight matrix - - def __init__( - self, in_features, out_features, bias=True, is_first=False, omega_0=30 - ): - super().__init__() - self.omega_0 = omega_0 - self.is_first = is_first - - self.in_features = in_features - self.linear = nn.Linear(in_features, out_features, bias_attr=bias) - - self.init_weights() - - def init_weights(self): - with paddle.no_grad(): - if self.is_first: - self.linear.weight.uniform_(-1 / self.in_features, 1 / self.in_features) - else: - self.linear.weight.uniform_( - -np.sqrt(6 / self.in_features) / self.omega_0, - np.sqrt(6 / self.in_features) / self.omega_0, - ) - - def forward(self, input): - return paddle.sin(self.omega_0 * self.linear(input)) - - def forward_with_intermediate(self, input): - # For visualization of activation distributions - intermediate = self.omega_0 * self.linear(input) - return paddle.sin(intermediate), intermediate - - -class IFM_DNN(nn.Layer): - """ - Args: - inputs (int): Input dim. - hidden_units (List[int]): Units num in hidden layers. - outputs (int): Output dim. - dp_ratio (float): Dropout ratio. - first_omega_0 (float): Frequency factor used in first layer. - hidden_omega_0 (float): Frequency factor used in hidden layer. - reg (bool): Regularization flag. - """ - - def __init__( - self, - inputs, - hidden_units, - outputs, - dp_ratio, - first_omega_0, - hidden_omega_0, - reg, - ): - super(IFM_DNN, self).__init__() - # parameters - self.reg = reg - # layers - self.hidden1 = SineLayer( - inputs, hidden_units[0], is_first=True, omega_0=first_omega_0 - ) - self.dropout1 = nn.Dropout(dp_ratio) - - self.hidden2 = SineLayer( - hidden_units[0], hidden_units[1], is_first=False, omega_0=hidden_omega_0 - ) - self.dropout2 = nn.Dropout(dp_ratio) - - self.hidden3 = SineLayer( - hidden_units[1], hidden_units[2], is_first=False, omega_0=hidden_omega_0 - ) - self.dropout3 = nn.Dropout(dp_ratio) - - if reg: - self.output = nn.Linear(hidden_units[2], 1) - with paddle.no_grad(): - self.output.weight.uniform_( - -np.sqrt(6 / hidden_units[2]) / hidden_omega_0, - np.sqrt(6 / hidden_units[2]) / hidden_omega_0, - ) - else: - self.output = nn.Linear(hidden_units[2], outputs) - with paddle.no_grad(): - self.output.weight.uniform_( - -np.sqrt(6 / hidden_units[2]) / hidden_omega_0, - np.sqrt(6 / hidden_units[2]) / hidden_omega_0, - ) - - def forward(self, x): - - x = self.hidden1(x) - x = F.relu(self.dropout1(x)) - # x = self.dropout1(x) - - x = self.hidden2(x) - x = F.relu(self.dropout2(x)) - # x = self.dropout2(x) - - x = self.hidden3(x) - x = F.relu(self.dropout3(x)) - # x = self.dropout3(x) - - return self.output(x) - - -class SIM_encoding(nn.Layer): - def __init__(self, n_num_features: int, d_out: int, sigma: float): - super().__init__() - self.d_out = d_out - self.sigma = sigma - self.n_num_features = n_num_features - self.coeffs = 2 * np.pi * sigma ** (paddle.arange(d_out) / d_out) - - def forward(self, x: paddle.Tensor) -> paddle.Tensor: - """ - x: (n_batch, n_features) - returns: (n_batch, n_features * 2 * d_out) - """ - xp = self.coeffs.to(x.device) * paddle.unsqueeze(x, -1) - xp_cat = paddle.concat((paddle.cos(xp), paddle.sin(xp)), axis=-1) - return xp_cat.flatten(-2, -1) - - -class SIM_DNN(nn.Layer): - """ - Args: - inputs (int): Input dim. - hidden_units (List[int]): Units num in hidden layers. - outputs (int): Output dim. - d_out (int): Embedding output dim for some architecture. - sigma (float): Hyper parameter for some architecture. - dp_ratio (float): Dropout ratio. - reg (bool): Regularization flag. - """ - - def __init__(self, inputs, hidden_units, outputs, d_out, sigma, dp_ratio, reg): - super(SIM_DNN, self).__init__() - # parameters - self.reg = reg - self.d_out = d_out - self.sigma = sigma - # layers - self.hidden1 = nn.Linear(d_out * 2 * inputs, hidden_units[0]) - self.dropout1 = nn.Dropout(dp_ratio) - - self.hidden2 = nn.Linear(hidden_units[0], hidden_units[1]) - self.dropout2 = nn.Dropout(dp_ratio) - - self.hidden3 = nn.Linear(hidden_units[1], hidden_units[2]) - self.dropout3 = nn.Dropout(dp_ratio) - - if reg: - self.output = nn.Linear(hidden_units[2], 1) - else: - self.output = nn.Linear(hidden_units[2], outputs) - self.embedding = SIM_encoding(inputs, d_out, sigma) - - def forward(self, x): - x = self.embedding(x) - x = self.hidden1(x) - x = F.relu(self.dropout1(x)) - - x = self.hidden2(x) - x = F.relu(self.dropout2(x)) - - x = self.hidden3(x) - x = F.relu(self.dropout3(x)) - - return self.output(x) diff --git a/examples/smc_reac/ppsci/arch/kan.py b/examples/smc_reac/ppsci/arch/kan.py deleted file mode 100644 index 2c46d60c64..0000000000 --- a/examples/smc_reac/ppsci/arch/kan.py +++ /dev/null @@ -1,385 +0,0 @@ -# Copyright (c) 2025 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. - -import math -from typing import Callable -from typing import Tuple - -import paddle - -from ppsci.arch import base -from ppsci.utils import initializer - -""" -This is the paddle implementation of Korogonov-Arnold-Network (KAN) -which is based on the torch implementation [efficient-kan] by Blealtan and akkashdash -please refer to their work (https://github.com/Blealtan/efficient-kan) -Authors: guhaohao0991(guhaohao@baidu.com) -Date: 2025/04/ -""" - - -class KANLinear(paddle.nn.Layer): - def __init__( - self, - in_features: int, - out_features: int, - grid_size: int = 5, - spline_order: int = 3, - scale_noise: float = 0.1, - scale_base: float = 1.0, - scale_spline: float = 1.0, - enable_standalone_scale_spline: bool = True, - base_activation: Callable[[paddle.Tensor], paddle.Tensor] = paddle.nn.Silu, - grid_eps: float = 0.02, - grid_range: Tuple[float, float] = (-1, 1), - ): - super().__init__() - self.in_features = in_features - self.out_features = out_features - self.grid_size = grid_size - self.spline_order = spline_order - - h = (grid_range[1] - grid_range[0]) / grid_size - grid = ( - ( - paddle.arange(start=-spline_order, end=grid_size + spline_order + 1) * h - + grid_range[0] - ) - .expand(shape=[in_features, -1]) - .contiguous() - ) - self.register_buffer(name="grid", tensor=grid) - - self.base_weight = self.create_parameter( - shape=[out_features, in_features], - default_initializer=paddle.nn.initializer.Assign( - paddle.empty(shape=[out_features, in_features]) - ), - ) - self.spline_weight = self.create_parameter( - shape=[out_features, in_features, grid_size + spline_order], - default_initializer=paddle.nn.initializer.Assign( - paddle.empty( - shape=[out_features, in_features, grid_size + spline_order] - ) - ), - ) - - if enable_standalone_scale_spline: - self.spline_scaler = self.create_parameter( - shape=[out_features, in_features], - default_initializer=paddle.nn.initializer.Assign( - paddle.empty(shape=[out_features, in_features]) - ), - ) - - self.scale_noise = scale_noise - self.scale_base = scale_base - self.scale_spline = scale_spline - self.enable_standalone_scale_spline = enable_standalone_scale_spline - self.base_activation = base_activation() - self.grid_eps = grid_eps - - self.reset_parameters() - - def reset_parameters(self): - self.base_weight = initializer.kaiming_uniform_( - tensor=self.base_weight, - a=math.sqrt(5) * self.scale_base, - nonlinearity="leaky_relu", - ) - with paddle.no_grad(): - noise = ( - ( - paddle.rand( - shape=[self.grid_size + 1, self.in_features, self.out_features] - ) - - 1 / 2 - ) - * self.scale_noise - / self.grid_size - ) - - paddle.assign( - (self.scale_spline if not self.enable_standalone_scale_spline else 1.0) - * self.curve2coeff( - self.grid.T[self.spline_order : -self.spline_order], noise - ), - output=self.spline_weight.data, - ) - - if self.enable_standalone_scale_spline: - self.spline_scaler = initializer.kaiming_uniform_( - tensor=self.spline_scaler, - a=math.sqrt(5) * self.scale_spline, - nonlinearity="leaky_relu", - ) - - def b_splines(self, x: paddle.Tensor): - """ - Compute the B-spline bases for the given input tensor. - - Args: - x (paddle.Tensor): Input tensor of shape (batch_size, in_features). - - Returns: - paddle.Tensor: B-spline bases tensor of shape (batch_size, in_features, grid_size + spline_order). - """ - assert x.dim() == 2 and x.shape[1] == self.in_features - grid: paddle.Tensor = self.grid - x = x.unsqueeze(axis=-1) - bases = ((x >= grid[:, :-1]) & (x < grid[:, 1:])).to(x.dtype) - - for k in range(1, self.spline_order + 1): - bases = (x - grid[:, : -(k + 1)]) / ( - grid[:, k:-1] - grid[:, : -(k + 1)] - ) * bases[:, :, :-1] + (grid[:, k + 1 :] - x) / ( - grid[:, k + 1 :] - grid[:, 1:-k] - ) * bases[ - :, :, 1: - ] - - assert tuple(bases.shape) == ( - x.shape[0], - self.in_features, - self.grid_size + self.spline_order, - ) - - return bases.contiguous() - - def curve2coeff(self, x: paddle.Tensor, y: paddle.Tensor): - """ - Compute the coefficients of the curve that interpolates the given points. - - Args: - x (paddle.Tensor): Input tensor of shape (batch_size, in_features). - y (paddle.Tensor): Output tensor of shape (batch_size, in_features, out_features). - - Returns: - paddle.Tensor: Coefficients tensor of shape (out_features, in_features, grid_size + spline_order). - """ - assert x.dim() == 2 and x.shape[1] == self.in_features - assert tuple(y.shape) == (x.shape[0], self.in_features, self.out_features) - - A = self.b_splines(x).transpose( - perm=dim2perm(self.b_splines(x).ndim, 0, 1) - ) # [in_features, batch_size, grid_size + spline_order] - B = y.transpose( - perm=dim2perm(y.ndim, 0, 1) - ) # [in_features, batch_size, out_features] - solution = paddle.linalg.lstsq(x=A, y=B)[ - 0 - ] # [in_features, grid_size + spline_order, out_features] - if A.shape[0] == 1: - solution = solution.unsqueeze(axis=0) - # print("A shape: ", A.shape, "B shape: ", B.shape, "Solution shape: ", solution.shape) - result = solution.transpose([2, 0, 1]) - assert tuple(result.shape) == ( - self.out_features, - self.in_features, - self.grid_size + self.spline_order, - ) - - return result.contiguous() - - @property - def scaled_spline_weight(self): - return self.spline_weight * ( - self.spline_scaler.unsqueeze(axis=-1) - if self.enable_standalone_scale_spline - else 1.0 - ) - - def forward(self, x: paddle.Tensor): - assert x.dim() == 2 and x.shape[1] == self.in_features - - base_output = paddle.nn.functional.linear( - x=self.base_activation(x), weight=self.base_weight.T - ) - - spline_output = paddle.nn.functional.linear( - x=self.b_splines(x).reshape([x.shape[0], -1]).contiguous(), - weight=self.scaled_spline_weight.reshape( - [self.out_features, -1] - ).T.contiguous(), - ) - # cant calculate 1st order derivation using view - # spline_output = paddle.nn.functional.linear( - # x=self.b_splines(x).view(x.shape[0], -1), - # weight=self.scaled_spline_weight.view(self.out_features, -1).T) - - return base_output + spline_output - - @paddle.no_grad() - def update_grid(self, x: paddle.Tensor, margin=0.01): - assert x.dim() == 2 and x.shape[1] == self.in_features - batch = x.shape[0] - - splines = self.b_splines(x) # [batch, in, coeff] - splines = splines.transpose(perm=[1, 0, 2]) # [in, batch, coeff] - orig_coeff = self.scaled_spline_weight # [out, in, coeff] - orig_coeff = orig_coeff.transpose(perm=[1, 2, 0]) # [in, coeff, out] - unreduced_spline_output = paddle.bmm( - x=splines, y=orig_coeff - ) # [in, batch, out] - unreduced_spline_output = unreduced_spline_output.transpose( - perm=[1, 0, 2] - ) # [batch, in, out] - - # sort each channel individually to collect data distribution - x_sorted = (paddle.sort(x=x, axis=0), paddle.argsort(x=x, axis=0))[0] - grid_adaptive = x_sorted[ - paddle.linspace( - start=0, stop=batch - 1, num=self.grid_size + 1, dtype="int64" - ) - ] - uniform_step = (x_sorted[-1] - x_sorted[0] + 2 * margin) / self.grid_size - grid_uniform = ( - paddle.arange(dtype="float32", end=self.grid_size + 1).unsqueeze(axis=1) - * uniform_step - + x_sorted[0] - - margin - ) - - grid = self.grid_eps * grid_uniform + (1 - self.grid_eps) * grid_adaptive - grid = paddle.concat( - x=[ - grid[:1] - - uniform_step - * paddle.arange( - start=self.spline_order, end=0, step=-1, dtype="float32" - ).unsqueeze(axis=1), - grid, - grid[-1:] - + uniform_step - * paddle.arange( - start=1, end=self.spline_order + 1, dtype="float32" - ).unsqueeze(axis=1), - ], - axis=0, - ) - - paddle.assign(grid.T, output=self.grid) - paddle.assign( - self.curve2coeff(x, unreduced_spline_output), output=self.spline_weight.data - ) - - def regularization_loss(self, regularize_activation=1.0, regularize_entropy=1.0): - """ - Compute the regularization loss. - - L1 and the entropy loss is for the feature selection, i.e., let the weight of the activation function be small. - """ - l1_fake = self.spline_weight.abs().mean(axis=-1) - regularization_loss_activation = l1_fake.sum() - p = l1_fake / regularization_loss_activation - regularization_loss_entropy = -paddle.sum(x=p * p.log()) - return ( - regularize_activation * regularization_loss_activation - + regularize_entropy * regularization_loss_entropy - ) - - -class KAN(base.Arch): - """Kolmogorov-Arnold Network (KAN). - - Args: - layers_hidden (Tuple[int, ...]): The number of hidden neurons in each layer. - input_keys (Tuple[str, ...]): The keys of the input dictionary. - output_keys (Tuple[str, ...]): The keys of the output dictionary. - grid_size (int): The size of the grid used by the spline basis functions. Default: 5. - spline_order (int): The order of the spline basis functions. Default: 3. - scale_noise (float): The scaling factor for the noise added to the weights of the KAN-linear layers. Default: 0.1. - scale_base (float): The scaling factor for the base activation output. Default: 1.0. - scale_spline (float): The scaling factor for the b-spline output. Default: 1.0. - base_activation (Callable[[paddle.Tensor], paddle.Tensor]): The base activation function. Default: paddle.nn.Silu. - grid_eps (float): The epsilon value used to initialize the grid. Default: 0.02. - grid_range (Tuple[float, float]): The domain range of the grid for b-spline interpolation. Default: (-1, 1). - - Examples: - >>> import paddle - >>> import ppsci - >>> model = ppsci.arch.KAN( - ... layers_hidden=(2, 5, 5, 1), - ... input_keys=("x", "y"), - ... output_keys=("z"), - ... grid_size=5, - ... spline_order=3 - >>> ) - >>> input_dict = {"x": paddle.rand([64, 1]), - ... "y": paddle.rand([64, 1])} - >>> output_dict = model(input_dict) - >>> print(output_dict["z"].shape) - [64, 1] - """ - - def __init__( - self, - layers_hidden: Tuple[int, ...], - input_keys: Tuple[str, ...], - output_keys: Tuple[str, ...], - grid_size: int = 5, - spline_order: int = 3, - scale_noise: float = 0.1, - scale_base: float = 1.0, - scale_spline: float = 1.0, - base_activation: Callable[[paddle.Tensor], paddle.Tensor] = paddle.nn.Silu, - grid_eps: float = 0.02, - grid_range: Tuple[float, float] = (-1, 1), - ): - super().__init__() - self.input_keys = input_keys - self.output_keys = output_keys - self.grid_size = grid_size - self.spline_order = spline_order - self.layers = paddle.nn.LayerList() - for in_features, out_features in zip(layers_hidden, layers_hidden[1:]): - self.layers.append( - KANLinear( - in_features, - out_features, - grid_size=grid_size, - spline_order=spline_order, - scale_noise=scale_noise, - scale_base=scale_base, - scale_spline=scale_spline, - base_activation=base_activation, - grid_eps=grid_eps, - grid_range=grid_range, - ) - ) - - def forward(self, x_dict, update_grid=False): - x = self.concat_to_tensor(x_dict, self.input_keys, axis=-1) - for index, layer in enumerate(self.layers): - if update_grid: - layer.update_grid(x) - x = layer(x) - if index < len(self.layers) - 1: - x = paddle.nn.functional.tanh(x=x) - out_dic = self.split_to_dict(x, self.output_keys, axis=-1) - return out_dic - - def regularization_loss(self, regularize_activation=1.0, regularize_entropy=1.0): - return sum( - layer.regularization_loss(regularize_activation, regularize_entropy) - for layer in self.layers - ) - - -def dim2perm(ndim, dim0, dim1): - perm = list(range(ndim)) - perm[dim0], perm[dim1] = perm[dim1], perm[dim0] - return perm diff --git a/examples/smc_reac/ppsci/arch/lno.py b/examples/smc_reac/ppsci/arch/lno.py deleted file mode 100644 index d600c5d028..0000000000 --- a/examples/smc_reac/ppsci/arch/lno.py +++ /dev/null @@ -1,312 +0,0 @@ -# Copyright (c) 2024 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. - -from __future__ import annotations - -import operator -from functools import reduce -from typing import Optional -from typing import Tuple - -import numpy as np -import paddle -import paddle.nn as nn - -from ppsci.arch import activation as act_mod -from ppsci.arch import base -from ppsci.utils import initializer - - -class Laplace(nn.Layer): - """Generic N-Dimensional Laplace Operator with Pole-Residue Method. - - Args: - in_channels (int): Number of input channels of the first layer. - out_channels (int): Number of output channels of the last layer. - modes (Tuple[int, ...]): Number of modes to use for contraction in Laplace domain during training. - T (paddle.Tensor): Linspace of time dimension. - data (Tuple[paddle.Tensor, ...]): Linspaces of other dimensions. - """ - - def __init__( - self, - in_channels: int, - out_channels: int, - modes: Tuple[int, ...], - T: paddle.Tensor, - data: Tuple[paddle.Tensor, ...], - ): - super().__init__() - self.char1 = "pqr" - self.char2 = "mnk" - self.modes = modes - self.scale = 1 / (in_channels * out_channels) - self.dims = len(modes) - - self.weights_pole_real = nn.ParameterList() - self.weights_pole_imag = nn.ParameterList() - for i in range(self.dims): - weight_real = self._init_weights( - self.create_parameter((in_channels, out_channels, modes[i], 1)) - ) - weight_imag = self._init_weights( - self.create_parameter((in_channels, out_channels, modes[i], 1)) - ) - self.weights_pole_real.append(weight_real) - self.weights_pole_imag.append(weight_imag) - - residues_shape = (in_channels, out_channels) + modes + (1,) - self.weights_residue_real = self._init_weights( - self.create_parameter(residues_shape) - ) - self.weights_residue_imag = self._init_weights( - self.create_parameter(residues_shape) - ) - - self.initialize_lambdas(T, data) - self.get_einsum_eqs() - - def _init_weights(self, weight) -> paddle.Tensor: - return initializer.uniform_(weight, a=0, b=self.scale) - - def initialize_lambdas(self, T, data) -> None: - self.t_lst = (T,) + data - self.lambdas = [] - for i in range(self.dims): - t_i = self.t_lst[i] - self.register_buffer(f"t_{i}", t_i) - dt = (t_i[0, 1] - t_i[0, 0]).item() - omega = paddle.fft.fftfreq(n=tuple(t_i.shape)[1], d=dt) * 2 * np.pi * 1.0j - lambda_ = omega.reshape([*omega.shape, 1, 1, 1]) - self.register_buffer(f"lambda_{i}", lambda_) - self.lambdas.append(lambda_) - - def get_einsum_eqs(self) -> None: - terms_eq = [] - terms_x2_eq = [] - for i in range(self.dims): - term_eq = self.char1[i] + "io" + self.char2[i] - terms_eq.append(term_eq) - term_x2_eq = "io" + self.char2[i] + self.char1[i] - terms_x2_eq.append(term_x2_eq) - self.eq1 = ( - "bi" - + "".join(self.char1) - + "," - + "io" - + "".join(self.char2) - + "," - + ",".join(terms_eq) - + "->" - + "bo" - + "".join(self.char1) - ) - self.eq2 = ( - "bi" - + "".join(self.char1) - + "," - + "io" - + "".join(self.char2) - + "," - + ",".join(terms_eq) - + "->" - + "bo" - + "".join(self.char2) - ) - self.eq_x2 = ( - "bi" - + "".join(self.char2) - + "," - + ",".join(terms_x2_eq) - + "->bo" - + "".join(self.char1) - ) - - def output_PR(self, alpha) -> Tuple[paddle.Tensor, paddle.Tensor]: - weights_residue = paddle.as_complex( - paddle.concat( - [self.weights_residue_real, self.weights_residue_imag], axis=-1 - ) - ) - self.weights_pole = [] - terms = [] - for i in range(self.dims): - weights_pole = paddle.as_complex( - paddle.concat( - [self.weights_pole_real[i], self.weights_pole_imag[i]], axis=-1 - ) - ) - self.weights_pole.append(weights_pole) - sub = paddle.subtract(self.lambdas[i], weights_pole) - terms.append(paddle.divide(paddle.to_tensor(1, dtype=sub.dtype), sub)) - - output_residue1 = paddle.einsum(self.eq1, alpha, weights_residue, *terms) - output_residue2 = (-1) ** self.dims * paddle.einsum( - self.eq2, alpha, weights_residue, *terms - ) - return output_residue1, output_residue2 - - def forward(self, x): - alpha = paddle.fft.fftn(x=x, axes=[-3, -2, -1]) - output_residue1, output_residue2 = self.output_PR(alpha) - - x1 = paddle.fft.ifftn( - x=output_residue1, s=(x.shape[-3], x.shape[-2], x.shape[-1]) - ) - x1 = paddle.real(x=x1) - - exp_terms = [] - for i in range(self.dims): - term = paddle.einsum( - "io" - + self.char2[i] - + ",d" - + self.char1[i] - + "->io" - + self.char2[i] - + self.char1[i], - self.weights_pole[i], - self.t_lst[i].astype(paddle.complex64).reshape([1, -1]), - ) - exp_terms.append(paddle.exp(term)) - - x2 = paddle.einsum(self.eq_x2, output_residue2, *exp_terms) - x2 = paddle.real(x2) - x2 = x2 / reduce(operator.mul, x.shape[-3:], 1) - return x1 + x2 - - -class LNO(base.Arch): - """Laplace Neural Operator net. - - Args: - input_keys (Tuple[str, ...]): Name of input keys, such as ("input1", "input2"). - output_keys (Tuple[str, ...]): Name of output keys, such as ("output1", "output2"). - width (int): Tensor width of Laplace Layer. - modes (Tuple[int, ...]): Number of modes to use for contraction in Laplace domain during training. - T (paddle.Tensor): Linspace of time dimension. - data (Tuple[paddle.Tensor, ...]): Linspaces of other dimensions. - in_features (int, optional): Number of input channels of the first layer.. Defaults to 1. - hidden_features (int, optional): Number of channels of the fully-connected layer. Defaults to 64. - activation (str, optional): The activation function. Defaults to "sin". - use_norm (bool, optional): Whether to use normalization layers. Defaults to True. - use_grid (bool, optional): Whether to create grid. Defaults to False. - """ - - def __init__( - self, - input_keys: Tuple[str, ...], - output_keys: Tuple[str, ...], - width: int, - modes: Tuple[int, ...], - T: paddle.Tensor, - data: Optional[Tuple[paddle.Tensor, ...]] = None, - in_features: int = 1, - hidden_features: int = 64, - activation: str = "sin", - use_norm: bool = True, - use_grid: bool = False, - ): - super().__init__() - self.input_keys = input_keys - self.output_keys = output_keys - self.width = width - self.modes = modes - self.dims = len(modes) - assert self.dims <= 3, "Only 3 dims and lower of modes are supported now." - - if data is None: - data = () - assert ( - self.dims == len(data) + 1 - ), f"Dims of modes is {self.dims} but only {len(data)} dims(except T) of data received." - - self.fc0 = nn.Linear(in_features=in_features, out_features=self.width) - self.laplace = Laplace(self.width, self.width, self.modes, T, data) - self.conv = getattr(nn, f"Conv{self.dims}D")( - in_channels=self.width, - out_channels=self.width, - kernel_size=1, - data_format="NCDHW", - ) - if use_norm: - self.norm = getattr(nn, f"InstanceNorm{self.dims}D")( - num_features=self.width, - weight_attr=False, - bias_attr=False, - ) - self.fc1 = nn.Linear(in_features=self.width, out_features=hidden_features) - self.fc2 = nn.Linear(in_features=hidden_features, out_features=1) - self.act = act_mod.get_activation(activation) - - self.use_norm = use_norm - self.use_grid = use_grid - - def get_grid(self, shape): - batchsize, size_t, size_x, size_y = shape[0], shape[1], shape[2], shape[3] - gridt = paddle.linspace(0, 1, size_t) - gridt = gridt.reshape([1, size_t, 1, 1, 1]).tile( - [batchsize, 1, size_x, size_y, 1] - ) - gridx = paddle.linspace(0, 1, size_x) - gridx = gridx.reshape([1, 1, size_x, 1, 1]).tile( - [batchsize, size_t, 1, size_y, 1] - ) - gridy = paddle.linspace(0, 1, size_y) - gridy = gridy.reshape([1, 1, 1, size_y, 1]).tile( - [batchsize, size_t, size_x, 1, 1] - ) - return paddle.concat([gridt, gridx, gridy], axis=-1) - - def transpoe_to_NCDHW(self, x): - perm = [0, self.dims + 1] + list(range(1, self.dims + 1)) - return paddle.transpose(x, perm=perm) - - def transpoe_to_NDHWC(self, x): - perm = [0] + list(range(2, self.dims + 2)) + [1] - return paddle.transpose(x, perm=perm) - - def forward_tensor(self, x): - if self.use_grid: - grid = self.get_grid(x.shape) - x = paddle.concat([x, grid], axis=-1) - x = self.fc0(x) - x = self.transpoe_to_NCDHW(x) - - if self.use_norm: - x1 = self.norm(self.laplace(self.norm(x))) - else: - x1 = self.laplace(x) - - x2 = self.conv(x) - x = x1 + x2 - - x = self.transpoe_to_NDHWC(x) - - x = self.fc1(x) - x = self.act(x) - x = self.fc2(x) - return x - - def forward(self, x): - if self._input_transform is not None: - x = self._input_transform(x) - - y = self.concat_to_tensor(x, self.input_keys, axis=-1) - y = self.forward_tensor(y) - y = self.split_to_dict(y, self.output_keys, axis=-1) - - if self._output_transform is not None: - y = self._output_transform(x, y) - return y diff --git a/examples/smc_reac/ppsci/arch/mlp.py b/examples/smc_reac/ppsci/arch/mlp.py deleted file mode 100644 index ef2d6e9e44..0000000000 --- a/examples/smc_reac/ppsci/arch/mlp.py +++ /dev/null @@ -1,828 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import Dict -from typing import Optional -from typing import Tuple -from typing import Union - -import numpy as np -import paddle -import paddle.nn as nn - -from ppsci.arch import activation as act_mod -from ppsci.arch import base -from ppsci.utils import initializer - - -class WeightNormLinear(nn.Layer): - def __init__(self, in_features: int, out_features: int, bias: bool = True) -> None: - super().__init__() - self.in_features = in_features - self.out_features = out_features - self.weight_v = self.create_parameter((in_features, out_features)) - self.weight_g = self.create_parameter((out_features,)) - if bias: - self.bias = self.create_parameter((out_features,)) - else: - self.bias = None - self._init_weights() - - def _init_weights(self) -> None: - initializer.xavier_uniform_(self.weight_v) - initializer.constant_(self.weight_g, 1.0) - if self.bias is not None: - initializer.constant_(self.bias, 0.0) - - def forward(self, input): - norm = self.weight_v.norm(p=2, axis=0, keepdim=True) - weight = self.weight_g * self.weight_v / norm - return nn.functional.linear(input, weight, self.bias) - - -class RandomWeightFactorization(nn.Layer): - def __init__( - self, - in_features: int, - out_features: int, - bias: bool = True, - mean: float = 0.5, - std: float = 0.1, - ): - super().__init__() - self.in_features = in_features - self.out_features = out_features - self.weight_v = self.create_parameter((in_features, out_features)) - self.weight_g = self.create_parameter((out_features,)) - if bias: - self.bias = self.create_parameter((out_features,)) - else: - self.bias = None - - self._init_weights(mean, std) - - def _init_weights(self, mean, std): - with paddle.no_grad(): - initializer.glorot_normal_(self.weight_v) - - nn.initializer.Normal(mean, std)(self.weight_g) - paddle.assign(paddle.exp(self.weight_g), self.weight_g) - paddle.assign(self.weight_v / self.weight_g, self.weight_v) - if self.bias is not None: - initializer.constant_(self.bias, 0.0) - - self.weight_g.stop_gradient = False - self.weight_v.stop_gradient = False - self.bias.stop_gradient = False - - def forward(self, input): - return nn.functional.linear(input, self.weight_g * self.weight_v, self.bias) - - -class PeriodEmbedding(nn.Layer): - def __init__(self, periods: Dict[str, Tuple[float, bool]]): - super().__init__() - self.freqs_dict = { - k: self.create_parameter( - [], - attr=paddle.ParamAttr(trainable=trainable), - default_initializer=nn.initializer.Constant(2 * np.pi / float(p)), - ) # mu = 2*pi / period for sin/cos function - for k, (p, trainable) in periods.items() - } - self.freqs = nn.ParameterList(list(self.freqs_dict.values())) - - def forward(self, x: Dict[str, paddle.Tensor]): - y = {k: v for k, v in x.items()} # shallow copy to avoid modifying input dict - - for k, w in self.freqs_dict.items(): - y[k] = paddle.concat([paddle.cos(w * x[k]), paddle.sin(w * x[k])], axis=-1) - - return y - - -class FourierEmbedding(nn.Layer): - def __init__(self, in_features, out_features, scale): - super().__init__() - if out_features % 2 != 0: - raise ValueError(f"out_features must be even, but got {out_features}.") - - self.kernel = self.create_parameter( - [in_features, out_features // 2], - default_initializer=nn.initializer.Normal(std=scale), - ) - - def forward(self, x: paddle.Tensor): - y = paddle.concat( - [ - paddle.cos(x @ self.kernel), - paddle.sin(x @ self.kernel), - ], - axis=-1, - ) - return y - - -class MLP(base.Arch): - """Multi layer perceptron network. - - Args: - input_keys (Tuple[str, ...]): Name of input keys, such as ("x", "y", "z"). - output_keys (Tuple[str, ...]): Name of output keys, such as ("u", "v", "w"). - num_layers (int): Number of hidden layers. - hidden_size (Union[int, Tuple[int, ...]]): Number of hidden size. - An integer for all layers, or list of integer specify each layer's size. - activation (str, optional): Name of activation function. Defaults to "tanh". - skip_connection (bool, optional): Whether to use skip connection. Defaults to False. - weight_norm (bool, optional): Whether to apply weight norm on parameter(s). Defaults to False. - input_dim (Optional[int]): Number of input's dimension. Defaults to None. - output_dim (Optional[int]): Number of output's dimension. Defaults to None. - periods (Optional[Dict[int, Tuple[float, bool]]]): Period of each input key, - input in given channel will be period embedded if specified, each tuple of - periods list is [period, trainable]. Defaults to None. - fourier (Optional[Dict[str, Union[float, int]]]): Random fourier feature embedding, - e.g. {'dim': 256, 'scale': 1.0}. Defaults to None. - random_weight (Optional[Dict[str, float]]): Mean and std of random weight - factorization layer, e.g. {"mean": 0.5, "std: 0.1"}. Defaults to None. - - Examples: - >>> import paddle - >>> import ppsci - >>> model = ppsci.arch.MLP( - ... input_keys=("x", "y"), - ... output_keys=("u", "v"), - ... num_layers=5, - ... hidden_size=128 - ... ) - >>> input_dict = {"x": paddle.rand([64, 1]), - ... "y": paddle.rand([64, 1])} - >>> output_dict = model(input_dict) - >>> print(output_dict["u"].shape) - [64, 1] - >>> print(output_dict["v"].shape) - [64, 1] - """ - - def __init__( - self, - input_keys: Tuple[str, ...], - output_keys: Tuple[str, ...], - num_layers: int, - hidden_size: Union[int, Tuple[int, ...]], - activation: str = "tanh", - skip_connection: bool = False, - weight_norm: bool = False, - input_dim: Optional[int] = None, - output_dim: Optional[int] = None, - periods: Optional[Dict[int, Tuple[float, bool]]] = None, - fourier: Optional[Dict[str, Union[float, int]]] = None, - random_weight: Optional[Dict[str, float]] = None, - ): - super().__init__() - self.input_keys = input_keys - self.output_keys = output_keys - self.linears = [] - self.acts = [] - self.periods = periods - self.fourier = fourier - if periods: - self.period_emb = PeriodEmbedding(periods) - - if isinstance(hidden_size, (tuple, list)): - if num_layers is not None: - raise ValueError( - "num_layers should be None when hidden_size is specified" - ) - elif isinstance(hidden_size, int): - if not isinstance(num_layers, int): - raise ValueError( - "num_layers should be an int when hidden_size is an int" - ) - hidden_size = [hidden_size] * num_layers - else: - raise ValueError( - f"hidden_size should be list of int or int, but got {type(hidden_size)}" - ) - - # initialize FC layer(s) - cur_size = len(self.input_keys) if input_dim is None else input_dim - if input_dim is None and periods: - # period embedded channel(s) will be doubled automatically - # if input_dim is not specified - cur_size += len(periods) - - if fourier: - self.fourier_emb = FourierEmbedding( - cur_size, fourier["dim"], fourier["scale"] - ) - cur_size = fourier["dim"] - - for i, _size in enumerate(hidden_size): - if weight_norm: - self.linears.append(WeightNormLinear(cur_size, _size)) - elif random_weight: - self.linears.append( - RandomWeightFactorization( - cur_size, - _size, - mean=random_weight["mean"], - std=random_weight["std"], - ) - ) - else: - self.linears.append(nn.Linear(cur_size, _size)) - - # initialize activation function - self.acts.append( - act_mod.get_activation(activation) - if activation != "stan" - else act_mod.get_activation(activation)(_size) - ) - # special initialization for certain activation - # TODO: Adapt code below to a more elegant style - if activation == "siren": - if i == 0: - act_mod.Siren.init_for_first_layer(self.linears[-1]) - else: - act_mod.Siren.init_for_hidden_layer(self.linears[-1]) - - cur_size = _size - - self.linears = nn.LayerList(self.linears) - self.acts = nn.LayerList(self.acts) - if random_weight: - self.last_fc = RandomWeightFactorization( - cur_size, - len(self.output_keys) if output_dim is None else output_dim, - mean=random_weight["mean"], - std=random_weight["std"], - ) - else: - self.last_fc = nn.Linear( - cur_size, - len(self.output_keys) if output_dim is None else output_dim, - ) - - self.skip_connection = skip_connection - - def forward_tensor(self, x): - y = x - skip = None - for i, linear in enumerate(self.linears): - y = linear(y) - if self.skip_connection and i % 2 == 0: - if skip is not None: - skip = y - y = y + skip - else: - skip = y - y = self.acts[i](y) - - y = self.last_fc(y) - - return y - - def forward(self, x): - if self._input_transform is not None: - x = self._input_transform(x) - - if self.periods: - x = self.period_emb(x) - - y = self.concat_to_tensor(x, self.input_keys, axis=-1) - - if self.fourier: - y = self.fourier_emb(y) - - y = self.forward_tensor(y) - y = self.split_to_dict(y, self.output_keys, axis=-1) - - if self._output_transform is not None: - y = self._output_transform(x, y) - return y - - -class ModifiedMLP(base.Arch): - """Modified Multi layer perceptron network. - - Understanding and mitigating gradient pathologies in physics-informed - neural networks. https://arxiv.org/pdf/2001.04536.pdf. - - Args: - input_keys (Tuple[str, ...]): Name of input keys, such as ("x", "y", "z"). - output_keys (Tuple[str, ...]): Name of output keys, such as ("u", "v", "w"). - num_layers (int): Number of hidden layers. - hidden_size (int): Number of hidden size, an integer for all layers. - activation (str, optional): Name of activation function. Defaults to "tanh". - skip_connection (bool, optional): Whether to use skip connection. Defaults to False. - weight_norm (bool, optional): Whether to apply weight norm on parameter(s). Defaults to False. - input_dim (Optional[int]): Number of input's dimension. Defaults to None. - output_dim (Optional[int]): Number of output's dimension. Defaults to None. - - Examples: - >>> import paddle - >>> import ppsci - >>> model = ppsci.arch.ModifiedMLP( - ... input_keys=("x", "y"), - ... output_keys=("u", "v"), - ... num_layers=5, - ... hidden_size=128 - ... ) - >>> input_dict = {"x": paddle.rand([64, 1]), - ... "y": paddle.rand([64, 1])} - >>> output_dict = model(input_dict) - >>> print(output_dict["u"].shape) - [64, 1] - >>> print(output_dict["v"].shape) - [64, 1] - """ - - def __init__( - self, - input_keys: Tuple[str, ...], - output_keys: Tuple[str, ...], - num_layers: int, - hidden_size: int, - activation: str = "tanh", - skip_connection: bool = False, - weight_norm: bool = False, - input_dim: Optional[int] = None, - output_dim: Optional[int] = None, - periods: Optional[Dict[int, Tuple[float, bool]]] = None, - fourier: Optional[Dict[str, Union[float, int]]] = None, - random_weight: Optional[Dict[str, float]] = None, - ): - super().__init__() - self.input_keys = input_keys - self.output_keys = output_keys - self.linears = [] - self.acts = [] - self.periods = periods - self.fourier = fourier - if periods: - self.period_emb = PeriodEmbedding(periods) - if isinstance(hidden_size, int): - if not isinstance(num_layers, int): - raise ValueError("num_layers should be an int") - hidden_size = [hidden_size] * num_layers - else: - raise ValueError(f"hidden_size should be int, but got {type(hidden_size)}") - - # initialize FC layer(s) - cur_size = len(self.input_keys) if input_dim is None else input_dim - if input_dim is None and periods: - # period embedded channel(s) will be doubled automatically - # if input_dim is not specified - cur_size += len(periods) - - if fourier: - self.fourier_emb = FourierEmbedding( - cur_size, fourier["dim"], fourier["scale"] - ) - cur_size = fourier["dim"] - - self.embed_u = nn.Sequential( - ( - WeightNormLinear(cur_size, hidden_size[0]) - if weight_norm - else ( - nn.Linear(cur_size, hidden_size[0]) - if random_weight is None - else RandomWeightFactorization( - cur_size, - hidden_size[0], - mean=random_weight["mean"], - std=random_weight["std"], - ) - ) - ), - ( - act_mod.get_activation(activation) - if activation != "stan" - else act_mod.get_activation(activation)(hidden_size[0]) - ), - ) - self.embed_v = nn.Sequential( - ( - WeightNormLinear(cur_size, hidden_size[0]) - if weight_norm - else ( - nn.Linear(cur_size, hidden_size[0]) - if random_weight is None - else RandomWeightFactorization( - cur_size, - hidden_size[0], - mean=random_weight["mean"], - std=random_weight["std"], - ) - ) - ), - ( - act_mod.get_activation(activation) - if activation != "stan" - else act_mod.get_activation(activation)(hidden_size[0]) - ), - ) - - for i, _size in enumerate(hidden_size): - if weight_norm: - self.linears.append(WeightNormLinear(cur_size, _size)) - elif random_weight: - self.linears.append( - RandomWeightFactorization( - cur_size, - _size, - mean=random_weight["mean"], - std=random_weight["std"], - ) - ) - else: - self.linears.append(nn.Linear(cur_size, _size)) - - # initialize activation function - self.acts.append( - act_mod.get_activation(activation) - if activation != "stan" - else act_mod.get_activation(activation)(_size) - ) - # special initialization for certain activation - # TODO: Adapt code below to a more elegant style - if activation == "siren": - if i == 0: - act_mod.Siren.init_for_first_layer(self.linears[-1]) - else: - act_mod.Siren.init_for_hidden_layer(self.linears[-1]) - - cur_size = _size - - self.linears = nn.LayerList(self.linears) - self.acts = nn.LayerList(self.acts) - if random_weight: - self.last_fc = RandomWeightFactorization( - cur_size, - len(self.output_keys) if output_dim is None else output_dim, - mean=random_weight["mean"], - std=random_weight["std"], - ) - else: - self.last_fc = nn.Linear( - cur_size, - len(self.output_keys) if output_dim is None else output_dim, - ) - - self.skip_connection = skip_connection - - def forward_tensor(self, x): - u = self.embed_u(x) - v = self.embed_v(x) - - y = x - skip = None - for i, linear in enumerate(self.linears): - y = linear(y) - y = self.acts[i](y) - y = y * u + (1 - y) * v - if self.skip_connection and i % 2 == 0: - if skip is not None: - skip = y - y = y + skip - else: - skip = y - - y = self.last_fc(y) - - return y - - def forward(self, x): - x_identity = x - if self._input_transform is not None: - x = self._input_transform(x) - - if self.periods: - x = self.period_emb(x) - - y = self.concat_to_tensor(x, self.input_keys, axis=-1) - - if self.fourier: - y = self.fourier_emb(y) - - y = self.forward_tensor(y) - y = self.split_to_dict(y, self.output_keys, axis=-1) - - if self._output_transform is not None: - y = self._output_transform(x_identity, y) - return y - - -class PirateNetBlock(nn.Layer): - r"""Basic block of PirateNet. - - $$ - \begin{align*} - \Phi(\mathbf{x})=\left[\begin{array}{l} - \cos (\mathbf{B} \mathbf{x}) \\ - \sin (\mathbf{B} \mathbf{x}) - \end{array}\right] \\ - \mathbf{f}^{(l)} & =\sigma\left(\mathbf{W}_1^{(l)} \mathbf{x}^{(l)}+\mathbf{b}_1^{(l)}\right) \\ - \mathbf{z}_1^{(l)} & =\mathbf{f}^{(l)} \odot \mathbf{U}+\left(1-\mathbf{f}^{(l)}\right) \odot \mathbf{V} \\ - \mathbf{g}^{(l)} & =\sigma\left(\mathbf{W}_2^{(l)} \mathbf{z}_1^{(l)}+\mathbf{b}_2^{(l)}\right) \\ - \mathbf{z}_2^{(l)} & =\mathbf{g}^{(l)} \odot \mathbf{U}+\left(1-\mathbf{g}^{(l)}\right) \odot \mathbf{V} \\ - \mathbf{h}^{(l)} & =\sigma\left(\mathbf{W}_3^{(l)} \mathbf{z}_2^{(l)}+\mathbf{b}_3^{(l)}\right) \\ - \mathbf{x}^{(l+1)} & =\alpha^{(l)} \cdot \mathbf{h}^{(l)}+\left(1-\alpha^{(l)}\right) \cdot \mathbf{x}^{(l)} - \end{align*} - $$ - - Args: - input_dim (int): Input dimension. - embed_dim (int): Embedding dimension. - activation (str, optional): Name of activation function. Defaults to "tanh". - random_weight (Optional[Dict[str, float]]): Mean and std of random weight - factorization layer, e.g. {"mean": 0.5, "std: 0.1"}. Defaults to None. - """ - - def __init__( - self, - input_dim: int, - embed_dim: int, - activation: str = "tanh", - random_weight: Optional[Dict[str, float]] = None, - ): - super().__init__() - self.linear1 = ( - nn.Linear(input_dim, embed_dim) - if random_weight is None - else RandomWeightFactorization( - input_dim, - embed_dim, - mean=random_weight["mean"], - std=random_weight["std"], - ) - ) - self.linear2 = ( - nn.Linear(embed_dim, embed_dim) - if random_weight is None - else RandomWeightFactorization( - embed_dim, - embed_dim, - mean=random_weight["mean"], - std=random_weight["std"], - ) - ) - self.linear3 = ( - nn.Linear(embed_dim, embed_dim) - if random_weight is None - else RandomWeightFactorization( - embed_dim, - embed_dim, - mean=random_weight["mean"], - std=random_weight["std"], - ) - ) - self.alpha = self.create_parameter( - [ - 1, - ], - default_initializer=nn.initializer.Constant(0), - ) - self.act1 = ( - act_mod.get_activation(activation) - if activation != "stan" - else act_mod.get_activation(activation)(embed_dim) - ) - self.act2 = ( - act_mod.get_activation(activation) - if activation != "stan" - else act_mod.get_activation(activation)(embed_dim) - ) - self.act3 = ( - act_mod.get_activation(activation) - if activation != "stan" - else act_mod.get_activation(activation)(embed_dim) - ) - - def forward(self, x, u, v): - f = self.act1(self.linear1(x)) - z1 = f * u + (1 - f) * v - g = self.act2(self.linear2(z1)) - z2 = g * u + (1 - g) * v - h = self.act3(self.linear3(z2)) - out = self.alpha * h + (1 - self.alpha) * x - return out - - -class PirateNet(base.Arch): - r"""PirateNet. - - [PIRATENETS: PHYSICS-INFORMED DEEP LEARNING WITHRESIDUAL ADAPTIVE NETWORKS](https://arxiv.org/pdf/2402.00326.pdf) - - $$ - \begin{align*} - \Phi(\mathbf{x}) &= \left[\begin{array}{l} - \cos (\mathbf{B} \mathbf{x}) \\ - \sin (\mathbf{B} \mathbf{x}) - \end{array}\right] \\ - \mathbf{f}^{(l)} &= \sigma\left(\mathbf{W}_1^{(l)} \mathbf{x}^{(l)}+\mathbf{b}_1^{(l)}\right) \\ - \mathbf{z}_1^{(l)} &= \mathbf{f}^{(l)} \odot \mathbf{U}+\left(1-\mathbf{f}^{(l)}\right) \odot \mathbf{V} \\ - \mathbf{g}^{(l)} &= \sigma\left(\mathbf{W}_2^{(l)} \mathbf{z}_1^{(l)}+\mathbf{b}_2^{(l)}\right) \\ - \mathbf{z}_2^{(l)} &= \mathbf{g}^{(l)} \odot \mathbf{U}+\left(1-\mathbf{g}^{(l)}\right) \odot \mathbf{V} \\ - \mathbf{h}^{(l)} &= \sigma\left(\mathbf{W}_3^{(l)} \mathbf{z}_2^{(l)}+\mathbf{b}_3^{(l)}\right) \\ - \mathbf{x}^{(l+1)} &= \text{PirateBlock}^{(l)}\left(\mathbf{x}^{(l)}\right), l=1...L-1\\ - \mathbf{u}_\theta &= \mathbf{W}^{(L+1)} \mathbf{x}^{(L)} - \end{align*} - $$ - - Args: - input_keys (Tuple[str, ...]): Name of input keys, such as ("x", "y", "z"). - output_keys (Tuple[str, ...]): Name of output keys, such as ("u", "v", "w"). - num_blocks (int): Number of PirateBlocks. - hidden_size (Union[int, Tuple[int, ...]]): Number of hidden size. - An integer for all layers, or list of integer specify each layer's size. - activation (str, optional): Name of activation function. Defaults to "tanh". - weight_norm (bool, optional): Whether to apply weight norm on parameter(s). Defaults to False. - input_dim (Optional[int]): Number of input's dimension. Defaults to None. - output_dim (Optional[int]): Number of output's dimension. Defaults to None. - periods (Optional[Dict[int, Tuple[float, bool]]]): Period of each input key, - input in given channel will be period embedded if specified, each tuple of - periods list is [period, trainable]. Defaults to None. - fourier (Optional[Dict[str, Union[float, int]]]): Random fourier feature embedding, - e.g. {'dim': 256, 'scale': 1.0}. Defaults to None. - random_weight (Optional[Dict[str, float]]): Mean and std of random weight - factorization layer, e.g. {"mean": 0.5, "std: 0.1"}. Defaults to None. - - Examples: - >>> import paddle - >>> import ppsci - >>> model = ppsci.arch.PirateNet( - ... input_keys=("x", "y"), - ... output_keys=("u", "v"), - ... num_blocks=3, - ... hidden_size=256, - ... fourier={'dim': 256, 'scale': 1.0}, - ... ) - >>> input_dict = {"x": paddle.rand([64, 1]), - ... "y": paddle.rand([64, 1])} - >>> output_dict = model(input_dict) - >>> print(output_dict["u"].shape) - [64, 1] - >>> print(output_dict["v"].shape) - [64, 1] - """ - - def __init__( - self, - input_keys: Tuple[str, ...], - output_keys: Tuple[str, ...], - num_blocks: int, - hidden_size: int, - activation: str = "tanh", - weight_norm: bool = False, - input_dim: Optional[int] = None, - output_dim: Optional[int] = None, - periods: Optional[Dict[int, Tuple[float, bool]]] = None, - fourier: Optional[Dict[str, Union[float, int]]] = None, - random_weight: Optional[Dict[str, float]] = None, - ): - super().__init__() - self.input_keys = input_keys - self.output_keys = output_keys - self.blocks = [] - self.periods = periods - self.fourier = fourier - if periods: - self.period_emb = PeriodEmbedding(periods) - - if isinstance(hidden_size, int): - if not isinstance(num_blocks, int): - raise ValueError("num_blocks should be an int") - hidden_size = [hidden_size] * num_blocks - else: - raise ValueError(f"hidden_size should be int, but got {type(hidden_size)}") - - # initialize FC layer(s) - cur_size = len(self.input_keys) if input_dim is None else input_dim - if input_dim is None and periods: - # period embedded channel(s) will be doubled automatically - # if input_dim is not specified - cur_size += len(periods) - - if fourier: - self.fourier_emb = FourierEmbedding( - cur_size, fourier["dim"], fourier["scale"] - ) - cur_size = fourier["dim"] - else: - self.linear_emb = nn.Linear(cur_size, hidden_size[0]) - cur_size = hidden_size[0] - - self.embed_u = nn.Sequential( - ( - WeightNormLinear(cur_size, hidden_size[0]) - if weight_norm - else ( - nn.Linear(cur_size, hidden_size[0]) - if random_weight is None - else RandomWeightFactorization( - cur_size, - hidden_size[0], - mean=random_weight["mean"], - std=random_weight["std"], - ) - ) - ), - ( - act_mod.get_activation(activation) - if activation != "stan" - else act_mod.get_activation(activation)(hidden_size[0]) - ), - ) - self.embed_v = nn.Sequential( - ( - WeightNormLinear(cur_size, hidden_size[0]) - if weight_norm - else ( - nn.Linear(cur_size, hidden_size[0]) - if random_weight is None - else RandomWeightFactorization( - cur_size, - hidden_size[0], - mean=random_weight["mean"], - std=random_weight["std"], - ) - ) - ), - ( - act_mod.get_activation(activation) - if activation != "stan" - else act_mod.get_activation(activation)(hidden_size[0]) - ), - ) - - for i, _size in enumerate(hidden_size): - self.blocks.append( - PirateNetBlock( - cur_size, - _size, - activation=activation, - random_weight=random_weight, - ) - ) - cur_size = _size - - self.blocks = nn.LayerList(self.blocks) - if random_weight: - self.last_fc = RandomWeightFactorization( - cur_size, - len(self.output_keys) if output_dim is None else output_dim, - mean=random_weight["mean"], - std=random_weight["std"], - ) - else: - self.last_fc = nn.Linear( - cur_size, - len(self.output_keys) if output_dim is None else output_dim, - ) - - def forward_tensor(self, x): - u = self.embed_u(x) - v = self.embed_v(x) - - y = x - for i, block in enumerate(self.blocks): - y = block(y, u, v) - - y = self.last_fc(y) - return y - - def forward(self, x): - if self._input_transform is not None: - x = self._input_transform(x) - - if self.periods: - x = self.period_emb(x) - - y = self.concat_to_tensor(x, self.input_keys, axis=-1) - - if self.fourier: - y = self.fourier_emb(y) - else: - y = self.linear_emb(y) - - y = self.forward_tensor(y) - y = self.split_to_dict(y, self.output_keys, axis=-1) - - if self._output_transform is not None: - y = self._output_transform(x, y) - return y diff --git a/examples/smc_reac/ppsci/arch/model_list.py b/examples/smc_reac/ppsci/arch/model_list.py deleted file mode 100644 index f5f7feeb8b..0000000000 --- a/examples/smc_reac/ppsci/arch/model_list.py +++ /dev/null @@ -1,72 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import Tuple - -from paddle import nn - -from ppsci.arch import base - - -class ModelList(base.Arch): - """ModelList layer which wrap more than one model that shares inputs. - - Args: - model_list (Tuple[base.Arch, ...]): Model(s) nested in tuple. - - Examples: - >>> import paddle - >>> import ppsci - >>> model1 = ppsci.arch.MLP(("x", "y"), ("u", "v"), 10, 128) - >>> model2 = ppsci.arch.MLP(("x", "y"), ("w", "p"), 5, 128) - >>> model = ppsci.arch.ModelList((model1, model2)) - >>> input_dict = {"x": paddle.rand([64, 64, 1]),"y": paddle.rand([64, 64, 1])} - >>> output_dict = model(input_dict) - >>> for k, v in output_dict.items(): - ... print(k, v.shape) - u [64, 64, 1] - v [64, 64, 1] - w [64, 64, 1] - p [64, 64, 1] - """ - - def __init__( - self, - model_list: Tuple[base.Arch, ...], - ): - super().__init__() - self.input_keys = sum([model.input_keys for model in model_list], ()) - self.input_keys = set(self.input_keys) - - output_keys_set = set() - for model in model_list: - if len(output_keys_set & set(model.output_keys)): - raise ValueError( - "output_keys of model from model_list should be unique," - f"but got duplicate keys: {output_keys_set & set(model.output_keys)}" - ) - output_keys_set = output_keys_set | set(model.output_keys) - self.output_keys = tuple(output_keys_set) - - self.model_list = nn.LayerList(model_list) - - def forward(self, x): - y_all = {} - for model in self.model_list: - y = model(x) - y_all.update(y) - - return y_all diff --git a/examples/smc_reac/ppsci/arch/moflow_basic.py b/examples/smc_reac/ppsci/arch/moflow_basic.py deleted file mode 100644 index 68f10efafb..0000000000 --- a/examples/smc_reac/ppsci/arch/moflow_basic.py +++ /dev/null @@ -1,297 +0,0 @@ -# Copyright (c) 2024 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. - -# Copyright 2020 Chengxi Zang - -import numpy as np -import paddle -import paddle.nn as nn -from scipy import linalg as la - - -# logabs = lambda x: paddle.log(x=paddle.abs(x=x)) -def logabs(x): - return paddle.log(paddle.abs(x)) - - -class ActNorm(nn.Layer): - def __init__(self, in_channel, logdet=True): - super().__init__() - self.loc = self.create_parameter( - [1, in_channel, 1, 1], - default_initializer=nn.initializer.Constant(value=0.0), - ) - - self.scale = self.create_parameter( - [1, in_channel, 1, 1], - default_initializer=nn.initializer.Constant(value=1.0), - ) - - self.register_buffer( - name="initialized", tensor=paddle.to_tensor(data=0, dtype="uint8") - ) - self.logdet = logdet - - def initialize(self, input): - with paddle.no_grad(): - flatten = input.transpose(perm=[1, 0, 2, 3]).reshape( - [tuple(input.shape)[1], -1] - ) - mean = ( - flatten.mean(axis=1) - .unsqueeze(axis=1) - .unsqueeze(axis=2) - .unsqueeze(axis=3) - .transpose(perm=[1, 0, 2, 3]) - ) - std = ( - flatten.std(axis=1) - .unsqueeze(axis=1) - .unsqueeze(axis=2) - .unsqueeze(axis=3) - .transpose(perm=[1, 0, 2, 3]) - ) - paddle.assign(-mean, output=self.loc.data) - paddle.assign(1 / (std + 1e-06), output=self.scale.data) - - def forward(self, input): - _, _, height, width = tuple(input.shape) - if self.initialized.item() == 0: - self.initialize(input) - self.initialized.fill_(value=1) - log_abs = logabs(self.scale) - logdet = height * width * paddle.sum(x=log_abs) - if self.logdet: - return self.scale * (input + self.loc), logdet - else: - return self.scale * (input + self.loc) - - def reverse(self, output): - return output / self.scale - self.loc - - -class ActNorm2D(nn.Layer): - def __init__(self, in_dim, logdet=True): - super().__init__() - self.loc = self.create_parameter( - [1, in_dim, 1], - default_initializer=nn.initializer.Constant(value=0.0), - ) - - self.scale = self.create_parameter( - [1, in_dim, 1], - default_initializer=nn.initializer.Constant(value=1.0), - ) - - self.register_buffer( - name="initialized", tensor=paddle.to_tensor(data=0, dtype="uint8") - ) - self.logdet = logdet - - def initialize(self, input): - with paddle.no_grad(): - flatten = input.transpose(perm=[1, 0, 2]).reshape( - [tuple(input.shape)[1], -1] - ) - mean = ( - flatten.mean(axis=1) - .unsqueeze(axis=1) - .unsqueeze(axis=2) - .transpose(perm=[1, 0, 2]) - ) - std = ( - flatten.std(axis=1) - .unsqueeze(axis=1) - .unsqueeze(axis=2) - .transpose(perm=[1, 0, 2]) - ) - paddle.assign(-mean, output=self.loc.data) - paddle.assign(1 / (std + 1e-06), output=self.scale.data) - - def forward(self, input): - _, _, height = tuple(input.shape) - if self.initialized.item() == 0: - self.initialize(input) - self.initialized.fill_(value=1) - log_abs = logabs(self.scale) - logdet = height * paddle.sum(x=log_abs) - if self.logdet: - return self.scale * (input + self.loc), logdet - else: - return self.scale * (input + self.loc) - - def reverse(self, output): - return output / self.scale - self.loc - - -class InvConv2d(nn.Layer): - def __init__(self, in_channel): - super().__init__() - weight = paddle.randn([in_channel, in_channel]) - q, _ = paddle.linalg.qr(weight) - weight = q.unsqueeze(2).unsqueeze(3) - self.weight = paddle.create_parameter( - weight.shape, - weight.numpy().dtype, - default_initializer=nn.initializer.Assign(weight), - ) - - def forward(self, input): - _, _, height, width = tuple(input.shape) - out = nn.functional.conv2d(x=input, weight=self.weight) - res = paddle.linalg.slogdet(self.weight.squeeze().astype(dtype="float64")) - logdet = height * width * (res[0], res[1])[1].astype(dtype="float32") - return out, logdet - - def reverse(self, output): - return nn.functional.conv2d( - x=output, - weight=self.weight.squeeze().inverse().unsqueeze(axis=2).unsqueeze(axis=3), - ) - - -class InvConv2dLU(nn.Layer): - def __init__(self, in_channel): - super().__init__() - weight = np.random.randn(in_channel, in_channel) - q, _ = la.qr(weight) - w_p, w_l, w_u = la.lu(q.astype(np.float32)) - w_s = np.diag(w_u) - w_u = np.triu(w_u, 1) - u_mask = np.triu(np.ones_like(w_u), 1) - l_mask = u_mask.T - w_p = paddle.to_tensor(data=w_p) - w_l = paddle.to_tensor(data=w_l) - w_s = paddle.to_tensor(data=w_s) - w_u = paddle.to_tensor(data=w_u) - self.register_buffer(name="w_p", tensor=w_p) - self.register_buffer(name="u_mask", tensor=paddle.to_tensor(data=u_mask)) - self.register_buffer(name="l_mask", tensor=paddle.to_tensor(data=l_mask)) - self.register_buffer(name="s_sign", tensor=paddle.sign(x=w_s)) - self.register_buffer( - name="l_eye", tensor=paddle.eye(num_rows=tuple(l_mask.shape)[0]) - ) - self.w_l = paddle.create_parameter( - w_l.shape, - w_l.numpy().dtype, - default_initializer=nn.initializer.Assign(w_l), - ) - - self.w_s = paddle.create_parameter( - logabs(w_s).shape, - logabs(w_s).numpy().dtype, - default_initializer=nn.initializer.Assign(logabs(w_s)), - ) - - self.w_u = paddle.create_parameter( - w_u.shape, - w_u.numpy().dtype, - default_initializer=nn.initializer.Assign(w_u), - ) - - def forward(self, input): - _, _, height, width = tuple(input.shape) - weight = self.calc_weight() - out = nn.functional.conv2d(x=input, weight=weight) - logdet = height * width * paddle.sum(x=self.w_s) - return out, logdet - - def calc_weight(self): - weight = ( - self.w_p - @ (self.w_l * self.l_mask + self.l_eye) - @ ( - self.w_u * self.u_mask - + paddle.diag(x=self.s_sign * paddle.exp(x=self.w_s)) - ) - ) - return weight.unsqueeze(axis=2).unsqueeze(axis=3) - - def reverse(self, output): - weight = self.calc_weight() - return nn.functional.conv2d( - x=output, - weight=weight.squeeze().inverse().unsqueeze(axis=2).unsqueeze(axis=3), - ) - - -class GraphLinear(nn.Layer): - """Graph Linear layer. - This function assumes its input is 3-dimensional. Or 4-dim or whatever, only last dim are changed - Differently from :class:`nn.Linear`, it applies an affine - transformation to the third axis of input `x`. - Warning: original Chainer.link.Link use i.i.d. Gaussian initialization as default, - while default nn.Linear initialization using init.kaiming_uniform_ - """ - - def __init__(self, in_size, out_size, bias=True): - super(GraphLinear, self).__init__() - self.in_size = in_size - self.out_size = out_size - self.linear = nn.Linear( - in_features=in_size, out_features=out_size, bias_attr=bias - ) - - def forward(self, x): - """Forward propagation. - Args: - x (:class:`chainer.Variable`, or :class:`numpy.ndarray` ): - Input array that should be a float array whose ``ndim`` is 3. - - It represents a minibatch of atoms, each of which consists - of a sequence of molecules. Each molecule is represented - by integer IDs. The first axis is an index of atoms - (i.e. minibatch dimension) and the second one an index - of molecules. - - Returns: - class:`chainer.Variable`: - A 3-dimeisional array. - - """ - h = x - h = h.reshape([-1, tuple(x.shape)[-1]]) - h = self.linear(h) - h = h.reshape(tuple(tuple(x.shape)[:-1] + (self.out_size,))) - return h - - -class GraphConv(nn.Layer): - """ - graph convolution over batch and multi-graphs - Args: - in_channels: e.g. 8 - out_channels: e.g. 64 - num_edge_type (types of edges/bonds): e.g. 4 - return: - class:`chainer.Variable`: - """ - - def __init__(self, in_channels, out_channels, num_edge_type=4): - super(GraphConv, self).__init__() - self.graph_linear_self = GraphLinear(in_channels, out_channels) - self.graph_linear_edge = GraphLinear(in_channels, out_channels * num_edge_type) - self.num_edge_type = num_edge_type - self.in_ch = in_channels - self.out_ch = out_channels - - def forward(self, adj, h): - mb, node, ch = tuple(h.shape) - hs = self.graph_linear_self(h) - m = self.graph_linear_edge(h) - m = m.reshape([mb, node, self.out_ch, self.num_edge_type]) - m = m.transpose(perm=[0, 3, 1, 2]) - hr = paddle.matmul(x=adj, y=m) - hr = hr.sum(axis=1) - return hs + hr diff --git a/examples/smc_reac/ppsci/arch/moflow_glow.py b/examples/smc_reac/ppsci/arch/moflow_glow.py deleted file mode 100644 index 5fbeb71520..0000000000 --- a/examples/smc_reac/ppsci/arch/moflow_glow.py +++ /dev/null @@ -1,477 +0,0 @@ -# Copyright (c) 2024 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. - -# Copyright 2020 Chengxi Zang - -import warnings - -import paddle -import paddle.nn as nn - -from ppsci.arch.moflow_basic import ActNorm -from ppsci.arch.moflow_basic import ActNorm2D -from ppsci.arch.moflow_basic import GraphConv -from ppsci.arch.moflow_basic import GraphLinear -from ppsci.arch.moflow_basic import InvConv2d -from ppsci.arch.moflow_basic import InvConv2dLU - -warnings.filterwarnings( - "ignore", message="when training, we now always track global mean and variance." -) - - -class AffineCoupling(nn.Layer): - def __init__(self, in_channel, hidden_channels, affine=True, mask_swap=False): - super(AffineCoupling, self).__init__() - self.affine = affine - self.layers = nn.LayerList() - self.norms = nn.LayerList() - self.mask_swap = mask_swap - last_h = in_channel // 2 - if affine: - vh = tuple(hidden_channels) + (in_channel,) - else: - vh = tuple(hidden_channels) + (in_channel // 2,) - for h in vh: - self.layers.append( - nn.Conv2D(in_channels=last_h, out_channels=h, kernel_size=3, padding=1) - ) - self.norms.append(nn.BatchNorm2D(num_features=h)) - last_h = h - - def forward(self, input): - in_a, in_b = input.chunk(chunks=2, axis=1) - if self.mask_swap: - in_a, in_b = in_b, in_a - if self.affine: - s, t = self._s_t_function(in_a) - out_b = (in_b + t) * s - logdet = paddle.sum( - x=paddle.log(x=paddle.abs(x=s)).reshape([tuple(input.shape)[0], -1]), - axis=1, - ) - else: - _, t = self._s_t_function(in_a) - out_b = in_b + t - logdet = None - if self.mask_swap: - result = paddle.concat(x=[out_b, in_a], axis=1) - else: - result = paddle.concat(x=[in_a, out_b], axis=1) - return result, logdet - - def reverse(self, output): - out_a, out_b = output.chunk(chunks=2, axis=1) - if self.mask_swap: - out_a, out_b = out_b, out_a - if self.affine: - s, t = self._s_t_function(out_a) - in_b = out_b / s - t - else: - _, t = self._s_t_function(out_a) - in_b = out_b - t - if self.mask_swap: - result = paddle.concat(x=[in_b, out_a], axis=1) - else: - result = paddle.concat(x=[out_a, in_b], axis=1) - return result - - def _s_t_function(self, x): - h = x - for i in range(len(self.layers) - 1): - h = self.layers[i](h) - h = self.norms[i](h) - h = nn.functional.relu(x=h) - h = self.layers[-1](h) - s = None - if self.affine: - log_s, t = h.chunk(chunks=2, axis=1) - s = nn.functional.sigmoid(x=log_s) - else: - t = h - return s, t - - -class GraphAffineCoupling(nn.Layer): - def __init__(self, n_node, in_dim, hidden_dim_dict, masked_row, affine=True): - super(GraphAffineCoupling, self).__init__() - self.n_node = n_node - self.in_dim = in_dim - self.hidden_dim_dict = hidden_dim_dict - self.masked_row = masked_row - self.affine = affine - self.hidden_dim_gnn = hidden_dim_dict["gnn"] - self.hidden_dim_linear = hidden_dim_dict["linear"] - self.net = nn.LayerList() - self.norm = nn.LayerList() - last_dim = in_dim - for out_dim in self.hidden_dim_gnn: - self.net.append(GraphConv(last_dim, out_dim)) - self.norm.append(nn.BatchNorm1D(num_features=n_node)) - last_dim = out_dim - self.net_lin = nn.LayerList() - self.norm_lin = nn.LayerList() - for out_dim in self.hidden_dim_linear: - self.net_lin.append(GraphLinear(last_dim, out_dim)) - self.norm_lin.append(nn.BatchNorm1D(num_features=n_node)) - last_dim = out_dim - if affine: - self.net_lin.append(GraphLinear(last_dim, in_dim * 2)) - else: - self.net_lin.append(GraphLinear(last_dim, in_dim)) - self.scale = paddle.create_parameter( - paddle.zeros(shape=[1]).shape, - paddle.zeros(shape=[1]).numpy().dtype, - default_initializer=nn.initializer.Assign(paddle.zeros(shape=[1])), - ) - - mask = paddle.ones(shape=[n_node, in_dim]) - mask[masked_row, :] = 0 - self.register_buffer(name="mask", tensor=mask) - - def forward(self, adj, input): - masked_x = self.mask * input - s, t = self._s_t_function(adj, masked_x) - if self.affine: - out = masked_x + (1 - self.mask) * (input + t) * s - logdet = paddle.sum( - x=paddle.log(x=paddle.abs(x=s)).reshape([tuple(input.shape)[0], -1]), - axis=1, - ) - else: - out = masked_x + t * (1 - self.mask) - logdet = None - return out, logdet - - def reverse(self, adj, output): - masked_y = self.mask * output - s, t = self._s_t_function(adj, masked_y) - if self.affine: - input = masked_y + (1 - self.mask) * (output / s - t) - else: - input = masked_y + (1 - self.mask) * (output - t) - return input - - def _s_t_function(self, adj, x): - s = None - h = x - for i in range(len(self.net)): - h = self.net[i](adj, h) - h = self.norm[i](h) - h = nn.functional.relu(x=h) - for i in range(len(self.net_lin) - 1): - h = self.net_lin[i](h) - h = self.norm_lin[i](h) - h = nn.functional.relu(x=h) - h = self.net_lin[-1](h) - if self.affine: - log_s, t = h.chunk(chunks=2, axis=-1) - s = nn.functional.sigmoid(x=log_s) - else: - t = h - return s, t - - -class Flow(nn.Layer): - def __init__( - self, in_channel, hidden_channels, affine=True, conv_lu=2, mask_swap=False - ): - super(Flow, self).__init__() - self.actnorm = ActNorm(in_channel) - if conv_lu == 0: - self.invconv = InvConv2d(in_channel) - elif conv_lu == 1: - self.invconv = InvConv2dLU(in_channel) - elif conv_lu == 2: - self.invconv = None - else: - raise ValueError( - "conv_lu in {0,1,2}, 0:InvConv2d, 1:InvConv2dLU, 2:none-just swap to update in coupling" - ) - self.coupling = AffineCoupling( - in_channel, hidden_channels, affine=affine, mask_swap=mask_swap - ) - - def forward(self, input): - out, logdet = self.actnorm(input) - if self.invconv: - out, det1 = self.invconv(out) - else: - det1 = 0 - out, det2 = self.coupling(out) - logdet = logdet + det1 - if det2 is not None: - logdet = logdet + det2 - return out, logdet - - def reverse(self, output): - input = self.coupling.reverse(output) - if self.invconv: - input = self.invconv.reverse(input) - input = self.actnorm.reverse(input) - return input - - -class FlowOnGraph(nn.Layer): - def __init__(self, n_node, in_dim, hidden_dim_dict, masked_row, affine=True): - super(FlowOnGraph, self).__init__() - self.n_node = n_node - self.in_dim = in_dim - self.hidden_dim_dict = hidden_dim_dict - self.masked_row = masked_row - self.affine = affine - self.actnorm = ActNorm2D(in_dim=n_node) - self.coupling = GraphAffineCoupling( - n_node, in_dim, hidden_dim_dict, masked_row, affine=affine - ) - - def forward(self, adj, input): - out, logdet = self.actnorm(input) - det1 = 0 - out, det2 = self.coupling(adj, out) - logdet = logdet + det1 - if det2 is not None: - logdet = logdet + det2 - return out, logdet - - def reverse(self, adj, output): - input = self.coupling.reverse(adj, output) - input = self.actnorm.reverse(input) - return input - - -class Block(nn.Layer): - def __init__( - self, in_channel, n_flow, squeeze_fold, hidden_channels, affine=True, conv_lu=2 - ): - super(Block, self).__init__() - self.squeeze_fold = squeeze_fold - squeeze_dim = in_channel * self.squeeze_fold * self.squeeze_fold - self.flows = nn.LayerList() - for i in range(n_flow): - if conv_lu in (0, 1): - self.flows.append( - Flow( - squeeze_dim, - hidden_channels, - affine=affine, - conv_lu=conv_lu, - mask_swap=False, - ) - ) - else: - self.flows.append( - Flow( - squeeze_dim, - hidden_channels, - affine=affine, - conv_lu=2, - mask_swap=bool(i % 2), - ) - ) - - def forward(self, input): - out = self._squeeze(input) - logdet = 0 - for flow in self.flows: - out, det = flow(out) - logdet = logdet + det - out = self._unsqueeze(out) - return out, logdet - - def reverse(self, output): - input = self._squeeze(output) - for flow in self.flows[::-1]: - input = flow.reverse(input) - unsqueezed = self._unsqueeze(input) - return unsqueezed - - def _squeeze(self, x): - """Trade spatial extent for channels. In forward direction, convert each - 1x4x4 volume of input into a 4x1x1 volume of output. - - Args: - x (paddle.Tensor): Input to squeeze or unsqueeze. - reverse (bool): Reverse the operation, i.e., unsqueeze. - - Returns: - x (paddle.Tensor): Squeezed or unsqueezed tensor. - """ - assert len(tuple(x.shape)) == 4 - b_size, n_channel, height, width = tuple(x.shape) - fold = self.squeeze_fold - squeezed = x.reshape( - [b_size, n_channel, height // fold, fold, width // fold, fold] - ) - squeezed = squeezed.transpose(perm=[0, 1, 3, 5, 2, 4]).contiguous() - out = squeezed.reshape( - [b_size, n_channel * fold * fold, height // fold, width // fold] - ) - return out - - def _unsqueeze(self, x): - assert len(tuple(x.shape)) == 4 - b_size, n_channel, height, width = tuple(x.shape) - fold = self.squeeze_fold - unsqueezed = x.reshape( - [b_size, n_channel // (fold * fold), fold, fold, height, width] - ) - unsqueezed = unsqueezed.transpose(perm=[0, 1, 4, 2, 5, 3]).contiguous() - out = unsqueezed.reshape( - [b_size, n_channel // (fold * fold), height * fold, width * fold] - ) - return out - - -class BlockOnGraph(nn.Layer): - def __init__( - self, - n_node, - in_dim, - hidden_dim_dict, - n_flow, - mask_row_size=1, - mask_row_stride=1, - affine=True, - ): - """ - - :param n_node: - :param in_dim: - :param hidden_dim: - :param n_flow: - :param mask_row_size: number of rows to be masked for update - :param mask_row_stride: number of steps between two masks' firs row - :param affine: - """ - super(BlockOnGraph, self).__init__() - assert 0 < mask_row_size < n_node - self.flows = nn.LayerList() - for i in range(n_flow): - start = i * mask_row_stride - masked_row = [(r % n_node) for r in range(start, start + mask_row_size)] - self.flows.append( - FlowOnGraph( - n_node, - in_dim, - hidden_dim_dict, - masked_row=masked_row, - affine=affine, - ) - ) - - def forward(self, adj, input): - out = input - logdet = 0 - for flow in self.flows: - out, det = flow(adj, out) - logdet = logdet + det - return out, logdet - - def reverse(self, adj, output): - input = output - for flow in self.flows[::-1]: - input = flow.reverse(adj, input) - return input - - -class Glow(nn.Layer): - def __init__( - self, - in_channel, - n_flow, - n_block, - squeeze_fold, - hidden_channel, - affine=True, - conv_lu=2, - ): - super(Glow, self).__init__() - self.blocks = nn.LayerList() - n_channel = in_channel - for i in range(n_block): - self.blocks.append( - Block( - n_channel, - n_flow, - squeeze_fold, - hidden_channel, - affine=affine, - conv_lu=conv_lu, - ) - ) - - def forward(self, input): - logdet = 0 - out = input - for block in self.blocks: - out, det = block(out) - logdet = logdet + det - return out, logdet - - def reverse(self, z): - h = z - for i, block in enumerate(self.blocks[::-1]): - h = block.reverse(h) - return h - - -class GlowOnGraph(nn.Layer): - def __init__( - self, - n_node, - in_dim, - hidden_dim_dict, - n_flow, - n_block, - mask_row_size_list=[2], - mask_row_stride_list=[1], - affine=True, - ): - super(GlowOnGraph, self).__init__() - assert len(mask_row_size_list) == n_block or len(mask_row_size_list) == 1 - assert len(mask_row_stride_list) == n_block or len(mask_row_stride_list) == 1 - if len(mask_row_size_list) == 1: - mask_row_size_list = mask_row_size_list * n_block - if len(mask_row_stride_list) == 1: - mask_row_stride_list = mask_row_stride_list * n_block - self.blocks = nn.LayerList() - for i in range(n_block): - mask_row_size = mask_row_size_list[i] - mask_row_stride = mask_row_stride_list[i] - self.blocks.append( - BlockOnGraph( - n_node, - in_dim, - hidden_dim_dict, - n_flow, - mask_row_size, - mask_row_stride, - affine=affine, - ) - ) - - def forward(self, adj, x): - logdet = 0 - out = x - for block in self.blocks: - out, det = block(adj, out) - logdet = logdet + det - return out, logdet - - def reverse(self, adj, z): - input = z - for i, block in enumerate(self.blocks[::-1]): - input = block.reverse(adj, input) - return input diff --git a/examples/smc_reac/ppsci/arch/moflow_net.py b/examples/smc_reac/ppsci/arch/moflow_net.py deleted file mode 100644 index c2f88607bc..0000000000 --- a/examples/smc_reac/ppsci/arch/moflow_net.py +++ /dev/null @@ -1,335 +0,0 @@ -# Copyright (c) 2024 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. - -# Copyright 2020 Chengxi Zang -from __future__ import annotations - -import math -from typing import Dict -from typing import Tuple - -import paddle - -from ppsci.arch import base -from ppsci.arch.moflow_glow import Glow -from ppsci.arch.moflow_glow import GlowOnGraph - - -def gaussian_nll(x, mean, ln_var, reduce="sum"): - """Computes the negative log-likelihood of a Gaussian distribution. - - Given two variable ``mean`` representing :math:`\\mu` and ``ln_var`` - representing :math:`\\log(\\sigma^2)`, this function computes in - elementwise manner the negative log-likelihood of :math:`x` on a - Gaussian distribution :math:`N(\\mu, S)`, - - .. math:: - - -\\log N(x; \\mu, \\sigma^2) = - \\log\\left(\\sqrt{(2\\pi)^D |S|}\\right) + - \\frac{1}{2}(x - \\mu)^\\top S^{-1}(x - \\mu), - - where :math:`D` is a dimension of :math:`x` and :math:`S` is a diagonal - matrix where :math:`S_{ii} = \\sigma_i^2`. - - The output is a variable whose value depends on the value of - the option ``reduce``. If it is ``'no'``, it holds the elementwise - loss values. If it is ``'sum'``, loss values are summed up. - - Args: - x (:class:`~chainer.Variable` or :ref:`ndarray`): Input variable. - mean (:class:`~chainer.Variable` or :ref:`ndarray`): A variable - representing mean of a Gaussian distribution, :math:`\\mu`. - ln_var (:class:`~chainer.Variable` or :ref:`ndarray`): A variable - representing logarithm of variance of a Gaussian distribution, - :math:`\\log(\\sigma^2)`. - reduce (str): Reduction option. Its value must be either - ``'sum'`` or ``'no'``. Otherwise, :class:`ValueError` is raised. - - Returns: - ~chainer.Variable: - A variable representing the negative log-likelihood. - If ``reduce`` is ``'no'``, the output variable holds array - whose shape is same as one of (hence both of) input variables. - If it is ``'sum'``, the output variable holds a scalar value. - - """ - if reduce not in ("sum", "no"): - raise ValueError( - "only 'sum' and 'no' are valid for 'reduce', but '%s' is given" % reduce - ) - x_prec = paddle.exp(x=-ln_var) - x_diff = x - mean - x_power = x_diff * x_diff * x_prec * -0.5 - loss = (ln_var + math.log(2 * math.pi)) / 2 - x_power - if reduce == "sum": - return loss.sum() - else: - return loss - - -def rescale_adj(adj, type="all"): - if type == "view": - out_degree = adj.sum(axis=-1) - out_degree_sqrt_inv = out_degree.pow(y=-1) - out_degree_sqrt_inv[out_degree_sqrt_inv == float("inf")] = 0 - adj_prime = out_degree_sqrt_inv.unsqueeze(axis=-1) * adj - else: - num_neighbors = adj.sum(axis=(1, 2)).astype(dtype="float32") - num_neighbors_inv = num_neighbors.pow(y=-1) - num_neighbors_inv[num_neighbors_inv == float("inf")] = 0 - adj_prime = num_neighbors_inv[:, None, None, :] * adj - return adj_prime - - -class MoFlowNet(base.Arch): - """ - Args: - input_keys (Tuple[str, ...]): Name of input keys, such as ("nodes","edges",). - output_keys (Tuple[str, ...]): Name of output keys, such as ("output","sum_log_det"). - hyper_params (object): More parameters derived from hyper_params for easy use. - """ - - def __init__( - self, - input_keys: Tuple[str, ...], - output_keys: Tuple[str, ...], - hyper_params: None, - ): - super(MoFlowNet, self).__init__() - self.input_keys = input_keys - self.output_keys = output_keys - self.hyper_params = hyper_params - self.b_n_type = hyper_params.b_n_type - self.a_n_node = hyper_params.a_n_node - self.a_n_type = hyper_params.a_n_type - self.b_size = self.a_n_node * self.a_n_node * self.b_n_type - self.a_size = self.a_n_node * self.a_n_type - self.noise_scale = hyper_params.noise_scale - if hyper_params.learn_dist: - self.ln_var = paddle.create_parameter( - paddle.zeros(shape=[1]).shape, - paddle.zeros(shape=[1]).numpy().dtype, - default_initializer=paddle.nn.initializer.Assign( - paddle.zeros(shape=[1]) - ), - ) - - else: - self.register_buffer(name="ln_var", tensor=paddle.zeros(shape=[1])) - self.bond_model = Glow( - in_channel=hyper_params.b_n_type, - n_flow=hyper_params.b_n_flow, - n_block=hyper_params.b_n_block, - squeeze_fold=hyper_params.b_n_squeeze, - hidden_channel=hyper_params.b_hidden_ch, - affine=hyper_params.b_affine, - conv_lu=hyper_params.b_conv_lu, - ) - self.atom_model = GlowOnGraph( - n_node=hyper_params.a_n_node, - in_dim=hyper_params.a_n_type, - hidden_dim_dict={ - "gnn": hyper_params.a_hidden_gnn, - "linear": hyper_params.a_hidden_lin, - }, - n_flow=hyper_params.a_n_flow, - n_block=hyper_params.a_n_block, - mask_row_size_list=hyper_params.mask_row_size_list, - mask_row_stride_list=hyper_params.mask_row_stride_list, - affine=hyper_params.a_affine, - ) - - def forward(self, x): - h = x[self.input_keys[0]] - adj = x[self.input_keys[1]] - adj_normalized = rescale_adj(adj).to(adj) - - if self.training: - if self.noise_scale == 0: - h = h / 2.0 - 0.5 + paddle.rand(shape=h.shape, dtype=h.dtype) * 0.4 - else: - h = h + paddle.rand(shape=h.shape, dtype=h.dtype) * self.noise_scale - h, sum_log_det_jacs_x = self.atom_model(adj_normalized, h) - if self.training: - if self.noise_scale == 0: - adj = ( - adj / 2.0 - - 0.5 - + paddle.rand(shape=adj.shape, dtype=adj.dtype) * 0.4 - ) - else: - adj = ( - adj - + paddle.rand(shape=adj.shape, dtype=adj.dtype) * self.noise_scale - ) - adj_h, sum_log_det_jacs_adj = self.bond_model(adj) - out = [h, adj_h] - result_dict = { - self.output_keys[0]: out, - self.output_keys[1]: [sum_log_det_jacs_x, sum_log_det_jacs_adj], - } - - return result_dict - - def reverse(self, z, true_adj=None): - """ - Returns a molecule, given its latent vector. - - Args: - z: latent vector. Shape: [B, N*N*M + N*T] (100,369) 369=9*9 * 4 + 9*5 - B = Batch size, N = number of atoms, M = number of bond types, - T = number of atom types (Carbon, Oxygen etc.) - true_adj: used for testing. An adjacency matrix of a real molecule - - return: - adjacency matrix and feature matrix of a molecule - """ - batch_size = tuple(z.shape)[0] - with paddle.no_grad(): - z_x = z[:, : self.a_size] - z_adj = z[:, self.a_size :] - if true_adj is None: - h_adj = z_adj.reshape( - [batch_size, self.b_n_type, self.a_n_node, self.a_n_node] - ) - h_adj = self.bond_model.reverse(h_adj) - if self.noise_scale == 0: - h_adj = (h_adj + 0.5) * 2 - adj = h_adj - adj = adj + adj.transpose(perm=[0, 1, 3, 2]) - adj = adj / 2 - adj = paddle.nn.functional.softmax(adj, axis=1) - max_bond = adj.max(axis=1).reshape( - [batch_size, -1, self.a_n_node, self.a_n_node] - ) - adj = paddle.floor(x=adj / max_bond) - else: - adj = true_adj - h_x = z_x.reshape([batch_size, self.a_n_node, self.a_n_type]) - adj_normalized = rescale_adj(adj).to(h_x) - h_x = self.atom_model.reverse(adj_normalized, h_x) - if self.noise_scale == 0: - h_x = (h_x + 0.5) * 2 - return adj, h_x - - def log_prob_loss(self, output_dict: Dict, *args): - losses = 0 - z = output_dict[self.output_keys[0]] - logdet = output_dict[self.output_keys[1]] - z[0] = z[0].reshape([tuple(z[0].shape)[0], -1]) - z[1] = z[1].reshape([tuple(z[1].shape)[0], -1]) - logdet[0] = logdet[0] - self.a_size * math.log(2.0) - logdet[1] = logdet[1] - self.b_size * math.log(2.0) - if len(self.ln_var) == 1: - ln_var_adj = self.ln_var * paddle.ones(shape=[self.b_size]).to(z[0]) - ln_var_x = self.ln_var * paddle.ones(shape=[self.a_size]).to(z[0]) - else: - ln_var_adj = self.ln_var[0] * paddle.ones(shape=[self.b_size]).to(z[0]) - ln_var_x = self.ln_var[1] * paddle.ones(shape=[self.a_size]).to(z[0]) - nll_adj = paddle.mean( - paddle.sum( - gaussian_nll( - z[1], - paddle.zeros(shape=self.b_size).to(z[0]), - ln_var_adj, - reduce="no", - ), - axis=1, - ) - - logdet[1] - ) - nll_adj = nll_adj / (self.b_size * math.log(2.0)) - nll_x = paddle.mean( - paddle.sum( - gaussian_nll( - z[0], - paddle.zeros(shape=self.a_size).to(z[0]), - ln_var_x, - reduce="no", - ), - axis=1, - ) - - logdet[0] - ) - nll_x = nll_x / (self.a_size * math.log(2.0)) - if nll_x.item() < 0: - print(f"nll_x: {nll_x.item()}") - losses = nll_x + nll_adj - return {"total_loss": losses} - - def save_hyperparams(self, path): - self.hyper_params.save(path) - - -class MoFlowProp(base.Arch): - """ - Args: - input_keys (Tuple[str, ...]): Name of input keys, such as ("nodes","edges",). - output_keys (Tuple[str, ...]): Name of output keys, such as ("output","sum_log_det"). - model (MoFlowNet): pre-trained model. - hidden_size (int): Hidden dimension list for output regression. - """ - - def __init__( - self, - input_keys: Tuple[str, ...], - output_keys: Tuple[str, ...], - model: MoFlowNet, - hidden_size, - ): - super(MoFlowProp, self).__init__() - self.input_keys = input_keys - self.output_keys = output_keys - self.model = model - self.latent_size = model.b_size + model.a_size - self.hidden_size = hidden_size - vh = (self.latent_size,) + tuple(hidden_size) + (1,) - modules = [] - for i in range(len(vh) - 1): - modules.append(paddle.nn.Linear(in_features=vh[i], out_features=vh[i + 1])) - if i < len(vh) - 2: - modules.append(paddle.nn.Tanh()) - self.propNN = paddle.nn.Sequential(*modules) - - def encode(self, x): - with paddle.no_grad(): - self.model.eval() - output_dict = self.model(x) - z = output_dict["output"] - sum_log_det_jacs = output_dict["sum_log_det"] - h = paddle.concat( - [ - z[0].reshape([tuple(z[0].shape)[0], -1]), - z[1].reshape([tuple(z[1].shape)[0], -1]), - ], - axis=1, - ) - return h, sum_log_det_jacs - - def reverse(self, z): - with paddle.no_grad(): - self.model.eval() - adj, x = self.model.reverse(z, true_adj=None) - return adj, x - - def forward(self, x): - h, sum_log_det_jacs = self.encode(x) - output = self.propNN(h) - result_dict = { - self.output_keys[0]: [h, output], - self.output_keys[1]: sum_log_det_jacs, - } - - return result_dict diff --git a/examples/smc_reac/ppsci/arch/nowcastnet.py b/examples/smc_reac/ppsci/arch/nowcastnet.py deleted file mode 100644 index bc7538ad91..0000000000 --- a/examples/smc_reac/ppsci/arch/nowcastnet.py +++ /dev/null @@ -1,639 +0,0 @@ -# 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. - -import collections -from typing import Tuple - -import paddle -from paddle import nn - -from ppsci.arch import base - - -class NowcastNet(base.Arch): - """The NowcastNet model. - - Args: - input_keys (Tuple[str, ...]): Name of input keys, such as ("input",). - output_keys (Tuple[str, ...]): Name of output keys, such as ("output",). - input_length (int, optional): Input length. Defaults to 9. - total_length (int, optional): Total length. Defaults to 29. - image_height (int, optional): Image height. Defaults to 512. - image_width (int, optional): Image width. Defaults to 512. - image_ch (int, optional): Image channel. Defaults to 2. - ngf (int, optional): Noise Projector input length. Defaults to 32. - - Examples: - >>> import ppsci - >>> model = ppsci.arch.NowcastNet(("input", ), ("output", )) - >>> input_data = paddle.rand([1, 9, 512, 512, 2]) - >>> input_dict = {"input": input_data} - >>> output_dict = model(input_dict) - >>> print(output_dict["output"].shape) - [1, 20, 512, 512, 1] - """ - - def __init__( - self, - input_keys: Tuple[str, ...], - output_keys: Tuple[str, ...], - input_length: int = 9, - total_length: int = 29, - image_height: int = 512, - image_width: int = 512, - image_ch: int = 2, - ngf: int = 32, - ): - super().__init__() - self.input_keys = input_keys - self.output_keys = output_keys - - self.input_length = input_length - self.total_length = total_length - self.image_height = image_height - self.image_width = image_width - self.image_ch = image_ch - self.ngf = ngf - - configs = collections.namedtuple( - "Object", ["ngf", "evo_ic", "gen_oc", "ic_feature"] - ) - configs.ngf = self.ngf - configs.evo_ic = self.total_length - self.input_length - configs.gen_oc = self.total_length - self.input_length - configs.ic_feature = self.ngf * 10 - - self.pred_length = self.total_length - self.input_length - self.evo_net = Evolution_Network(self.input_length, self.pred_length, base_c=32) - self.gen_enc = Generative_Encoder(self.total_length, base_c=self.ngf) - self.gen_dec = Generative_Decoder(configs) - self.proj = Noise_Projector(self.ngf) - sample_tensor = paddle.zeros(shape=[1, 1, self.image_height, self.image_width]) - self.grid = make_grid(sample_tensor) - - @staticmethod - def split_to_dict(data_tensors: Tuple[paddle.Tensor, ...], keys: Tuple[str, ...]): - return {key: data_tensors[i] for i, key in enumerate(keys)} - - def forward(self, x): - if self._input_transform is not None: - x = self._input_transform(x) - - x_tensor = self.concat_to_tensor(x, self.input_keys) - - y = [] - out = self.forward_tensor(x_tensor) - y.append(out) - y = self.split_to_dict(y, self.output_keys) - - if self._output_transform is not None: - y = self._output_transform(x, y) - return y - - def forward_tensor(self, x): - all_frames = x[:, :, :, :, :1] - frames = all_frames.transpose(perm=[0, 1, 4, 2, 3]) - batch = frames.shape[0] - height = frames.shape[3] - width = frames.shape[4] - # Input Frames - input_frames = frames[:, : self.input_length] - input_frames = input_frames.reshape((batch, self.input_length, height, width)) - # Evolution Network - intensity, motion = self.evo_net(input_frames) - motion_ = motion.reshape((batch, self.pred_length, 2, height, width)) - intensity_ = intensity.reshape((batch, self.pred_length, 1, height, width)) - series = [] - last_frames = all_frames[:, self.input_length - 1 : self.input_length, :, :, 0] - grid = self.grid.tile((batch, 1, 1, 1)) - for i in range(self.pred_length): - last_frames = warp( - last_frames, motion_[:, i], grid, mode="nearest", padding_mode="border" - ) - last_frames = last_frames + intensity_[:, i] - series.append(last_frames) - evo_result = paddle.concat(x=series, axis=1) - evo_result = evo_result / 128 - # Generative Network - evo_feature = self.gen_enc(paddle.concat(x=[input_frames, evo_result], axis=1)) - noise = paddle.randn(shape=[batch, self.ngf, height // 32, width // 32]) - noise = self.proj(noise) - ngf = noise.shape[1] - noise_feature = ( - noise.reshape((batch, -1, 4, 4, 8, 8)) - .transpose(perm=[0, 1, 4, 5, 2, 3]) - .reshape((batch, ngf // 16, height // 8, width // 8)) - ) - feature = paddle.concat(x=[evo_feature, noise_feature], axis=1) - gen_result = self.gen_dec(feature, evo_result) - return gen_result.unsqueeze(axis=-1) - - -class Evolution_Network(nn.Layer): - def __init__(self, n_channels, n_classes, base_c=64, bilinear=True): - super().__init__() - self.n_channels = n_channels - self.n_classes = n_classes - self.bilinear = bilinear - base_c = base_c - self.inc = DoubleConv(n_channels, base_c) - self.down1 = Down(base_c * 1, base_c * 2) - self.down2 = Down(base_c * 2, base_c * 4) - self.down3 = Down(base_c * 4, base_c * 8) - factor = 2 if bilinear else 1 - self.down4 = Down(base_c * 8, base_c * 16 // factor) - self.up1 = Up(base_c * 16, base_c * 8 // factor, bilinear) - self.up2 = Up(base_c * 8, base_c * 4 // factor, bilinear) - self.up3 = Up(base_c * 4, base_c * 2 // factor, bilinear) - self.up4 = Up(base_c * 2, base_c * 1, bilinear) - self.outc = OutConv(base_c * 1, n_classes) - param1 = paddle.zeros(shape=[1, n_classes, 1, 1]) - gamma = self.create_parameter( - shape=param1.shape, - dtype=param1.dtype, - default_initializer=nn.initializer.Assign(param1), - ) - gamma.stop_gradient = False - self.gamma = gamma - self.up1_v = Up(base_c * 16, base_c * 8 // factor, bilinear) - self.up2_v = Up(base_c * 8, base_c * 4 // factor, bilinear) - self.up3_v = Up(base_c * 4, base_c * 2 // factor, bilinear) - self.up4_v = Up(base_c * 2, base_c * 1, bilinear) - self.outc_v = OutConv(base_c * 1, n_classes * 2) - - def forward(self, x): - x1 = self.inc(x) - x2 = self.down1(x1) - x3 = self.down2(x2) - x4 = self.down3(x3) - x5 = self.down4(x4) - x = self.up1(x5, x4) - x = self.up2(x, x3) - x = self.up3(x, x2) - x = self.up4(x, x1) - x = self.outc(x) * self.gamma - v = self.up1_v(x5, x4) - v = self.up2_v(v, x3) - v = self.up3_v(v, x2) - v = self.up4_v(v, x1) - v = self.outc_v(v) - return x, v - - -class DoubleConv(nn.Layer): - def __init__(self, in_channels, out_channels, kernel=3, mid_channels=None): - super().__init__() - if not mid_channels: - mid_channels = out_channels - self.double_conv = nn.Sequential( - nn.BatchNorm2D(num_features=in_channels), - nn.ReLU(), - nn.utils.spectral_norm( - layer=nn.Conv2D( - in_channels=in_channels, - out_channels=mid_channels, - kernel_size=kernel, - padding=kernel // 2, - ) - ), - nn.BatchNorm2D(num_features=mid_channels), - nn.ReLU(), - nn.utils.spectral_norm( - layer=nn.Conv2D( - in_channels=mid_channels, - out_channels=out_channels, - kernel_size=kernel, - padding=kernel // 2, - ) - ), - ) - self.single_conv = nn.Sequential( - nn.BatchNorm2D(num_features=in_channels), - nn.utils.spectral_norm( - layer=nn.Conv2D( - in_channels=in_channels, - out_channels=out_channels, - kernel_size=kernel, - padding=kernel // 2, - ) - ), - ) - - def forward(self, x): - shortcut = self.single_conv(x) - x = self.double_conv(x) - x = x + shortcut - return x - - -class Down(nn.Layer): - def __init__(self, in_channels, out_channels, kernel=3): - super().__init__() - self.maxpool_conv = nn.Sequential( - nn.MaxPool2D(kernel_size=2), - DoubleConv(in_channels, out_channels, kernel), - ) - - def forward(self, x): - x = self.maxpool_conv(x) - return x - - -class Up(nn.Layer): - def __init__(self, in_channels, out_channels, bilinear=True, kernel=3): - super().__init__() - if bilinear: - self.up = nn.Upsample(scale_factor=2, mode="bilinear", align_corners=True) - self.conv = DoubleConv( - in_channels, out_channels, kernel=kernel, mid_channels=in_channels // 2 - ) - else: - self.up = nn.Conv2DTranspose( - in_channels=in_channels, - out_channels=in_channels // 2, - kernel_size=2, - stride=2, - ) - self.conv = DoubleConv(in_channels, out_channels, kernel) - - def forward(self, x1, x2): - x1 = self.up(x1) - # input is CHW - diffY = x2.shape[2] - x1.shape[2] - diffX = x2.shape[3] - x1.shape[3] - x1 = nn.functional.pad( - x1, [diffX // 2, diffX - diffX // 2, diffY // 2, diffY - diffY // 2] - ) - x = paddle.concat(x=[x2, x1], axis=1) - return self.conv(x) - - -class Up_S(nn.Layer): - def __init__(self, in_channels, out_channels, bilinear=True, kernel=3): - super().__init__() - if bilinear: - self.up = nn.Upsample(scale_factor=2, mode="bilinear", align_corners=True) - self.conv = DoubleConv( - in_channels, out_channels, kernel=kernel, mid_channels=in_channels - ) - else: - self.up = nn.Conv2DTranspose( - in_channels=in_channels, - out_channels=in_channels, - kernel_size=2, - stride=2, - ) - self.conv = DoubleConv(in_channels, out_channels, kernel) - - def forward(self, x): - x = self.up(x) - return self.conv(x) - - -class OutConv(nn.Layer): - def __init__(self, in_channels, out_channels): - super().__init__() - self.conv = nn.Conv2D( - in_channels=in_channels, out_channels=out_channels, kernel_size=1 - ) - - def forward(self, x): - return self.conv(x) - - -class Generative_Encoder(nn.Layer): - def __init__(self, n_channels, base_c=64): - super().__init__() - base_c = base_c - self.inc = DoubleConv(n_channels, base_c, kernel=3) - self.down1 = Down(base_c * 1, base_c * 2, 3) - self.down2 = Down(base_c * 2, base_c * 4, 3) - self.down3 = Down(base_c * 4, base_c * 8, 3) - - def forward(self, x): - x = self.inc(x) - x = self.down1(x) - x = self.down2(x) - x = self.down3(x) - return x - - -class Generative_Decoder(nn.Layer): - def __init__(self, opt): - super().__init__() - self.opt = opt - nf = opt.ngf - ic = opt.ic_feature - self.fc = nn.Conv2D( - in_channels=ic, out_channels=8 * nf, kernel_size=3, padding=1 - ) - self.head_0 = GenBlock(8 * nf, 8 * nf, opt) - self.G_middle_0 = GenBlock(8 * nf, 4 * nf, opt, double_conv=True) - self.G_middle_1 = GenBlock(4 * nf, 4 * nf, opt, double_conv=True) - self.up_0 = GenBlock(4 * nf, 2 * nf, opt) - self.up_1 = GenBlock(2 * nf, 1 * nf, opt, double_conv=True) - self.up_2 = GenBlock(1 * nf, 1 * nf, opt, double_conv=True) - final_nc = nf * 1 - self.conv_img = nn.Conv2D( - in_channels=final_nc, out_channels=self.opt.gen_oc, kernel_size=3, padding=1 - ) - self.up = nn.Upsample(scale_factor=2) - - def forward(self, x, evo): - x = self.fc(x) - x = self.head_0(x, evo) - x = self.up(x) - x = self.G_middle_0(x, evo) - x = self.G_middle_1(x, evo) - x = self.up(x) - x = self.up_0(x, evo) - x = self.up(x) - x = self.up_1(x, evo) - x = self.up_2(x, evo) - x = self.conv_img(nn.functional.leaky_relu(x=x, negative_slope=0.2)) - return x - - -class GenBlock(nn.Layer): - def __init__(self, fin, fout, opt, use_se=False, dilation=1, double_conv=False): - super().__init__() - self.learned_shortcut = fin != fout - fmiddle = min(fin, fout) - self.opt = opt - self.double_conv = double_conv - self.pad = nn.Pad2D(padding=dilation, mode="reflect") - self.conv_0 = nn.Conv2D( - in_channels=fin, - out_channels=fmiddle, - kernel_size=3, - padding=0, - dilation=dilation, - ) - self.conv_1 = nn.Conv2D( - in_channels=fmiddle, - out_channels=fout, - kernel_size=3, - padding=0, - dilation=dilation, - ) - if self.learned_shortcut: - self.conv_s = nn.Conv2D( - in_channels=fin, out_channels=fout, kernel_size=1, bias_attr=False - ) - self.conv_0 = nn.utils.spectral_norm(layer=self.conv_0) - self.conv_1 = nn.utils.spectral_norm(layer=self.conv_1) - if self.learned_shortcut: - self.conv_s = nn.utils.spectral_norm(layer=self.conv_s) - ic = opt.evo_ic - self.norm_0 = SPADE(fin, ic) - self.norm_1 = SPADE(fmiddle, ic) - if self.learned_shortcut: - self.norm_s = SPADE(fin, ic) - - def forward(self, x, evo): - x_s = self.shortcut(x, evo) - dx = self.conv_0(self.pad(self.actvn(self.norm_0(x, evo)))) - if self.double_conv: - dx = self.conv_1(self.pad(self.actvn(self.norm_1(dx, evo)))) - out = x_s + dx - return out - - def shortcut(self, x, evo): - if self.learned_shortcut: - x_s = self.conv_s(self.norm_s(x, evo)) - else: - x_s = x - return x_s - - def actvn(self, x): - return nn.functional.leaky_relu(x=x, negative_slope=0.2) - - -class SPADE(nn.Layer): - def __init__(self, norm_nc, label_nc): - super().__init__() - ks = 3 - self.param_free_norm = nn.InstanceNorm2D( - num_features=norm_nc, weight_attr=False, bias_attr=False, momentum=1 - 0.1 - ) - nhidden = 64 - ks = 3 - pw = ks // 2 - self.mlp_shared = nn.Sequential( - nn.Pad2D(padding=pw, mode="reflect"), - nn.Conv2D( - in_channels=label_nc, out_channels=nhidden, kernel_size=ks, padding=0 - ), - nn.ReLU(), - ) - self.pad = nn.Pad2D(padding=pw, mode="reflect") - self.mlp_gamma = nn.Conv2D( - in_channels=nhidden, out_channels=norm_nc, kernel_size=ks, padding=0 - ) - self.mlp_beta = nn.Conv2D( - in_channels=nhidden, out_channels=norm_nc, kernel_size=ks, padding=0 - ) - - def forward(self, x, evo): - normalized = self.param_free_norm(x) - evo = nn.functional.adaptive_avg_pool2d(x=evo, output_size=x.shape[2:]) - actv = self.mlp_shared(evo) - gamma = self.mlp_gamma(self.pad(actv)) - beta = self.mlp_beta(self.pad(actv)) - out = normalized * (1 + gamma) + beta - return out - - -class Noise_Projector(nn.Layer): - def __init__(self, input_length): - super().__init__() - self.input_length = input_length - self.conv_first = nn.utils.spectral_norm( - nn.Conv2D( - in_channels=self.input_length, - out_channels=self.input_length * 2, - kernel_size=3, - padding=1, - ) - ) - self.L1 = ProjBlock(self.input_length * 2, self.input_length * 4) - self.L2 = ProjBlock(self.input_length * 4, self.input_length * 8) - self.L3 = ProjBlock(self.input_length * 8, self.input_length * 16) - self.L4 = ProjBlock(self.input_length * 16, self.input_length * 32) - - def forward(self, x): - x = self.conv_first(x) - x = self.L1(x) - x = self.L2(x) - x = self.L3(x) - x = self.L4(x) - return x - - -class ProjBlock(nn.Layer): - def __init__(self, in_channel, out_channel): - super().__init__() - self.one_conv = nn.utils.spectral_norm( - nn.Conv2D( - in_channels=in_channel, - out_channels=out_channel - in_channel, - kernel_size=1, - padding=0, - ) - ) - self.double_conv = nn.Sequential( - nn.utils.spectral_norm( - nn.Conv2D( - in_channels=in_channel, - out_channels=out_channel, - kernel_size=3, - padding=1, - ) - ), - nn.ReLU(), - nn.utils.spectral_norm( - nn.Conv2D( - in_channels=out_channel, - out_channels=out_channel, - kernel_size=3, - padding=1, - ) - ), - ) - - def forward(self, x): - x1 = paddle.concat(x=[x, self.one_conv(x)], axis=1) - x2 = self.double_conv(x) - output = x1 + x2 - return output - - -def make_grid(input): - B, C, H, W = input.shape - xx = paddle.arange(start=0, end=W).reshape((1, -1)).tile((H, 1)) - yy = paddle.arange(start=0, end=H).reshape((-1, 1)).tile((1, W)) - xx = xx.reshape((1, 1, H, W)).tile((B, 1, 1, 1)) - yy = yy.reshape((1, 1, H, W)).tile((B, 1, 1, 1)) - grid = paddle.concat(x=(xx, yy), axis=1).astype(dtype=paddle.get_default_dtype()) - return grid - - -def warp(input, flow, grid, mode="bilinear", padding_mode="zeros"): - B, C, H, W = input.shape - vgrid = grid + flow - vgrid[:, 0, :, :] = 2.0 * vgrid[:, 0, :, :].clone() / max(W - 1, 1) - 1.0 - vgrid[:, 1, :, :] = 2.0 * vgrid[:, 1, :, :].clone() / max(H - 1, 1) - 1.0 - vgrid = vgrid.transpose(perm=[0, 2, 3, 1]) - output = nn.functional.grid_sample( - x=input.cpu(), - grid=vgrid.cpu(), - padding_mode=padding_mode, - mode=mode, - align_corners=True, - ) - return output.cuda() - - -def l2normalize(v, eps=1e-12): - return v / (v.norm() + eps) - - -class spectral_norm(nn.Layer): - def __init__(self, module, name="weight", power_iterations=1): - super().__init__() - self.module = module - self.name = name - self.power_iterations = power_iterations - if not self._made_params(): - self._make_params() - - def _update_u_v(self): - u = getattr(self.module, self.name + "_u") - v = getattr(self.module, self.name + "_v") - w = getattr(self.module, self.name + "_bar") - height = w.detach().shape[0] - for _ in range(self.power_iterations): - v = l2normalize( - paddle.mv( - x=paddle.t(input=w.reshape((height, -1)).detach()), vec=u.detach() - ) - ) - u = l2normalize( - paddle.mv(x=w.reshape((height, -1)).detach(), vec=v.detach()) - ) - sigma = u.dot(y=w.reshape((height, -1)).mv(vec=v)) - setattr(self.module, self.name, w / sigma.expand_as(y=w)) - - def _made_params(self): - try: - _ = getattr(self.module, self.name + "_u") - _ = getattr(self.module, self.name + "_v") - _ = getattr(self.module, self.name + "_bar") - return True - except AttributeError: - return False - - def _make_params(self): - w = getattr(self.module, self.name) - height = w.detach().shape[0] - width = w.reshape((height, -1)).detach().shape[1] - - tmp_w = paddle.normal(shape=[height]) - out_0 = paddle.create_parameter( - shape=tmp_w.shape, - dtype=tmp_w.numpy().dtype, - default_initializer=nn.initializer.Assign(tmp_w), - ) - out_0.stop_gradient = True - u = out_0 - - tmp_w = paddle.normal(shape=[width]) - out_1 = paddle.create_parameter( - shape=tmp_w.shape, - dtype=tmp_w.numpy().dtype, - default_initializer=nn.initializer.Assign(tmp_w), - ) - out_1.stop_gradient = True - v = out_1 - u = l2normalize(u) - v = l2normalize(v) - tmp_w = w.detach() - out_2 = paddle.create_parameter( - shape=tmp_w.shape, - dtype=tmp_w.numpy().dtype, - default_initializer=nn.initializer.Assign(tmp_w), - ) - out_2.stop_gradient = False - w_bar = out_2 - del self.module._parameters[self.name] - - u = create_param(u) - v = create_param(v) - self.module.add_parameter(name=self.name + "_u", parameter=u) - self.module.add_parameter(name=self.name + "_v", parameter=v) - self.module.add_parameter(name=self.name + "_bar", parameter=w_bar) - - def forward(self, *args): - self._update_u_v() - return self.module.forward(*args) - - -def create_param(x): - param = paddle.create_parameter( - shape=x.shape, - dtype=x.dtype, - default_initializer=nn.initializer.Assign(x), - ) - param.stop_gradient = x.stop_gradient - return param diff --git a/examples/smc_reac/ppsci/arch/paddle_harmonics/legendre.py b/examples/smc_reac/ppsci/arch/paddle_harmonics/legendre.py deleted file mode 100644 index 376599719e..0000000000 --- a/examples/smc_reac/ppsci/arch/paddle_harmonics/legendre.py +++ /dev/null @@ -1,176 +0,0 @@ -# coding=utf-8 - -# SPDX-FileCopyrightText: Copyright (c) 2022 The torch-harmonics Authors. All rights reserved. -# SPDX-License-Identifier: BSD-3-Clause -# -# Redistribution and use in source and binary forms, with or without -# modification, are permitted provided that the following conditions are met: -# -# 1. Redistributions of source code must retain the above copyright notice, this -# list of conditions and the following disclaimer. -# -# 2. Redistributions in binary form must reproduce the above copyright notice, -# this list of conditions and the following disclaimer in the documentation -# and/or other materials provided with the distribution. -# -# 3. Neither the name of the copyright holder nor the names of its -# contributors may be used to endorse or promote products derived from -# this software without specific prior written permission. -# -# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" -# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE -# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE -# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE -# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL -# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR -# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER -# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, -# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE -# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. - -""" -Code below is heavily based on [torch-harmonics](https://github.com/NVIDIA/torch-harmonics/blob/main/torch_harmonics/legendre.py) -""" - -import numpy as np - - -def clm(l, m): - """ - Defines the normalization factor to orthonormalize the Spherical Harmonics - """ - return np.sqrt((2 * l + 1) / 4 / np.pi) * np.sqrt( - np.math.factorial(l - m) / np.math.factorial(l + m) - ) - - -def legpoly(mmax, lmax, x, norm="ortho", inverse=False, csphase=True): - r""" - Computes the values of (-1)^m c^l_m P^l_m(x) at the positions specified by x. - The resulting tensor has shape (mmax, lmax, len(x)). The Condon-Shortley Phase (-1)^m - can be turned off optionally. - - method of computation follows - [1] Schaeffer, N.; Efficient spherical harmonic transforms aimed at pseudospectral numerical simulations, G3: Geochemistry, Geophysics, Geosystems. - [2] Rapp, R.H.; A Fortran Program for the Computation of Gravimetric Quantities from High Degree Spherical Harmonic Expansions, Ohio State University Columbus; report; 1982; - https://apps.dtic.mil/sti/citations/ADA123406 - [3] Schrama, E.; Orbit integration based upon interpolated gravitational gradients - """ - # compute the tensor P^m_n: - nmax = max(mmax, lmax) - vdm = np.zeros((nmax, nmax, len(x)), dtype=np.float64) - - norm_factor = 1.0 if norm == "ortho" else np.sqrt(4 * np.pi) - norm_factor = 1.0 / norm_factor if inverse else norm_factor - - # initial values to start the recursion - vdm[0, 0, :] = norm_factor / np.sqrt(4 * np.pi) - - # fill the diagonal and the lower diagonal - for l in range(1, nmax): - vdm[l - 1, l, :] = np.sqrt(2 * l + 1) * x * vdm[l - 1, l - 1, :] - vdm[l, l, :] = ( - np.sqrt((2 * l + 1) * (1 + x) * (1 - x) / 2 / l) * vdm[l - 1, l - 1, :] - ) - - # fill the remaining values on the upper triangle and multiply b - for l in range(2, nmax): - for m in range(0, l - 1): - vdm[m, l, :] = ( - x - * np.sqrt((2 * l - 1) / (l - m) * (2 * l + 1) / (l + m)) - * vdm[m, l - 1, :] - - np.sqrt( - (l + m - 1) - / (l - m) - * (2 * l + 1) - / (2 * l - 3) - * (l - m - 1) - / (l + m) - ) - * vdm[m, l - 2, :] - ) - - if norm == "schmidt": - for l in range(0, nmax): - if inverse: - vdm[:, l, :] = vdm[:, l, :] * np.sqrt(2 * l + 1) - else: - vdm[:, l, :] = vdm[:, l, :] / np.sqrt(2 * l + 1) - - vdm = vdm[:mmax, :lmax] - - if csphase: - for m in range(1, mmax, 2): - vdm[m] *= -1 - - return vdm - - -def _precompute_legpoly(mmax, lmax, t, norm="ortho", inverse=False, csphase=True): - r""" - Computes the values of (-1)^m c^l_m P^l_m(\cos \theta) at the positions specified by t (theta). - The resulting tensor has shape (mmax, lmax, len(x)). The Condon-Shortley Phase (-1)^m - can be turned off optionally. - - method of computation follows - [1] Schaeffer, N.; Efficient spherical harmonic transforms aimed at pseudospectral numerical simulations, G3: Geochemistry, Geophysics, Geosystems. - [2] Rapp, R.H.; A Fortran Program for the Computation of Gravimetric Quantities from High Degree Spherical Harmonic Expansions, Ohio State University Columbus; report; 1982; - https://apps.dtic.mil/sti/citations/ADA123406 - [3] Schrama, E.; Orbit integration based upon interpolated gravitational gradients - """ - return legpoly(mmax, lmax, np.cos(t), norm=norm, inverse=inverse, csphase=csphase) - - -def _precompute_dlegpoly(mmax, lmax, t, norm="ortho", inverse=False, csphase=True): - r""" - Computes the values of the derivatives $\frac{d}{d \theta} P^m_l(\cos \theta)$ - at the positions specified by t (theta), as well as $\frac{1}{\sin \theta} P^m_l(\cos \theta)$, - needed for the computation of the vector spherical harmonics. The resulting tensor has shape - (2, mmax, lmax, len(t)). - - computation follows - [2] Wang, B., Wang, L., Xie, Z.; Accurate calculation of spherical and vector spherical harmonic expansions via spectral element grids; Adv Comput Math. - """ - pct = _precompute_legpoly( - mmax + 1, lmax + 1, t, norm=norm, inverse=inverse, csphase=False - ) - - dpct = np.zeros((2, mmax, lmax, len(t)), dtype=np.float64) - - # fill the derivative terms wrt theta - for l in range(0, lmax): - - # m = 0 - dpct[0, 0, l] = -np.sqrt(l * (l + 1)) * pct[1, l] - - # 0 < m < l - for m in range(1, min(l, mmax)): - dpct[0, m, l] = 0.5 * ( - np.sqrt((l + m) * (l - m + 1)) * pct[m - 1, l] - - np.sqrt((l - m) * (l + m + 1)) * pct[m + 1, l] - ) - - # m == l - if mmax > l: - dpct[0, l, l] = np.sqrt(l / 2) * pct[l - 1, l] - - # fill the - 1j m P^m_l / sin(phi). as this component is purely imaginary, - # we won't store it explicitly in a complex array - for m in range(1, min(l + 1, mmax)): - # this component is implicitly complex - # we do not divide by m here as this cancels with the derivative of the exponential - dpct[1, m, l] = ( - 0.5 - * np.sqrt((2 * l + 1) / (2 * l + 3)) - * ( - np.sqrt((l - m + 1) * (l - m + 2)) * pct[m - 1, l + 1] - + np.sqrt((l + m + 1) * (l + m + 2)) * pct[m + 1, l + 1] - ) - ) - - if csphase: - for m in range(1, mmax, 2): - dpct[:, m] *= -1 - - return dpct diff --git a/examples/smc_reac/ppsci/arch/paddle_harmonics/quadrature.py b/examples/smc_reac/ppsci/arch/paddle_harmonics/quadrature.py deleted file mode 100644 index 25de0e1436..0000000000 --- a/examples/smc_reac/ppsci/arch/paddle_harmonics/quadrature.py +++ /dev/null @@ -1,156 +0,0 @@ -# coding=utf-8 - -# SPDX-FileCopyrightText: Copyright (c) 2022 The torch-harmonics Authors. All rights reserved. -# SPDX-License-Identifier: BSD-3-Clause -# -# Redistribution and use in source and binary forms, with or without -# modification, are permitted provided that the following conditions are met: -# -# 1. Redistributions of source code must retain the above copyright notice, this -# list of conditions and the following disclaimer. -# -# 2. Redistributions in binary form must reproduce the above copyright notice, -# this list of conditions and the following disclaimer in the documentation -# and/or other materials provided with the distribution. -# -# 3. Neither the name of the copyright holder nor the names of its -# contributors may be used to endorse or promote products derived from -# this software without specific prior written permission. -# -# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" -# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE -# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE -# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE -# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL -# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR -# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER -# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, -# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE -# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. - -""" -Code below is heavily based on [torch-harmonics](https://github.com/NVIDIA/torch-harmonics/blob/main/torch_harmonics/quadrature.py) -""" - -import numpy as np - - -def legendre_gauss_weights(n, a=-1.0, b=1.0): - r""" - Helper routine which returns the Legendre-Gauss nodes and weights - on the interval [a, b] - """ - xlg, wlg = np.polynomial.legendre.leggauss(n) - xlg = (b - a) * 0.5 * xlg + (b + a) * 0.5 - wlg = wlg * (b - a) * 0.5 - - return xlg, wlg - - -def lobatto_weights(n, a=-1.0, b=1.0, tol=1e-16, maxiter=100): - r""" - Helper routine which returns the Legendre-Gauss-Lobatto nodes and weights - on the interval [a, b] - """ - wlg = np.zeros((n,)) - tlg = np.zeros((n,)) - tmp = np.zeros((n,)) - - # Vandermonde Matrix - vdm = np.zeros((n, n)) - - # initialize Chebyshev nodes as first guess - for i in range(n): - tlg[i] = -np.cos(np.pi * i / (n - 1)) - - tmp = 2.0 - - for i in range(maxiter): - tmp = tlg - - vdm[:, 0] = 1.0 - vdm[:, 1] = tlg - - for k in range(2, n): - vdm[:, k] = ( - (2 * k - 1) * tlg * vdm[:, k - 1] - (k - 1) * vdm[:, k - 2] - ) / k - - tlg = tmp - (tlg * vdm[:, n - 1] - vdm[:, n - 2]) / (n * vdm[:, n - 1]) - - if max(abs(tlg - tmp).flatten()) < tol: - break - - wlg = 2.0 / ((n * (n - 1)) * (vdm[:, n - 1] ** 2)) - - # rescale - tlg = (b - a) * 0.5 * tlg + (b + a) * 0.5 - wlg = wlg * (b - a) * 0.5 - - return tlg, wlg - - -def clenshaw_curtiss_weights(n, a=-1.0, b=1.0): - r""" - Computation of the Clenshaw-Curtis quadrature nodes and weights. - This implementation follows - - [1] Joerg Waldvogel, Fast Construction of the Fejer and Clenshaw-Curtis Quadrature Rules; BIT Numerical Mathematics, Vol. 43, No. 1, pp. 001–018. - """ - assert n > 1 - - tcc = np.cos(np.linspace(np.pi, 0, n)) - - if n == 2: - wcc = np.array([1.0, 1.0]) - else: - - n1 = n - 1 - N = np.arange(1, n1, 2) - l = len(N) - m = n1 - l - - v = np.concatenate([2 / N / (N - 2), 1 / N[-1:], np.zeros(m)]) - v = 0 - v[:-1] - v[-1:0:-1] - - g0 = -np.ones(n1) - g0[l] = g0[l] + n1 - g0[m] = g0[m] + n1 - g = g0 / (n1**2 - 1 + (n1 % 2)) - wcc = np.fft.ifft(v + g).real - wcc = np.concatenate((wcc, wcc[:1])) - - # rescale - tcc = (b - a) * 0.5 * tcc + (b + a) * 0.5 - wcc = wcc * (b - a) * 0.5 - - return tcc, wcc - - -def fejer2_weights(n, a=-1.0, b=1.0): - r""" - Computation of the Fejer quadrature nodes and weights. - This implementation follows - - [1] Joerg Waldvogel, Fast Construction of the Fejer and Clenshaw-Curtis Quadrature Rules; BIT Numerical Mathematics, Vol. 43, No. 1, pp. 001–018. - """ - assert n > 2 - - tcc = np.cos(np.linspace(np.pi, 0, n)) - - n1 = n - 1 - N = np.arange(1, n1, 2) - l = len(N) - m = n1 - l - - v = np.concatenate([2 / N / (N - 2), 1 / N[-1:], np.zeros(m)]) - v = 0 - v[:-1] - v[-1:0:-1] - - wcc = np.fft.ifft(v).real - wcc = np.concatenate((wcc, wcc[:1])) - - # rescale - tcc = (b - a) * 0.5 * tcc + (b + a) * 0.5 - wcc = wcc * (b - a) * 0.5 - - return tcc, wcc diff --git a/examples/smc_reac/ppsci/arch/paddle_harmonics/random_fields.py b/examples/smc_reac/ppsci/arch/paddle_harmonics/random_fields.py deleted file mode 100644 index 8fad28cf26..0000000000 --- a/examples/smc_reac/ppsci/arch/paddle_harmonics/random_fields.py +++ /dev/null @@ -1,148 +0,0 @@ -# coding=utf-8 - -# SPDX-FileCopyrightText: Copyright (c) 2022 The paddle-harmonics Authors. All rights reserved. -# SPDX-License-Identifier: BSD-3-Clause -# -# Redistribution and use in source and binary forms, with or without -# modification, are permitted provided that the following conditions are met: -# -# 1. Redistributions of source code must retain the above copyright notice, this -# list of conditions and the following disclaimer. -# -# 2. Redistributions in binary form must reproduce the above copyright notice, -# this list of conditions and the following disclaimer in the documentation -# and/or other materials provided with the distribution. -# -# 3. Neither the name of the copyright holder nor the names of its -# contributors may be used to endorse or promote products derived from -# this software without specific prior written permission. -# -# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" -# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE -# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE -# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE -# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL -# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR -# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER -# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, -# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE -# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. - -""" -Code below is heavily based on [torch-harmonics](https://github.com/NVIDIA/torch-harmonics/blob/main/torch_harmonics/random_fields.py) -""" - -import paddle -from paddle import nn - -from ppsci.arch.paddle_harmonics.sht import InverseRealSHT - - -class GaussianRandomFieldS2(nn.Layer): - r""" - A mean-zero Gaussian Random Field on the sphere with Matern covariance: - C = sigma^2 (-Lap + tau^2 I)^(-alpha). - - Lap is the Laplacian on the sphere, I the identity operator, - and sigma, tau, alpha are scalar parameters. - - Note: C is trace-class on L^2 if and only if alpha > 1. - - Args: - nlat (int): Number of latitudinal modes.longitudinal modes are 2*nlat. - alpha (float, optional): Regularity parameter. Larger means smoother. Defaults to 2.0. - tau (float, optional): Lenght-scale parameter. Larger means more scales. Defaults to 3.0. - sigma (float, optional): Scale parameter. Larger means bigger. - If None, sigma = tau**(0.5*(2*alpha - 2.0)). Defaults to None. - radius (float, optional): Radius of the sphere. Defaults to 1.0. - grid (str, optional): Grid type. Currently supports "equiangular" and - "legendre-gauss". Defaults to "equiangular". - dtype (paddle.dtype, optional): Numerical type for the calculations. Defaults to paddle.float32. - """ - - def __init__( - self, - nlat, - alpha: float = 2.0, - tau: float = 3.0, - sigma: float = None, - radius: float = 1.0, - grid: str = "equiangular", - dtype: paddle.dtype = paddle.float32, - ): - - super().__init__() - - # Number of latitudinal modes. - self.nlat = nlat - - # Default value of sigma if None is given. - if sigma is None: - assert alpha > 1.0, f"Alpha must be greater than one, got {alpha}." - sigma = tau ** (0.5 * (2 * alpha - 2.0)) - - # Inverse SHT - self.isht = InverseRealSHT( - self.nlat, 2 * self.nlat, grid=grid, norm="backward" - ).to(dtype=dtype) - - # Square root of the eigenvalues of C. - sqrt_eig = ( - paddle.to_tensor([j * (j + 1) for j in range(self.nlat)]) - .reshape([self.nlat, 1]) - .tile([1, self.nlat + 1]) - ) - sqrt_eig = paddle.tril( - sigma * (((sqrt_eig / radius**2) + tau**2) ** (-alpha / 2.0)) - ) - sqrt_eig[0, 0] = 0.0 - sqrt_eig = sqrt_eig.unsqueeze(0) - self.register_buffer("sqrt_eig", sqrt_eig) - - # Save mean and var of the standard Gaussian. - # Need these to re-initialize distribution on a new device. - mean = paddle.to_tensor([0.0]).astype(dtype) - var = paddle.to_tensor([1.0]).astype(dtype) - self.register_buffer("mean", mean) - self.register_buffer("var", var) - - # Standard normal noise sampler. - self.gaussian_noise = paddle.distribution.Normal(self.mean, self.var) - - def forward(self, N, xi=None): - """Sample random functions from a spherical GRF. - - Args: - N (int): Number of functions to sample. - xi (paddle.Tensor, optional): Noise is a complex tensor of size (N, nlat, nlat+1). - If None, new Gaussian noise is sampled. - If xi is provided, N is ignored.. Defaults to None. - - Returns: - u (paddle.Tensor): N random samples from the GRF returned as a - tensor of size (N, nlat, 2*nlat) on a equiangular grid. - """ - - # Sample Gaussian noise. - if xi is None: - xi = self.gaussian_noise.sample((N, self.nlat, self.nlat + 1, 2)).squeeze() - xi = paddle.as_complex(xi) - - # Karhunen-Loeve expansion. - u = self.isht(xi * self.sqrt_eig) - - return u - - # Override cuda and to methods so sampler gets initialized with mean - # and variance on the correct device. - def cuda(self, *args, **kwargs): - super().cuda(*args, **kwargs) - self.gaussian_noise = paddle.distribution.Normal(self.mean, self.var) - - return self - - def to(self, *args, **kwargs): - super().to(*args, **kwargs) - self.gaussian_noise = paddle.distribution.Normal(self.mean, self.var) - - return self diff --git a/examples/smc_reac/ppsci/arch/paddle_harmonics/sht.py b/examples/smc_reac/ppsci/arch/paddle_harmonics/sht.py deleted file mode 100644 index bf5e685a04..0000000000 --- a/examples/smc_reac/ppsci/arch/paddle_harmonics/sht.py +++ /dev/null @@ -1,461 +0,0 @@ -# coding=utf-8 - -# SPDX-FileCopyrightText: Copyright (c) 2022 The torch-harmonics Authors. All rights reserved. -# SPDX-License-Identifier: BSD-3-Clause -# -# Redistribution and use in source and binary forms, with or without -# modification, are permitted provided that the following conditions are met: -# -# 1. Redistributions of source code must retain the above copyright notice, this -# list of conditions and the following disclaimer. -# -# 2. Redistributions in binary form must reproduce the above copyright notice, -# this list of conditions and the following disclaimer in the documentation -# and/or other materials provided with the distribution. -# -# 3. Neither the name of the copyright holder nor the names of its -# contributors may be used to endorse or promote products derived from -# this software without specific prior written permission. -# -# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" -# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE -# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE -# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE -# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL -# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR -# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER -# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, -# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE -# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. - -""" -Code below is heavily based on [torch-harmonics](https://github.com/NVIDIA/torch-harmonics/blob/main/torch_harmonics/sht.py) -""" - -import math - -import numpy as np -import paddle -import paddle.fft -import paddle.nn as nn - -from ppsci.arch.paddle_harmonics.legendre import _precompute_dlegpoly -from ppsci.arch.paddle_harmonics.legendre import _precompute_legpoly -from ppsci.arch.paddle_harmonics.quadrature import clenshaw_curtiss_weights -from ppsci.arch.paddle_harmonics.quadrature import legendre_gauss_weights -from ppsci.arch.paddle_harmonics.quadrature import lobatto_weights - - -class RealSHT(nn.Layer): - """ - Defines a module for computing the forward (real-valued) SHT. - Precomputes Legendre Gauss nodes, weights and associated Legendre polynomials on these nodes. - The SHT is applied to the last two dimensions of the input - - [1] Schaeffer, N. Efficient spherical harmonic transforms aimed at pseudospectral numerical simulations, G3: Geochemistry, Geophysics, Geosystems. - [2] Wang, B., Wang, L., Xie, Z.; Accurate calculation of spherical and vector spherical harmonic expansions via spectral element grids; Adv Comput Math. - - Initializes the SHT Layer, precomputing the necessary quadrature weights. - - Args: - nlat (int): Input grid resolution in the latitudinal direction. - nlon (int): Input grid resolution in the longitudinal direction. - lmax (int, optional): The max input grid resolution in the latitudinal direction. Defaults to None. - mmax (int, optional): The max input grid resolution in the longitudinal direction. Defaults to None. - grid (str, optional): Grid in the latitude direction (for now only tensor product grids are supported). - Defaults to "lobatto". - norm (str, optional): The type of normalization to use. Defaults to "ortho". - csphase (bool, optional): Whether to apply the complex-conjugate symmetry phase factor. Defaults to True. - """ - - def __init__( - self, - nlat, - nlon, - lmax=None, - mmax=None, - grid="lobatto", - norm="ortho", - csphase=True, - ): - super().__init__() - - self.nlat = nlat - self.nlon = nlon - self.grid = grid - self.norm = norm - self.csphase = csphase - - # compute quadrature points - if self.grid == "legendre-gauss": - cost, w = legendre_gauss_weights(nlat, -1, 1) - self.lmax = lmax or self.nlat - elif self.grid == "lobatto": - cost, w = lobatto_weights(nlat, -1, 1) - self.lmax = lmax or self.nlat - 1 - elif self.grid == "equiangular": - cost, w = clenshaw_curtiss_weights(nlat, -1, 1) - # cost, w = fejer2_weights(nlat, -1, 1) - self.lmax = lmax or self.nlat - else: - raise (ValueError("Unknown quadrature mode")) - - # apply cosine transform and flip them - tq = np.flip(np.arccos(cost)) - - # determine the dimensions - self.mmax = mmax or self.nlon // 2 + 1 - - # combine quadrature weights with the legendre weights - weights = paddle.to_tensor(w) - pct = _precompute_legpoly( - self.mmax, self.lmax, tq, norm=self.norm, csphase=self.csphase - ) - pct = paddle.to_tensor(pct) - self.weights = paddle.einsum("mlk,k->mlk", pct, weights) - - def extra_repr(self): - return f"nlat={self.nlat}, nlon={self.nlon},\n lmax={self.lmax}, mmax={self.mmax},\n grid={self.grid}, csphase={self.csphase}" - - def forward(self, x: paddle.Tensor): - - assert x.shape[-2] == self.nlat - assert x.shape[-1] == self.nlon - - # apply real fft in the longitudinal direction - x = ( - 2.0 - * paddle.to_tensor(math.pi) - * paddle.fft.rfft(x, axis=-1, norm="forward") - ) - - # do the Legendre-Gauss quadrature - x = paddle.as_real(x) - # distributed contraction: fork - out_shape = list(x.shape) - out_shape[-3] = self.lmax - out_shape[-2] = self.mmax - xout = paddle.zeros(out_shape, dtype=x.dtype) - - # contraction - xout[..., 0] = paddle.einsum( - "...km,mlk->...lm", x[..., : self.mmax, 0], self.weights.astype(x.dtype) - ) - xout[..., 1] = paddle.einsum( - "...km,mlk->...lm", x[..., : self.mmax, 1], self.weights.astype(x.dtype) - ) - x = paddle.as_complex(xout) - - return x - - -class InverseRealSHT(nn.Layer): - """ - Defines a module for computing the inverse (real-valued) SHT. - Precomputes Legendre Gauss nodes, weights and associated Legendre polynomials on these nodes. - nlat, nlon: Output dimensions - lmax, mmax: Input dimensions (spherical coefficients). For convenience, these are inferred from the output dimensions - - [1] Schaeffer, N. Efficient spherical harmonic transforms aimed at pseudospectral numerical simulations, G3: Geochemistry, Geophysics, Geosystems. - [2] Wang, B., Wang, L., Xie, Z.; Accurate calculation of spherical and vector spherical harmonic expansions via spectral element grids; Adv Comput Math. - """ - - def __init__( - self, - nlat, - nlon, - lmax=None, - mmax=None, - grid="lobatto", - norm="ortho", - csphase=True, - ): - - super().__init__() - - self.nlat = nlat - self.nlon = nlon - self.grid = grid - self.norm = norm - self.csphase = csphase - - # compute quadrature points - if self.grid == "legendre-gauss": - cost, _ = legendre_gauss_weights(nlat, -1, 1) - self.lmax = lmax or self.nlat - elif self.grid == "lobatto": - cost, _ = lobatto_weights(nlat, -1, 1) - self.lmax = lmax or self.nlat - 1 - elif self.grid == "equiangular": - cost, _ = clenshaw_curtiss_weights(nlat, -1, 1) - self.lmax = lmax or self.nlat - else: - raise (ValueError("Unknown quadrature mode")) - - # apply cosine transform and flip them - t = np.flip(np.arccos(cost)) - - # determine the dimensions - self.mmax = mmax or self.nlon // 2 + 1 - - pct = _precompute_legpoly( - self.mmax, self.lmax, t, norm=self.norm, inverse=True, csphase=self.csphase - ) - self.pct = paddle.to_tensor(pct) - - def extra_repr(self): - return f"nlat={self.nlat}, nlon={self.nlon},\n lmax={self.lmax}, mmax={self.mmax},\n grid={self.grid}, csphase={self.csphase}" - - def forward(self, x: paddle.Tensor): - - assert x.shape[-2] == self.lmax - assert x.shape[-1] == self.mmax - - # Evaluate associated Legendre functions on the output nodes - x = paddle.as_real(x) - - rl = paddle.einsum("...lm, mlk->...km", x[..., 0], self.pct.astype(x.dtype)) - im = paddle.einsum("...lm, mlk->...km", x[..., 1], self.pct.astype(x.dtype)) - xs = paddle.stack((rl, im), -1) - - # apply the inverse (real) FFT - x = paddle.as_complex(xs) - x = paddle.fft.irfft(x, n=self.nlon, axis=-1, norm="forward") - - return x - - -class RealVectorSHT(nn.Layer): - """ - Defines a module for computing the forward (real) vector SHT. - Precomputes Legendre Gauss nodes, weights and associated Legendre polynomials on these nodes. - The SHT is applied to the last three dimensions of the input. - - [1] Schaeffer, N. Efficient spherical harmonic transforms aimed at pseudospectral numerical simulations, G3: Geochemistry, Geophysics, Geosystems. - [2] Wang, B., Wang, L., Xie, Z.; Accurate calculation of spherical and vector spherical harmonic expansions via spectral element grids; Adv Comput Math. - - Initializes the vector SHT Layer, precomputing the necessary quadrature weights. - """ - - def __init__( - self, - nlat, - nlon, - lmax=None, - mmax=None, - grid="lobatto", - norm="ortho", - csphase=True, - ): - super().__init__() - - self.nlat = nlat - self.nlon = nlon - self.grid = grid - self.norm = norm - self.csphase = csphase - - # compute quadrature points - if self.grid == "legendre-gauss": - cost, w = legendre_gauss_weights(nlat, -1, 1) - self.lmax = lmax or self.nlat - elif self.grid == "lobatto": - cost, w = lobatto_weights(nlat, -1, 1) - self.lmax = lmax or self.nlat - 1 - elif self.grid == "equiangular": - cost, w = clenshaw_curtiss_weights(nlat, -1, 1) - # cost, w = fejer2_weights(nlat, -1, 1) - self.lmax = lmax or self.nlat - else: - raise (ValueError("Unknown quadrature mode")) - - # apply cosine transform and flip them - tq = np.flip(np.arccos(cost)) - - # determine the dimensions - self.mmax = mmax or self.nlon // 2 + 1 - - weights = paddle.to_tensor(w) - dpct = _precompute_dlegpoly( - self.mmax, self.lmax, tq, norm=self.norm, csphase=self.csphase - ) - dpct = paddle.to_tensor(dpct) - - # combine integration weights, normalization factor in to one: - l = paddle.arange(0, self.lmax) - norm_factor = 1.0 / l / (l + 1) - norm_factor[0] = 1.0 - weights = paddle.einsum("dmlk,k,l->dmlk", dpct, weights, norm_factor) - # since the second component is imaginary, we need to take complex conjugation into account - weights[1] = -1 * weights[1] - - self.weights = weights - - def extra_repr(self): - return f"nlat={self.nlat}, nlon={self.nlon},\n lmax={self.lmax}, mmax={self.mmax},\n grid={self.grid}, csphase={self.csphase}" - - def forward(self, x: paddle.Tensor): - - assert len(x.shape) >= 3 - - # apply real fft in the longitudinal direction - x = 2.0 * paddle.to_tensor(np.pi) * paddle.fft.rfft(x, axis=-1, norm="forward") - - # do the Legendre-Gauss quadrature - x = paddle.as_real(x) - - # distributed contraction: fork - out_shape = list(x.shape) - out_shape[-3] = self.lmax - out_shape[-2] = self.mmax - xout = paddle.zeros(out_shape, dtype=x.dtype) - - # contraction - spheroidal component - # real component - xout[..., 0, :, :, 0] = paddle.einsum( - "...km,mlk->...lm", - x[..., 0, :, : self.mmax, 0], - self.weights[0].astype(x.dtype), - ) - paddle.einsum( - "...km,mlk->...lm", - x[..., 1, :, : self.mmax, 1], - self.weights[1].astype(x.dtype), - ) - - # iamg component - xout[..., 0, :, :, 1] = paddle.einsum( - "...km,mlk->...lm", - x[..., 0, :, : self.mmax, 1], - self.weights[0].astype(x.dtype), - ) + paddle.einsum( - "...km,mlk->...lm", - x[..., 1, :, : self.mmax, 0], - self.weights[1].astype(x.dtype), - ) - - # contraction - toroidal component - # real component - xout[..., 1, :, :, 0] = -paddle.einsum( - "...km,mlk->...lm", - x[..., 0, :, : self.mmax, 1], - self.weights[1].astype(x.dtype), - ) - paddle.einsum( - "...km,mlk->...lm", - x[..., 1, :, : self.mmax, 0], - self.weights[0].astype(x.dtype), - ) - # imag component - xout[..., 1, :, :, 1] = paddle.einsum( - "...km,mlk->...lm", - x[..., 0, :, : self.mmax, 0], - self.weights[1].astype(x.dtype), - ) - paddle.einsum( - "...km,mlk->...lm", - x[..., 1, :, : self.mmax, 1], - self.weights[0].astype(x.dtype), - ) - - return paddle.as_complex(xout) - - -class InverseRealVectorSHT(nn.Layer): - """ - Defines a module for computing the inverse (real-valued) vector SHT. - Precomputes Legendre Gauss nodes, weights and associated Legendre polynomials on these nodes. - - [1] Schaeffer, N. Efficient spherical harmonic transforms aimed at pseudospectral numerical simulations, G3: Geochemistry, Geophysics, Geosystems. - [2] Wang, B., Wang, L., Xie, Z.; Accurate calculation of spherical and vector spherical harmonic expansions via spectral element grids; Adv Comput Math. - """ - - def __init__( - self, - nlat, - nlon, - lmax=None, - mmax=None, - grid="lobatto", - norm="ortho", - csphase=True, - ): - - super().__init__() - - self.nlat = nlat - self.nlon = nlon - self.grid = grid - self.norm = norm - self.csphase = csphase - - # compute quadrature points - if self.grid == "legendre-gauss": - cost, _ = legendre_gauss_weights(nlat, -1, 1) - self.lmax = lmax or self.nlat - elif self.grid == "lobatto": - cost, _ = lobatto_weights(nlat, -1, 1) - self.lmax = lmax or self.nlat - 1 - elif self.grid == "equiangular": - cost, _ = clenshaw_curtiss_weights(nlat, -1, 1) - self.lmax = lmax or self.nlat - else: - raise (ValueError("Unknown quadrature mode")) - - # apply cosine transform and flip them - t = np.flip(np.arccos(cost)) - - # determine the dimensions - self.mmax = mmax or self.nlon // 2 + 1 - - dpct = _precompute_dlegpoly( - self.mmax, self.lmax, t, norm=self.norm, inverse=True, csphase=self.csphase - ) - self.dpct = paddle.to_tensor(dpct) - - def extra_repr(self): - return f"nlat={self.nlat}, nlon={self.nlon},\n lmax={self.lmax}, mmax={self.mmax},\n grid={self.grid}, csphase={self.csphase}" - - def forward(self, x: paddle.Tensor): - - assert x.shape[-2] == self.lmax - assert x.shape[-1] == self.mmax - - # Evaluate associated Legendre functions on the output nodes - x = paddle.as_real(x) - - # contraction - spheroidal component - # real component - srl = paddle.einsum( - "...lm,mlk->...km", x[..., 0, :, :, 0], self.dpct[0].astype(x.dtype) - ) - paddle.einsum( - "...lm,mlk->...km", x[..., 1, :, :, 1], self.dpct[1].astype(x.dtype) - ) - # iamg component - sim = paddle.einsum( - "...lm,mlk->...km", x[..., 0, :, :, 1], self.dpct[0].astype(x.dtype) - ) + paddle.einsum( - "...lm,mlk->...km", x[..., 1, :, :, 0], self.dpct[1].astype(x.dtype) - ) - - # contraction - toroidal component - # real component - trl = -paddle.einsum( - "...lm,mlk->...km", x[..., 0, :, :, 1], self.dpct[1].astype(x.dtype) - ) - paddle.einsum( - "...lm,mlk->...km", x[..., 1, :, :, 0], self.dpct[0].astype(x.dtype) - ) - # imag component - tim = paddle.einsum( - "...lm,mlk->...km", x[..., 0, :, :, 0], self.dpct[1].astype(x.dtype) - ) - paddle.einsum( - "...lm,mlk->...km", x[..., 1, :, :, 1], self.dpct[0].astype(x.dtype) - ) - - # reassemble - s = paddle.stack((srl, sim), -1) - t = paddle.stack((trl, tim), -1) - xs = paddle.stack((s, t), -4) - - # apply the inverse (real) FFT - x = paddle.as_complex(xs) - x = paddle.fft.irfft(x, n=self.nlon, axis=-1, norm="forward") - - return x diff --git a/examples/smc_reac/ppsci/arch/phycrnet.py b/examples/smc_reac/ppsci/arch/phycrnet.py deleted file mode 100644 index c72583ebf9..0000000000 --- a/examples/smc_reac/ppsci/arch/phycrnet.py +++ /dev/null @@ -1,540 +0,0 @@ -# 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. - -from typing import Tuple - -import numpy as np -import paddle -import paddle.nn as nn -from paddle.nn import utils - -from ppsci.arch import base - -# define the high-order finite difference kernels -LALP_OP = [ - [ - [ - [0, 0, -1 / 12, 0, 0], - [0, 0, 4 / 3, 0, 0], - [-1 / 12, 4 / 3, -5, 4 / 3, -1 / 12], - [0, 0, 4 / 3, 0, 0], - [0, 0, -1 / 12, 0, 0], - ] - ] -] - -PARTIAL_Y = [ - [ - [ - [0, 0, 0, 0, 0], - [0, 0, 0, 0, 0], - [1 / 12, -8 / 12, 0, 8 / 12, -1 / 12], - [0, 0, 0, 0, 0], - [0, 0, 0, 0, 0], - ] - ] -] - -PARTIAL_X = [ - [ - [ - [0, 0, 1 / 12, 0, 0], - [0, 0, -8 / 12, 0, 0], - [0, 0, 0, 0, 0], - [0, 0, 8 / 12, 0, 0], - [0, 0, -1 / 12, 0, 0], - ] - ] -] - - -# specific parameters for burgers equation -def _initialize_weights(module): - if isinstance(module, nn.Conv2D): - c = 1.0 # 0.5 - initializer = nn.initializer.Uniform( - -c * np.sqrt(1 / (3 * 3 * 320)), c * np.sqrt(1 / (3 * 3 * 320)) - ) - initializer(module.weight) - elif isinstance(module, nn.Linear): - initializer = nn.initializer.Constant(0.0) - initializer(module.bias) - - -class PhyCRNet(base.Arch): - """Physics-informed convolutional-recurrent neural networks. - - Args: - input_channels (int): The input channels. - hidden_channels (Tuple[int, ...]): The hidden channels. - input_kernel_size (Tuple[int, ...]): The input kernel size(s). - input_stride (Tuple[int, ...]): The input stride(s). - input_padding (Tuple[int, ...]): The input padding(s). - dt (float): The dt parameter. - num_layers (Tuple[int, ...]): The number of layers. - upscale_factor (int): The upscale factor. - step (int, optional): The step(s). Defaults to 1. - effective_step (Tuple[int, ...], optional): The effective step. Defaults to (1, ). - - Examples: - >>> import ppsci - >>> model = ppsci.arch.PhyCRNet( - ... input_channels=2, - ... hidden_channels=[8, 32, 128, 128], - ... input_kernel_size=[4, 4, 4, 3], - ... input_stride=[2, 2, 2, 1], - ... input_padding=[1, 1, 1, 1], - ... dt=0.002, - ... num_layers=[3, 1], - ... upscale_factor=8 - ... ) - """ - - def __init__( - self, - input_channels: int, - hidden_channels: Tuple[int, ...], - input_kernel_size: Tuple[int, ...], - input_stride: Tuple[int, ...], - input_padding: Tuple[int, ...], - dt: float, - num_layers: Tuple[int, ...], - upscale_factor: int, - step: int = 1, - effective_step: Tuple[int, ...] = (1,), - ): - super(PhyCRNet, self).__init__() - - # input channels of layer includes input_channels and hidden_channels of cells - self.input_channels = [input_channels] + hidden_channels - self.hidden_channels = hidden_channels - self.input_kernel_size = input_kernel_size - self.input_stride = input_stride - self.input_padding = input_padding - self.step = step - self.effective_step = effective_step - self._all_layers = [] - self.dt = dt - self.upscale_factor = upscale_factor - - # number of layers - self.num_encoder = num_layers[0] - self.num_convlstm = num_layers[1] - - # encoder - downsampling - self.encoder = nn.LayerList( - [ - encoder_block( - input_channels=self.input_channels[i], - hidden_channels=self.hidden_channels[i], - input_kernel_size=self.input_kernel_size[i], - input_stride=self.input_stride[i], - input_padding=self.input_padding[i], - ) - for i in range(self.num_encoder) - ] - ) - - # ConvLSTM - self.convlstm = nn.LayerList( - [ - ConvLSTMCell( - input_channels=self.input_channels[i], - hidden_channels=self.hidden_channels[i], - input_kernel_size=self.input_kernel_size[i], - input_stride=self.input_stride[i], - input_padding=self.input_padding[i], - ) - for i in range(self.num_encoder, self.num_encoder + self.num_convlstm) - ] - ) - - # output layer - self.output_layer = nn.Conv2D( - 2, 2, kernel_size=5, stride=1, padding=2, padding_mode="circular" - ) - - # pixelshuffle - upscale - self.pixelshuffle = nn.PixelShuffle(self.upscale_factor) - - # initialize weights - self.apply(_initialize_weights) - initializer_0 = nn.initializer.Constant(0.0) - initializer_0(self.output_layer.bias) - self.enable_transform = True - - def forward(self, x): - if self.enable_transform: - if self._input_transform is not None: - x = self._input_transform(x) - output_x = x - - self.initial_state = x["initial_state"] - x = x["input"] - internal_state = [] - outputs = [] - second_last_state = [] - - for step in range(self.step): - xt = x - - # encoder - for encoder in self.encoder: - x = encoder(x) - - # convlstm - for i, lstm in enumerate(self.convlstm, self.num_encoder): - if step == 0: - (h, c) = lstm.init_hidden_tensor( - prev_state=self.initial_state[i - self.num_encoder] - ) - internal_state.append((h, c)) - - # one-step forward - (h, c) = internal_state[i - self.num_encoder] - x, new_c = lstm(x, h, c) - internal_state[i - self.num_encoder] = (x, new_c) - - # output - x = self.pixelshuffle(x) - x = self.output_layer(x) - - # residual connection - x = xt + self.dt * x - - if step == (self.step - 2): - second_last_state = internal_state.copy() - - if step in self.effective_step: - outputs.append(x) - - result_dict = {"outputs": outputs, "second_last_state": second_last_state} - if self.enable_transform: - if self._output_transform is not None: - result_dict = self._output_transform(output_x, result_dict) - return result_dict - - -class ConvLSTMCell(nn.Layer): - """Convolutional LSTM""" - - def __init__( - self, - input_channels, - hidden_channels, - input_kernel_size, - input_stride, - input_padding, - hidden_kernel_size=3, - num_features=4, - ): - super(ConvLSTMCell, self).__init__() - - self.input_channels = input_channels - self.hidden_channels = hidden_channels - self.hidden_kernel_size = hidden_kernel_size # Page 9, The convolutional operations in ConvLSTM have 3x3 kernels. - self.input_kernel_size = input_kernel_size - self.input_stride = input_stride - self.input_padding = input_padding - self.num_features = ( - num_features # Page 10, block of different dense layers {4, 3, 4} - ) - - # padding for hidden state - self.padding = int((self.hidden_kernel_size - 1) / 2) - - self.Wxi = nn.Conv2D( - self.input_channels, - self.hidden_channels, - self.input_kernel_size, - self.input_stride, - self.input_padding, - bias_attr=None, - padding_mode="circular", - ) - - self.Whi = nn.Conv2D( - self.hidden_channels, - self.hidden_channels, - self.hidden_kernel_size, - 1, - padding=1, - bias_attr=False, - padding_mode="circular", - ) - - self.Wxf = nn.Conv2D( - self.input_channels, - self.hidden_channels, - self.input_kernel_size, - self.input_stride, - self.input_padding, - bias_attr=None, - padding_mode="circular", - ) - - self.Whf = nn.Conv2D( - self.hidden_channels, - self.hidden_channels, - self.hidden_kernel_size, - 1, - padding=1, - bias_attr=False, - padding_mode="circular", - ) - - self.Wxc = nn.Conv2D( - self.input_channels, - self.hidden_channels, - self.input_kernel_size, - self.input_stride, - self.input_padding, - bias_attr=None, - padding_mode="circular", - ) - - self.Whc = nn.Conv2D( - self.hidden_channels, - self.hidden_channels, - self.hidden_kernel_size, - 1, - padding=1, - bias_attr=False, - padding_mode="circular", - ) - - self.Wxo = nn.Conv2D( - self.input_channels, - self.hidden_channels, - self.input_kernel_size, - self.input_stride, - self.input_padding, - bias_attr=None, - padding_mode="circular", - ) - - self.Who = nn.Conv2D( - self.hidden_channels, - self.hidden_channels, - self.hidden_kernel_size, - 1, - padding=1, - bias_attr=False, - padding_mode="circular", - ) - - initializer_0 = nn.initializer.Constant(0.0) - initializer_1 = nn.initializer.Constant(1.0) - - initializer_0(self.Wxi.bias) - initializer_0(self.Wxf.bias) - initializer_0(self.Wxc.bias) - initializer_1(self.Wxo.bias) - - def forward(self, x, h, c): - ci = nn.functional.sigmoid(self.Wxi(x) + self.Whi(h)) - cf = nn.functional.sigmoid(self.Wxf(x) + self.Whf(h)) - cc = cf * c + ci * paddle.tanh(self.Wxc(x) + self.Whc(h)) - co = nn.functional.sigmoid(self.Wxo(x) + self.Who(h)) - ch = co * paddle.tanh(cc) - return ch, cc - - def init_hidden_tensor(self, prev_state): - return ((prev_state[0]).cuda(), (prev_state[1]).cuda()) - - -class encoder_block(nn.Layer): - """Encoder with CNN""" - - def __init__( - self, - input_channels, - hidden_channels, - input_kernel_size, - input_stride, - input_padding, - ): - super(encoder_block, self).__init__() - - self.input_channels = input_channels - self.hidden_channels = hidden_channels - self.input_kernel_size = input_kernel_size - self.input_stride = input_stride - self.input_padding = input_padding - - self.conv = utils.weight_norm( - nn.Conv2D( - self.input_channels, - self.hidden_channels, - self.input_kernel_size, - self.input_stride, - self.input_padding, - bias_attr=None, - padding_mode="circular", - ) - ) - - self.act = nn.ReLU() - - initializer_0 = nn.initializer.Constant(0.0) - initializer_0(self.conv.bias) - - def forward(self, x): - return self.act(self.conv(x)) - - -class Conv2DDerivative(nn.Layer): - def __init__(self, der_filter, resol, kernel_size=3, name=""): - super(Conv2DDerivative, self).__init__() - - self.resol = resol # constant in the finite difference - self.name = name - self.input_channels = 1 - self.output_channels = 1 - self.kernel_size = kernel_size - - self.padding = int((kernel_size - 1) / 2) - self.filter = nn.Conv2D( - self.input_channels, - self.output_channels, - self.kernel_size, - 1, - padding=0, - bias_attr=False, - ) - - # Fixed gradient operator - self.filter.weight = self.create_parameter( - shape=self.filter.weight.shape, - dtype=self.filter.weight.dtype, - default_initializer=nn.initializer.Assign( - paddle.to_tensor( - der_filter, dtype=paddle.get_default_dtype(), stop_gradient=True - ) - ), - ) - self.filter.weight.stop_gradient = True - - def forward(self, input): - derivative = self.filter(input) - return derivative / self.resol - - -class Conv1DDerivative(nn.Layer): - def __init__(self, der_filter, resol, kernel_size=3, name=""): - super(Conv1DDerivative, self).__init__() - - self.resol = resol # $\delta$*constant in the finite difference - self.name = name - self.input_channels = 1 - self.output_channels = 1 - self.kernel_size = kernel_size - - self.padding = int((kernel_size - 1) / 2) - self.filter = nn.Conv1D( - self.input_channels, - self.output_channels, - self.kernel_size, - 1, - padding=0, - bias_attr=False, - ) - - # Fixed gradient operator - self.filter.weight = self.create_parameter( - shape=self.filter.weight.shape, - dtype=self.filter.weight.dtype, - default_initializer=nn.initializer.Assign( - paddle.to_tensor( - der_filter, dtype=paddle.get_default_dtype(), stop_gradient=True - ) - ), - ) - self.filter.weight.stop_gradient = True - - def forward(self, input): - derivative = self.filter(input) - return derivative / self.resol - - -class loss_generator(nn.Layer): - """Loss generator for physics loss""" - - def __init__(self, dt, dx): - """Construct the derivatives, X = Width, Y = Height""" - super(loss_generator, self).__init__() - - # spatial derivative operator - self.laplace = Conv2DDerivative( - der_filter=LALP_OP, resol=(dx**2), kernel_size=5, name="laplace_operator" - ) - - self.dx = Conv2DDerivative( - der_filter=PARTIAL_X, resol=(dx * 1), kernel_size=5, name="dx_operator" - ) - - self.dy = Conv2DDerivative( - der_filter=PARTIAL_Y, resol=(dx * 1), kernel_size=5, name="dy_operator" - ) - - # temporal derivative operator - self.dt = Conv1DDerivative( - der_filter=[[[-1, 0, 1]]], resol=(dt * 2), kernel_size=3, name="partial_t" - ) - - def get_phy_Loss(self, output): - # spatial derivatives - laplace_u = self.laplace(output[1:-1, 0:1, :, :]) # [t,c,h,w] - laplace_v = self.laplace(output[1:-1, 1:2, :, :]) - - u_x = self.dx(output[1:-1, 0:1, :, :]) - u_y = self.dy(output[1:-1, 0:1, :, :]) - v_x = self.dx(output[1:-1, 1:2, :, :]) - v_y = self.dy(output[1:-1, 1:2, :, :]) - - # temporal derivative - u - u = output[:, 0:1, 2:-2, 2:-2] - lent = u.shape[0] - lenx = u.shape[3] - leny = u.shape[2] - u_conv1d = u.transpose((2, 3, 1, 0)) # [height(Y), width(X), c, step] - u_conv1d = u_conv1d.reshape((lenx * leny, 1, lent)) - u_t = self.dt(u_conv1d) # lent-2 due to no-padding - u_t = u_t.reshape((leny, lenx, 1, lent - 2)) - u_t = u_t.transpose((3, 2, 0, 1)) # [step-2, c, height(Y), width(X)] - - # temporal derivative - v - v = output[:, 1:2, 2:-2, 2:-2] - v_conv1d = v.transpose((2, 3, 1, 0)) # [height(Y), width(X), c, step] - v_conv1d = v_conv1d.reshape((lenx * leny, 1, lent)) - v_t = self.dt(v_conv1d) # lent-2 due to no-padding - v_t = v_t.reshape((leny, lenx, 1, lent - 2)) - v_t = v_t.transpose((3, 2, 0, 1)) # [step-2, c, height(Y), width(X)] - - u = output[1:-1, 0:1, 2:-2, 2:-2] # [t, c, height(Y), width(X)] - v = output[1:-1, 1:2, 2:-2, 2:-2] # [t, c, height(Y), width(X)] - - assert laplace_u.shape == u_t.shape - assert u_t.shape == v_t.shape - assert laplace_u.shape == u.shape - assert laplace_v.shape == v.shape - - # Reynolds number - R = 200.0 - - # 2D burgers eqn - f_u = u_t + u * u_x + v * u_y - (1 / R) * laplace_u - f_v = v_t + u * v_x + v * v_y - (1 / R) * laplace_v - - return f_u, f_v diff --git a/examples/smc_reac/ppsci/arch/phylstm.py b/examples/smc_reac/ppsci/arch/phylstm.py deleted file mode 100644 index a840d3f7ad..0000000000 --- a/examples/smc_reac/ppsci/arch/phylstm.py +++ /dev/null @@ -1,239 +0,0 @@ -# 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. - -import paddle -import paddle.nn as nn - -from ppsci.arch import base - - -class DeepPhyLSTM(base.Arch): - """DeepPhyLSTM init function. - - Args: - input_size (int): The input size. - output_size (int): The output size. - hidden_size (int, optional): The hidden size. Defaults to 100. - model_type (int, optional): The model type, value is 2 or 3, 2 indicates having two sub-models, 3 indicates having three submodels. Defaults to 2. - - Examples: - >>> import paddle - >>> import ppsci - >>> # model_type is `2` - >>> model = ppsci.arch.DeepPhyLSTM( - ... input_size=16, - ... output_size=1, - ... hidden_size=100, - ... model_type=2) - >>> out = model( - ... {"ag":paddle.rand([64, 16, 16]), - ... "ag_c":paddle.rand([64, 16, 16]), - ... "phi":paddle.rand([1, 16, 16])}) - >>> for k, v in out.items(): - ... print(f"{k} {v.dtype} {v.shape}") - eta_pred paddle.float32 [64, 16, 1] - eta_dot_pred paddle.float32 [64, 16, 1] - g_pred paddle.float32 [64, 16, 1] - eta_t_pred_c paddle.float32 [64, 16, 1] - eta_dot_pred_c paddle.float32 [64, 16, 1] - lift_pred_c paddle.float32 [64, 16, 1] - >>> # model_type is `3` - >>> model = ppsci.arch.DeepPhyLSTM( - ... input_size=16, - ... output_size=1, - ... hidden_size=100, - ... model_type=3) - >>> out = model( - ... {"ag":paddle.rand([64, 16, 1]), - ... "ag_c":paddle.rand([64, 16, 1]), - ... "phi":paddle.rand([1, 16, 16])}) - >>> for k, v in out.items(): - ... print(f"{k} {v.dtype} {v.shape}") - eta_pred paddle.float32 [64, 16, 1] - eta_dot_pred paddle.float32 [64, 16, 1] - g_pred paddle.float32 [64, 16, 1] - eta_t_pred_c paddle.float32 [64, 16, 1] - eta_dot_pred_c paddle.float32 [64, 16, 1] - lift_pred_c paddle.float32 [64, 16, 1] - g_t_pred_c paddle.float32 [64, 16, 1] - g_dot_pred_c paddle.float32 [64, 16, 1] - """ - - def __init__(self, input_size, output_size, hidden_size=100, model_type=2): - super().__init__() - self.input_size = input_size - self.output_size = output_size - self.hidden_size = hidden_size - self.model_type = model_type - - if self.model_type == 2: - self.lstm_model = nn.Sequential( - nn.LSTM(input_size, hidden_size), - nn.ReLU(), - nn.LSTM(hidden_size, hidden_size), - nn.ReLU(), - nn.LSTM(hidden_size, hidden_size), - nn.ReLU(), - nn.Linear(hidden_size, hidden_size), - nn.Linear(hidden_size, 3 * output_size), - ) - - self.lstm_model_f = nn.Sequential( - nn.LSTM(3 * output_size, hidden_size), - nn.ReLU(), - nn.LSTM(hidden_size, hidden_size), - nn.ReLU(), - nn.LSTM(hidden_size, hidden_size), - nn.ReLU(), - nn.Linear(hidden_size, hidden_size), - nn.Linear(hidden_size, output_size), - ) - elif self.model_type == 3: - self.lstm_model = nn.Sequential( - nn.LSTM(1, hidden_size), - nn.ReLU(), - nn.LSTM(hidden_size, hidden_size), - nn.ReLU(), - nn.LSTM(hidden_size, hidden_size), - nn.ReLU(), - nn.Linear(hidden_size, 3 * output_size), - ) - - self.lstm_model_f = nn.Sequential( - nn.LSTM(3 * output_size, hidden_size), - nn.ReLU(), - nn.LSTM(hidden_size, hidden_size), - nn.ReLU(), - nn.LSTM(hidden_size, hidden_size), - nn.ReLU(), - nn.Linear(hidden_size, output_size), - ) - - self.lstm_model_g = nn.Sequential( - nn.LSTM(2 * output_size, hidden_size), - nn.ReLU(), - nn.LSTM(hidden_size, hidden_size), - nn.ReLU(), - nn.LSTM(hidden_size, hidden_size), - nn.ReLU(), - nn.Linear(hidden_size, output_size), - ) - else: - raise ValueError(f"model_type should be 2 or 3, but got {model_type}") - - def forward(self, x): - if self._input_transform is not None: - x = self._input_transform(x) - - if self.model_type == 2: - result_dict = self._forward_type_2(x) - elif self.model_type == 3: - result_dict = self._forward_type_3(x) - if self._output_transform is not None: - result_dict = self._output_transform(x, result_dict) - return result_dict - - def _forward_type_2(self, x): - output = x["ag"] - for layer in self.lstm_model: - output = layer(output) - if isinstance(output, tuple): - output = output[0] - - eta_pred = output[:, :, 0 : self.output_size] - eta_dot_pred = output[:, :, self.output_size : 2 * self.output_size] - g_pred = output[:, :, 2 * self.output_size :] - - # for ag_c - output_c = x["ag_c"] - for layer in self.lstm_model: - output_c = layer(output_c) - if isinstance(output_c, tuple): - output_c = output_c[0] - - eta_pred_c = output_c[:, :, 0 : self.output_size] - eta_dot_pred_c = output_c[:, :, self.output_size : 2 * self.output_size] - g_pred_c = output_c[:, :, 2 * self.output_size :] - eta_t_pred_c = paddle.matmul(x["phi"], eta_pred_c) - eta_tt_pred_c = paddle.matmul(x["phi"], eta_dot_pred_c) - eta_dot1_pred_c = eta_dot_pred_c[:, :, 0:1] - tmp = paddle.concat([eta_pred_c, eta_dot1_pred_c, g_pred_c], 2) - f = tmp - for layer in self.lstm_model_f: - f = layer(f) - if isinstance(f, tuple): - f = f[0] - - lift_pred_c = eta_tt_pred_c + f - - return { - "eta_pred": eta_pred, - "eta_dot_pred": eta_dot_pred, - "g_pred": g_pred, - "eta_t_pred_c": eta_t_pred_c, - "eta_dot_pred_c": eta_dot_pred_c, - "lift_pred_c": lift_pred_c, - } - - def _forward_type_3(self, x): - # physics informed neural networks - output = x["ag"] - for layer in self.lstm_model: - output = layer(output) - if isinstance(output, tuple): - output = output[0] - - eta_pred = output[:, :, 0 : self.output_size] - eta_dot_pred = output[:, :, self.output_size : 2 * self.output_size] - g_pred = output[:, :, 2 * self.output_size :] - - output_c = x["ag_c"] - for layer in self.lstm_model: - output_c = layer(output_c) - if isinstance(output_c, tuple): - output_c = output_c[0] - - eta_pred_c = output_c[:, :, 0 : self.output_size] - eta_dot_pred_c = output_c[:, :, self.output_size : 2 * self.output_size] - g_pred_c = output_c[:, :, 2 * self.output_size :] - - eta_t_pred_c = paddle.matmul(x["phi"], eta_pred_c) - eta_tt_pred_c = paddle.matmul(x["phi"], eta_dot_pred_c) - g_t_pred_c = paddle.matmul(x["phi"], g_pred_c) - - f = paddle.concat([eta_pred_c, eta_dot_pred_c, g_pred_c], 2) - for layer in self.lstm_model_f: - f = layer(f) - if isinstance(f, tuple): - f = f[0] - - lift_pred_c = eta_tt_pred_c + f - - eta_dot1_pred_c = eta_dot_pred_c[:, :, 0:1] - g_dot_pred_c = paddle.concat([eta_dot1_pred_c, g_pred_c], 2) - for layer in self.lstm_model_g: - g_dot_pred_c = layer(g_dot_pred_c) - if isinstance(g_dot_pred_c, tuple): - g_dot_pred_c = g_dot_pred_c[0] - - return { - "eta_pred": eta_pred, - "eta_dot_pred": eta_dot_pred, - "g_pred": g_pred, - "eta_t_pred_c": eta_t_pred_c, - "eta_dot_pred_c": eta_dot_pred_c, - "lift_pred_c": lift_pred_c, - "g_t_pred_c": g_t_pred_c, - "g_dot_pred_c": g_dot_pred_c, - } diff --git a/examples/smc_reac/ppsci/arch/physx_transformer.py b/examples/smc_reac/ppsci/arch/physx_transformer.py deleted file mode 100644 index 267fb458c6..0000000000 --- a/examples/smc_reac/ppsci/arch/physx_transformer.py +++ /dev/null @@ -1,407 +0,0 @@ -# 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. - -""" -Code below is heavily based on [transformer-physx](https://github.com/zabaras/transformer-physx) -""" - -from __future__ import annotations - -from typing import Optional -from typing import Tuple - -import paddle -import paddle.nn.functional as F -from paddle import nn -from paddle.nn.initializer import Constant -from paddle.nn.initializer import Normal - -from ppsci.arch import base -from ppsci.arch.embedding_koopman import CylinderEmbedding - -zeros_ = Constant(value=0.0) -ones_ = Constant(value=1.0) - - -class MaskedAttention(nn.Layer): - """Masked self-attention module. - - Args: - embed_dim (int): The expected feature size in the input and output. - num_ctx (int): Context length of block. - num_heads (int): The number of heads in multi-head attention. - attn_drop (float, optional): The dropout probability used on attention - weights to drop some attention targets. Defaults to 0. - proj_drop (float, optional): The dropout probability used on output. Defaults to 0. - scale (bool, optional): Whether to scale attention weights. Defaults to False. - """ - - def __init__( - self, - embed_dim: int, - num_ctx: int, - num_heads: int, - attn_drop: float = 0.0, - proj_drop: float = 0.0, - scale: bool = False, - ): - super().__init__() - self.register_buffer( - "bias", - paddle.tril(paddle.ones((num_ctx, num_ctx), dtype="int32")).reshape( - [1, 1, num_ctx, num_ctx] - ), - ) - - self.register_buffer("masked_bias", paddle.to_tensor(-1e4)) - self.num_heads = num_heads - self.split_size = embed_dim - self.scale = scale - - self.qkv_proj = nn.Linear(embed_dim, embed_dim * 3) - self.out_proj = nn.Linear(embed_dim, embed_dim) - self.attn_drop = nn.Dropout(attn_drop) - self.proj_drop = nn.Dropout(proj_drop) - - def _attn( - self, - query, - key, - value, - attention_mask=None, - head_mask=None, - output_attentions=False, - ): - attn = paddle.matmul(query, key) - if self.scale: - attn = attn / (float(value.shape[-1]) ** 0.5) - - nd, ns = attn.shape[-2], attn.shape[-1] - mask = self.bias[:, :, ns - nd : ns, :ns] - attn = paddle.where(mask > 0, attn, self.masked_bias.cast(attn.dtype)) - - if attention_mask is not None: - attn = attn + attention_mask - - attn = F.softmax(attn, axis=-1) - attn = self.attn_drop(attn) - - if head_mask is not None: - attn = attn * head_mask - - outputs = [paddle.matmul(attn, value)] - if output_attentions: - outputs.append(attn) - return outputs - - def merge_heads(self, x): - x = x.transpose([0, 2, 1, 3]) - new_x_shape = x.shape[:-2] + [ - x.shape[-2] * x.shape[-1], - ] - return x.reshape(new_x_shape) - - def split_heads(self, x, k=False): - new_x_shape = x.shape[:-1] + [self.num_heads, x.shape[-1] // self.num_heads] - x = x.reshape(new_x_shape) - if k: - return x.transpose([0, 2, 3, 1]) - return x.transpose([0, 2, 1, 3]) - - def forward( - self, - x, - layer_past=None, - attention_mask=None, - head_mask=None, - output_attentions=False, - ): - x = self.qkv_proj(x) - query, key, value = x.split(x.shape[2] // self.split_size, axis=2) - query = self.split_heads(query) - key = self.split_heads(key, k=True) - value = self.split_heads(value) - # Concat previous key and value tensors - if layer_past is not None: - past_key, past_value = layer_past[0].transpose([0, 1, 3, 2]), layer_past[1] - key = paddle.concat((past_key, key), axis=-1) - value = paddle.concat((past_value, value), axis=-2) - - attn_outputs = self._attn( - query, key, value, attention_mask, head_mask, output_attentions - ) - output = attn_outputs[0] - output = self.merge_heads(output) - output = self.out_proj(output) - output = self.proj_drop(output) - - outputs = [output] + attn_outputs[1:] - return outputs - - -class MLP(nn.Layer): - """Multi layer perceptron module used in Transformer. - - Args: - in_features (int): Number of the input features. - hidden_features (Optional[int]): Number of the hidden size. Defaults to None. - out_features (Optional[int]): Number of the output features. Defaults to None. - drop (float, optional): Probability of dropout the units. Defaults to 0. - """ - - def __init__( - self, - in_features: int, - hidden_features: Optional[int] = None, - out_features: Optional[int] = None, - drop: float = 0.0, - ): - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features - - self.fc1 = nn.Linear(in_features, hidden_features) - self.act = nn.GELU(approximate=True) - self.fc2 = nn.Linear(hidden_features, out_features) - self.drop = nn.Dropout(drop) - - def forward(self, x): - x = self.fc1(x) - x = self.act(x) - x = self.fc2(x) - x = self.drop(x) - return x - - -class Block(nn.Layer): - """Transformer decoder block consisting of layer norm, - masked self-attention, layer norm and fully connected layer. - - Args: - num_ctx (int): Context length of block - embed_size (int): The number of embedding size. - num_heads (int): The number of heads in multi-head attention. - attn_pdrop (float): The dropout probability used on attention - weights to drop some attention targets. - resid_pdrop (float): The dropout probability used on output. - scale (bool, optional): Scaled self-attention calculation. Defaults to False. - """ - - def __init__( - self, - num_ctx: int, - embed_size: int, - num_heads: int, - attn_pdrop: float, - resid_pdrop: float, - scale: bool = False, - ): - super().__init__() - self.ln_1 = nn.LayerNorm(embed_size) - self.attn = MaskedAttention( - embed_size, num_ctx, num_heads, attn_pdrop, resid_pdrop, scale - ) - self.ln_2 = nn.LayerNorm(embed_size) - self.mlp = MLP(embed_size, 4 * embed_size, resid_pdrop) - - def forward( - self, - x, - layer_past=None, - attention_mask=None, - head_mask=None, - output_attentions=False, - ): - # Evaluate attention heads - output_attn = self.attn.forward( - self.ln_1(x), - layer_past=layer_past, - attention_mask=attention_mask, - head_mask=head_mask, - output_attentions=output_attentions, - ) - x = x + output_attn[0] - m = self.mlp(self.ln_2(x)) - x = x + m - outputs = [x] + output_attn[1:] - return outputs - - -class PhysformerGPT2(base.Arch): - """Transformer decoder model for modeling physics. - - Args: - input_keys (Tuple[str, ...]): Input keys, such as ("embeds",). - output_keys (Tuple[str, ...]): Output keys, such as ("pred_embeds",). - num_layers (int): Number of transformer layers. - num_ctx (int): Context length of block. - embed_size (int): The number of embedding size. - num_heads (int): The number of heads in multi-head attention. - embd_pdrop (float, optional): The dropout probability used on embedding features. Defaults to 0.0. - attn_pdrop (float, optional): The dropout probability used on attention weights. Defaults to 0.0. - resid_pdrop (float, optional): The dropout probability used on block outputs. Defaults to 0.0. - initializer_range (float, optional): Initializer range of linear layer. Defaults to 0.05. - embedding_model (Optional[base.Arch]): Embedding model, If this parameter is set, - the embedding model will map the input data to the embedding space and the - output data to the physical space. Defaults to None. - - Examples: - >>> import paddle - >>> import ppsci - >>> model = ppsci.arch.PhysformerGPT2(("embeds", ), ("pred_embeds", ), 6, 16, 128, 4) - >>> data = paddle.to_tensor(paddle.randn([10, 16, 128])) - >>> inputs = {"embeds": data} - >>> outputs = model(inputs) - >>> print(outputs["pred_embeds"].shape) - [10, 16, 128] - """ - - def __init__( - self, - input_keys: Tuple[str, ...], - output_keys: Tuple[str, ...], - num_layers: int, - num_ctx: int, - embed_size: int, - num_heads: int, - embd_pdrop: float = 0.0, - attn_pdrop: float = 0.0, - resid_pdrop: float = 0.0, - initializer_range: float = 0.05, - embedding_model: Optional[base.Arch] = None, - ): - super().__init__() - self.input_keys = input_keys - self.output_keys = output_keys - - self.num_layers = num_layers - self.num_ctx = num_ctx - self.embed_size = embed_size - self.num_heads = num_heads - self.embd_pdrop = embd_pdrop - self.attn_pdrop = attn_pdrop - self.resid_pdrop = resid_pdrop - self.initializer_range = initializer_range - - self.drop = nn.Dropout(embd_pdrop) - self.blocks = nn.LayerList( - [ - Block( - num_ctx, embed_size, num_heads, attn_pdrop, resid_pdrop, scale=True - ) - for _ in range(num_layers) - ] - ) - self.ln = nn.LayerNorm(embed_size) - self.linear = nn.Linear(embed_size, embed_size) - - self.apply(self._init_weights) - self.embedding_model = embedding_model - - def _init_weights(self, module): - if isinstance(module, nn.Linear): - normal_ = Normal(mean=0.0, std=self.initializer_range) - normal_(module.weight) - if module.bias is not None: - zeros_(module.bias) - elif isinstance(module, nn.LayerNorm): - zeros_(module.bias) - ones_(module.weight) - - def get_position_embed(self, x): - B, N, _ = x.shape - position_ids = paddle.arange(0, N, dtype=paddle.get_default_dtype()).reshape( - [1, N, 1] - ) - position_ids = position_ids.repeat_interleave(B, axis=0) - - position_embeds = paddle.zeros_like(x) - i = paddle.arange(0, self.embed_size // 2).unsqueeze(0).unsqueeze(0) - position_embeds[:, :, ::2] = paddle.sin( - position_ids / 10000 ** (2 * i / self.embed_size) - ) - position_embeds[:, :, 1::2] = paddle.cos( - position_ids / 10000 ** (2 * i / self.embed_size) - ) - return position_embeds - - def _generate_time_series(self, x, max_length): - cur_len = x.shape[1] - if cur_len >= max_length: - raise ValueError( - f"max_length({max_length}) should be larger than " - f"the length of input context({cur_len})" - ) - - while cur_len < max_length: - model_inputs = x[:, -1:] - outputs = self.forward_tensor(model_inputs) - next_output = outputs[0][:, -1:] - x = paddle.concat([x, next_output], axis=1) - cur_len = cur_len + 1 - return x - - @paddle.no_grad() - def generate(self, x, max_length=256): - if max_length <= 0: - raise ValueError( - f"max_length({max_length}) should be a strictly positive integer." - ) - outputs = self._generate_time_series(x, max_length) - return outputs - - def forward_tensor(self, x): - position_embeds = self.get_position_embed(x) - # Combine input embedding, position embedding - hidden_states = x + position_embeds - hidden_states = self.drop(hidden_states) - - # Loop through transformer self-attention layers - for block in self.blocks: - block_outputs = block(hidden_states) - hidden_states = block_outputs[0] - outputs = self.linear(self.ln(hidden_states)) - return (outputs,) - - def forward_eval(self, x): - input_embeds = x[:, :1] - outputs = self.generate(input_embeds) - return (outputs[:, 1:],) - - @staticmethod - def split_to_dict(data_tensors, keys): - return {key: data_tensors[i] for i, key in enumerate(keys)} - - def forward(self, x): - if self._input_transform is not None: - x = self._input_transform(x) - x_tensor = self.concat_to_tensor(x, self.input_keys, axis=-1) - if self.embedding_model is not None: - if isinstance(self.embedding_model, CylinderEmbedding): - x_tensor = self.embedding_model.encoder(x_tensor, x["visc"]) - else: - x_tensor = self.embedding_model.encoder(x_tensor) - - if self.training: - y = self.forward_tensor(x_tensor) - else: - y = self.forward_eval(x_tensor) - - if self.embedding_model is not None: - y = (self.embedding_model.decoder(y[0]),) - - y = self.split_to_dict(y, self.output_keys) - if self._output_transform is not None: - y = self._output_transform(x, y) - return y diff --git a/examples/smc_reac/ppsci/arch/regdgcnn.py b/examples/smc_reac/ppsci/arch/regdgcnn.py deleted file mode 100644 index c39a5b0540..0000000000 --- a/examples/smc_reac/ppsci/arch/regdgcnn.py +++ /dev/null @@ -1,250 +0,0 @@ -# Copyright (c) 2024 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. - -""" -Created on Mon May 29 22:18:28 2023 - -@author: Mohamed Elrefaie, mohamed.elrefaie@mit.edu, mohamed.elrefaie@tum.de - -This module is part of the research presented in the paper -"DrivAerNet: A Parametric Car Dataset for Data-driven Aerodynamic Design and Graph-Based Drag Prediction". -It extends the work by introducing a Deep Graph Convolutional Neural Network (RegDGCNN) model for Regression Tasks, -specifically designed for processing 3D point cloud data of car models from the DrivAerNet dataset. - -The RegDGCNN model utilizes a series of graph-based convolutional layers to effectively capture the complex geometric -and topological structure of 3D car models, facilitating advanced aerodynamic analyses and predictions. -The model architecture incorporates several techniques, including dynamic graph construction, -EdgeConv operations, and global feature aggregation, to robustly learn from graphs and point cloud data. - -Parts of this code are modified from the original version authored by Yue Wang -""" - -from __future__ import annotations - -from typing import Dict -from typing import Tuple - -import paddle - - -def transpose_aux_func(dims, dim0, dim1): - perm = list(range(dims)) - perm[dim0], perm[dim1] = perm[dim1], perm[dim0] - return perm - - -def knn(x, k): - """ - Computes the k-nearest neighbors for each point in x. - - Args: - x (paddle.Tensor): The input tensor of shape (batch_size, num_dims, num_points). - k (int): The number of nearest neighbors to find. - - Returns: - paddle.Tensor: Indices of the k-nearest neighbors for each point, shape (batch_size, num_points, k). - """ - inner = -2 * paddle.matmul( - x=x.transpose(perm=transpose_aux_func(x.ndim, 2, 1)), y=x - ) - xx = paddle.sum(x=x**2, axis=1, keepdim=True) - pairwise_distance = ( - -xx - inner - xx.transpose(perm=transpose_aux_func(xx.ndim, 2, 1)) - ) - idx = pairwise_distance.topk(k=k, axis=-1)[1] - return idx - - -def get_graph_feature(x, k=20, idx=None): - """ - Constructs local graph features for each point by finding its k-nearest neighbors and - concatenating the relative position vectors. - - Args: - x (paddle.Tensor): The input tensor of shape (batch_size, num_dims, num_points). - k (int): The number of neighbors to consider for graph construction. - idx (paddle.Tensor, optional): Precomputed k-nearest neighbor indices. - - Returns: - paddle.Tensor: The constructed graph features of shape (batch_size, 2*num_dims, num_points, k). - """ - batch_size = x.shape[0] - num_points = x.shape[2] - x = x.reshape([batch_size, -1, num_points]) - if idx is None: - idx = knn(x, k=k) - idx_base = paddle.arange(start=0, end=batch_size).reshape([-1, 1, 1]) * num_points - idx = idx + idx_base - idx = idx.reshape([-1]) - _, num_dims, _ = tuple(x.shape) - x = x.transpose(perm=transpose_aux_func(x.ndim, 2, 1)).contiguous() - feature = x.reshape([batch_size * num_points, -1])[idx, :] - feature = feature.reshape([batch_size, num_points, k, num_dims]) - x = x.reshape([batch_size, num_points, 1, num_dims]).tile(repeat_times=[1, 1, k, 1]) - feature = ( - paddle.concat(x=(feature - x, x), axis=3) - .transpose(perm=[0, 3, 1, 2]) - .contiguous() - ) - del x, idx, idx_base - paddle.device.cuda.empty_cache() - return feature - - -class RegDGCNN(paddle.nn.Layer): - """Deep Graph Convolutional Neural Network for Regression Tasks (RegDGCNN) designed to process 3D point cloud data. - - This network architecture extracts hierarchical features from point clouds using graph-based convolutions, - enabling effective learning of spatial structures. - - Args: - input_keys (Tuple[str, ...]): Keys for input data fields. - label_keys (Tuple[str, ...]): Keys for label data fields. - weight_keys (Tuple[str, ...]): Keys for weight data fields. - args (dict): Configuration parameters including: - - 'k' (int): Number of neighbors for graph convolution. - - 'emb_dims' (int): Embedding dimensions for feature aggregation. - - 'dropout' (float): Dropout rate for regularization. - output_channels (int, optional): Number of output channels. Defaults to 1. - """ - - def __init__( - self, - input_keys: Tuple[str, ...], - label_keys: Tuple[str, ...], - weight_keys: Tuple[str, ...], - args: dict, - output_channels=1, - ): - - super(RegDGCNN, self).__init__() - self.input_keys = input_keys - self.label_keys = label_keys - self.weight_keys = weight_keys - self.args = args - self.k = args["k"] - self.bn1 = paddle.nn.BatchNorm2D(num_features=256) - self.bn2 = paddle.nn.BatchNorm2D(num_features=512) - self.bn3 = paddle.nn.BatchNorm2D(num_features=512) - self.bn4 = paddle.nn.BatchNorm2D(num_features=1024) - self.bn5 = paddle.nn.BatchNorm1D(num_features=args["emb_dims"]) - self.conv1 = paddle.nn.Sequential( - paddle.nn.Conv2D( - in_channels=6, out_channels=256, kernel_size=1, bias_attr=False - ), - self.bn1, - paddle.nn.LeakyReLU(negative_slope=0.2), - ) - self.conv2 = paddle.nn.Sequential( - paddle.nn.Conv2D( - in_channels=256 * 2, out_channels=512, kernel_size=1, bias_attr=False - ), - self.bn2, - paddle.nn.LeakyReLU(negative_slope=0.2), - ) - self.conv3 = paddle.nn.Sequential( - paddle.nn.Conv2D( - in_channels=512 * 2, out_channels=512, kernel_size=1, bias_attr=False - ), - self.bn3, - paddle.nn.LeakyReLU(negative_slope=0.2), - ) - self.conv4 = paddle.nn.Sequential( - paddle.nn.Conv2D( - in_channels=512 * 2, out_channels=1024, kernel_size=1, bias_attr=False - ), - self.bn4, - paddle.nn.LeakyReLU(negative_slope=0.2), - ) - self.conv5 = paddle.nn.Sequential( - paddle.nn.Conv1D( - in_channels=2304, - out_channels=args["emb_dims"], - kernel_size=1, - bias_attr=False, - ), - self.bn5, - paddle.nn.LeakyReLU(negative_slope=0.2), - ) - self.linear1 = paddle.nn.Linear( - in_features=args["emb_dims"] * 2, out_features=128, bias_attr=False - ) - self.bn6 = paddle.nn.BatchNorm1D(num_features=128) - self.dp1 = paddle.nn.Dropout(p=args["dropout"]) - self.linear2 = paddle.nn.Linear(in_features=128, out_features=64) - self.bn7 = paddle.nn.BatchNorm1D(num_features=64) - self.dp2 = paddle.nn.Dropout(p=args["dropout"]) - self.linear3 = paddle.nn.Linear(in_features=64, out_features=32) - self.bn8 = paddle.nn.BatchNorm1D(num_features=32) - self.dp3 = paddle.nn.Dropout(p=args["dropout"]) - self.linear4 = paddle.nn.Linear(in_features=32, out_features=16) - self.bn9 = paddle.nn.BatchNorm1D(num_features=16) - self.dp4 = paddle.nn.Dropout(p=args["dropout"]) - self.linear5 = paddle.nn.Linear(in_features=16, out_features=output_channels) - - def forward(self, x: paddle.Tensor) -> Dict[str, paddle.Tensor]: - """ - Forward pass of the model to process input data and predict outputs. - - Args: - x (paddle.Tensor): Input tensor representing a batch of point clouds. - - Returns: - Dict[str, paddle.Tensor]: Model predictions for the input batch. - - """ - - x = x[self.input_keys[0]] - batch_size = x.shape[0] - x = x.transpose(perm=[0, 2, 1]) - - x = get_graph_feature(x, k=self.k) - x = self.conv1(x) - x1 = x.max(axis=-1, keepdim=False) - x = get_graph_feature(x1, k=self.k) - x = self.conv2(x) - x2 = x.max(axis=-1, keepdim=False) - x = get_graph_feature(x2, k=self.k) - x = self.conv3(x) - x3 = x.max(axis=-1, keepdim=False) - x = get_graph_feature(x3, k=self.k) - x = self.conv4(x) - x4 = x.max(axis=-1, keepdim=False) - x = paddle.concat(x=(x1, x2, x3, x4), axis=1) - x = self.conv5(x) - x1 = paddle.nn.functional.adaptive_max_pool1d(x=x, output_size=1).reshape( - [batch_size, -1] - ) - x2 = paddle.nn.functional.adaptive_avg_pool1d(x=x, output_size=1).reshape( - [batch_size, -1] - ) - x = paddle.concat(x=(x1, x2), axis=1) - x = paddle.nn.functional.leaky_relu( - x=self.bn6(self.linear1(x)), negative_slope=0.2 - ) - x = self.dp1(x) - x = paddle.nn.functional.leaky_relu( - x=self.bn7(self.linear2(x)), negative_slope=0.2 - ) - x = self.dp2(x) - x = paddle.nn.functional.leaky_relu( - x=self.bn8(self.linear3(x)), negative_slope=0.2 - ) - x = self.dp3(x) - x = paddle.nn.functional.leaky_relu( - x=self.bn9(self.linear4(x)), negative_slope=0.2 - ) - x = self.dp4(x) - x = self.linear5(x) - return {self.label_keys[0]: x} diff --git a/examples/smc_reac/ppsci/arch/regpointnet.py b/examples/smc_reac/ppsci/arch/regpointnet.py deleted file mode 100644 index 2dee522326..0000000000 --- a/examples/smc_reac/ppsci/arch/regpointnet.py +++ /dev/null @@ -1,146 +0,0 @@ -# Copyright (c) 2024 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. - -# Copyright 2024 Mohamed Elrefaie -""" -@author: Mohamed Elrefaie, mohamed.elrefaie@mit.edu mohamed.elrefaie@tum.de - -This module is part of the research presented in the paper: -"DrivAerNet++: A Large-Scale Multimodal Car Dataset with Computational Fluid Dynamics Simulations and Deep Learning Benchmarks". - -This module is used to define point-cloud models, includingPointNet -for the task of surrogate modeling of the aerodynamic drag. -""" - -from __future__ import annotations - -from typing import Dict -from typing import Tuple - -import paddle - - -class RegPointNet(paddle.nn.Layer): - """ - PointNet-based regression model for 3D point cloud data. - - This network architecture is designed to process 3D point cloud data using a series of convolutional layers, - followed by fully connected layers, enabling effective learning of spatial structures and features. - - Args: - input_keys (Tuple[str, ...]): Keys for input data fields. - output_keys (Tuple[str, ...]): Keys for output data fields. - weight_keys (Tuple[str, ...]): Keys for weight data fields. - args (dict): Configuration parameters including: - - 'emb_dims' (int): Dimensionality of the embedding space. - - 'dropout' (float): Dropout probability. - """ - - def __init__( - self, - input_keys: Tuple[str, ...], - output_keys: Tuple[str, ...], - weight_keys: Tuple[str, ...], - args, - ): - super().__init__() - self.input_keys = input_keys - self.output_keys = output_keys - self.weight_keys = weight_keys - self.args = args - self.conv1 = paddle.nn.Conv1D( - in_channels=3, out_channels=512, kernel_size=1, bias_attr=False - ) - self.conv2 = paddle.nn.Conv1D( - in_channels=512, out_channels=1024, kernel_size=1, bias_attr=False - ) - self.conv3 = paddle.nn.Conv1D( - in_channels=1024, out_channels=1024, kernel_size=1, bias_attr=False - ) - self.conv4 = paddle.nn.Conv1D( - in_channels=1024, out_channels=1024, kernel_size=1, bias_attr=False - ) - self.conv5 = paddle.nn.Conv1D( - in_channels=1024, out_channels=1024, kernel_size=1, bias_attr=False - ) - self.conv6 = paddle.nn.Conv1D( - in_channels=1024, - out_channels=args["emb_dims"], - kernel_size=1, - bias_attr=False, - ) - self.bn1 = paddle.nn.BatchNorm1D(num_features=512) - self.bn2 = paddle.nn.BatchNorm1D(num_features=1024) - self.bn3 = paddle.nn.BatchNorm1D(num_features=1024) - self.bn4 = paddle.nn.BatchNorm1D(num_features=1024) - self.bn5 = paddle.nn.BatchNorm1D(num_features=1024) - self.bn6 = paddle.nn.BatchNorm1D(num_features=args["emb_dims"]) - self.dropout_conv = paddle.nn.Dropout(p=args["dropout"]) - self.dropout_linear = paddle.nn.Dropout(p=args["dropout"]) - self.conv_shortcut = paddle.nn.Conv1D( - in_channels=3, out_channels=args["emb_dims"], kernel_size=1, bias_attr=False - ) - self.bn_shortcut = paddle.nn.BatchNorm1D(num_features=args["emb_dims"]) - self.linear1 = paddle.nn.Linear( - in_features=args["emb_dims"], out_features=512, bias_attr=False - ) - self.bn7 = paddle.nn.BatchNorm1D(num_features=512) - self.linear2 = paddle.nn.Linear( - in_features=512, out_features=256, bias_attr=False - ) - self.bn8 = paddle.nn.BatchNorm1D(num_features=256) - self.linear3 = paddle.nn.Linear(in_features=256, out_features=128) - self.bn9 = paddle.nn.BatchNorm1D(num_features=128) - self.linear4 = paddle.nn.Linear(in_features=128, out_features=64) - self.bn10 = paddle.nn.BatchNorm1D(num_features=64) - self.final_linear = paddle.nn.Linear(in_features=64, out_features=1) - - def forward(self, x: Dict[str, paddle.Tensor]) -> Dict[str, paddle.Tensor]: - """ - Forward pass of the network. - - Args: - x (Dict[str, paddle.Tensor]): Input tensor of shape (batch_size, 3, num_points). - - Returns: - Dict[str, paddle.Tensor]: A dictionary where the key is the first element of `self.output_keys` - and the value is the output tensor of the predicted scalar value. - """ - - x: paddle.Tensor = x[self.input_keys[0]] - - x_processed = x.transpose(perm=[0, 2, 1]) - - shortcut = self.bn_shortcut(self.conv_shortcut(x_processed)) - x = paddle.nn.functional.relu(x=self.bn1(self.conv1(x_processed))) - x = self.dropout_conv(x) - x = paddle.nn.functional.relu(x=self.bn2(self.conv2(x))) - x = self.dropout_conv(x) - x = paddle.nn.functional.relu(x=self.bn3(self.conv3(x))) - x = self.dropout_conv(x) - x = paddle.nn.functional.relu(x=self.bn4(self.conv4(x))) - x = self.dropout_conv(x) - x = paddle.nn.functional.relu(x=self.bn5(self.conv5(x))) - x = self.dropout_conv(x) - x = paddle.nn.functional.relu(x=self.bn6(self.conv6(x))) - x = x + shortcut - x = paddle.nn.functional.adaptive_max_pool1d(x=x, output_size=1).squeeze( - axis=-1 - ) - x = paddle.nn.functional.relu(x=self.bn7(self.linear1(x))) - x = paddle.nn.functional.relu(x=self.bn8(self.linear2(x))) - x = paddle.nn.functional.relu(x=self.bn9(self.linear3(x))) - x = paddle.nn.functional.relu(x=self.bn10(self.linear4(x))) - x = self.final_linear(x) - return {self.output_keys[0]: x} diff --git a/examples/smc_reac/ppsci/arch/sfnonet.py b/examples/smc_reac/ppsci/arch/sfnonet.py deleted file mode 100644 index c7695fb02d..0000000000 --- a/examples/smc_reac/ppsci/arch/sfnonet.py +++ /dev/null @@ -1,568 +0,0 @@ -from typing import Dict -from typing import List -from typing import Optional -from typing import Tuple -from typing import Union - -import paddle -import paddle.nn.functional as F -from paddle import nn - -from ppsci.arch import base -from ppsci.arch import fno_block -from ppsci.arch.paddle_harmonics import sht as paddle_sht -from ppsci.utils import initializer - -einsum_symbols = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ" - - -def _contract_dense(x, weight, separable=False, dhconv=True): - order = len(x.shape) - x_syms = list(einsum_symbols[:order]) - - # in_channels, out_channels, x, y... - weight_syms = list(x_syms[1:]) # no batch-size - - # batch-size, out_channels, x, y... - if separable: - out_syms = [x_syms[0]] + list(weight_syms) - else: - weight_syms.insert(1, einsum_symbols[order]) # outputs - out_syms = list(weight_syms) - out_syms[0] = x_syms[0] - - if dhconv: - weight_syms.pop() - - eq = "".join(x_syms) + "," + "".join(weight_syms) + "->" + "".join(out_syms) - # For the darcy flow, the only einsum is abcd,becd->aecd, where x and weights are shaped [32,32,8,8] - if not isinstance(weight, paddle.Tensor): - weight = paddle.to_tensor(weight) - - return paddle.einsum(eq, x, weight) - - -def _contract_dense_trick(x, weight_real, weight_imag, separable=False, dhconv=True): - # the same as above function, but do the complex multiplication manually to avoid the einsum bug in paddle - order = len(x.shape) - # batch-size, in_channels, x, y... - x_syms = list(einsum_symbols[:order]) - - # in_channels, out_channels, x, y... - weight_syms = list(x_syms[1:]) # no batch-size - - # batch-size, out_channels, x, y... - if separable: - out_syms = [x_syms[0]] + list(weight_syms) - else: - weight_syms.insert(1, einsum_symbols[order]) # outputs - out_syms = list(weight_syms) - out_syms[0] = x_syms[0] - - if dhconv: - weight_syms.pop() - - eq = "".join(x_syms) + "," + "".join(weight_syms) + "->" + "".join(out_syms) - - o1_real = paddle.einsum(eq, x.real(), weight_real) - paddle.einsum( - eq, x.imag(), weight_imag - ) - o1_imag = paddle.einsum(eq, x.imag(), weight_real) + paddle.einsum( - eq, x.real(), weight_imag - ) - x = paddle.complex(o1_real, o1_imag) - return x - - -def _contract_dense_separable(x, weight, separable=True): - if not separable: - raise ValueError("This function is only for separable=True") - return x * weight - - -def get_contract_fun(weight, implementation="reconstructed", separable=False): - """Generic ND implementation of Fourier Spectral Conv contraction. - - Args: - weight (FactorizedTensor): The factoriz Tensor. - implementation (str, optional): Whether to reconstruct the weight and do a forward pass (reconstructed) - or contract directly the factors of the factorized weight with the input (factorized). - {'reconstructed', 'factorized'} Defaults to "reconstructed". - separable (bool, optional): Whether to use the separable implementation of contraction. This arg is - only checked when `implementation=reconstructed`. Defaults to False. - """ - - if implementation == "reconstructed": - if separable: - return _contract_dense_separable - else: - return _contract_dense_trick - elif implementation == "factorized": - if isinstance(weight, paddle.Tensor): - return _contract_dense_trick - - else: - raise ValueError( - f'Got implementation={implementation}, expected "reconstructed" or "factorized"' - ) - - -class SHT(nn.Layer): - """A wrapper for the Spherical Harmonics transform - - Allows to call it with an interface similar to that of FFT - """ - - def __init__(self, dtype=paddle.float32): - super().__init__() - self.dtype = dtype - self._SHT_cache = nn.LayerDict() - self._iSHT_cache = nn.LayerDict() - - def sht(self, x, s=None, norm="ortho", grid="equiangular"): - *_, height, width = x.shape # height = latitude, width = longitude - if s is None: - if grid == "equiangular": - modes_width = height // 2 - else: - modes_width = height - modes_height = height - else: - modes_height, modes_width = s - - cache_key = f"{height}_{width}_{modes_height}_{modes_width}_{norm}_{grid}" - - try: - sht = self._SHT_cache[cache_key] - except KeyError: - sht = paddle_sht.RealSHT( - nlat=height, - nlon=width, - lmax=modes_height, - mmax=modes_width, - grid=grid, - norm=norm, - ).astype(dtype=self.dtype) - - self._SHT_cache[cache_key] = sht - - return sht(x) - - def isht(self, x, s=None, norm="ortho", grid="equiangular"): - *_, modes_height, modes_width = x.shape # height = latitude, width = longitude - if s is None: - if grid == "equiangular": - width = modes_width * 2 - else: - width = modes_width - height = modes_height - else: - height, width = s - - cache_key = f"{height}_{width}_{modes_height}_{modes_width}_{norm}_{grid}" - - try: - isht = self._iSHT_cache[cache_key] - except KeyError: - isht = paddle_sht.InverseRealSHT( - nlat=height, - nlon=width, - lmax=modes_height, - mmax=modes_width, - grid=grid, - norm=norm, - ).astype(dtype=self.dtype) - self._iSHT_cache[cache_key] = isht - - return isht(x) - - -Number = Union[int, float] - - -class SphericalConv(nn.Layer): - """Spherical Convolution, base class for the SFNO [1]. - .. [1] Spherical Fourier Neural Operators: Learning Stable Dynamics on the Sphere, - Boris Bonev, Thorsten Kurth, Christian Hundt, Jaideep Pathak, Maximilian Baust, Karthik Kashinath, Anima Anandkumar, - ICML 2023. - - Args: - in_channels (int): Number of input channels. - out_channels (int): Number of output channels. - n_modes (Tuple[int, ...]): Number of modes to use for contraction in Fourier domain during - training. - max_n_modes (int, optional): The maximum number of modes to use for contraction in Fourier domain during - training. Defaults to None. - bias (bool, optional): Whether to use bias in the layers. Defaults to True. - n_layers (int, optional): Number of Fourier Layers. Defaults to 1. - separable (bool, optional): Whether to use separable Fourier Conv. Defaults to False. - output_scaling_factor (Optional[Union[Number, List[Number]]], optional): Scaling factor for the - output. Defaults to None. - rank (float, optional): Rank of the tensor factorization of the Fourier weights. Defaults to 0.5. - factorization (str, optional): Tensor factorization of the parameters weight to use. Defaults to "dense". - implementation (str, optional): If factorization is not None, forward mode to use. Defaults to "reconstructed". - joint_factorization (bool, optional): Whether all the Fourier Layers should be parametrized by a - single tensor. Defaults to False. - init_std (str, optional): The std to use for the init. Defaults to "auto". - sht_norm (str, optional): The normalization mode of the SHT. Defaults to "ortho". - sht_grids (str, optional): The grid of the SHT. Defaults to "equiangular". - dtype (paddle.float32, optional): The data type. Defaults to paddle.float32. - """ - - def __init__( - self, - in_channels: int, - out_channels: int, - n_modes: Tuple[int, ...], - max_n_modes: int = None, - bias: bool = True, - n_layers: int = 1, - separable: bool = False, - output_scaling_factor: Optional[Union[Number, List[Number]]] = None, - rank: float = 0.5, - factorization: str = "dense", - implementation: str = "reconstructed", - joint_factorization: bool = False, - init_std: str = "auto", - sht_norm: str = "ortho", - sht_grids: str = "equiangular", - dtype: paddle.dtype = paddle.float32, - ): - super().__init__() - self.in_channels = in_channels - self.out_channels = out_channels - - self.dtype = dtype - - self.joint_factorization = joint_factorization - - if isinstance(n_modes, int): - n_modes = [n_modes] - self._n_modes = n_modes - self.order = len(n_modes) - - if max_n_modes is None: - max_n_modes = self.n_modes - elif isinstance(max_n_modes, int): - max_n_modes = [max_n_modes] - self.max_n_modes = max_n_modes - - self.rank = rank - self.factorization = factorization - self.n_layers = n_layers - self.implementation = implementation - - self.output_scaling_factor: Union[ - None, List[List[float]] - ] = fno_block.validate_scaling_factor( - output_scaling_factor, self.order, n_layers - ) - - if init_std == "auto": - init_std = (2 / (in_channels + out_channels)) ** 0.5 - else: - init_std = init_std - - if separable: - if in_channels != out_channels: - raise ValueError( - f"To use separable Fourier Conv, in_channels must be equal to out_channels, but got in_channels={in_channels} and out_channels={out_channels}" - ) - weight_shape = (in_channels, *self.n_modes[:-1]) - else: - weight_shape = (in_channels, out_channels, *self.n_modes[:-1]) - self.separable = separable - - if joint_factorization: - self.weight = paddle.create_parameter( - shape=(n_layers, *weight_shape), - dtype="float32", - ) - self.weight = initializer.normal_(self.weight, 0, init_std) - else: - self.weight = nn.LayerList( - [ - fno_block.FactorizedTensor(weight_shape, init_scale=init_std) - for _ in range(n_layers) - ] - ) - self._contract = get_contract_fun( - self.weight[0].data, implementation=implementation, separable=separable - ) - if bias: - shape = (n_layers, self.out_channels) + (1,) * self.order - init_bias = init_std * paddle.randn(shape) - self.bias = paddle.create_parameter( - shape=shape, - dtype=(init_bias.dtype), - default_initializer=nn.initializer.Assign(init_bias), - ) - self.bias.stop_gradient = False - else: - self.bias = None - - self.sht_norm = sht_norm - if isinstance(sht_grids, str): - sht_grids = [sht_grids] * (self.n_layers + 1) - self.sht_grids = sht_grids - self.sht_handle = SHT(dtype=self.dtype) - - @property - def n_modes(self): - return self._n_modes - - @n_modes.setter - def n_modes(self, n_modes): - if isinstance(n_modes, int): # Should happen for 1D FNO only - n_modes = [n_modes] - else: - n_modes = list(n_modes) - self._n_modes = n_modes - - def forward(self, x, indices=0, output_shape=None): - batchsize, channels, height, width = x.shape - - if self.output_scaling_factor is not None and output_shape is None: - scaling_factors = self.output_scaling_factor[indices] - height = round(height * scaling_factors[0]) - width = round(width * scaling_factors[1]) - elif output_shape is not None: - height, width = output_shape[0], output_shape[1] - - out_fft = self.sht_handle.sht( - x, - s=(self.n_modes[0], self.n_modes[1] // 2), - norm=self.sht_norm, - grid=self.sht_grids[indices], - ) - - w_real = self.weight[indices].real[:, :, : self.n_modes[0]] - w_imag = self.weight[indices].imag[:, :, : self.n_modes[0]] - - out_fft = self._contract( - out_fft[:, :, : self.n_modes[0], : self.n_modes[1] // 2], - w_real, - w_imag, - separable=self.separable, - dhconv=True, - ) - - x = self.sht_handle.isht( - out_fft, - s=(height, width), - norm=self.sht_norm, - grid=self.sht_grids[indices + 1], - ) - - if self.bias is not None: - x = x + self.bias[indices, ...] - - return x - - def transform(self, x, layer_index=0, output_shape=None): - *_, in_height, in_width = x.shape - - if self.output_scaling_factor is not None and output_shape is None: - height = round(in_height * self.output_scaling_factor[layer_index][0]) - width = round(in_width * self.output_scaling_factor[layer_index][1]) - elif output_shape is not None: - height, width = output_shape[0], output_shape[1] - else: - height, width = in_height, in_width - - # Return the identity if the resolution and grid of the input and output are the same - if ((in_height, in_width) == (height, width)) and ( - self.sht_grids[layer_index] == self.sht_grids[layer_index + 1] - ): - return x - else: - coefs = self.sht_handle.sht( - x, s=self.n_modes, norm=self.sht_norm, grid=self.sht_grids[layer_index] - ) - return self.sht_handle.isht( - coefs, - s=(height, width), - norm=self.sht_norm, - grid=self.sht_grids[layer_index + 1], - ) - - -class SFNONet(base.Arch): - """N-Dimensional Tensorized Fourier Neural Operator. - - Args: - input_keys (Tuple[str, ...]): Name of input keys, such as ("input",). - output_keys (Tuple[str, ...]): Name of output keys, such as ("output",). - n_modes (Tuple[int, ...]): Number of modes to keep in Fourier Layer, along each dimension - The dimensionality of the SFNO is inferred from ``len(n_modes)` - hidden_channels (int): Width of the FNO (i.e. number of channels) - in_channels (int, optional): Number of input channels. Defaults to 3. - out_channels (int, optional): Number of output channels. Defaults to 1. - lifting_channels (int, optional): Number of hidden channels of the lifting block of the FNO. - Defaults to 256. - projection_channels (int, optional): Number of hidden channels of the projection block of the FNO. - Defaults to 256. - n_layers (int, optional): Number of Fourier Layers. Defaults to 4. - use_mlp (bool, optional): Whether to use an MLP layer after each FNO block. Defaults to False. - mlp (Dict[str, float], optional): Parameters of the MLP. {'expansion': float, 'dropout': float}. - Defaults to None. - non_linearity (nn.functional, optional): Non-Linearity module to use. Defaults to F.gelu. - norm (str, optional): Normalization layer to use. Defaults to None. - ada_in_features (int,optional): The input channels of the adaptive normalization.Defaults to None. - preactivation (bool, optional): Whether to use resnet-style preactivation. Defaults to False. - fno_skip (str, optional): Type of skip connection to use,{'linear', 'identity', 'soft-gating'}. - Defaults to "soft-gating". - separable (bool, optional): Whether to use a depthwise separable spectral convolution. - Defaults to False. - factorization (str, optional): Tensor factorization of the parameters weight to use. - * If None, a dense tensor parametrizes the Spectral convolutions. - * Otherwise, the specified tensor factorization is used. Defaults to "Tucker". - rank (float, optional): Rank of the tensor factorization of the Fourier weights. Defaults to 1.0. - joint_factorization (bool, optional): Whether all the Fourier Layers should be parametrized by a - single tensor (vs one per layer). Defaults to False. - implementation (str, optional): {'factorized', 'reconstructed'}, optional. Defaults to "factorized". - If factorization is not None, forward mode to use:: - * `reconstructed` : the full weight tensor is reconstructed from the factorization and used for the forward pass. - * `factorized` : the input is directly contracted with the factors of the decomposition. - domain_padding (Optional[list], optional): Whether to use percentage of padding. Defaults to None. - domain_padding_mode (str, optional): {'symmetric', 'one-sided'}, optional - How to perform domain padding, by default 'one-sided'. Defaults to "one-sided". - fft_norm (str, optional): The normalization mode for the FFT. Defaults to "forward". - patching_levels (int, optional): Number of patching levels to use. Defaults to 0. - """ - - def __init__( - self, - input_keys: Tuple[str, ...], - output_keys: Tuple[str, ...], - n_modes: Tuple[int, ...], - hidden_channels: int, - in_channels: int = 3, - out_channels: int = 1, - lifting_channels: int = 256, - projection_channels: int = 256, - n_layers: int = 4, - use_mlp: bool = False, - mlp: Optional[Dict[str, float]] = None, - max_n_modes: int = None, - non_linearity: nn.functional = F.gelu, - stabilizer: str = None, - norm: str = None, - ada_in_features: Optional[int] = None, - preactivation: bool = False, - fno_skip: str = "linear", - mlp_skip: str = "soft-gating", - separable: bool = False, - factorization: str = None, - rank: float = 1.0, - joint_factorization: bool = False, - implementation: str = "factorized", - domain_padding: Optional[list] = None, - domain_padding_mode: str = "one-sided", - fft_norm: str = "forward", - patching_levels: int = 0, - **kwargs, - ): - super().__init__() - self.input_keys = input_keys - self.output_keys = output_keys - - self.n_dim = len(n_modes) - self.n_modes = n_modes - self.hidden_channels = hidden_channels - self.lifting_channels = lifting_channels - self.projection_channels = projection_channels - self.in_channels = in_channels - if patching_levels: - self.in_channels = self.in_channels * patching_levels + 1 - self.out_channels = out_channels - self.n_layers = n_layers - self.joint_factorization = joint_factorization - self.non_linearity = non_linearity - self.rank = rank - self.factorization = factorization - self.fno_skip = (fno_skip,) - self.mlp_skip = (mlp_skip,) - self.fft_norm = fft_norm - self.implementation = implementation - self.separable = separable - self.preactivation = preactivation - self.stabilizer = stabilizer - if domain_padding is not None and ( - (isinstance(domain_padding, list) and sum(domain_padding) > 0) - or (isinstance(domain_padding, (float, int)) and domain_padding > 0) - ): - self.domain_padding = fno_block.DomainPadding( - domain_padding=domain_padding, padding_mode=domain_padding_mode - ) - else: - self.domain_padding = None - self.domain_padding_mode = domain_padding_mode - - self.fno_blocks = fno_block.FNOBlocks( - in_channels=hidden_channels, - out_channels=hidden_channels, - n_modes=self.n_modes, - n_layers=n_layers, - max_n_modes=max_n_modes, - use_mlp=use_mlp, - mlp=mlp, - non_linearity=non_linearity, - stabilizer=stabilizer, - norm=norm, - ada_in_features=ada_in_features, - preactivation=preactivation, - fno_skip=fno_skip, - mlp_skip=mlp_skip, - separable=separable, - factorization=factorization, - rank=rank, - SpectralConv=SphericalConv, - joint_factorization=joint_factorization, - implementation=implementation, - fft_norm=fft_norm, - ) - # if lifting_channels is passed, make lifting an MLP - # with a hidden layer of size lifting_channels - if self.lifting_channels: - self.lifting = fno_block.MLP( - in_channels=in_channels, - out_channels=self.hidden_channels, - hidden_channels=self.lifting_channels, - n_layers=2, - n_dim=self.n_dim, - ) - # otherwise, make it a linear layer - else: - self.lifting = fno_block.MLP( - in_channels=in_channels, - out_channels=self.hidden_channels, - hidden_channels=self.hidden_channels, - n_layers=1, - n_dim=self.n_dim, - ) - self.projection = fno_block.MLP( - in_channels=self.hidden_channels, - out_channels=out_channels, - hidden_channels=self.projection_channels, - n_layers=2, - n_dim=self.n_dim, - non_linearity=non_linearity, - ) - - def forward(self, x): - """SFNO's forward pass""" - x = self.concat_to_tensor(x, self.input_keys) - - x = self.lifting(x) - if self.domain_padding is not None: - x = self.domain_padding.pad(x) - # x is 0.4 * [1, 32, 16, 16], passed - for index in range(self.n_layers): - x = self.fno_blocks(x, index) - - if self.domain_padding is not None: - x = self.domain_padding.unpad(x) - out = self.projection(x) - - return {self.output_keys[0]: out} diff --git a/examples/smc_reac/ppsci/arch/smc_reac.py b/examples/smc_reac/ppsci/arch/smc_reac.py deleted file mode 100644 index 9e4d2595db..0000000000 --- a/examples/smc_reac/ppsci/arch/smc_reac.py +++ /dev/null @@ -1,107 +0,0 @@ -import paddle -from paddle import nn - -from ppsci.arch import base - - -class SuzukiMiyauraModel(base.Arch): - def __init__( - self, input_dim, hidden_dim, hidden_dim2, hidden_dim3, hidden_dim4, output_dim - ): - super().__init__() - - self.r1_fc = nn.Sequential( - nn.Linear(input_dim, hidden_dim), - nn.ReLU(), - nn.Linear(hidden_dim, hidden_dim2), - nn.ReLU(), - nn.Linear(hidden_dim2, hidden_dim3), - ) - - self.r2_fc = nn.Sequential( - nn.Linear(input_dim, hidden_dim), - nn.ReLU(), - nn.Linear(hidden_dim, hidden_dim2), - nn.ReLU(), - nn.Linear(hidden_dim2, hidden_dim3), - ) - - self.ligand_fc = nn.Sequential( - nn.Linear(input_dim, hidden_dim), - nn.ReLU(), - nn.Linear(hidden_dim, hidden_dim2), - nn.ReLU(), - nn.Linear(hidden_dim2, hidden_dim3), - ) - - self.base_fc = nn.Sequential( - nn.Linear(input_dim, hidden_dim), - nn.ReLU(), - nn.Linear(hidden_dim, hidden_dim2), - nn.ReLU(), - nn.Linear(hidden_dim2, hidden_dim3), - ) - - self.solvent_fc = nn.Sequential( - nn.Linear(input_dim, hidden_dim), - nn.ReLU(), - nn.Linear(hidden_dim, hidden_dim2), - nn.ReLU(), - nn.Linear(hidden_dim2, hidden_dim3), - nn.ReLU(), - ) - - self.weights = paddle.create_parameter( - shape=[5], - dtype="float32", - default_initializer=paddle.nn.initializer.Assign( - paddle.to_tensor([0.2, 0.2, 0.2, 0.2, 0.2]) - ), - ) - - self.fc_combined = nn.Sequential( - nn.Linear(hidden_dim3, hidden_dim4), - nn.ReLU(), - nn.Linear(hidden_dim4, output_dim), - ) - - def weighted_average(self, features, weights): - - weights = weights.clone().detach() - - weighted_sum = sum(f * w for f, w in zip(features, weights)) - - total_weight = weights.sum() - - return weighted_sum / total_weight - - def forward(self, x): - x = self.concat_to_tensor(x, ("v"), axis=-1) - - input_splits = paddle.split(x, num_or_sections=5, axis=1) - - r1_input, r2_input, ligand_input, base_input, solvent_input = input_splits - - r1_features = self.r1_fc(r1_input) - - r2_features = self.r2_fc(r2_input) - - ligand_features = self.ligand_fc(ligand_input) - - base_features = self.base_fc(base_input) - - solvent_features = self.solvent_fc(solvent_input) - - features = [ - r1_features, - r2_features, - ligand_features, - base_features, - solvent_features, - ] - - combined_features = self.weighted_average(features, self.weights) - - output = self.fc_combined(combined_features) - output = self.split_to_dict(output, ("u"), axis=-1) - return output diff --git a/examples/smc_reac/ppsci/arch/spinn.py b/examples/smc_reac/ppsci/arch/spinn.py deleted file mode 100644 index 014446941f..0000000000 --- a/examples/smc_reac/ppsci/arch/spinn.py +++ /dev/null @@ -1,180 +0,0 @@ -# Copyright (c) 2024 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. - -from __future__ import annotations - -from typing import Dict -from typing import List -from typing import Optional -from typing import Tuple -from typing import Union - -import paddle -import paddle.nn as nn - -from ppsci.arch import base -from ppsci.arch.mlp import ModifiedMLP -from ppsci.utils import initializer - - -class SPINN(base.Arch): - """Separable Physics-Informed Neural Networks. - - Args: - input_keys (Tuple[str, ...]): Keys of input variables. - output_keys (Tuple[str, ...]): Keys of output variables. - r (int): Number of features for each output dimension. - num_layers (int): Number of layers. - hidden_size (Union[int, Tuple[int, ...]]): Size of hidden layer. - activation (str, optional): Name of activation function. - skip_connection (bool, optional): Whether to use skip connection. - weight_norm (bool, optional): Whether to use weight normalization. - periods (Optional[Dict[int, Tuple[float, bool]]], optional): Periodicity of input variables. - fourier (Optional[Dict[str, Union[float, int]]], optional): Frequency of input variables. - random_weight (Optional[Dict[str, float]], optional): Random weight of linear layer. - - Examples: - >>> from ppsci.arch import SPINN - >>> model = SPINN( - ... input_keys=('x', 'y', 'z'), - ... output_keys=('u', 'v'), - ... r=32, - ... num_layers=4, - ... hidden_size=32, - ... ) - >>> input_dict = {"x": paddle.rand([3, 1]), - ... "y": paddle.rand([4, 1]), - ... "z": paddle.rand([5, 1])} - >>> output_dict = model(input_dict) - >>> print(output_dict["u"].shape) - [3, 4, 5, 1] - >>> print(output_dict["v"].shape) - [3, 4, 5, 1] - """ - - def __init__( - self, - input_keys: Tuple[str, ...], - output_keys: Tuple[str, ...], - r: int, - num_layers: int, - hidden_size: Union[int, Tuple[int, ...]], - activation: str = "tanh", - skip_connection: bool = False, - weight_norm: bool = False, - periods: Optional[Dict[int, Tuple[float, bool]]] = None, - fourier: Optional[Dict[str, Union[float, int]]] = None, - random_weight: Optional[Dict[str, float]] = None, - ): - - super().__init__() - self.input_keys = input_keys - self.output_keys = output_keys - self.r = r - input_dim = len(self.input_keys) - - self.branch_nets = nn.LayerList() - for i in range(input_dim): - self.branch_nets.append( - ModifiedMLP( - input_keys=(input_keys[i],), - output_keys=("f",), - num_layers=num_layers, - hidden_size=hidden_size, - activation=activation, - skip_connection=skip_connection, - weight_norm=weight_norm, - output_dim=r * len(output_keys), - periods=periods, - fourier=fourier, - random_weight=random_weight, - ) - ) - - self._init_weights() - - def _init_weights(self): - for m in self.sublayers(True): - if isinstance(m, nn.Linear): - initializer.glorot_normal_(m.weight) - initializer.zeros_(m.bias) - - def _tensor_contraction(self, x: paddle.Tensor, y: paddle.Tensor) -> paddle.Tensor: - """Tensor contraction between two tensors along the last channel. - - Args: - x (Tensor): Input tensor with shape [*N, C]. - y (Tensor): Input tensor with shape [*M, C] - - Returns: - Tensor: Output tensor with shape [*N, *M, C]. - """ - x_ndim = x.ndim - y_ndim = y.ndim - out_dim = x_ndim + y_ndim - 1 - - # Align the dimensions of x and y to out_dim - if x_ndim < out_dim: - # Add singleton dimensions to x at the end of dimensions - x = x.unsqueeze([-2] * (out_dim - x_ndim)) - if y_ndim < out_dim: - # Add singleton dimensions to y at the begin of dimensions - y = y.unsqueeze([0] * (out_dim - y_ndim)) - - # Multiply x and y with implicit broadcasting - out = x * y - - return out - - def forward_tensor(self, *xs) -> List[paddle.Tensor]: - # forward each dim branch - feature_f = [] - for i, input_var in enumerate(xs): - input_i = {self.input_keys[i]: input_var} - output_f_i = self.branch_nets[i](input_i) - feature_f.append(output_f_i["f"]) # [B, r*output_dim] - - output = [] - for i, key in enumerate(self.output_keys): - st, ed = i * self.r, (i + 1) * self.r - # do tensor contraction and sum over all branch outputs - if ed - st == self.r: - output_i = feature_f[0] - else: - output_i = feature_f[0][:, st:ed] - - for j in range(1, len(self.input_keys)): - if ed - st == self.r: - output_ii = feature_f[j] - else: - output_ii = feature_f[j][:, st:ed] - output_i = self._tensor_contraction(output_i, output_ii) - - output_i = output_i.sum(-1, keepdim=True) - output.append(output_i) - - return output - - def forward(self, x): - if self._input_transform is not None: - x = self._input_transform(x) - - output = self.forward_tensor(*[x[key] for key in self.input_keys]) - - output = {key: output[i] for i, key in enumerate(self.output_keys)} - - if self._output_transform is not None: - output = self._output_transform(x, output) - - return output diff --git a/examples/smc_reac/ppsci/arch/tfnonet.py b/examples/smc_reac/ppsci/arch/tfnonet.py deleted file mode 100644 index 91bcfd6f5c..0000000000 --- a/examples/smc_reac/ppsci/arch/tfnonet.py +++ /dev/null @@ -1,514 +0,0 @@ -from typing import Dict -from typing import Optional -from typing import Tuple -from typing import Union - -import paddle.nn.functional as F -from paddle import nn - -from ppsci.arch import base -from ppsci.arch import fno_block - - -class FNONet(base.Arch): - """N-Dimensional Tensorized Fourier Neural Operator. - - Args: - input_keys (Tuple[str, ...]): Name of input keys, such as ("input",). - output_keys (Tuple[str, ...]): Name of output keys, such as ("output",). - n_modes (Tuple[int, ...]): Number of modes to keep in Fourier Layer, along each dimension - The dimensionality of the TFNO is inferred from ``len(n_modes)` - hidden_channels (int): Width of the FNO (i.e. number of channels) - in_channels (int, optional): Number of input channels. Defaults to 3. - out_channels (int, optional): Number of output channels. Defaults to 1. - lifting_channels (int, optional): Number of hidden channels of the lifting block of the FNO. - Defaults to 256. - projection_channels (int, optional): Number of hidden channels of the projection block of the FNO. - Defaults to 256. - n_layers (int, optional): Number of Fourier Layers. Defaults to 4. - use_mlp (bool, optional): Whether to use an MLP layer after each FNO block. Defaults to False. - mlp (Dict[str, float], optional): Parameters of the MLP. {'expansion': float, 'dropout': float}. - Defaults to None. - non_linearity (nn.functional, optional): Non-Linearity module to use. Defaults to F.gelu. - norm (str, optional): Normalization layer to use. Defaults to None. - ada_in_features (int,optional): The input channels of the adaptive normalization.Defaults to None.s - preactivation (bool, optional): Whether to use resnet-style preactivation. Defaults to False. - skip (str, optional): Type of skip connection to use,{'linear', 'identity', 'soft-gating'}. - Defaults to "soft-gating". - separable (bool, optional): Whether to use a depthwise separable spectral convolution. - Defaults to False. - factorization (str, optional): Tensor factorization of the parameters weight to use. - * If None, a dense tensor parametrizes the Spectral convolutions. - * Otherwise, the specified tensor factorization is used. Defaults to "Tucker". - rank (float, optional): Rank of the tensor factorization of the Fourier weights. Defaults to 1.0. - joint_factorization (bool, optional): Whether all the Fourier Layers should be parametrized by a - single tensor (vs one per layer). Defaults to False. - implementation (str, optional): {'factorized', 'reconstructed'}, optional. Defaults to "factorized". - If factorization is not None, forward mode to use:: - * `reconstructed` : the full weight tensor is reconstructed from the factorization and used for the forward pass. - * `factorized` : the input is directly contracted with the factors of the decomposition. - domain_padding (Optional[Union[list,float,int]]): Whether to use percentage of padding. Defaults to - None. - domain_padding_mode (str, optional): {'symmetric', 'one-sided'}, optional - How to perform domain padding, by default 'one-sided'. Defaults to "one-sided". - fft_norm (str, optional): The normalization mode for the FFT. Defaults to "forward". - patching_levels (int, optional): Number of patching levels to use. Defaults to 0. - SpectralConv (nn.layer, optional): Spectral convolution layer to use. - Defaults to fno_block.FactorizedSpectralConv. - """ - - def __init__( - self, - input_keys: Tuple[str, ...], - output_keys: Tuple[str, ...], - n_modes: Tuple[int, ...], - hidden_channels: int, - in_channels: int = 3, - out_channels: int = 1, - lifting_channels: int = 256, - projection_channels: int = 256, - n_layers: int = 4, - use_mlp: bool = False, - mlp: Optional[Dict[str, float]] = None, - max_n_modes: int = None, - non_linearity: nn.functional = F.gelu, - stabilizer: str = None, - norm: str = None, - ada_in_features: Optional[int] = None, - preactivation: bool = False, - fno_skip: str = "linear", - mlp_skip: str = "soft-gating", - separable: bool = False, - factorization: str = None, - rank: float = 1.0, - joint_factorization: bool = False, - implementation: str = "factorized", - domain_padding: Optional[Union[list, float, int]] = None, - domain_padding_mode: str = "one-sided", - fft_norm: str = "forward", - patching_levels: int = 0, - SpectralConv: nn.Layer = fno_block.FactorizedSpectralConv, - **kwargs, - ): - super().__init__() - self.input_keys = input_keys - self.output_keys = output_keys - - self.n_dim = len(n_modes) - self.n_modes = n_modes - self.hidden_channels = hidden_channels - self.lifting_channels = lifting_channels - self.projection_channels = projection_channels - self.in_channels = in_channels - if patching_levels: - self.in_channels = self.in_channels * patching_levels + 1 - self.out_channels = out_channels - self.n_layers = n_layers - self.joint_factorization = joint_factorization - self.non_linearity = non_linearity - self.rank = rank - self.factorization = factorization - self.fno_skip = (fno_skip,) - self.mlp_skip = (mlp_skip,) - self.fft_norm = fft_norm - self.implementation = implementation - self.separable = separable - self.preactivation = preactivation - self.stabilizer = stabilizer - if domain_padding is not None and ( - (isinstance(domain_padding, list) and sum(domain_padding) > 0) - or (isinstance(domain_padding, (float, int)) and domain_padding > 0) - ): - self.domain_padding = fno_block.DomainPadding( - domain_padding=domain_padding, padding_mode=domain_padding_mode - ) - else: - self.domain_padding = None - self.domain_padding_mode = domain_padding_mode - self.fno_blocks = fno_block.FNOBlocks( - in_channels=hidden_channels, - out_channels=hidden_channels, - n_modes=self.n_modes, - n_layers=n_layers, - max_n_modes=max_n_modes, - use_mlp=use_mlp, - mlp=mlp, - non_linearity=non_linearity, - stabilizer=stabilizer, - norm=norm, - ada_in_features=ada_in_features, - preactivation=preactivation, - fno_skip=fno_skip, - mlp_skip=mlp_skip, - separable=separable, - factorization=factorization, - rank=rank, - SpectralConv=SpectralConv, - joint_factorization=joint_factorization, - implementation=implementation, - fft_norm=fft_norm, - ) - # if lifting_channels is passed, make lifting an MLP - # with a hidden layer of size lifting_channels - if self.lifting_channels: - self.lifting = fno_block.MLP( - in_channels=in_channels, - out_channels=self.hidden_channels, - hidden_channels=self.lifting_channels, - n_layers=2, - n_dim=self.n_dim, - ) - # otherwise, make it a linear layer - else: - self.lifting = fno_block.MLP( - in_channels=in_channels, - out_channels=self.hidden_channels, - hidden_channels=self.hidden_channels, - n_layers=1, - n_dim=self.n_dim, - ) - self.projection = fno_block.MLP( - in_channels=self.hidden_channels, - out_channels=out_channels, - hidden_channels=self.projection_channels, - n_layers=2, - n_dim=self.n_dim, - non_linearity=non_linearity, - ) - - def forward(self, x): - """TFNO's forward pass""" - x = self.concat_to_tensor(x, self.input_keys) - - x = self.lifting(x) - if self.domain_padding is not None: - x = self.domain_padding.pad(x) - # x is 0.4 * [1, 32, 16, 16], passed - for index in range(self.n_layers): - x = self.fno_blocks(x, index) - - if self.domain_padding is not None: - x = self.domain_padding.unpad(x) - out = self.projection(x) - return {self.output_keys[0]: out} - - -class TFNO1dNet(FNONet): - """1D Fourier Neural Operator. - - Args: - input_keys (Tuple[str, ...]): Name of input keys, such as ("input",). - output_keys (Tuple[str, ...]): Name of output keys, such as ("output",). - n_modes_height (Tuple[int, ...]): Number of Fourier modes to keep along the height, along each - dimension. - hidden_channels (int): Width of the FNO (i.e. number of channels). - in_channels (int, optional): Number of input channels. Defaults to 3. - out_channels (int, optional): Number of output channels. Defaults to 1. - lifting_channels (int, optional): Number of hidden channels of the lifting block of the FNO. - Defaults to 256. - projection_channels (int, optional): Number of hidden channels of the projection block of the FNO. - Defaults to 256. - n_layers (int, optional): Number of Fourier Layers. Defaults to 4. - use_mlp (bool, optional): Whether to use an MLP layer after each FNO block. Defaults to False. - mlp (dict[str, float], optional): Parameters of the MLP. {'expansion': float, 'dropout': float}. - Defaults to None. - non_linearity (nn.functional, optional): Non-Linearity module to use. Defaults to F.gelu. - norm (F.module, optional): Normalization layer to use. Defaults to None. - preactivation (bool, optional): Whether to use resnet-style preactivation. Defaults to False. - skip (str, optional): Type of skip connection to use,{'linear', 'identity', 'soft-gating'}. - Defaults to "soft-gating". - separable (bool, optional): Whether to use a depthwise separable spectral convolution. - Defaults to False. - factorization (str, optional): Tensor factorization of the parameters weight to use. - * If None, a dense tensor parametrizes the Spectral convolutions. - * Otherwise, the specified tensor factorization is used. Defaults to "Tucker". - rank (float, optional): Rank of the tensor factorization of the Fourier weights. Defaults to 1.0. - joint_factorization (bool, optional): Whether all the Fourier Layers should be parametrized by a - single tensor (vs one per layer). Defaults to False. - implementation (str, optional): {'factorized', 'reconstructed'}, optional. Defaults to "factorized". - If factorization is not None, forward mode to use:: - * `reconstructed` : the full weight tensor is reconstructed from the factorization and used for the forward pass. - * `factorized` : the input is directly contracted with the factors of the decomposition. - domain_padding (Optional[Union[list, float, int]], optional): Whether to use percentage of padding. - Defaults to None. - domain_padding_mode (str, optional): {'symmetric', 'one-sided'}, optional - How to perform domain padding, by default 'one-sided'. Defaults to "one-sided". - fft_norm (str, optional): The normalization mode for the FFT. Defaults to "forward". - patching_levels (int, optional): Number of patching levels to use. Defaults to 0. - SpectralConv (nn.layer, optional): Spectral convolution layer to use. - Defaults to fno_block.FactorizedSpectralConv. - """ - - def __init__( - self, - input_keys: Tuple[str, ...], - output_keys: Tuple[str, ...], - n_modes_height: Tuple[int, ...], - hidden_channels: int, - in_channels: int = 3, - out_channels: int = 1, - lifting_channels: int = 256, - projection_channels: int = 256, - n_layers: int = 4, - non_linearity: nn.functional = F.gelu, - use_mlp: bool = False, - mlp: Optional[Dict[str, float]] = None, - norm: str = None, - skip: str = "soft-gating", - separable: bool = False, - preactivation: bool = False, - factorization: str = "Tucker", - rank: float = 1.0, - joint_factorization: bool = False, - implementation: str = "factorized", - domain_padding: Optional[Union[list, float, int]] = None, - domain_padding_mode: str = "one-sided", - fft_norm: str = "forward", - patching_levels: int = 0, - SpectralConv: nn.Layer = fno_block.FactorizedSpectralConv, - **kwargs, - ): - super().__init__( - input_keys=input_keys, - output_keys=output_keys, - n_modes=(n_modes_height,), - hidden_channels=hidden_channels, - in_channels=in_channels, - out_channels=out_channels, - lifting_channels=lifting_channels, - projection_channels=projection_channels, - n_layers=n_layers, - non_linearity=non_linearity, - use_mlp=use_mlp, - mlp=mlp, - norm=norm, - skip=skip, - separable=separable, - preactivation=preactivation, - factorization=factorization, - rank=rank, - joint_factorization=joint_factorization, - implementation=implementation, - domain_padding=domain_padding, - domain_padding_mode=domain_padding_mode, - fft_norm=fft_norm, - patching_levels=patching_levels, - SpectralConv=SpectralConv, - ) - self.n_modes_height = n_modes_height - - -class TFNO2dNet(FNONet): - """2D Fourier Neural Operator. - - Args: - input_keys (Tuple[str, ...]): Name of input keys, such as ("input",). - output_keys (Tuple[str, ...]): Name of output keys, such as ("output",). - n_modes_height (int): Number of Fourier modes to keep along the height. - n_modes_width (int): Number of modes to keep in Fourier Layer, along the width. - hidden_channels (int): Width of the FNO (i.e. number of channels). - in_channels (int, optional): Number of input channels. Defaults to 3. - out_channels (int, optional): Number of output channels. Defaults to 1. - lifting_channels (int, optional): Number of hidden channels of the lifting block of the FNO. - Defaults to 256. - projection_channels (int, optional): Number of hidden channels of the projection block of the FNO. - Defaults to 256. - n_layers (int, optional): Number of Fourier Layers. Defaults to 4. - use_mlp (bool, optional): Whether to use an MLP layer after each FNO block. Defaults to False. - mlp (Dict[str, float], optional): Parameters of the MLP. {'expansion': float, 'dropout': float}. - Defaults to None. - non_linearity (nn.Layer, optional): Non-Linearity module to use. Defaults to F.gelu. - norm (F.module, optional): Normalization layer to use. Defaults to None. - preactivation (bool, optional): Whether to use resnet-style preactivation. Defaults to False. - skip (str, optional): Type of skip connection to use,{'linear', 'identity', 'soft-gating'}. - Defaults to "soft-gating". - separable (bool, optional): Whether to use a depthwise separable spectral convolution. - Defaults to False. - factorization (str, optional): Tensor factorization of the parameters weight to use. - * If None, a dense tensor parametrizes the Spectral convolutions. - * Otherwise, the specified tensor factorization is used. Defaults to "Tucker". - rank (float, optional): Rank of the tensor factorization of the Fourier weights. Defaults to 1.0. - joint_factorization (bool, optional): Whether all the Fourier Layers should be parametrized by a - single tensor (vs one per layer). Defaults to False. - implementation (str, optional): {'factorized', 'reconstructed'}, optional. Defaults to "factorized". - If factorization is not None, forward mode to use:: - * `reconstructed` : the full weight tensor is reconstructed from the factorization and used for the forward pass. - * `factorized` : the input is directly contracted with the factors of the decomposition. - domain_padding (Union[list,float,int], optional): Whether to use percentage of padding. Defaults to - None. - domain_padding_mode (str, optional): {'symmetric', 'one-sided'}, optional - How to perform domain padding, by default 'one-sided'. Defaults to "one-sided". - fft_norm (str, optional): The normalization mode for the FFT. Defaults to "forward". - patching_levels (int, optional): Number of patching levels to use. Defaults to 0. - SpectralConv (nn.layer, optional): Spectral convolution layer to use. - Defaults to fno_block.FactorizedSpectralConv. - """ - - def __init__( - self, - input_keys: Tuple[str, ...], - output_keys: Tuple[str, ...], - n_modes_height: int, - n_modes_width: int, - hidden_channels: int, - in_channels: int = 3, - out_channels: int = 1, - lifting_channels: int = 256, - projection_channels: int = 256, - n_layers: int = 4, - non_linearity: nn.functional = F.gelu, - use_mlp: bool = False, - mlp: Optional[Dict[str, float]] = None, - norm: str = None, - skip: str = "soft-gating", - separable: bool = False, - preactivation: bool = False, - factorization: str = "Tucker", - rank: float = 1.0, - joint_factorization: bool = False, - implementation: str = "factorized", - domain_padding: Optional[Union[list, float, int]] = None, - domain_padding_mode: str = "one-sided", - fft_norm: str = "forward", - patching_levels: int = 0, - SpectralConv: nn.layer = fno_block.FactorizedSpectralConv, - **kwargs, - ): - super().__init__( - input_keys=input_keys, - output_keys=output_keys, - n_modes=(n_modes_height, n_modes_width), - hidden_channels=hidden_channels, - in_channels=in_channels, - out_channels=out_channels, - lifting_channels=lifting_channels, - projection_channels=projection_channels, - n_layers=n_layers, - non_linearity=non_linearity, - use_mlp=use_mlp, - mlp=mlp, - norm=norm, - skip=skip, - separable=separable, - preactivation=preactivation, - factorization=factorization, - rank=rank, - joint_factorization=joint_factorization, - implementation=implementation, - domain_padding=domain_padding, - domain_padding_mode=domain_padding_mode, - fft_norm=fft_norm, - patching_levels=patching_levels, - SpectralConv=SpectralConv, - ) - self.n_modes_height = n_modes_height - self.n_modes_width = n_modes_width - - -class TFNO3dNet(FNONet): - """3D Fourier Neural Operator. - - Args: - input_keys (Tuple[str, ...]): Name of input keys, such as ("input",). - output_keys (Tuple[str, ...]): Name of output keys, such as ("output",). - n_modes_height (int): Number of Fourier modes to keep along the height. - n_modes_width (int): Number of modes to keep in Fourier Layer, along the width. - n_modes_depth (int): Number of Fourier modes to keep along the depth. - hidden_channels (int): Width of the FNO (i.e. number of channels). - in_channels (int, optional): Number of input channels. Defaults to 3. - out_channels (int, optional): Number of output channels. Defaults to 1. - lifting_channels (int, optional): Number of hidden channels of the lifting block of the FNO. - Defaults to 256. - projection_channels (int, optional): Number of hidden channels of the projection block of the FNO. - Defaults to 256. - n_layers (int, optional): Number of Fourier Layers. Defaults to 4. - use_mlp (bool, optional): Whether to use an MLP layer after each FNO block. Defaults to False. - mlp (Dict[str, float], optional): Parameters of the MLP. {'expansion': float, 'dropout': float}. - Defaults to None. - non_linearity (nn.Layer, optional): Non-Linearity module to use. Defaults to F.gelu. - norm (F.module, optional): Normalization layer to use. Defaults to None. - preactivation (bool, optional): Whether to use resnet-style preactivation. Defaults to False. - skip (str, optional): Type of skip connection to use,{'linear', 'identity', 'soft-gating'}. - Defaults to "soft-gating". - separable (bool, optional): Whether to use a depthwise separable spectral convolution. - Defaults to False. - factorization (str, optional): Tensor factorization of the parameters weight to use. - * If None, a dense tensor parametrizes the Spectral convolutions. - * Otherwise, the specified tensor factorization is used. Defaults to "Tucker". - rank (float, optional): Rank of the tensor factorization of the Fourier weights. Defaults to 1.0. - joint_factorization (bool, optional): Whether all the Fourier Layers should be parametrized by a - single tensor (vs one per layer). Defaults to False. - implementation (str, optional): {'factorized', 'reconstructed'}, optional. Defaults to "factorized". - If factorization is not None, forward mode to use:: - * `reconstructed` : the full weight tensor is reconstructed from the factorization and used for the forward pass. - * `factorized` : the input is directly contracted with the factors of the decomposition. - domain_padding (str, optional): Whether to use percentage of padding. Defaults to None. - domain_padding_mode (str, optional): {'symmetric', 'one-sided'}, optional - How to perform domain padding, by default 'one-sided'. Defaults to "one-sided". - fft_norm (str, optional): The normalization mode for the FFT. Defaults to "forward". - patching_levels (int, optional): Number of patching levels to use. Defaults to 0. - SpectralConv (nn.layer, optional): Spectral convolution layer to use. Defaults to fno_block. - FactorizedSpectralConv. - """ - - def __init__( - self, - input_keys: Tuple[str, ...], - output_keys: Tuple[str, ...], - n_modes_height: int, - n_modes_width: int, - n_modes_depth: int, - hidden_channels: int, - in_channels: int = 3, - out_channels: int = 1, - lifting_channels: int = 256, - projection_channels: int = 256, - n_layers: int = 4, - non_linearity: nn.functional = F.gelu, - use_mlp: bool = False, - mlp: Optional[Dict[str, float]] = None, - norm: str = None, - skip: str = "soft-gating", - separable: bool = False, - preactivation: bool = False, - factorization: str = "Tucker", - rank: float = 1.0, - joint_factorization: bool = False, - implementation: str = "factorized", - domain_padding: Optional[Union[list, float, int]] = None, - domain_padding_mode: str = "one-sided", - fft_norm: str = "forward", - patching_levels: int = 0, - SpectralConv: nn.layer = fno_block.FactorizedSpectralConv, - **kwargs, - ): - super().__init__( - input_keys=input_keys, - output_keys=output_keys, - n_modes=(n_modes_height, n_modes_width, n_modes_depth), - hidden_channels=hidden_channels, - in_channels=in_channels, - out_channels=out_channels, - lifting_channels=lifting_channels, - projection_channels=projection_channels, - n_layers=n_layers, - non_linearity=non_linearity, - use_mlp=use_mlp, - mlp=mlp, - norm=norm, - skip=skip, - separable=separable, - preactivation=preactivation, - factorization=factorization, - rank=rank, - joint_factorization=joint_factorization, - implementation=implementation, - domain_padding=domain_padding, - domain_padding_mode=domain_padding_mode, - fft_norm=fft_norm, - patching_levels=patching_levels, - SpectralConv=SpectralConv, - ) - self.n_modes_height = n_modes_height - self.n_modes_width = n_modes_width - self.n_modes_depth = n_modes_depth diff --git a/examples/smc_reac/ppsci/arch/tgcn.py b/examples/smc_reac/ppsci/arch/tgcn.py deleted file mode 100644 index 5cf1ebc3eb..0000000000 --- a/examples/smc_reac/ppsci/arch/tgcn.py +++ /dev/null @@ -1,200 +0,0 @@ -from typing import Tuple - -import paddle as pp -import paddle.nn.functional as F -from numpy import ndarray -from paddle import nn -from paddle.nn.initializer import KaimingNormal - -from ppsci.arch.base import Arch - - -class graph_conv(nn.Layer): - def __init__(self, in_dim, out_dim, dropout, num_layer=2): - super(graph_conv, self).__init__() - self.mlp = nn.Conv2D( - (num_layer + 1) * in_dim, - out_dim, - kernel_size=(1, 1), - weight_attr=KaimingNormal(), - ) - self.dropout = dropout - self.num_layer = num_layer - - def forward(self, x, adj): - # B C N T - out = [x] - for _ in range(self.num_layer): - new_x = pp.matmul(adj, x) - out.append(new_x) - x = new_x - - h = pp.concat(out, axis=1) - h = self.mlp(h) - h = F.dropout(h, self.dropout, training=self.training) - return h - - -class tempol_conv(nn.Layer): - def __init__(self, in_dim, out_dim, hidden, num_layer=3, k_s=3, alpha=0.1): - super(tempol_conv, self).__init__() - self.leakyrelu = nn.LeakyReLU(alpha) - self.tc_convs = nn.LayerList() - self.num_layer = num_layer - for i in range(num_layer): - in_channels = in_dim if i == 0 else hidden - self.tc_convs.append( - nn.Conv2D( - in_channels=in_channels, - out_channels=hidden, - kernel_size=(1, k_s), - padding=(0, i + 1), - dilation=i + 1, - weight_attr=KaimingNormal(), - ) - ) - - self.mlp = nn.Conv2D( - in_channels=in_dim + hidden * num_layer, - out_channels=out_dim, - kernel_size=(1, 1), - weight_attr=KaimingNormal(), - ) - - def forward(self, x): - # B C N T - x_cat = [x] - for i in range(self.num_layer): - x = self.leakyrelu(self.tc_convs[i](x)) - x_cat.append(x) - tc_out = self.mlp(pp.concat(x_cat, axis=1)) - return tc_out - - -class TGCN(Arch): - """ - TGCN is a class that represents an Temporal Graph Convolutional Network model. - - Args: - input_keys (Tuple[str, ...]): A tuple of input keys. - output_keys (Tuple[str, ...]): A tuple of output keys. - adj (ndarray): The adjacency matrix of the graph. - in_dim (int): The dimension of the input data. - emb_dim (int): The dimension of the embedded space. - hidden (int): The dimension of the latent space. - gc_layer (int): The number of the graph convolutional layer. - tc_layer (int): The number of the temporal convolutional layer. - k_s (int): The kernel size of the temporal convolutional layer. - dropout (float): The dropout rate. - alpha (float): The negative slope of LeakyReLU. - input_len (int): The input timesteps. - label_len (int): The output timesteps. - - Examples: - >>> import paddle - >>> import ppsci - >>> model = ppsci.arch.TGCN( - ... input_keys=("input",), - ... output_keys=("label",), - ... adj=numpy.ones((307, 307), dtype=numpy.float32), - ... in_dim=1, - ... emb_dim=32 - ... hidden=64, - ... gc_layer=2, - ... tc_layer=2 - ... k_s=3, - ... dropout=0.25, - ... alpha=0.1, - ... input_len=12, - ... label_len=12, - ... ) - >>> input_dict = {"input": paddle.rand([64, 12, 307, 1]),} - >>> label_dict = model(input_dict) - >>> print(label_dict["label"].shape) - [64, 12, 307, 1] - """ - - def __init__( - self, - input_keys: Tuple[str, ...], - output_keys: Tuple[str, ...], - adj: ndarray, - in_dim: int, - emb_dim: int, - hidden: int, - gc_layer: int, - tc_layer: int, - k_s: int, - dropout: float, - alpha: float, - input_len: int, - label_len: int, - ): - super(TGCN, self).__init__() - - self.input_keys = input_keys - self.output_keys = output_keys - - self.register_buffer("adj", pp.to_tensor(data=adj)) - - self.emb_conv = nn.Conv2D( - in_channels=in_dim, - out_channels=emb_dim, - kernel_size=(1, 1), - weight_attr=KaimingNormal(), - ) - - self.tc1_conv = tempol_conv( - emb_dim, hidden, hidden, num_layer=tc_layer, k_s=k_s, alpha=alpha - ) - self.sc1_conv = graph_conv(hidden, hidden, dropout, num_layer=gc_layer) - self.bn1 = nn.BatchNorm2D(hidden) - - self.tc2_conv = tempol_conv( - hidden, hidden, hidden, num_layer=tc_layer, k_s=k_s, alpha=alpha - ) - self.sc2_conv = graph_conv(hidden, hidden, dropout, num_layer=gc_layer) - self.bn2 = nn.BatchNorm2D(hidden) - - self.end_conv_1 = nn.Conv2D( - in_channels=emb_dim + hidden + hidden, - out_channels=2 * hidden, - kernel_size=(1, 1), - weight_attr=KaimingNormal(), - ) - self.end_conv_2 = nn.Conv2D( - in_channels=2 * hidden, - out_channels=label_len, - kernel_size=(1, input_len), - weight_attr=KaimingNormal(), - ) - - def forward(self, raw): - # emb block - x = raw[self.input_keys[0]] - x = x.transpose(perm=[0, 3, 2, 1]) # B in_dim N T - emb_x = self.emb_conv(x) # B emd_dim N T - - # TC1 - tc1_out = self.tc1_conv(emb_x) # B hidden N T - - # SC1 - sc1_out = self.sc1_conv(tc1_out, self.adj) # B hidden N T - sc1_out = sc1_out + tc1_out - sc1_out = self.bn1(sc1_out) - - # TC2 - tc2_out = self.tc2_conv(sc1_out) # B hidden N T - - # SC2 - sc2_out = self.sc2_conv(tc2_out, self.adj) # B hidden N T - sc2_out = sc2_out + tc2_out - sc2_out = self.bn2(sc2_out) - - # readout block - x_out = F.relu(pp.concat((emb_x, sc1_out, sc2_out), axis=1)) - x_out = F.relu(self.end_conv_1(x_out)) - # transform - x_out = self.end_conv_2(x_out) # B T N 1 - - return {self.output_keys[0]: x_out} diff --git a/examples/smc_reac/ppsci/arch/transformer.py b/examples/smc_reac/ppsci/arch/transformer.py deleted file mode 100644 index a80eb46dcf..0000000000 --- a/examples/smc_reac/ppsci/arch/transformer.py +++ /dev/null @@ -1,417 +0,0 @@ -# Copyright (c) 2024 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. - -""" -Reference: https://github.com/omron-sinicx/transformer4sr -""" - -from __future__ import annotations - -import math -from typing import Callable -from typing import Tuple - -import paddle -import paddle.nn as nn - -from ppsci.arch import activation as act_mod -from ppsci.arch import base - - -def transpose_aux_func(dims, dim0, dim1): - perm = list(range(dims)) - perm[dim0], perm[dim1] = perm[dim1], perm[dim0] - return perm - - -class MultiHeadAttention(nn.Layer): - def __init__(self, heads, d_model): - super().__init__() - self.heads = heads - self.d_model = d_model - assert d_model % heads == 0 - self.d_k = d_model // heads - self.W_Q = nn.Linear(in_features=d_model, out_features=d_model) - self.W_K = nn.Linear(in_features=d_model, out_features=d_model) - self.W_V = nn.Linear(in_features=d_model, out_features=d_model) - self.W_O = nn.Linear(in_features=d_model, out_features=d_model) - - def scaled_dot_product_attention(self, Q, K, V, mask=None): - scores = paddle.matmul( - x=Q, y=K.transpose(perm=transpose_aux_func(K.ndim, -1, -2)) - ) / math.sqrt(self.d_k) - if mask is not None: - scores = paddle.where( - condition=mask, - x=paddle.to_tensor(data=[-1e9], dtype="float32"), - y=scores, - ) - weights = nn.functional.softmax(x=scores, axis=-1) - return paddle.matmul(x=weights, y=V) - - def forward(self, Q, K, V, mask=None): - Q_temp = paddle.reshape( - x=self.W_Q(Q), - shape=[i for i in tuple(Q.shape)[:-1]] + [self.heads] + [self.d_k], - ).transpose( - perm=transpose_aux_func( - paddle.reshape( - x=self.W_Q(Q), - shape=[i for i in tuple(Q.shape)[:-1]] + [self.heads] + [self.d_k], - ).ndim, - 1, - 2, - ) - ) - K_temp = paddle.reshape( - x=self.W_K(K), - shape=[i for i in tuple(K.shape)[:-1]] + [self.heads] + [self.d_k], - ).transpose( - perm=transpose_aux_func( - paddle.reshape( - x=self.W_K(K), - shape=[i for i in tuple(K.shape)[:-1]] + [self.heads] + [self.d_k], - ).ndim, - 1, - 2, - ) - ) - V_temp = paddle.reshape( - x=self.W_V(V), - shape=[i for i in tuple(V.shape)[:-1]] + [self.heads] + [self.d_k], - ).transpose( - perm=transpose_aux_func( - paddle.reshape( - x=self.W_V(V), - shape=[i for i in tuple(V.shape)[:-1]] + [self.heads] + [self.d_k], - ).ndim, - 1, - 2, - ) - ) - sdpa = self.scaled_dot_product_attention( - Q_temp, K_temp, V_temp, mask - ).transpose( - perm=transpose_aux_func( - self.scaled_dot_product_attention(Q_temp, K_temp, V_temp, mask).ndim, - 1, - 2, - ) - ) - sdpa = paddle.reshape( - x=sdpa, shape=[i for i in tuple(sdpa.shape)[:-2]] + [self.d_model] - ) - y_mha = self.W_O(sdpa) - return y_mha - - -class MLP(nn.Layer): - def __init__(self, list_dims, act="relu", dropout=0.0): - super().__init__() - self.layers = nn.LayerList() - for i in range(len(list_dims) - 1): - self.layers.append( - nn.Linear(in_features=list_dims[i], out_features=list_dims[i + 1]) - ) - self.layers.append(act_mod.get_activation(act) if act else None) - self.layers.append(nn.Dropout(p=dropout)) - - def forward(self, x): - y = x - for layer in self.layers: - y = layer(y) - return y - - -class EncoderLayerMix(nn.Layer): - def __init__(self, in_features, d_model, heads, act="relu", dropout=0.0): - super().__init__() - self.mlp = MLP([in_features, d_model, d_model], act="relu", dropout=dropout) - self.multihead_attention = MultiHeadAttention(heads, d_model) - self.dropout = nn.Dropout(p=dropout) - self.norm = nn.LayerNorm(normalized_shape=d_model) - - def forward(self, x): - y = x - y = paddle.flatten(y, start_axis=2) - y = self.mlp(y) - y = self.multihead_attention(y, y, y, mask=None) - y = self.dropout(y) - y = paddle.unsqueeze(y, axis=2) - y = x + y - y = self.norm(y) - return y - - -class Encoder(nn.Layer): - def __init__( - self, num_layers, num_var_max, d_model, heads, act="relu", dropout=0.0 - ): - super().__init__() - self.first_mlp = MLP([1, d_model, d_model], act="relu", dropout=dropout) - self.layers = nn.LayerList( - sublayers=[ - EncoderLayerMix( - d_model * num_var_max, d_model, heads, act="relu", dropout=dropout - ) - for _ in range(num_layers) - ] - ) - self.last_mlp = MLP([d_model, d_model], act="relu", dropout=dropout) - - def forward(self, x): - y = x - y = self.first_mlp(y) - for layer in self.layers: - y = layer(y) - y = self.last_mlp(y) - y = paddle.max(y, axis=1) - return y - - -class TokenEmbeddings(nn.Layer): - def __init__(self, vocab_size, seq_length, d_model, dropout=0.0): - super().__init__() - self.embed = nn.Embedding(num_embeddings=vocab_size, embedding_dim=d_model) - self.seq_length = seq_length - self.d_model = d_model - self.dropout = nn.Dropout(dropout) - self.get_pe_num() - - def get_pe_num(self): - self.pe = paddle.zeros(shape=[self.seq_length, self.d_model]) - numerator = paddle.arange( - self.seq_length, dtype=paddle.get_default_dtype() - ).unsqueeze(axis=1) - denominator = paddle.pow( - x=paddle.to_tensor(10e4, dtype=paddle.get_default_dtype()), - y=paddle.arange(self.d_model, step=2) / self.d_model, - ).unsqueeze(axis=0) - self.pe[:, 0::2] = paddle.sin(x=numerator / denominator) - self.pe[:, 1::2] = paddle.cos(x=numerator / denominator) - self.pe.stop_gradient = True - - def forward(self, x): - # embedding - y = x - y = self.embed(y) * math.sqrt(self.d_model) - # position encoding - y = self.dropout(y + self.pe) - return y - - -class DecoderLayer(nn.Layer): - def __init__(self, heads, d_model, act="relu", dropout=0.0): - super().__init__() - self.multihead_attention_1 = MultiHeadAttention(heads, d_model) - self.dropout_1 = nn.Dropout(p=dropout) - self.norm_1 = nn.LayerNorm(d_model) - - self.multihead_attention_2 = MultiHeadAttention(heads, d_model) - self.dropout_2 = nn.Dropout(p=dropout) - self.norm_2 = nn.LayerNorm(d_model) - - self.mlp = MLP([d_model, 2 * d_model, d_model], act="relu", dropout=dropout) - self.norm_3 = nn.LayerNorm(d_model) - - def forward(self, x_emb, x_enc, mask): - y_mha_1 = self.multihead_attention_1(x_emb, x_emb, x_emb, mask=mask) - y_mha_1 = self.dropout_1(y_mha_1) - y = y_mha_1 + x_emb - y = self.norm_1(y) - y_mha_2 = self.multihead_attention_2(y, x_enc, x_enc, mask=None) - y_mha_2 = self.dropout_2(y_mha_2) - y = y + y_mha_2 - y = self.norm_2(y) - y_mlp = self.mlp(y) - y = y + y_mlp - y = self.norm_3(y) - return y - - -class Decoder(nn.Layer): - def __init__( - self, - num_layers, - vocab_size, - seq_length, - d_model, - heads, - act="relu", - dropout=0.0, - ): - super().__init__() - self.token_embeddings = TokenEmbeddings( - vocab_size, seq_length, d_model, dropout - ) - self.dropout = nn.Dropout(p=dropout) - self.layers = nn.LayerList( - sublayers=[ - DecoderLayer(heads, d_model, act="relu", dropout=dropout) - for _ in range(num_layers) - ] - ) - - def forward(self, x_target, x_enc, mask): - y = x_target - y = self.token_embeddings(y) - y = self.dropout(y) - for layer in self.layers: - y = layer(y, x_enc, mask) - return y - - -class Transformer(base.Arch): - """A Kind of Transformer Model. - - Args: - input_keys (Tuple[str, ...]): Name of input keys, such as ("x", "y", "z"). - output_keys (Tuple[str, ...]): Name of output keys, such as ("u", "v", "w"). - num_var_max (int): Maximum number of variables. - vocab_size (int): Size of vocab. Size of unary operators = 1, binary operators = 2. - seq_length (int): Length of sequance. - d_model (int, optional): The innermost dimension of model. Defaults to 256. - heads (int, optional): The number of independent heads for the multi-head attention layers. Defaults to 4. - num_layers_enc (int, optional): The number of encoders. Defaults to 4. - num_layers_dec (int, optional): The number of decoders. Defaults to 8. - dropout (float, optional): Dropout regularization. Defaults to 0.0. - - Examples: - >>> import paddle - >>> import ppsci - >>> model = ppsci.arch.Transformer( - ... input_keys=("input", "target_seq"), - ... output_keys=("output",), - ... num_var_max=7, - ... vocab_size=20, - ... seq_length=30, - ... ) - >>> input_dict = {"input": paddle.rand([512, 50, 7, 1]), - ... "target_seq": paddle.rand([512, 30])} - >>> output_dict = model(input_dict) - >>> print(output_dict["output"].shape) - [512, 30, 20] - """ - - def __init__( - self, - input_keys: Tuple[str, ...], - output_keys: Tuple[str, ...], - num_var_max: int, - vocab_size: int, - seq_length: int, - d_model: int = 256, - heads: int = 4, - num_layers_enc: int = 4, - num_layers_dec: int = 8, - act: str = "relu", - dropout: float = 0.0, - ): - super().__init__() - self.input_keys = input_keys - self.output_keys = output_keys - self.num_var_max = num_var_max - self.vocab_size = vocab_size - self.seq_length = seq_length - self.d_model = d_model - self.heads = heads - self.num_layers_enc = num_layers_enc - self.num_layers_dec = num_layers_dec - self.act = act - self.dropout = dropout - - self.encoder = Encoder( - num_layers_enc, num_var_max, d_model, heads, act="relu", dropout=dropout - ) - self.decoder = Decoder( - num_layers_dec, - vocab_size, - seq_length, - d_model, - heads, - act="relu", - dropout=dropout, - ) - self.last_layer = paddle.nn.Linear(in_features=d_model, out_features=vocab_size) - - def get_mask(self, target_seq): - padding_mask = paddle.equal(target_seq, 0).unsqueeze(axis=1).unsqueeze(axis=1) - future_mask = paddle.triu( - paddle.ones(shape=[target_seq.shape[1], target_seq.shape[1]]), - diagonal=1, - ).astype(dtype="bool") - mask = paddle.logical_or(x=padding_mask, y=future_mask) - return mask - - def forward_tensor(self, x_lst): - y, target_seq = x_lst[0], x_lst[1] - mask = self.get_mask(target_seq) - y_enc = self.encoder(y) - y = self.decoder(target_seq, y_enc, mask) - y = self.last_layer(y) - return y - - def forward(self, x): - if self._input_transform is not None: - x = self._input_transform(x) - - x_lst = [x[key] for key in self.input_keys] # input, target_seq - y = self.forward_tensor(x_lst) - y = self.split_to_dict(y, self.output_keys, axis=-1) - - if self._output_transform is not None: - y = self._output_transform(x, y) - return y - - @paddle.no_grad() - def decode_process( - self, dataset: paddle.Tensor, complete_func: Callable - ) -> paddle.Tensor: - """Greedy decode with the Transformer model, decode until the equation tree is completed. - - Args: - dataset (paddle.Tensor): Tabular dataset. - complete_func (Callable): Function used to calculate whether inference is complete. - """ - encoder_output = self.encoder(dataset) - decoder_output = paddle.zeros( - shape=(dataset.shape[0], self.seq_length + 1), dtype=paddle.int64 - ) - decoder_output[:, 0] = 1 - is_complete = paddle.zeros(shape=dataset.shape[0], dtype=paddle.bool) - for n1 in range(self.seq_length): - padding_mask = ( - paddle.equal(x=decoder_output[:, :-1], y=0) - .unsqueeze(axis=1) - .unsqueeze(axis=1) - ) - future_mask = paddle.triu( - x=paddle.ones(shape=[self.seq_length, self.seq_length]), diagonal=1 - ).astype(dtype=paddle.bool) - mask_dec = paddle.logical_or(x=padding_mask, y=future_mask) - y_dec = self.decoder( - x_target=decoder_output[:, :-1], - x_enc=encoder_output, - mask=mask_dec, - ) - y_mlp = self.last_layer(y_dec) - # set value depending on complete condition - decoder_output[:, n1 + 1] = paddle.where( - is_complete, 0, paddle.argmax(y_mlp[:, n1], axis=-1) - ) - # set complete condition - for n2 in range(dataset.shape[0]): - if complete_func(decoder_output[n2, 1:]): - is_complete[n2] = True - return decoder_output diff --git a/examples/smc_reac/ppsci/arch/unetex.py b/examples/smc_reac/ppsci/arch/unetex.py deleted file mode 100644 index d0ba170464..0000000000 --- a/examples/smc_reac/ppsci/arch/unetex.py +++ /dev/null @@ -1,290 +0,0 @@ -# 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. - -from typing import Optional -from typing import Tuple -from typing import Type - -import paddle -from paddle import nn - -from ppsci.arch import base - - -def create_layer( - in_channel, - out_channel, - kernel_size, - weight_norm=True, - batch_norm=True, - activation=nn.ReLU, - convolution=nn.Conv2D, -): - if kernel_size % 2 == 0: - raise ValueError("kernel_size should even number") - conv = convolution(in_channel, out_channel, kernel_size, padding=kernel_size // 2) - if weight_norm: - conv = nn.util.weight_norm(conv) - layer = [] - layer.append(conv) - if activation is not None: - layer.append(activation()) - if batch_norm: - layer.append(nn.BatchNorm2D(out_channel)) - return nn.Sequential(*layer) - - -def create_encoder_block( - in_channel, - out_channel, - kernel_size, - weight_norm=True, - batch_norm=True, - activation=nn.ReLU, - layers=2, -): - encoder = [] - encoder.append( - create_layer( - in_channel, - out_channel, - kernel_size, - weight_norm, - batch_norm, - activation, - nn.Conv2D, - ) - ) - for i in range(layers - 1): - encoder.append( - create_layer( - out_channel, - out_channel, - kernel_size, - weight_norm, - batch_norm, - activation, - nn.Conv2D, - ) - ) - return nn.Sequential(*encoder) - - -def create_decoder_block( - in_channel, - out_channel, - kernel_size, - weight_norm=True, - batch_norm=True, - activation=nn.ReLU, - layers=2, - final_layer=False, -): - decoder = [] - for i in range(layers): - _in = in_channel - _out = in_channel - _batch_norm = batch_norm - _activation = activation - if i == 0: - _in = in_channel * 2 - if i == layers - 1: - _out = out_channel - if final_layer: - _batch_norm = False - _activation = None - decoder.append( - create_layer( - _in, - _out, - kernel_size, - weight_norm, - _batch_norm, - _activation, - nn.Conv2DTranspose, - ) - ) - return nn.Sequential(*decoder) - - -def create_encoder( - in_channel, filters, kernel_size, wn=True, bn=True, activation=nn.ReLU, layers=2 -): - encoder = [] - for i in range(len(filters)): - encoder_layer = create_encoder_block( - in_channel if i == 0 else filters[i - 1], - filters[i], - kernel_size, - wn, - bn, - activation, - layers, - ) - encoder = encoder + [encoder_layer] - return nn.Sequential(*encoder) - - -def create_decoder( - out_channel, - filters, - kernel_size, - weight_norm=True, - batch_norm=True, - activation=nn.ReLU, - layers=2, -): - decoder = [] - for i in range(len(filters)): - if i == 0: - decoder_layer = create_decoder_block( - filters[i], - out_channel, - kernel_size, - weight_norm, - batch_norm, - activation, - layers, - final_layer=True, - ) - else: - decoder_layer = create_decoder_block( - filters[i], - filters[i - 1], - kernel_size, - weight_norm, - batch_norm, - activation, - layers, - final_layer=False, - ) - decoder = [decoder_layer] + decoder - return nn.Sequential(*decoder) - - -class UNetEx(base.Arch): - """U-Net Extension for CFD. - - Reference: [Ribeiro M D, Rehman A, Ahmed S, et al. DeepCFD: Efficient steady-state laminar flow approximation with deep convolutional neural networks[J]. arXiv preprint arXiv:2004.08826, 2020.](https://arxiv.org/abs/2004.08826) - - Args: - input_key (str): Name of function data for input. - output_key (str): Name of function data for output. - in_channel (int): Number of channels of input. - out_channel (int): Number of channels of output. - kernel_size (int, optional): Size of kernel of convolution layer. Defaults to 3. - filters (Tuple[int, ...], optional): Number of filters. Defaults to (16, 32, 64). - layers (int, optional): Number of encoders or decoders. Defaults to 3. - weight_norm (bool, optional): Whether use weight normalization layer. Defaults to True. - batch_norm (bool, optional): Whether add batch normalization layer. Defaults to True. - activation (Type[nn.Layer], optional): Name of activation function. Defaults to nn.ReLU. - final_activation (Optional[Type[nn.Layer]]): Name of final activation function. Defaults to None. - - Examples: - >>> import ppsci - >>> model = ppsci.arch.UNetEx( - ... input_key="input", - ... output_key="output", - ... in_channel=3, - ... out_channel=3, - ... kernel_size=5, - ... filters=(4, 4, 4, 4), - ... layers=3, - ... weight_norm=False, - ... batch_norm=False, - ... activation=None, - ... final_activation=None, - ... ) - >>> input_dict = {'input': paddle.rand([4, 3, 4, 4])} - >>> output_dict = model(input_dict) - >>> print(output_dict['output']) # doctest: +SKIP - >>> print(output_dict['output'].shape) - [4, 3, 4, 4] - """ - - def __init__( - self, - input_key: str, - output_key: str, - in_channel: int, - out_channel: int, - kernel_size: int = 3, - filters: Tuple[int, ...] = (16, 32, 64), - layers: int = 3, - weight_norm: bool = True, - batch_norm: bool = True, - activation: Type[nn.Layer] = nn.ReLU, - final_activation: Optional[Type[nn.Layer]] = None, - ): - if len(filters) == 0: - raise ValueError("The filters shouldn't be empty ") - - super().__init__() - self.input_keys = (input_key,) - self.output_keys = (output_key,) - self.final_activation = final_activation - self.encoder = create_encoder( - in_channel, - filters, - kernel_size, - weight_norm, - batch_norm, - activation, - layers, - ) - decoders = [ - create_decoder( - 1, filters, kernel_size, weight_norm, batch_norm, activation, layers - ) - for i in range(out_channel) - ] - self.decoders = nn.Sequential(*decoders) - - def encode(self, x): - tensors = [] - indices = [] - sizes = [] - for encoder in self.encoder: - x = encoder(x) - sizes.append(x.shape) - tensors.append(x) - x, ind = nn.functional.max_pool2d(x, 2, 2, return_mask=True) - indices.append(ind) - return x, tensors, indices, sizes - - def decode(self, x, tensors, indices, sizes): - y = [] - for _decoder in self.decoders: - _x = x - _tensors = tensors[:] - _indices = indices[:] - _sizes = sizes[:] - for decoder in _decoder: - tensor = _tensors.pop() - size = _sizes.pop() - indice = _indices.pop() - # upsample operations - _x = nn.functional.max_unpool2d(_x, indice, 2, 2, output_size=size) - _x = paddle.concat([tensor, _x], axis=1) - _x = decoder(_x) - y.append(_x) - return paddle.concat(y, axis=1) - - def forward(self, x): - x = x[self.input_keys[0]] - x, tensors, indices, sizes = self.encode(x) - x = self.decode(x, tensors, indices, sizes) - if self.final_activation is not None: - x = self.final_activation(x) - return {self.output_keys[0]: x} diff --git a/examples/smc_reac/ppsci/arch/unonet.py b/examples/smc_reac/ppsci/arch/unonet.py deleted file mode 100644 index c238a55be1..0000000000 --- a/examples/smc_reac/ppsci/arch/unonet.py +++ /dev/null @@ -1,289 +0,0 @@ -from typing import Dict -from typing import Optional -from typing import Tuple -from typing import Union - -import paddle -import paddle.nn as nn -import paddle.nn.functional as F - -from ppsci.arch import base -from ppsci.arch import fno_block - - -class UNONet(base.Arch): - """N-Dimensional U-Shaped Neural Operator. - - Args: - input_keys (Tuple[str, ...]): Name of input keys, such as ("input",). - output_keys (Tuple[str, ...]): Name of output keys, such as ("output",). - in_channels (int, optional): Number of input channels. - out_channels (int, optional): Number of output channels. - hidden_channels (int): Width of the FNO (i.e. number of channels). - lifting_channels (int, optional): Number of hidden channels of the lifting block of the FNO. - Defaults to 256. - projection_channels (int, optional): Number of hidden channels of the projection block of the FNO. - Defaults to 256. - n_layers (int, optional): Number of Fourier Layers. Defaults to 4. - uno_out_channels (Tuple[int, ...], optional): Number of output channel of each Fourier Layers. - Eaxmple: For a Five layer UNO uno_out_channels can be [32,64,64,64,32].c - uno_n_modes (Tuple[Tuple[int, ...], ...]): Number of Fourier Modes to use in integral operation of each - Fourier Layers (along each dimension). - Example: For a five layer UNO with 2D input the uno_n_modes can be: [[5,5],[5,5],[5,5],[5,5],[5,5]]. Defaults to None. - uno_scalings (Tuple[Tuple[int, ...], ...]): Scaling Factors for each Fourier Layers. - Example: For a five layer UNO with 2D input, the uno_scalings can be : [[1.0,1.0],[0.5,0.5],[1,1],[1,1],[2,2]].Defaults to None. - horizontal_skips_map (Dict, optional): A map {...., b: a, ....} denoting horizontal skip connection - from a-th layer to b-th layer. If None default skip connection is applied. - Example: For a 5 layer UNO architecture, the skip connections can be horizontal_skips_map ={4:0,3:1}.Defaults to None. - incremental_n_modes (tuple[int],optional): Incremental number of modes to use in Fourier domain. - * If not None, this allows to incrementally increase the number of modes in Fourier domain - during training. Has to verify n <= N for (n, m) in zip(incremental_n_modes, n_modes). - * If None, all the n_modes are used. - This can be updated dynamically during training.Defaults to None. - use_mlp (bool, optional): Whether to use an MLP layer after each FNO block. Defaults to False. - mlp (Dict[str, float], optional): Parameters of the MLP. {'expansion': float, 'dropout': float}. - Defaults to None. - non_linearity (nn.functional, optional): Non-Linearity module to use. Defaults to F.gelu. - norm (str, optional): Normalization layer to use. Defaults to None. - ada_in_features (Optional[int],optional): The input channels of the adaptive normalization.Defaults to - None. - preactivation (bool, optional): Whether to use resnet-style preactivation. Defaults to False. - fno_skip (str, optional): Type of skip connection to use for fno_block. Defaults to "linear". - horizontal_skip (str, optional): Type of skip connection to use for horizontal skip. Defaults to - "linear". - mlp_skip (str, optional): Type of skip connection to use for mlp. Defaults to "soft-gating". - separable (bool, optional): Whether to use a depthwise separable spectral convolution. - Defaults to False. - factorization (str, optional): Tensor factorization of the parameters weight to use. - * If None, a dense tensor parametrizes the Spectral convolutions. - * Otherwise, the specified tensor factorization is used. Defaults to "Tucker". - rank (float, optional): Rank of the tensor factorization of the Fourier weights. Defaults to 1.0. - joint_factorization (bool, optional): Whether all the Fourier Layers should be parametrized by a - single tensor (vs one per layer). Defaults to False. - implementation (str, optional): {'factorized', 'reconstructed'}, optional. Defaults to "factorized". - If factorization is not None, forward mode to use:: - * `reconstructed` : the full weight tensor is reconstructed from the factorization and used for the forward pass. - * `factorized` : the input is directly contracted with the factors of the decomposition. - domain_padding (Optional[Union[list, float, int]], optional): Whether to use percentage of padding. - Defaults to None. - domain_padding_mode (str, optional): {'symmetric', 'one-sided'}, optional - How to perform domain padding, by default 'one-sided'. Defaults to "one-sided". - fft_norm (str, optional): The normalization mode for the FFT. Defaults to "forward". - patching_levels (int, optional): Number of patching levels to use. Defaults to 0. - """ - - def __init__( - self, - input_keys: Tuple[str, ...], - output_keys: Tuple[str, ...], - in_channels: int, - out_channels: int, - hidden_channels: int, - lifting_channels: int = 256, - projection_channels: int = 256, - n_layers: int = 4, - uno_out_channels: Tuple[int, ...] = None, - uno_n_modes: Tuple[Tuple[int, ...], ...] = None, - uno_scalings: Tuple[Tuple[int, ...], ...] = None, - horizontal_skips_map: Dict = None, - incremental_n_modes: Tuple[int, ...] = None, - use_mlp: bool = False, - mlp: Optional[Dict[str, float]] = None, - non_linearity: nn.functional = F.gelu, - norm: str = None, - ada_in_features: Optional[int] = None, - preactivation: bool = False, - fno_skip: str = "linear", - horizontal_skip: str = "linear", - mlp_skip: str = "soft-gating", - separable: bool = False, - factorization: str = None, - rank: float = 1.0, - joint_factorization: bool = False, - implementation: str = "factorized", - domain_padding: Optional[Union[list, float, int]] = None, - domain_padding_mode: str = "one-sided", - fft_norm: str = "forward", - patching_levels: int = 0, - **kwargs, - ): - super().__init__() - self.input_keys = input_keys - self.output_keys = output_keys - if uno_out_channels is None: - raise ValueError("uno_out_channels can not be None") - if uno_n_modes is None: - raise ValueError("uno_n_modes can not be None") - if uno_scalings is None: - raise ValueError("uno_scalings can not be None") - - if len(uno_out_channels) != n_layers: - raise ValueError("Output channels for all layers are not given") - - if len(uno_n_modes) != n_layers: - raise ValueError("Number of modes for all layers are not given") - - if len(uno_scalings) != n_layers: - raise ValueError("Scaling factor for all layers are not given") - - self.n_dim = len(uno_n_modes[0]) - self.uno_out_channels = uno_out_channels - self.uno_n_modes = uno_n_modes - self.uno_scalings = uno_scalings - - self.hidden_channels = hidden_channels - self.lifting_channels = lifting_channels - self.projection_channels = projection_channels - self.in_channels = in_channels - if patching_levels: - self.in_channels = self.in_channels * patching_levels + 1 - self.out_channels = out_channels - self.n_layers = n_layers - self.horizontal_skips_map = horizontal_skips_map - self.joint_factorization = joint_factorization - self.non_linearity = non_linearity - self.rank = rank - self.factorization = factorization - self.fno_skip = (fno_skip,) - self.mlp_skip = (mlp_skip,) - self.fft_norm = fft_norm - self.implementation = implementation - self.separable = separable - self.preactivation = preactivation - self._incremental_n_modes = incremental_n_modes - self.mlp = mlp - # constructing default skip maps - if self.horizontal_skips_map is None: - self.horizontal_skips_map = {} - for i in range( - 0, - n_layers // 2, - ): - # example, if n_layers = 5, then 4:0, 3:1 - self.horizontal_skips_map[n_layers - i - 1] = i - # self.uno_scalings may be a 1d list specifying uniform scaling factor at each layer - # or a 2d list, where each row specifies scaling factors along each dimension. - # To get the final (end to end) scaling factors we need to multiply - # the scaling factors (a list) of all layer. - - self.end_to_end_scaling_factor = [1] * len(self.uno_scalings[0]) - # multiplying scaling factors - for k in self.uno_scalings: - self.end_to_end_scaling_factor = [ - i * j for (i, j) in zip(self.end_to_end_scaling_factor, k) - ] - - # list with a single element is replaced by the scaler. - if len(self.end_to_end_scaling_factor) == 1: - self.end_to_end_scaling_factor = self.end_to_end_scaling_factor[0] - - if isinstance(self.end_to_end_scaling_factor, (float, int)): - self.end_to_end_scaling_factor = [ - self.end_to_end_scaling_factor - ] * self.n_dim - - if domain_padding is not None and ( - (isinstance(domain_padding, list) and sum(domain_padding) > 0) - or (isinstance(domain_padding, (float, int)) and domain_padding > 0) - ): - self.domain_padding = fno_block.DomainPadding( - domain_padding=domain_padding, padding_mode=domain_padding_mode - ) - else: - self.domain_padding = None - self.domain_padding_mode = domain_padding_mode - - self.lifting = fno_block.MLP( - in_channels=in_channels, - out_channels=self.hidden_channels, - hidden_channels=self.lifting_channels, - n_layers=2, - n_dim=self.n_dim, - ) - - self.fno_blocks = nn.LayerList([]) - self.horizontal_skips = nn.LayerDict({}) - prev_out = self.hidden_channels - for i in range(self.n_layers): - if i in self.horizontal_skips_map.keys(): - prev_out = ( - prev_out + self.uno_out_channels[self.horizontal_skips_map[i]] - ) - self.fno_blocks.append( - fno_block.FNOBlocks( - in_channels=prev_out, - out_channels=self.uno_out_channels[i], - n_modes=self.uno_n_modes[i], - use_mlp=use_mlp, - mlp=mlp, - output_scaling_factor=[self.uno_scalings[i]], - non_linearity=non_linearity, - norm=norm, - ada_in_features=ada_in_features, - preactivation=preactivation, - fno_skip=fno_skip, - mlp_skip=mlp_skip, - separable=separable, - incremental_n_modes=incremental_n_modes, - factorization=factorization, - rank=rank, - SpectralConv=fno_block.FactorizedSpectralConv, - joint_factorization=joint_factorization, - implementation=implementation, - fft_norm=fft_norm, - ) - ) - - if i in self.horizontal_skips_map.values(): - self.horizontal_skips[str(i)] = fno_block.skip_connection( - self.uno_out_channels[i], - self.uno_out_channels[i], - type=horizontal_skip, - n_dim=self.n_dim, - ) - prev_out = self.uno_out_channels[i] - - self.projection = fno_block.MLP( - in_channels=prev_out, - out_channels=out_channels, - hidden_channels=self.projection_channels, - n_layers=2, - n_dim=self.n_dim, - non_linearity=non_linearity, - ) - - def forward(self, x, **kwargs): - x = self.concat_to_tensor(x, self.input_keys) - x = self.lifting(x) - if self.domain_padding is not None: - x = self.domain_padding.pad(x) - output_shape = [ - int(round(i * j)) - for (i, j) in zip(x.shape[-self.n_dim :], self.end_to_end_scaling_factor) - ] - - skip_outputs = {} - cur_output = None - for layer_idx in range(self.n_layers): - if layer_idx in self.horizontal_skips_map.keys(): - skip_val = skip_outputs[self.horizontal_skips_map[layer_idx]] - output_scaling_factors = [ - m / n for (m, n) in zip(x.shape, skip_val.shape) - ] - output_scaling_factors = output_scaling_factors[-1 * self.n_dim :] - t = fno_block.resample( - skip_val, output_scaling_factors, list(range(-self.n_dim, 0)) - ) - x = paddle.concat([x, t], axis=1) - - if layer_idx == self.n_layers - 1: - cur_output = output_shape - x = self.fno_blocks[layer_idx](x, output_shape=cur_output) - if layer_idx in self.horizontal_skips_map.values(): - skip_outputs[layer_idx] = self.horizontal_skips[str(layer_idx)](x) - - if self.domain_padding is not None: - x = self.domain_padding.unpad(x) - - out = self.projection(x) - return {self.output_keys[0]: out} diff --git a/examples/smc_reac/ppsci/arch/uscnn.py b/examples/smc_reac/ppsci/arch/uscnn.py deleted file mode 100644 index 5da8440ca4..0000000000 --- a/examples/smc_reac/ppsci/arch/uscnn.py +++ /dev/null @@ -1,124 +0,0 @@ -# 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. - -from typing import Tuple -from typing import Union - -import numpy as np -from paddle import nn - -import ppsci -from ppsci.arch import base - - -class USCNN(base.Arch): - """Physics-informed convolutional neural networks. - - Args: - input_keys (Tuple[str, ...]): Name of input keys, such as ("coords"). - output_keys (Tuple[str, ...]):Name of output keys, such as ("outputV"). - hidden_size (Union[int, Tuple[int, ...]]): The hidden channel for convolutional layers - h (float): The spatial step - nx (int): the number of grids along x-axis - ny (int): The number of grids along y-axis - nvar_in (int, optional): input channel. Defaults to 1. - nvar_out (int, optional): Output channel. Defaults to 1. - pad_singleside (int, optional): Pad for hard boundary constraint. Defaults to 1. - k (int, optional): Kernel_size. Defaults to 5. - s (int, optional): Stride. Defaults to 1. - p (int, optional): Padding. Defaults to 2. - - Examples: - >>> import ppsci - >>> model = ppsci.arch.USCNN( - ... ["coords"], - ... ["outputV"], - ... [16, 32, 16], - ... h=0.01, - ... ny=19, - ... nx=84, - ... nvar_in=2, - ... nvar_out=1, - ... pad_singleside=1, - ... ) - """ - - def __init__( - self, - input_keys: Tuple[str, ...], - output_keys: Tuple[str, ...], - hidden_size: Union[int, Tuple[int, ...]], - h: float, - nx: int, - ny: int, - nvar_in: int = 1, - nvar_out: int = 1, - pad_singleside: int = 1, - k: int = 5, - s: int = 1, - p: int = 2, - ): - super().__init__() - self.input_keys = input_keys - self.output_keys = output_keys - self.nvar_in = nvar_in - self.nvar_out = nvar_out - self.k = k - self.s = s - self.p = p - self.deltaX = h - self.nx = nx - self.ny = ny - self.pad_singleside = pad_singleside - self.relu = nn.ReLU() - self.US = nn.Upsample(size=[self.ny - 2, self.nx - 2], mode="bicubic") - self.conv1 = nn.Conv2D( - self.nvar_in, hidden_size[0], kernel_size=k, stride=s, padding=p - ) - self.conv2 = nn.Conv2D( - hidden_size[0], hidden_size[1], kernel_size=k, stride=s, padding=p - ) - self.conv3 = nn.Conv2D( - hidden_size[1], hidden_size[2], kernel_size=k, stride=s, padding=p - ) - self.conv4 = nn.Conv2D( - hidden_size[2], self.nvar_out, kernel_size=k, stride=s, padding=p - ) - self.pixel_shuffle = nn.PixelShuffle(1) - self.apply(self.init_weights) - self.udfpad = nn.Pad2D( - [pad_singleside, pad_singleside, pad_singleside, pad_singleside], value=0 - ) - - def init_weights(self, m): - if isinstance(m, nn.Conv2D): - bound = 1 / np.sqrt(np.prod(m.weight.shape[1:])) - ppsci.utils.initializer.uniform_(m.weight, -bound, bound) - if m.bias is not None: - ppsci.utils.initializer.uniform_(m.bias, -bound, bound) - - def forward(self, x): - y = self.concat_to_tensor(x, self.input_keys, axis=-1) - y = self.US(y) - y = self.relu(self.conv1(y)) - y = self.relu(self.conv2(y)) - y = self.relu(self.conv3(y)) - y = self.pixel_shuffle(self.conv4(y)) - - y = self.udfpad(y) - y = y[:, 0, :, :].reshape([y.shape[0], 1, y.shape[2], y.shape[3]]) - y = self.split_to_dict(y, self.output_keys) - if self._output_transform is not None: - y = self._output_transform(x, y) - return y diff --git a/examples/smc_reac/ppsci/arch/vae.py b/examples/smc_reac/ppsci/arch/vae.py deleted file mode 100644 index 2a05f0d648..0000000000 --- a/examples/smc_reac/ppsci/arch/vae.py +++ /dev/null @@ -1,103 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import Tuple - -import paddle -import paddle.nn as nn - -from ppsci.arch import base - - -class AutoEncoder(base.Arch): - """ - AutoEncoder is a class that represents an autoencoder neural network model. - - Args: - input_keys (Tuple[str, ...]): A tuple of input keys. - output_keys (Tuple[str, ...]): A tuple of output keys. - input_dim (int): The dimension of the input data. - latent_dim (int): The dimension of the latent space. - hidden_dim (int): The dimension of the hidden layer. - - Examples: - >>> import paddle - >>> import ppsci - >>> model = ppsci.arch.AutoEncoder( - ... input_keys=("input1",), - ... output_keys=("mu", "log_sigma", "decoder_z",), - ... input_dim=100, - ... latent_dim=50, - ... hidden_dim=200 - ... ) - >>> input_dict = {"input1": paddle.rand([200, 100]),} - >>> output_dict = model(input_dict) - >>> print(output_dict["mu"].shape) - [200, 50] - >>> print(output_dict["log_sigma"].shape) - [200, 50] - >>> print(output_dict["decoder_z"].shape) - [200, 100] - """ - - def __init__( - self, - input_keys: Tuple[str, ...], - output_keys: Tuple[str, ...], - input_dim: int, - latent_dim: int, - hidden_dim: int, - ): - super(AutoEncoder, self).__init__() - self.input_keys = input_keys - self.output_keys = output_keys - # encoder - self._encoder_linear = nn.Sequential( - nn.Linear(input_dim, hidden_dim), - nn.Tanh(), - ) - self._encoder_mu = nn.Linear(hidden_dim, latent_dim) - self._encoder_log_sigma = nn.Linear(hidden_dim, latent_dim) - - self._decoder = nn.Sequential( - nn.Linear(latent_dim, hidden_dim), - nn.Tanh(), - nn.Linear(hidden_dim, input_dim), - ) - - def encoder(self, x): - h = self._encoder_linear(x) - mu = self._encoder_mu(h) - log_sigma = self._encoder_log_sigma(h) - return mu, log_sigma - - def decoder(self, x): - return self._decoder(x) - - def forward_tensor(self, x): - mu, log_sigma = self.encoder(x) - z = mu + paddle.randn(mu.shape) * paddle.exp(log_sigma) - return mu, log_sigma, self.decoder(z) - - def forward(self, x): - x = self.concat_to_tensor(x, self.input_keys, axis=-1) - mu, log_sigma, decoder_z = self.forward_tensor(x) - result_dict = { - self.output_keys[0]: mu, - self.output_keys[1]: log_sigma, - self.output_keys[2]: decoder_z, - } - return result_dict diff --git a/examples/smc_reac/ppsci/arch/velocitygan.py b/examples/smc_reac/ppsci/arch/velocitygan.py deleted file mode 100644 index 28f44d14f1..0000000000 --- a/examples/smc_reac/ppsci/arch/velocitygan.py +++ /dev/null @@ -1,354 +0,0 @@ -# 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. - -from math import ceil -from math import sqrt -from typing import Tuple - -import paddle - -import ppsci.utils.initializer as init -from ppsci.arch import base - - -class VelocityGenerator(base.Arch): - """The Generator Of VelocityGAN. - VelocityGAN is applied to full waveform inversion tasks, and the structure of the model - comes from https://arxiv.org/abs/2111.02926# - - Args: - input_keys (Tuple[str, ...]): Name of input keys, such as ("input",). - output_keys (Tuple[str, ...]): Name of output keys, such as ("output",). - dim1 (int, optional): Number of channels in the outermost layers of both encoder and decoder segments. Default is 32. - dim2 (int, optional): Number of channels in the second set of layers from the outermost in both encoder and decoder segments. Default is 64. - dim3 (int, optional): Number of channels in the intermediate layers. Default is 128. - dim4 (int, optional): Number of channels near the bottleneck, just before and after the deepest layer. Default is 256. - dim5 (int, optional): Number of channels at the bottleneck, the deepest layer in the network. Default is 512. - sample_spatial (float, optional): Spatial sampling rate of the input, used to dynamically calculate the kernel size in the last encoder layer. Default is 1.0. - - Examples: - >>> import ppsci - >>> import paddle - >>> model = ppsci.arch.VelocityGenerator(("input", ), ("output", )) - >>> input_dict = {"input": paddle.randn((1, 5, 1000, 70))} - >>> output_dict = model(input_dict) # doctest: +SKIP - >>> print(output_dict["output"].shape) # doctest: +SKIP - [1, 1, 70, 70] - """ - - def __init__( - self, - input_keys: Tuple[str, ...], - output_keys: Tuple[str, ...], - dim1: int = 32, - dim2: int = 64, - dim3: int = 128, - dim4: int = 256, - dim5: int = 512, - sample_spatial: float = 1.0, - ): - super().__init__() - self.input_keys = input_keys - self.output_keys = output_keys - self.generator = Generator( - dim1=dim1, - dim2=dim2, - dim3=dim3, - dim4=dim4, - dim5=dim5, - sample_spatial=sample_spatial, - ) - - def forward(self, x): - if self._input_transform is not None: - x = self._input_transform(x) - - y = self.concat_to_tensor(x, self.input_keys, axis=-1) - y = self.generator(y) - y = self.split_to_dict(y, self.output_keys, axis=-1) - - if self._output_transform is not None: - y = self._output_transform(x, y) - - return y - - -class VelocityDiscriminator(base.Arch): - """The Discriminator Of VelocityGAN. - VelocityGAN is applied to full waveform inversion tasks, and the structure of the model - comes from https://arxiv.org/abs/2111.02926# - - Args: - input_keys (Tuple[str, ...]): Name of input keys, such as ("input",). - output_keys (Tuple[str, ...]): Name of output keys, such as ("output",). - dim1 (int, optional): The number of output channels for convblock1_1 and convblock1_2. Default is 32. - dim2 (int, optional): The number of output channels for convblock2_1 and convblock2_2. Default is 64. - dim3 (int, optional): The number of output channels for convblock3_1 and convblock3_2. Default is 128. - dim4 (int, optional): The number of output channels for convblock4_1 and convblock4_2. Default is 256. - - Examples: - >>> import ppsci - >>> import paddle - >>> model = ppsci.arch.VelocityDiscriminator(("input", ), ("output", )) - >>> input_dict = {"input": paddle.randn((1, 1, 70, 70))} - >>> output_dict = model(input_dict) # doctest: +SKIP - >>> print(output_dict["output"].shape) # doctest: +SKIP - [1, 1] - """ - - def __init__( - self, - input_keys: Tuple[str, ...], - output_keys: Tuple[str, ...], - dim1: int = 32, - dim2: int = 64, - dim3: int = 128, - dim4: int = 256, - ): - super().__init__() - self.input_keys = input_keys - self.output_keys = output_keys - self.discriminator = Discriminator(dim1=dim1, dim2=dim2, dim3=dim3, dim4=dim4) - - def forward(self, x): - if self._input_transform is not None: - x = self._input_transform(x) - - y = self.concat_to_tensor(x, self.input_keys, axis=-1) - y = self.discriminator(y) - y = self.split_to_dict(y, self.output_keys, axis=-1) - - if self._output_transform is not None: - y = self._output_transform(x, y) - - return y - - -class Generator(paddle.nn.Layer): - """The specific implementation of the generator, which is encapsulated in the VelocityGenerator class. - - Args: - dim1 (int, optional): Number of channels in the outermost layers of both encoder and decoder segments. Default is 32. - dim2 (int, optional): Number of channels in the second set of layers from the outermost in both encoder and decoder segments. Default is 64. - dim3 (int, optional): Number of channels in the intermediate layers. Default is 128. - dim4 (int, optional): Number of channels near the bottleneck, just before and after the deepest layer. Default is 256. - dim5 (int, optional): Number of channels at the bottleneck, the deepest layer in the network. Default is 512. - sample_spatial (float, optional): Spatial sampling rate of the input, used to dynamically calculate the kernel size in the last encoder layer. Default is 1.0. - """ - - def __init__( - self, - dim1: int = 32, - dim2: int = 64, - dim3: int = 128, - dim4: int = 256, - dim5: int = 512, - sample_spatial: float = 1.0, - ): - super(Generator, self).__init__() - self.convblock1 = ConvBlock( - 5, dim1, kernel_size=(7, 1), stride=(2, 1), padding=(3, 0) - ) - self.convblock2_1 = ConvBlock( - dim1, dim2, kernel_size=(3, 1), stride=(2, 1), padding=(1, 0) - ) - self.convblock2_2 = ConvBlock(dim2, dim2, kernel_size=(3, 1), padding=(1, 0)) - self.convblock3_1 = ConvBlock( - dim2, dim2, kernel_size=(3, 1), stride=(2, 1), padding=(1, 0) - ) - self.convblock3_2 = ConvBlock(dim2, dim2, kernel_size=(3, 1), padding=(1, 0)) - self.convblock4_1 = ConvBlock( - dim2, dim3, kernel_size=(3, 1), stride=(2, 1), padding=(1, 0) - ) - self.convblock4_2 = ConvBlock(dim3, dim3, kernel_size=(3, 1), padding=(1, 0)) - self.convblock5_1 = ConvBlock(dim3, dim3, stride=2) - self.convblock5_2 = ConvBlock(dim3, dim3) - self.convblock6_1 = ConvBlock(dim3, dim4, stride=2) - self.convblock6_2 = ConvBlock(dim4, dim4) - self.convblock7_1 = ConvBlock(dim4, dim4, stride=2) - self.convblock7_2 = ConvBlock(dim4, dim4) - self.convblock8 = ConvBlock( - dim4, dim5, kernel_size=(8, ceil(70 * sample_spatial / 8)), padding=0 - ) - self.deconv1_1 = DeconvBlock(dim5, dim5, kernel_size=5) - self.deconv1_2 = ConvBlock(dim5, dim5) - self.deconv2_1 = DeconvBlock(dim5, dim4, kernel_size=4, stride=2, padding=1) - self.deconv2_2 = ConvBlock(dim4, dim4) - self.deconv3_1 = DeconvBlock(dim4, dim3, kernel_size=4, stride=2, padding=1) - self.deconv3_2 = ConvBlock(dim3, dim3) - self.deconv4_1 = DeconvBlock(dim3, dim2, kernel_size=4, stride=2, padding=1) - self.deconv4_2 = ConvBlock(dim2, dim2) - self.deconv5_1 = DeconvBlock(dim2, dim1, kernel_size=4, stride=2, padding=1) - self.deconv5_2 = ConvBlock(dim1, dim1) - self.deconv6 = ConvBlock_Tanh(dim1, 1) - self.initial_weight() - - def initial_weight(self): - for _, m in self.named_sublayers(): - if isinstance(m, paddle.nn.Conv2D) or isinstance( - m, paddle.nn.Conv2DTranspose - ): - init.kaiming_uniform_(m.weight, a=sqrt(5)) - if m.bias is not None: - fan_in, _ = init._calculate_fan_in_and_fan_out(m.weight) - bound = 1 / sqrt(fan_in) - init.uniform_(m.bias, -bound, bound) - - def forward(self, x): - x = self.convblock1(x) - x = self.convblock2_1(x) - x = self.convblock2_2(x) - x = self.convblock3_1(x) - x = self.convblock3_2(x) - x = self.convblock4_1(x) - x = self.convblock4_2(x) - x = self.convblock5_1(x) - x = self.convblock5_2(x) - x = self.convblock6_1(x) - x = self.convblock6_2(x) - x = self.convblock7_1(x) - x = self.convblock7_2(x) - x = self.convblock8(x) - x = self.deconv1_1(x) - x = self.deconv1_2(x) - x = self.deconv2_1(x) - x = self.deconv2_2(x) - x = self.deconv3_1(x) - x = self.deconv3_2(x) - x = self.deconv4_1(x) - x = self.deconv4_2(x) - x = self.deconv5_1(x) - x = self.deconv5_2(x) - x = paddle.nn.functional.pad(x, pad=[-5, -5, -5, -5], mode="constant", value=0) - x = self.deconv6(x) - return x - - -class Discriminator(paddle.nn.Layer): - """The specific implementation of the discriminator, which is encapsulated in the VelocityDiscriminator class. - - Args: - dim1 (int, optional): The number of output channels for convblock1_1 and convblock1_2. Default is 32. - dim2 (int, optional): The number of output channels for convblock2_1 and convblock2_2. Default is 64. - dim3 (int, optional): The number of output channels for convblock3_1 and convblock3_2. Default is 128. - dim4 (int, optional): The number of output channels for convblock4_1 and convblock4_2. Default is 256. - """ - - def __init__( - self, dim1: int = 32, dim2: int = 64, dim3: int = 128, dim4: int = 256 - ): - super(Discriminator, self).__init__() - self.convblock1_1 = ConvBlock(1, dim1, stride=2) - self.convblock1_2 = ConvBlock(dim1, dim1) - self.convblock2_1 = ConvBlock(dim1, dim2, stride=2) - self.convblock2_2 = ConvBlock(dim2, dim2) - self.convblock3_1 = ConvBlock(dim2, dim3, stride=2) - self.convblock3_2 = ConvBlock(dim3, dim3) - self.convblock4_1 = ConvBlock(dim3, dim4, stride=2) - self.convblock4_2 = ConvBlock(dim4, dim4) - self.convblock5 = ConvBlock(dim4, 1, kernel_size=5, padding=0) - self.initial_weight() - - def initial_weight(self): - for _, m in self.named_sublayers(): - if isinstance(m, paddle.nn.Conv2D): - init.kaiming_uniform_(m.weight, a=sqrt(5)) - if m.bias is not None: - fan_in, _ = init._calculate_fan_in_and_fan_out(m.weight) - bound = 1 / sqrt(fan_in) - init.uniform_(m.bias, -bound, bound) - - def forward(self, x): - x = self.convblock1_1(x) - x = self.convblock1_2(x) - x = self.convblock2_1(x) - x = self.convblock2_2(x) - x = self.convblock3_1(x) - x = self.convblock3_2(x) - x = self.convblock4_1(x) - x = self.convblock4_2(x) - x = self.convblock5(x) - x = x.reshape([x.shape[0], -1]) - return x - - -class ConvBlock(paddle.nn.Layer): - """A convolution block, including Conv2D, BatchNorm2D, and LeakyReLU""" - - def __init__( - self, in_fea, out_fea, kernel_size=3, stride=1, padding=1, relu_slop=0.2 - ): - super(ConvBlock, self).__init__() - layers = [ - paddle.nn.Conv2D( - in_channels=in_fea, - out_channels=out_fea, - kernel_size=kernel_size, - stride=stride, - padding=padding, - ), - paddle.nn.BatchNorm2D(out_fea), - paddle.nn.LeakyReLU(negative_slope=relu_slop), - ] - self.layers = paddle.nn.Sequential(*layers) - - def forward(self, x): - return self.layers(x) - - -class DeconvBlock(paddle.nn.Layer): - """A deconvolution block, including Conv2DTranspose, BatchNorm2D, and LeakyReLU""" - - def __init__( - self, in_fea, out_fea, kernel_size=2, stride=2, padding=0, output_padding=0 - ): - super(DeconvBlock, self).__init__() - layers = [ - paddle.nn.Conv2DTranspose( - in_channels=in_fea, - out_channels=out_fea, - kernel_size=kernel_size, - stride=stride, - padding=padding, - output_padding=output_padding, - ), - paddle.nn.BatchNorm2D(out_fea), - paddle.nn.LeakyReLU(negative_slope=0.2), - ] - self.layers = paddle.nn.Sequential(*layers) - - def forward(self, x): - return self.layers(x) - - -class ConvBlock_Tanh(paddle.nn.Layer): - """A convolution block, including Conv2D, BatchNorm2D, and Tanh""" - - def __init__(self, in_fea, out_fea, kernel_size=3, stride=1, padding=1): - super(ConvBlock_Tanh, self).__init__() - layers = [ - paddle.nn.Conv2D( - in_channels=in_fea, - out_channels=out_fea, - kernel_size=kernel_size, - stride=stride, - padding=padding, - ), - paddle.nn.BatchNorm2D(out_fea), - paddle.nn.Tanh(), - ] - self.layers = paddle.nn.Sequential(*layers) - - def forward(self, x): - return self.layers(x) diff --git a/examples/smc_reac/ppsci/autodiff/__init__.py b/examples/smc_reac/ppsci/autodiff/__init__.py deleted file mode 100644 index 68a0570c24..0000000000 --- a/examples/smc_reac/ppsci/autodiff/__init__.py +++ /dev/null @@ -1,17 +0,0 @@ -# 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. - -from ppsci.autodiff.ad import clear -from ppsci.autodiff.ad import hessian -from ppsci.autodiff.ad import jacobian diff --git a/examples/smc_reac/ppsci/autodiff/ad.py b/examples/smc_reac/ppsci/autodiff/ad.py deleted file mode 100644 index ba3afd14a4..0000000000 --- a/examples/smc_reac/ppsci/autodiff/ad.py +++ /dev/null @@ -1,341 +0,0 @@ -# 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. - -""" -This module is adapted from [https://github.com/lululxvi/deepxde](https://github.com/lululxvi/deepxde) -""" - -from __future__ import annotations - -from typing import Callable -from typing import Dict -from typing import List -from typing import Optional -from typing import Union - -import paddle - - -class _Jacobian: - """Compute Jacobian matrix J: J[i][j] = dy_i/dx_j, where i = 0, ..., dim_y-1 and - j = 0, ..., dim_x - 1. - - It is lazy evaluation, i.e., it only computes J[i][j] when needed, and will cache - by output tensor(row index in jacobian matrix). - - Args: - ys (paddle.Tensor): Output Tensor of shape [batch_size, dim_y]. - xs (paddle.Tensor): Input Tensor of shape [batch_size, dim_x]. - """ - - def __init__( - self, - ys: "paddle.Tensor", - xs: "paddle.Tensor", - J: Optional[Dict[int, paddle.Tensor]] = None, - ): - self.ys = ys - self.xs = xs - - self.dim_y = ys.shape[1] - self.dim_x = xs.shape[1] - - self.J: Dict[int, paddle.Tensor] = {} if J is None else J - - def __call__( - self, - i: int = 0, - j: Optional[int] = None, - retain_graph: Optional[bool] = None, - create_graph: bool = True, - ) -> "paddle.Tensor": - """ - Returns J[`i`][`j`]. If `j` is ``None``, returns the gradient of y_i, i.e. J[i]. - """ - if not 0 <= i < self.dim_y: - raise ValueError(f"i({i}) should in range [0, {self.dim_y}).") - if j is not None and not 0 <= j < self.dim_x: - raise ValueError(f"j({j}) should in range [0, {self.dim_x}).") - # Compute J[i] - if i not in self.J: - y = self.ys[:, i : i + 1] if self.dim_y > 1 else self.ys - self.J[i] = paddle.grad( - y, self.xs, retain_graph=retain_graph, create_graph=create_graph - )[0] - - return self.J[i] if (j is None or self.dim_x == 1) else self.J[i][:, j : j + 1] - - -class Jacobians: - r"""Compute multiple Jacobians. - - $$ - \rm Jacobian(ys, xs, i, j) = \dfrac{\partial ys_i}{\partial xs_j} - $$ - - A new instance will be created for a new pair of (output, input). For the (output, - input) pair that has been computed before, it will reuse the previous instance, - rather than creating a new one. - """ - - def __init__(self): - self.Js = {} - - def __call__( - self, - ys: "paddle.Tensor", - xs: Union["paddle.Tensor", List["paddle.Tensor"]], - i: int = 0, - j: Optional[int] = None, - retain_graph: Optional[bool] = None, - create_graph: bool = True, - ) -> Union["paddle.Tensor", List["paddle.Tensor"]]: - """Compute jacobians for given ys and xs. - - Args: - ys (paddle.Tensor): Output tensor. - xs (Union[paddle.Tensor, List[paddle.Tensor]]): Input tensor(s). - i (int, optional): i-th output variable. Defaults to 0. - j (Optional[int]): j-th input variable. Defaults to None. - retain_graph (Optional[bool]): Whether to retain the forward graph which - is used to calculate the gradient. When it is True, the graph would - be retained, in which way users can calculate backward twice for the - same graph. When it is False, the graph would be freed. Default None, - which means it is equal to `create_graph`. - create_graph (bool, optional): Whether to create the gradient graphs of - the computing process. When it is True, higher order derivatives are - supported to compute; when it is False, the gradient graphs of the - computing process would be discarded. Default False. - - Returns: - paddle.Tensor: Jacobian matrix of ys[i] to xs[j]. - - Examples: - >>> import paddle - >>> import ppsci - >>> x = paddle.randn([4, 1]) - >>> x.stop_gradient = False - >>> y = x * x - >>> dy_dx = ppsci.autodiff.jacobian(y, x) - >>> print(dy_dx.shape) - [4, 1] - """ - if not isinstance(xs, (list, tuple)): - key = (ys, xs) - if key not in self.Js: - self.Js[key] = _Jacobian(ys, xs) - return self.Js[key](i, j, retain_graph, create_graph) - else: - xs_require = [xs[i] for i in range(len(xs)) if (ys, xs[i]) not in self.Js] - grads_require = paddle.grad( - ys, - xs_require, - create_graph=create_graph, - retain_graph=retain_graph, - ) - - idx = 0 - Js_list = [] - for k, xs_ in enumerate(xs): - key = (ys, xs_) - assert xs_.shape[-1] == 1, ( - f"The last dim of each xs should be 1, but xs[{k}] has shape " - f"{xs_.shape}" - ) - if key not in self.Js: - self.Js[key] = _Jacobian(ys, xs_, {0: grads_require[idx]}) - idx += 1 - Js_list.append(self.Js[key](i, j, retain_graph, create_graph)) - return Js_list - - def _clear(self): - """Clear cached Jacobians.""" - self.Js = {} - - -# Use high-order differentiation with singleton pattern for convenient -jacobian: Callable[ - [ - "paddle.Tensor", - Union["paddle.Tensor", List["paddle.Tensor"]], - int, - Optional[int], - Optional[bool], - bool, - ], - Union["paddle.Tensor", List["paddle.Tensor"]], -] = Jacobians() - - -class _Hessian: - """Compute Hessian matrix H: H[i][j] = d^2y / dx_i dx_j, where i,j = 0,..., dim_x-1. - - It is lazy evaluation, i.e., it only computes H[i][j] when needed. - - Args: - ys: Output Tensor of shape (batch_size, 1) or (batch_size, dim_y > 1). - xs: Input Tensor of shape (batch_size, dim_x). - component: If `y` has the shape (batch_size, dim_y > 1), then `y[:, component]` - is used to compute the Hessian. Do not use if `y` has the shape (batch_size, - 1). - grad_y: The gradient of `y` w.r.t. `xs`. Provide `grad_y` if known to avoid - duplicate computation. `grad_y` can be computed from ``Jacobian``. - """ - - def __init__( - self, - ys: "paddle.Tensor", - xs: "paddle.Tensor", - component: Optional[int] = None, - grad_y: Optional["paddle.Tensor"] = None, - ): - dim_y = ys.shape[1] - - if dim_y > 1: - if component is None: - raise ValueError( - f"component({component}) can not be None when dim_y({dim_y})>1." - ) - if component >= dim_y: - raise ValueError( - f"component({component}) should be smaller than dim_y({dim_y})." - ) - else: - if component is not None: - raise ValueError( - f"component{component} should be set to None when dim_y({dim_y})=1." - ) - component = 0 - - if grad_y is None: - # `create_graph` of first order(jacobian) should be `True` in _Hessian. - grad_y = jacobian( - ys, xs, i=component, j=None, retain_graph=None, create_graph=True - ) - self.H = _Jacobian(grad_y, xs) - - def __call__( - self, - i: int = 0, - j: int = 0, - retain_graph: Optional[bool] = None, - create_graph: bool = True, - ): - """Returns H[`i`][`j`].""" - return self.H(i, j, retain_graph, create_graph) - - -class Hessians: - r"""Compute multiple Hessians. - - $$ - \rm Hessian(ys, xs, component, i, j) = \dfrac{\partial ys_{component}}{\partial xs_i \partial xs_j} - $$ - - A new instance will be created for a new pair of (output, input). For the (output, - input) pair that has been computed before, it will reuse the previous instance, - rather than creating a new one. - """ - - def __init__(self): - self.Hs = {} - - def __call__( - self, - ys: "paddle.Tensor", - xs: "paddle.Tensor", - component: Optional[int] = None, - i: int = 0, - j: int = 0, - grad_y: Optional["paddle.Tensor"] = None, - retain_graph: Optional[bool] = None, - create_graph: bool = True, - ) -> "paddle.Tensor": - """Compute hessian matrix for given ys and xs. - - Args: - ys (paddle.Tensor): Output tensor. - xs (paddle.Tensor): Input tensor. - component (Optional[int]): If `y` has the shape (batch_size, dim_y > 1), then `y[:, component]` - is used to compute the Hessian. Do not use if `y` has the shape (batch_size, - 1). Defaults to None. - i (int, optional): I-th input variable. Defaults to 0. - j (int, optional): J-th input variable. Defaults to 0. - grad_y (Optional[paddle.Tensor]): The gradient of `y` w.r.t. `xs`. Provide `grad_y` if known to avoid - duplicate computation. Defaults to None. - retain_graph (Optional[bool]): Whether to retain the forward graph which - is used to calculate the gradient. When it is True, the graph would - be retained, in which way users can calculate backward twice for the - same graph. When it is False, the graph would be freed. Default None, - which means it is equal to `create_graph`. - create_graph (bool, optional): Whether to create the gradient graphs of - the computing process. When it is True, higher order derivatives are - supported to compute; when it is False, the gradient graphs of the - computing process would be discarded. Default False. - - Returns: - paddle.Tensor: Hessian matrix. - - Examples: - >>> import paddle - >>> import ppsci - >>> x = paddle.randn([4, 3]) - >>> x.stop_gradient = False - >>> y = (x * x).sin() - >>> dy_dxx = ppsci.autodiff.hessian(y, x, component=0) - >>> print(dy_dxx.shape) - [4, 1] - """ - key = (ys, xs, component) - if key not in self.Hs: - self.Hs[key] = _Hessian(ys, xs, component=component, grad_y=grad_y) - return self.Hs[key](i, j, retain_graph, create_graph) - - def _clear(self): - """Clear cached Hessians.""" - self.Hs = {} - - -# Use high-order differentiation with singleton pattern for convenient -hessian: Callable[ - [ - "paddle.Tensor", - "paddle.Tensor", - Optional[int], - int, - int, - Optional["paddle.Tensor"], - Optional[bool], - bool, - ], - "paddle.Tensor", -] = Hessians() - - -def clear(): - """Clear cached Jacobians and Hessians. - - Examples: - >>> import paddle - >>> import ppsci - >>> x = paddle.randn([4, 3]) - >>> x.stop_gradient = False - >>> y = (x * x).sin() - >>> dy_dxx = ppsci.autodiff.hessian(y, x, component=0) - >>> ppsci.autodiff.clear() - >>> print(ppsci.autodiff.hessian.Hs) - {} - """ - jacobian._clear() - hessian._clear() diff --git a/examples/smc_reac/ppsci/constraint/__init__.py b/examples/smc_reac/ppsci/constraint/__init__.py deleted file mode 100644 index 9179439436..0000000000 --- a/examples/smc_reac/ppsci/constraint/__init__.py +++ /dev/null @@ -1,86 +0,0 @@ -# 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. - -from __future__ import annotations - -import copy - -from ppsci.constraint.base import Constraint -from ppsci.constraint.boundary_constraint import BoundaryConstraint -from ppsci.constraint.initial_constraint import InitialConstraint -from ppsci.constraint.integral_constraint import IntegralConstraint -from ppsci.constraint.interior_constraint import InteriorConstraint -from ppsci.constraint.periodic_constraint import PeriodicConstraint -from ppsci.constraint.supervised_constraint import SupervisedConstraint -from ppsci.loss import build_loss -from ppsci.utils import logger -from ppsci.utils import misc - -__all__ = [ - "Constraint", - "BoundaryConstraint", - "InitialConstraint", - "IntegralConstraint", - "InteriorConstraint", - "PeriodicConstraint", - "SupervisedConstraint", -] - - -def build_constraint(cfg, equation_dict, geom_dict): - """Build constraint(s). - - Args: - cfg (List[DictConfig]): Constraint config list. - equation_dict (Dct[str, Equation]): Equation(s) in dict. - geom_dict (Dct[str, Geometry]): Geometry(ies) in dict. - - Returns: - Dict[str, constraint]: Constraint(s) in dict. - """ - if cfg is None: - return None - cfg = copy.deepcopy(cfg) - global_dataloader_cfg = cfg["dataloader"] - constraint_cfg = cfg["content"] - - constraint_dict = misc.PrettyOrderedDict() - for _item in constraint_cfg: - constraint_cls = next(iter(_item.keys())) - _constraint_cfg = _item[constraint_cls] - constraint_name = _constraint_cfg.get("name", constraint_cls) - - # select equation - if isinstance(_constraint_cfg["output_expr"], str): - equation_name = _constraint_cfg.pop("output_expr") - _constraint_cfg["output_expr"] = equation_dict[equation_name].equations - - # select geometry - geom_name = _constraint_cfg.pop("geom") - _constraint_cfg["geom"] = geom_dict[geom_name] - - # update complete dataloader config - local_dataloader_cfg = _constraint_cfg["dataloader"] - local_dataloader_cfg.update(global_dataloader_cfg) - - # build loss - _constraint_cfg["loss"] = build_loss(_constraint_cfg["loss"]) - - # instantiate constraint - _constraint_cfg["dataloader_cfg"] = _constraint_cfg.pop("dataloader") - constraint_dict[constraint_name] = eval(constraint_cls)(**_constraint_cfg) - - logger.debug(str(constraint_dict[constraint_name])) - - return constraint_dict diff --git a/examples/smc_reac/ppsci/constraint/base.py b/examples/smc_reac/ppsci/constraint/base.py deleted file mode 100644 index c3b4c8a122..0000000000 --- a/examples/smc_reac/ppsci/constraint/base.py +++ /dev/null @@ -1,62 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import TYPE_CHECKING -from typing import Any -from typing import Dict - -from paddle import io - -from ppsci import data - -if TYPE_CHECKING: - from ppsci import loss - - -class Constraint: - """Base class for constraint. - - Args: - dataset (io.Dataset): Dataset. - dataloader_cfg (Dict[str, Any]): Dataloader config. - loss (loss.Loss): Loss functor. - name (str): Name of constraint. - """ - - def __init__( - self, - dataset: io.Dataset, - dataloader_cfg: Dict[str, Any], - loss: "loss.Loss", - name: str, - ): - self.data_loader = data.build_dataloader(dataset, dataloader_cfg) - self.data_iter = iter(self.data_loader) - self.loss = loss - self.name = name - - def __str__(self): - return ", ".join( - [ - self.__class__.__name__, - f"name = {self.name}", - f"input_keys = {self.input_keys}", - f"output_keys = {self.output_keys}", - f"output_expr = {self.output_expr}", - f"label_dict = {self.label_dict}", - f"loss = {self.loss}", - ] - ) diff --git a/examples/smc_reac/ppsci/constraint/boundary_constraint.py b/examples/smc_reac/ppsci/constraint/boundary_constraint.py deleted file mode 100644 index 8ac30fcb41..0000000000 --- a/examples/smc_reac/ppsci/constraint/boundary_constraint.py +++ /dev/null @@ -1,163 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import TYPE_CHECKING -from typing import Any -from typing import Callable -from typing import Dict -from typing import Optional -from typing import Union - -import numpy as np -import sympy -from typing_extensions import Literal - -from ppsci import geometry -from ppsci.constraint import base -from ppsci.data import dataset - -if TYPE_CHECKING: - from ppsci import loss - - -class BoundaryConstraint(base.Constraint): - """Class for boundary constraint. - - Args: - output_expr (Dict[str, Callable]): Function in dict for computing output. - e.g. {"u_mul_v": lambda out: out["u"] * out["v"]} means the model output u - will be multiplied by model output v and the result will be named "u_mul_v". - label_dict (Dict[str, Union[float, Callable]]): Function in dict for computing - label, which will be a reference value to participate in the loss calculation. - geom (geometry.Geometry): Geometry where data sampled from. - dataloader_cfg (Dict[str, Any]): Dataloader config. - loss (loss.Loss): Loss functor. - random (Literal["pseudo", "Halton", "LHS"], optional): Random method for sampling data in - geometry. Defaults to "pseudo". - criteria (Optional[Callable]): Criteria for refining specified boundaries. - Defaults to None. - evenly (bool, optional): Whether to use evenly distribution sampling. - Defaults to False. - weight_dict (Optional[Dict[str, Union[float, Callable]]]): Define the weight of each - constraint variable. Defaults to None. - name (str, optional): Name of constraint object. Defaults to "BC". - - Examples: - >>> import ppsci - >>> rect = ppsci.geometry.Rectangle((0, 0), (1, 1)) - >>> bc = ppsci.constraint.BoundaryConstraint( - ... {"u": lambda out: out["u"]}, - ... {"u": 0}, - ... rect, - ... { - ... "dataset": "IterableNamedArrayDataset", - ... "iters_per_epoch": 1, - ... "batch_size": 16, - ... }, - ... ppsci.loss.MSELoss("mean"), - ... name="BC", - ... ) # doctest: +SKIP - """ - - def __init__( - self, - output_expr: Dict[str, Callable], - label_dict: Dict[str, Union[float, Callable]], - geom: geometry.Geometry, - dataloader_cfg: Dict[str, Any], - loss: "loss.Loss", - random: Literal["pseudo", "Halton", "LHS"] = "pseudo", - criteria: Optional[Callable] = None, - evenly: bool = False, - weight_dict: Optional[Dict[str, Union[float, Callable]]] = None, - name: str = "BC", - ): - self.label_dict = label_dict - self.input_keys = geom.dim_keys - self.output_keys = tuple(label_dict.keys()) - self.output_expr = { - k: v for k, v in output_expr.items() if k in self.output_keys - } - - if isinstance(criteria, str): - criteria = eval(criteria) - - # prepare input - input = geom.sample_boundary( - dataloader_cfg["batch_size"] * dataloader_cfg["iters_per_epoch"], - random, - criteria, - evenly, - ) - if "area" in input: - input["area"] *= dataloader_cfg["iters_per_epoch"] - - # prepare label - label = {} - for key, value in label_dict.items(): - if isinstance(value, (int, float)): - label[key] = np.full_like(next(iter(input.values())), value) - elif isinstance(value, sympy.Basic): - func = sympy.lambdify( - sympy.symbols(geom.dim_keys), - value, - [{"amax": lambda xy, axis: np.maximum(xy[0], xy[1])}, "numpy"], - ) - label[key] = func( - **{k: v for k, v in input.items() if k in geom.dim_keys} - ) - elif callable(value): - func = value - label[key] = func(input) - if isinstance(label[key], (int, float)): - label[key] = np.full_like(next(iter(input.values())), label[key]) - else: - raise NotImplementedError(f"type of {type(value)} is invalid yet.") - - # prepare weight - weight = None - if weight_dict is not None: - weight = {key: np.ones_like(next(iter(label.values()))) for key in label} - for key, value in weight_dict.items(): - if isinstance(value, (int, float)): - weight[key] = np.full_like(next(iter(label.values())), value) - elif isinstance(value, sympy.Basic): - func = sympy.lambdify( - [sympy.Symbol(k) for k in geom.dim_keys], - value, - [{"amax": lambda xy, _: np.maximum(xy[0], xy[1])}, "numpy"], - ) - weight[key] = func(**{k: input[k] for k in geom.dim_keys}) - elif callable(value): - func = value - weight[key] = func(input) - if isinstance(weight[key], (int, float)): - weight[key] = np.full_like( - next(iter(input.values())), weight[key] - ) - else: - raise NotImplementedError(f"type of {type(value)} is invalid yet.") - - # wrap input, label, weight into a dataset - if isinstance(dataloader_cfg["dataset"], str): - dataloader_cfg["dataset"] = {"name": dataloader_cfg["dataset"]} - dataloader_cfg["dataset"].update( - {"input": input, "label": label, "weight": weight} - ) - _dataset = dataset.build_dataset(dataloader_cfg["dataset"]) - - # construct dataloader with dataset and dataloader_cfg - super().__init__(_dataset, dataloader_cfg, loss, name) diff --git a/examples/smc_reac/ppsci/constraint/initial_constraint.py b/examples/smc_reac/ppsci/constraint/initial_constraint.py deleted file mode 100644 index 63ae320993..0000000000 --- a/examples/smc_reac/ppsci/constraint/initial_constraint.py +++ /dev/null @@ -1,172 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import TYPE_CHECKING -from typing import Any -from typing import Callable -from typing import Dict -from typing import Optional -from typing import Union - -import numpy as np -import sympy -from typing_extensions import Literal - -from ppsci import geometry -from ppsci.constraint import base -from ppsci.data import dataset - -if TYPE_CHECKING: - from ppsci import loss - - -class InitialConstraint(base.Constraint): - """Class for initial interior constraint. - - Args: - output_expr (Dict[str, Callable]): Function in dict for computing output. - e.g. {"u_mul_v": lambda out: out["u"] * out["v"]} means the model output u - will be multiplied by model output v and the result will be named "u_mul_v". - label_dict (Dict[str, Union[float, Callable]]): Function in dict for computing - label, which will be a reference value to participate in the loss calculation. - geom (geometry.TimeXGeometry): Geometry where data sampled from. - dataloader_cfg (Dict[str, Any]): Dataloader config. - loss (loss.Loss): Loss functor. - random (Literal["pseudo", "Halton", "LHS"], optional): Random method for sampling data in - geometry. Defaults to "pseudo". - criteria (Optional[Callable]): Criteria for refining specified boundaries. - Defaults to None. - evenly (bool, optional): Whether to use evenly distribution sampling. - Defaults to False. - weight_dict (Optional[Dict[str, Callable]]): Define the weight of each - constraint variable. Defaults to None. - compute_sdf_derivatives (Optional[bool]): Whether compute derivatives for SDF. - Defaults to False. - name (str, optional): Name of constraint object. Defaults to "IC". - - Examples: - >>> import ppsci - >>> rect = ppsci.geometry.TimeXGeometry( - ... ppsci.geometry.TimeDomain(0, 1), - ... ppsci.geometry.Rectangle((0, 0), (1, 1)), - ... ) - >>> ic = ppsci.constraint.InitialConstraint( - ... {"u": lambda out: out["u"]}, - ... {"u": 0}, - ... rect, - ... { - ... "dataset": "IterableNamedArrayDataset", - ... "iters_per_epoch": 1, - ... "batch_size": 16, - ... }, - ... ppsci.loss.MSELoss("mean"), - ... name="IC", - ... ) # doctest: +SKIP - """ - - def __init__( - self, - output_expr: Dict[str, Callable], - label_dict: Dict[str, Union[float, Callable]], - geom: geometry.TimeXGeometry, - dataloader_cfg: Dict[str, Any], - loss: "loss.Loss", - random: Literal["pseudo", "Halton", "LHS"] = "pseudo", - criteria: Optional[Callable] = None, - evenly: bool = False, - weight_dict: Optional[Dict[str, Callable]] = None, - compute_sdf_derivatives: bool = False, - name: str = "IC", - ): - self.label_dict = label_dict - self.input_keys = geom.dim_keys - self.output_keys = tuple(label_dict.keys()) - self.output_expr = { - k: v for k, v in output_expr.items() if k in self.output_keys - } - - if isinstance(criteria, str): - criteria = eval(criteria) - - # prepare input - input = geom.sample_initial_interior( - dataloader_cfg["batch_size"] * dataloader_cfg["iters_per_epoch"], - random, - criteria, - evenly, - compute_sdf_derivatives, - ) - if "area" in input: - input["area"] *= dataloader_cfg["iters_per_epoch"] - - # prepare label - label = {} - for key, value in label_dict.items(): - if isinstance(value, (int, float)): - label[key] = np.full_like(next(iter(input.values())), value) - elif isinstance(value, sympy.Basic): - func = sympy.lambdify( - sympy.symbols(geom.dim_keys), - value, - [{"amax": lambda xy, _: np.maximum(xy[0], xy[1])}, "numpy"], - ) - label[key] = func( - **{k: v for k, v in input.items() if k in geom.dim_keys} - ) - elif callable(value): - func = value - label[key] = func(input) - if isinstance(label[key], (int, float)): - label[key] = np.full_like(next(iter(input.values())), label[key]) - else: - raise NotImplementedError(f"type of {type(value)} is invalid yet.") - - # prepare weight - weight = None - if weight_dict is not None: - weight = {key: np.ones_like(next(iter(label.values()))) for key in label} - for key, value in weight_dict.items(): - if isinstance(value, (int, float)): - weight[key] = np.full_like(next(iter(label.values())), value) - elif isinstance(value, sympy.Basic): - func = sympy.lambdify( - sympy.symbols(geom.dim_keys), - value, - [{"amax": lambda xy, _: np.maximum(xy[0], xy[1])}, "numpy"], - ) - weight[key] = func( - **{k: v for k, v in input.items() if k in geom.dim_keys} - ) - elif callable(value): - func = value - weight[key] = func(input) - if isinstance(weight[key], (int, float)): - weight[key] = np.full_like( - next(iter(input.values())), weight[key] - ) - else: - raise NotImplementedError(f"type of {type(value)} is invalid yet.") - - # wrap input, label, weight into a dataset - if isinstance(dataloader_cfg["dataset"], str): - dataloader_cfg["dataset"] = {"name": dataloader_cfg["dataset"]} - dataloader_cfg["dataset"].update( - {"input": input, "label": label, "weight": weight} - ) - _dataset = dataset.build_dataset(dataloader_cfg["dataset"]) - - # construct dataloader with dataset and dataloader_cfg - super().__init__(_dataset, dataloader_cfg, loss, name) diff --git a/examples/smc_reac/ppsci/constraint/integral_constraint.py b/examples/smc_reac/ppsci/constraint/integral_constraint.py deleted file mode 100644 index 19f6f8a1fe..0000000000 --- a/examples/smc_reac/ppsci/constraint/integral_constraint.py +++ /dev/null @@ -1,178 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import TYPE_CHECKING -from typing import Any -from typing import Callable -from typing import Dict -from typing import List -from typing import Optional -from typing import Union - -import numpy as np -import paddle -import sympy -from typing_extensions import Literal - -from ppsci import geometry -from ppsci.constraint import base -from ppsci.data import dataset -from ppsci.utils import misc - -if TYPE_CHECKING: - from ppsci import loss - - -class IntegralConstraint(base.Constraint): - """Class for integral constraint. - - Args: - output_expr (Dict[str, Callable]): Function in dict for computing output. - e.g. {"u_mul_v": lambda out: out["u"] * out["v"]} means the model output u - will be multiplied by model output v and the result will be named "u_mul_v". - label_dict (Dict[str, Union[float, Callable]]): Function in dict for computing - label, which will be a reference value to participate in the loss calculation. - geom (geometry.Geometry): Geometry where data sampled from. - dataloader_cfg (Dict[str, Any]): Dataloader config. - loss (loss.Loss): Loss functor. - random (Literal["pseudo", "Halton", "LHS"], optional): Random method for sampling data in - geometry. Defaults to "pseudo". - criteria (Optional[Callable]): Criteria for refining specified boundaries. - Defaults to None. - weight_dict (Optional[Dict[str, Callable]]): Define the weight of each - constraint variable. Defaults to None. - name (str, optional): Name of constraint object. Defaults to "IgC". - - Examples: - >>> import ppsci - >>> rect = ppsci.geometry.Rectangle((0, 0), (1, 1)) - >>> igc = ppsci.constraint.IntegralConstraint( - ... {"u": lambda out: out["u"]}, - ... {"u": 0}, - ... rect, - ... { - ... "dataset": "IterableNamedArrayDataset", - ... "iters_per_epoch": 1, - ... "batch_size": 16, - ... "integral_batch_size": 8, - ... }, - ... ppsci.loss.MSELoss("mean"), - ... name="IgC", - ... ) # doctest: +SKIP - """ - - def __init__( - self, - output_expr: Dict[str, Callable], - label_dict: Dict[str, Union[float, Callable]], - geom: geometry.Geometry, - dataloader_cfg: Dict[str, Any], - loss: "loss.Loss", - random: Literal["pseudo", "Halton", "LHS"] = "pseudo", - criteria: Optional[Callable] = None, - weight_dict: Optional[Dict[str, Callable]] = None, - name: str = "IgC", - ): - self.label_dict = label_dict - self.input_keys = geom.dim_keys - self.output_keys = tuple(label_dict.keys()) - self.output_expr = { - k: v for k, v in output_expr.items() if k in self.output_keys - } - - if isinstance(criteria, str): - criteria = eval(criteria) - - # prepare input - input_list: List[Dict[str, np.ndarray]] = [] - for _ in range( - dataloader_cfg["batch_size"] * dataloader_cfg["iters_per_epoch"] - ): - input = geom.sample_boundary( - dataloader_cfg["integral_batch_size"], random, criteria - ) - input_list.append(input) - input = misc.stack_dict_list(input_list) - # shape of each input is [batch_size, integral_batch_size, ndim] - - # prepare label - # shape of each label is [batch_size, ndim] - label = {} - for key, value in label_dict.items(): - if isinstance(value, (int, float)): - label[key] = np.full( - (next(iter(input.values())).shape[0], 1), - value, - paddle.get_default_dtype(), - ) - elif isinstance(value, sympy.Basic): - func = sympy.lambdify( - sympy.symbols(geom.dim_keys), - value, - [{"amax": lambda xy, _: np.maximum(xy[0], xy[1])}, "numpy"], - ) - label[key] = func( - **{k: v for k, v in input.items() if k in geom.dim_keys} - ) - elif callable(value): - func = value - label[key] = func(input) - if isinstance(label[key], (int, float)): - label[key] = np.full( - (next(iter(input.values())).shape[0], 1), - label[key], - paddle.get_default_dtype(), - ) - else: - raise NotImplementedError(f"type of {type(value)} is invalid yet.") - - # prepare weight - # shape of each weight is [batch_size, ndim] - weight = None - if weight_dict is not None: - weight = {key: np.ones_like(next(iter(label.values()))) for key in label} - for key, value in weight_dict.items(): - if isinstance(value, (int, float)): - weight[key] = np.full_like(next(iter(label.values())), value) - elif isinstance(value, sympy.Basic): - func = sympy.lambdify( - sympy.symbols(geom.dim_keys), - value, - [{"amax": lambda xy, _: np.maximum(xy[0], xy[1])}, "numpy"], - ) - weight[key] = func( - **{k: v for k, v in input.items() if k in geom.dim_keys} - ) - elif callable(value): - func = value - weight[key] = func(input) - if isinstance(weight[key], (int, float)): - weight[key] = np.full_like( - next(iter(input.values())), weight[key] - ) - else: - raise NotImplementedError(f"type of {type(value)} is invalid yet.") - - # wrap input, label, weight into a dataset - if isinstance(dataloader_cfg["dataset"], str): - dataloader_cfg["dataset"] = {"name": dataloader_cfg["dataset"]} - dataloader_cfg["dataset"].update( - {"input": input, "label": label, "weight": weight} - ) - _dataset = dataset.build_dataset(dataloader_cfg["dataset"]) - - # construct dataloader with dataset and dataloader_cfg - super().__init__(_dataset, dataloader_cfg, loss, name) diff --git a/examples/smc_reac/ppsci/constraint/interior_constraint.py b/examples/smc_reac/ppsci/constraint/interior_constraint.py deleted file mode 100644 index 3c1eb7ed3f..0000000000 --- a/examples/smc_reac/ppsci/constraint/interior_constraint.py +++ /dev/null @@ -1,174 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import TYPE_CHECKING -from typing import Any -from typing import Callable -from typing import Dict -from typing import Optional -from typing import Union - -import numpy as np -import sympy -from typing_extensions import Literal - -from ppsci import geometry -from ppsci.constraint import base -from ppsci.data import dataset - -if TYPE_CHECKING: - from ppsci import loss - - -class InteriorConstraint(base.Constraint): - """Class for interior constraint. - - Args: - output_expr (Dict[str, Callable]): Function in dict for computing output. - e.g. {"u_mul_v": lambda out: out["u"] * out["v"]} means the model output u - will be multiplied by model output v and the result will be named "u_mul_v". - label_dict (Dict[str, Union[float, Callable]]): Function in dict for computing - label, which will be a reference value to participate in the loss calculation. - geom (geometry.Geometry): Geometry where data sampled from. - dataloader_cfg (Dict[str, Any]): Dataloader config. - loss (loss.Loss): Loss functor. - random (Literal["pseudo", "Halton", "LHS"], optional): Random method for sampling data in - geometry. Defaults to "pseudo". - criteria (Optional[Callable]): Criteria for refining specified boundaries. - Defaults to None. - evenly (bool, optional): Whether to use evenly distribution sampling. - Defaults to False. - weight_dict (Optional[Dict[str, Union[Callable, float]]]): Define the - weight of each constraint variable. Defaults to None. - compute_sdf_derivatives (Optional[bool]): Whether compute derivatives for SDF. - Defaults to False. - name (str, optional): Name of constraint object. Defaults to "EQ". - - Examples: - >>> import ppsci - >>> rect = ppsci.geometry.Rectangle((0, 0), (1, 1)) - >>> pde_constraint = ppsci.constraint.InteriorConstraint( - ... {"u": lambda out: out["u"]}, - ... {"u": 0}, - ... rect, - ... { - ... "dataset": "IterableNamedArrayDataset", - ... "iters_per_epoch": 1, - ... "batch_size": 16, - ... }, - ... ppsci.loss.MSELoss("mean"), - ... name="EQ", - ... ) # doctest: +SKIP - """ - - def __init__( - self, - output_expr: Dict[str, Callable], - label_dict: Dict[str, Union[float, Callable]], - geom: geometry.Geometry, - dataloader_cfg: Dict[str, Any], - loss: "loss.Loss", - random: Literal["pseudo", "Halton", "LHS"] = "pseudo", - criteria: Optional[Callable] = None, - evenly: bool = False, - weight_dict: Optional[Dict[str, Union[Callable, float]]] = None, - compute_sdf_derivatives: bool = False, - name: str = "EQ", - ): - self.label_dict = label_dict - self.input_keys = geom.dim_keys - self.output_keys = tuple(label_dict.keys()) - self.output_expr = { - k: v for k, v in output_expr.items() if k in self.output_keys - } - - if isinstance(criteria, str): - criteria = eval(criteria) - - # prepare input - input = geom.sample_interior( - dataloader_cfg["batch_size"] * dataloader_cfg["iters_per_epoch"], - random, - criteria, - evenly, - compute_sdf_derivatives, - ) - if "area" in input: - input["area"] *= dataloader_cfg["iters_per_epoch"] - - # prepare label - label = {} - for key, value in label_dict.items(): - if isinstance(value, (int, float)): - label[key] = np.full_like(next(iter(input.values())), value) - elif isinstance(value, sympy.Basic): - func = sympy.lambdify( - sympy.symbols(geom.dim_keys), - value, - [{"amax": lambda xy, _: np.maximum(xy[0], xy[1])}, "numpy"], - ) - label[key] = func( - **{k: v for k, v in input.items() if k in geom.dim_keys} - ) - elif callable(value): - func = value - label[key] = func(input) - if isinstance(label[key], (int, float)): - label[key] = np.full_like(next(iter(input.values())), label[key]) - else: - raise NotImplementedError(f"type of {type(value)} is invalid yet.") - - # prepare weight - weight = None - if weight_dict is not None: - weight = {key: np.ones_like(next(iter(label.values()))) for key in label} - for key, value in weight_dict.items(): - if isinstance(value, str): - if value == "sdf": - weight[key] = input["sdf"] - else: - raise NotImplementedError(f"string {value} is invalid yet.") - elif isinstance(value, (int, float)): - weight[key] = np.full_like(next(iter(label.values())), float(value)) - elif isinstance(value, sympy.Basic): - func = sympy.lambdify( - sympy.symbols(geom.dim_keys), - value, - [{"amax": lambda xy, _: np.maximum(xy[0], xy[1])}, "numpy"], - ) - weight[key] = func( - **{k: v for k, v in input.items() if k in geom.dim_keys} - ) - elif callable(value): - func = value - weight[key] = func(input) - if isinstance(weight[key], (int, float)): - weight[key] = np.full_like( - next(iter(input.values())), weight[key] - ) - else: - raise NotImplementedError(f"type of {type(value)} is invalid yet.") - - # wrap input, label, weight into a dataset - if isinstance(dataloader_cfg["dataset"], str): - dataloader_cfg["dataset"] = {"name": dataloader_cfg["dataset"]} - dataloader_cfg["dataset"].update( - {"input": input, "label": label, "weight": weight} - ) - _dataset = dataset.build_dataset(dataloader_cfg["dataset"]) - - # construct dataloader with dataset and dataloader_cfg - super().__init__(_dataset, dataloader_cfg, loss, name) diff --git a/examples/smc_reac/ppsci/constraint/periodic_constraint.py b/examples/smc_reac/ppsci/constraint/periodic_constraint.py deleted file mode 100644 index cb5fc1a332..0000000000 --- a/examples/smc_reac/ppsci/constraint/periodic_constraint.py +++ /dev/null @@ -1,169 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import TYPE_CHECKING -from typing import Any -from typing import Callable -from typing import Dict -from typing import Optional -from typing import Union - -import numpy as np -import paddle -import sympy -from typing_extensions import Literal - -from ppsci import geometry -from ppsci.constraint import base -from ppsci.data import dataset - -if TYPE_CHECKING: - from ppsci import loss - - -class PeriodicConstraint(base.Constraint): - """Class for periodic constraint. - - Args: - output_expr (Dict[str, Callable]): Function in dict for computing output. - e.g. {"u_mul_v": lambda out: out["u"] * out["v"]} means the model output u - will be multiplied by model output v and the result will be named "u_mul_v". - label_dict (Dict[str, Union[float, Callable]]): Function in dict for computing - label, which will be a reference value to participate in the loss calculation. - geom (geometry.Geometry): Geometry where data sampled from. - dataloader_cfg (Dict[str, Any]): Dataloader config. - periodic_key (str): Name of dimension which periodic constraint applied to. - loss (loss.Loss): Loss functor. - random (Literal["pseudo", "Halton", "LHS"], optional): Random method for sampling data in - geometry. Defaults to "pseudo". - criteria (Optional[Callable]): Criteria for refining specified boundaries. - Defaults to None. - evenly (bool, optional): Whether to use evenly distribution sampling. - Defaults to False. - weight_dict (Optional[Dict[str, Callable]]): Define the weight of each - constraint variable. Defaults to None. - name (str, optional): Name of constraint object. Defaults to "PeriodicBC". - """ - - def __init__( - self, - output_expr: Dict[str, Callable], - label_dict: Dict[str, Union[float, Callable]], - geom: geometry.Geometry, - periodic_key: str, - dataloader_cfg: Dict[str, Any], - loss: "loss.Loss", - random: Literal["pseudo", "Halton", "LHS"] = "pseudo", - criteria: Optional[Callable] = None, - evenly: bool = False, - weight_dict: Optional[Dict[str, Callable]] = None, - name: str = "PeriodicBC", - ): - self.input_keys = geom.dim_keys - self.output_keys = tuple(output_expr.keys()) - self.output_expr = { - k: v for k, v in output_expr.items() if k in self.output_keys - } - - if isinstance(criteria, str): - criteria = eval(criteria) - - if dataloader_cfg["batch_size"] % 2 > 0: - raise ValueError( - f"batch_size({dataloader_cfg['sampler']['batch_size']}) " - "should be positive and even when using PeriodicConstraint" - ) - if dataloader_cfg.get("shuffle", False): - raise ValueError( - f"shuffle({dataloader_cfg['sampler']['batch_size']}) " - "should be False when using PeriodicConstraint" - ) - - # prepare input - _bs_half = dataloader_cfg["batch_size"] // 2 - input = geom.sample_boundary( - _bs_half * dataloader_cfg["iters_per_epoch"], - random, - criteria, - evenly, - ) - if "area" in input: - input["area"] *= dataloader_cfg["iters_per_epoch"] - - input_periodic = geom.periodic_point( - input, - geom.geometry.dim_keys.index(periodic_key) - if isinstance(geom, geometry.TimeXGeometry) - else geom.dim_keys.index(periodic_key), - ) - # concatenate original data next to periodic data, i.e. - # [orignal1, periodic1, orignal2, periodic2, ..., orignalN, periodicN] - mixed_input = {} - for key in input: - mixed_input[key] = [] - for iter_id in range(dataloader_cfg["iters_per_epoch"]): - mixed_input[key].append( - input[key][iter_id * _bs_half : (iter_id + 1) * _bs_half] - ) - mixed_input[key].append( - input_periodic[key][iter_id * _bs_half : (iter_id + 1) * _bs_half] - ) - mixed_input[key] = np.vstack(mixed_input[key]) - - # prepare label, keep label the same shape as input_periodic - label = {} - for key, value in label_dict.items(): - # set all label's to zero for dummy data. - label[key] = np.full( - (next(iter(mixed_input.values())).shape[0], 1), - 0, - paddle.get_default_dtype(), - ) - - # # prepare weight, keep weight the same shape as input_periodic - weight = None - if weight_dict is not None: - weight = {key: np.ones_like(next(iter(label.values()))) for key in label} - for key, value in weight_dict.items(): - if isinstance(value, (int, float)): - weight[key] = np.full_like(next(iter(label.values())), value) - elif isinstance(value, sympy.Basic): - func = sympy.lambdify( - [sympy.Symbol(k) for k in geom.dim_keys], - value, - [{"amax": lambda xy, _: np.maximum(xy[0], xy[1])}, "numpy"], - ) - weight[key] = func(**{k: mixed_input[k] for k in geom.dim_keys}) - elif callable(value): - func = value - weight[key] = func(mixed_input) - if isinstance(weight[key], (int, float)): - weight[key] = np.full_like( - next(iter(mixed_input.values())), weight[key] - ) - else: - raise NotImplementedError(f"type of {type(value)} is invalid yet.") - - # wrap input, label, weight into a dataset - if isinstance(dataloader_cfg["dataset"], str): - dataloader_cfg["dataset"] = {"name": dataloader_cfg["dataset"]} - dataloader_cfg["dataset"].update( - {"input": mixed_input, "label": label, "weight": weight} - ) - _dataset = dataset.build_dataset(dataloader_cfg["dataset"]) - - # construct dataloader with dataset and dataloader_cfg - super().__init__(_dataset, dataloader_cfg, loss, name) diff --git a/examples/smc_reac/ppsci/constraint/supervised_constraint.py b/examples/smc_reac/ppsci/constraint/supervised_constraint.py deleted file mode 100644 index 84b8816222..0000000000 --- a/examples/smc_reac/ppsci/constraint/supervised_constraint.py +++ /dev/null @@ -1,92 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import TYPE_CHECKING -from typing import Any -from typing import Callable -from typing import Dict -from typing import Optional - -from ppsci.constraint import base -from ppsci.data import dataset - -if TYPE_CHECKING: - from ppsci import loss - - -class SupervisedConstraint(base.Constraint): - """Class for supervised constraint. - - Args: - dataloader_cfg (Dict[str, Any]): Dataloader config. - loss (loss.Loss): Loss functor. - output_expr (Optional[Dict[str, Callable]]): List of label expression. - Defaults to None. - name (str, optional): Name of constraint object. Defaults to "Sup". - - Examples: - >>> import ppsci - >>> bc_sup = ppsci.constraint.SupervisedConstraint( - ... { - ... "dataset": { - ... "name": "IterableCSVDataset", - ... "file_path": "/path/to/file.csv", - ... "input_keys": ("x", "y"), - ... "label_keys": ("u", "v"), - ... }, - ... }, - ... ppsci.loss.MSELoss("mean"), - ... name="bc_sup", - ... ) # doctest: +SKIP - """ - - def __init__( - self, - dataloader_cfg: Dict[str, Any], - loss: "loss.Loss", - output_expr: Optional[Dict[str, Callable]] = None, - name: str = "Sup", - ): - # build dataset - _dataset = dataset.build_dataset(dataloader_cfg["dataset"]) - - self.input_keys = _dataset.input_keys - self.output_keys = ( - tuple(output_expr.keys()) - if output_expr is not None - else _dataset.label_keys - ) - - self.output_expr = output_expr - if self.output_expr is None: - self.output_expr = { - key: (lambda out, k=key: out[k]) for key in self.output_keys - } - - # construct dataloader with dataset and dataloader_cfg - super().__init__(_dataset, dataloader_cfg, loss, name) - - def __str__(self): - return ", ".join( - [ - self.__class__.__name__, - f"name = {self.name}", - f"input_keys = {self.input_keys}", - f"output_keys = {self.output_keys}", - f"output_expr = {self.output_expr}", - f"loss = {self.loss}", - ] - ) diff --git a/examples/smc_reac/ppsci/data/__init__.py b/examples/smc_reac/ppsci/data/__init__.py deleted file mode 100644 index 41ea77147d..0000000000 --- a/examples/smc_reac/ppsci/data/__init__.py +++ /dev/null @@ -1,205 +0,0 @@ -# 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. - -import copy -import random -from functools import partial -from typing import Callable -from typing import Optional - -import numpy as np -import paddle.distributed as dist -from paddle import device -from paddle import io - -from ppsci.data import dataloader -from ppsci.data import dataset -from ppsci.data import process -from ppsci.data.process import batch_transform -from ppsci.data.process import transform -from ppsci.utils import logger - -__all__ = [ - "dataset", - "process", - "dataloader", - "build_dataloader", - "transform", - "batch_transform", -] - - -def worker_init_fn(worker_id: int, num_workers: int, rank: int, base_seed: int) -> None: - """Callback function on each worker subprocess after seeding and before data loading. - - Args: - worker_id (int): Worker id in [0, num_workers - 1]. - num_workers (int): Number of subprocesses to use for data loading. - rank (int): Rank of process in distributed environment. If in non-distributed - environment, it is a constant number `0`. - base_seed (int): Base random seed. - """ - # The seed of each worker equals to: user_seed + num_worker * rank + worker_id - worker_seed = base_seed + num_workers * rank + worker_id - np.random.seed(worker_seed) - random.seed(worker_seed) - - -def build_dataloader(_dataset, cfg): - world_size = dist.get_world_size() - # just return IterableDataset as dataloader - if isinstance(_dataset, io.IterableDataset): - return _dataset - - cfg = copy.deepcopy(cfg) - - # build sampler - sampler_cfg = cfg.pop("sampler", None) - if sampler_cfg is not None: - batch_sampler_cls = sampler_cfg.pop("name") - - if batch_sampler_cls == "BatchSampler": - if world_size > 1: - batch_sampler_cls = "DistributedBatchSampler" - logger.warning( - f"Automatically use 'DistributedBatchSampler' instead of " - f"'BatchSampler' when world_size({world_size}) > 1." - ) - - sampler_cfg["batch_size"] = cfg["batch_size"] - batch_sampler = getattr(io, batch_sampler_cls)(_dataset, **sampler_cfg) - else: - batch_sampler_cls = "BatchSampler" - if world_size > 1: - batch_sampler_cls = "DistributedBatchSampler" - logger.warning( - f"Automatically use 'DistributedBatchSampler' instead of " - f"'BatchSampler' when world_size({world_size}) > 1." - ) - batch_sampler = getattr(io, batch_sampler_cls)( - _dataset, - batch_size=cfg["batch_size"], - shuffle=False, - drop_last=False, - ) - logger.message( - "'shuffle' and 'drop_last' are both set to False in default as sampler config is not specified." - ) - - # build collate_fn if specified - batch_transforms_cfg = cfg.pop("batch_transforms", None) - collate_fn: Optional[Callable] = cfg.pop("collate_fn", None) - if isinstance(batch_transforms_cfg, (list, tuple)): - collate_fn = batch_transform.build_batch_transforms( - batch_transforms_cfg, collate_fn - ) - - # build init function - _DEFAULT_NUM_WORKERS = 1 - _DEFAULT_SEED = 42 - init_fn = partial( - worker_init_fn, - num_workers=cfg.get("num_workers", _DEFAULT_NUM_WORKERS), - rank=dist.get_rank(), - base_seed=cfg.get("seed", _DEFAULT_SEED), - ) - - # build dataloader - if getattr(_dataset, "use_pgl", False): - # Use special dataloader from "Paddle Graph Learning" toolkit. - try: - from pgl.utils import data as pgl_data - except ModuleNotFoundError as e: - logger.error("Please install pgl with `pip install pgl`.") - raise ModuleNotFoundError(str(e)) - - if collate_fn is None: - collate_fn = batch_transform.default_collate_fn - dataloader_ = pgl_data.Dataloader( - dataset=_dataset, - batch_size=cfg["batch_size"], - drop_last=sampler_cfg.get("drop_last", False), - shuffle=sampler_cfg.get("shuffle", False), - num_workers=cfg.get("num_workers", _DEFAULT_NUM_WORKERS), - collate_fn=collate_fn, - ) - elif getattr(_dataset, "use_graph_grid_mesh", False): - # Use special dataloader `GridMeshAtmosphericDataset`. - - if collate_fn is None: - collate_fn = batch_transform.default_collate_fn - dataloader_ = io.DataLoader( - dataset=_dataset, - places=device.get_device(), - batch_sampler=batch_sampler, - collate_fn=collate_fn, - num_workers=cfg.get("num_workers", _DEFAULT_NUM_WORKERS), - use_shared_memory=cfg.get("use_shared_memory", False), - worker_init_fn=init_fn, - ) - else: - if ( - cfg.get("auto_collation", not getattr(_dataset, "batch_index", False)) - is False - ): - if "transforms" in cfg["dataset"] and "auto_collation" not in cfg: - logger.warning( - "'transforms' and batch indexing(auto_collation=False) are both " - "enabled. If you do want to apply transforms to the batch samples, " - "please explicitly set 'auto_collation' to False in dataloader_cfg;" - " otherwise, the 'transforms' will be retained, but batch indexing " - "will be disabled." - ) - else: - # 1. wrap batch_sampler again into BatchSampler for disabling auto collation, - # which can speed up the process of batch samples indexing from dataset. See - # details at: https://discuss.pytorch.org/t/efficiency-of-dataloader-and-collate-for-large-array-like-datasets/59569/8 - batch_sampler = io.BatchSampler(sampler=batch_sampler, batch_size=1) - if collate_fn is not None: - raise NotImplementedError( - "Detected collate_fn is not None for 'batch_transforms' might " - "be specified in 'dataloader_cfg', which is not supported yet " - "with 'auto_collation' is False at the same time" - ) - # 2. disable auto collation by given identity collate_fn which return the first - # (also the only) batch data in batch list, or there will be a redundant - # axis at the first dimension returned by dataloader. This step is necessary - # because paddle do not support 'sampler' as instantiation argument of 'io.DataLoader' - collate_fn = lambda batch: batch[0] # noqa: E731 - _DEFAULT_NUM_WORKERS = 0 - logger.info( - "Auto collation is disabled and set num_workers to " - f"{_DEFAULT_NUM_WORKERS} to speed up batch sampling." - ) - - dataloader_ = io.DataLoader( - dataset=_dataset, - places=device.get_device(), - batch_sampler=batch_sampler, - collate_fn=collate_fn, - num_workers=cfg.get("num_workers", _DEFAULT_NUM_WORKERS), - use_shared_memory=cfg.get("use_shared_memory", False), - worker_init_fn=init_fn, - # TODO: Do not enable 'persistent_workers' below for - # 'IndexError: pop from empty list ...' will be raised in certain cases - # persistent_workers=cfg.get("num_workers", _DEFAULT_NUM_WORKERS) > 0, - ) - - if len(dataloader_) == 0: - raise ValueError( - f"batch_size({sampler_cfg['batch_size']}) should not bigger than number of " - f"samples({len(_dataset)}) when drop_last is {sampler_cfg.get('drop_last', False)}." - ) - - return dataloader_ diff --git a/examples/smc_reac/ppsci/data/dataloader.py b/examples/smc_reac/ppsci/data/dataloader.py deleted file mode 100644 index 4c01da873e..0000000000 --- a/examples/smc_reac/ppsci/data/dataloader.py +++ /dev/null @@ -1,47 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import Union - -from paddle import io - - -class InfiniteDataLoader: - """A wrapper for infinite dataloader. - - Args: - dataloader (Union[io.DataLoader, io.IterableDataset]): A finite and iterable loader or iterable dataset to be wrapped. - """ - - def __init__(self, dataloader: Union[io.DataLoader, io.IterableDataset]): - self.dataloader = dataloader - if isinstance(dataloader, io.DataLoader): - self.dataset = dataloader.dataset - elif isinstance(dataloader, io.IterableDataset): - self.dataset = dataloader - else: - raise TypeError( - f"dataloader should be io.DataLoader or io.IterableDataset, but got {type(dataloader)}" - ) - - def __iter__(self): - while True: - dataloader_iter = iter(self.dataloader) - for batch in dataloader_iter: - yield batch - - def __len__(self): - return len(self.dataloader) diff --git a/examples/smc_reac/ppsci/data/dataset/__init__.py b/examples/smc_reac/ppsci/data/dataset/__init__.py deleted file mode 100644 index 9ece354700..0000000000 --- a/examples/smc_reac/ppsci/data/dataset/__init__.py +++ /dev/null @@ -1,118 +0,0 @@ -# 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. - -import copy -from typing import TYPE_CHECKING - -from ppsci.data.dataset.airfoil_dataset import MeshAirfoilDataset -from ppsci.data.dataset.array_dataset import ChipHeatDataset -from ppsci.data.dataset.array_dataset import ContinuousNamedArrayDataset -from ppsci.data.dataset.array_dataset import IterableNamedArrayDataset -from ppsci.data.dataset.array_dataset import NamedArrayDataset -from ppsci.data.dataset.atmospheric_dataset import GridMeshAtmosphericDataset -from ppsci.data.dataset.cgcnn_dataset import CIFData as CGCNNDataset -from ppsci.data.dataset.csv_dataset import CSVDataset -from ppsci.data.dataset.csv_dataset import IterableCSVDataset -from ppsci.data.dataset.cylinder_dataset import MeshCylinderDataset -from ppsci.data.dataset.darcyflow_dataset import DarcyFlowDataset -from ppsci.data.dataset.dgmr_dataset import DGMRDataset -from ppsci.data.dataset.drivaernet_dataset import DrivAerNetDataset -from ppsci.data.dataset.drivaernetplusplus_dataset import DrivAerNetPlusPlusDataset -from ppsci.data.dataset.enso_dataset import ENSODataset -from ppsci.data.dataset.era5_dataset import ERA5Dataset -from ppsci.data.dataset.era5_dataset import ERA5SampledDataset -from ppsci.data.dataset.ext_moe_enso_dataset import ExtMoEENSODataset -from ppsci.data.dataset.fwi_dataset import FWIDataset -from ppsci.data.dataset.ifm_moe_dataset import IFMMoeDataset -from ppsci.data.dataset.mat_dataset import IterableMatDataset -from ppsci.data.dataset.mat_dataset import MatDataset -from ppsci.data.dataset.moflow_dataset import MOlFLOWDataset -from ppsci.data.dataset.mrms_dataset import MRMSDataset -from ppsci.data.dataset.mrms_dataset import MRMSSampledDataset -from ppsci.data.dataset.npz_dataset import IterableNPZDataset -from ppsci.data.dataset.npz_dataset import NPZDataset -from ppsci.data.dataset.pems_dataset import PEMSDataset -from ppsci.data.dataset.radar_dataset import RadarDataset -from ppsci.data.dataset.sevir_dataset import SEVIRDataset -from ppsci.data.dataset.spherical_swe_dataset import SphericalSWEDataset -from ppsci.data.dataset.trphysx_dataset import CylinderDataset -from ppsci.data.dataset.trphysx_dataset import LorenzDataset -from ppsci.data.dataset.trphysx_dataset import RosslerDataset -from ppsci.data.dataset.vtu_dataset import VtuDataset -from ppsci.data.process import transform -from ppsci.utils import logger - -if TYPE_CHECKING: - from paddle import io - -__all__ = [ - "IterableNamedArrayDataset", - "NamedArrayDataset", - "ContinuousNamedArrayDataset", - "ChipHeatDataset", - "CSVDataset", - "IterableCSVDataset", - "ERA5Dataset", - "ERA5SampledDataset", - "GridMeshAtmosphericDataset", - "IterableMatDataset", - "MatDataset", - "MRMSDataset", - "MRMSSampledDataset", - "IterableNPZDataset", - "NPZDataset", - "PEMSDataset", - "CylinderDataset", - "LorenzDataset", - "RadarDataset", - "RosslerDataset", - "VtuDataset", - "DGMRDataset", - "MeshAirfoilDataset", - "MeshCylinderDataset", - "DarcyFlowDataset", - "SphericalSWEDataset", - "ENSODataset", - "ExtMoEENSODataset", - "SEVIRDataset", - "MOlFLOWDataset", - "build_dataset", - "CGCNNDataset", - "FWIDataset", - "DrivAerNetDataset", - "DrivAerNetPlusPlusDataset", - "IFMMoeDataset", -] - - -def build_dataset(cfg) -> "io.Dataset": - """Build dataset - - Args: - cfg (List[DictConfig]): Dataset config list. - - Returns: - Dict[str, io.Dataset]: dataset. - """ - cfg = copy.deepcopy(cfg) - - dataset_cls = cfg.pop("name") - if "transforms" in cfg: - cfg["transforms"] = transform.build_transforms(cfg.pop("transforms")) - - dataset = eval(dataset_cls)(**cfg) - - logger.debug(str(dataset)) - - return dataset diff --git a/examples/smc_reac/ppsci/data/dataset/airfoil_dataset.py b/examples/smc_reac/ppsci/data/dataset/airfoil_dataset.py deleted file mode 100644 index 2a249104f7..0000000000 --- a/examples/smc_reac/ppsci/data/dataset/airfoil_dataset.py +++ /dev/null @@ -1,241 +0,0 @@ -# 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. - -from __future__ import annotations - -import os -import pickle -from os import path as osp -from typing import Dict -from typing import List -from typing import Tuple - -import numpy as np -import paddle -from paddle import io - -try: - import pgl -except ModuleNotFoundError: - pass - -SU2_SHAPE_IDS = { - "line": 3, - "triangle": 5, - "quad": 9, -} - - -# HACK: Simplify code below -def _get_mesh_graph( - mesh_filename: str, dtype: np.dtype = np.float32 -) -> Tuple[np.ndarray, np.ndarray, List[List[List[int]]], Dict[str, List[List[int]]]]: - def get_rhs(s: str) -> str: - return s.split("=")[-1] - - marker_dict = {} - with open(mesh_filename) as f: - for line in f: - if line.startswith("NPOIN"): - num_points = int(get_rhs(line)) - mesh_points = [ - [float(p) for p in f.readline().split()[:2]] - for _ in range(num_points) - ] - nodes = np.array(mesh_points, dtype=dtype) - elif line.startswith("NMARK"): - num_markers = int(get_rhs(line)) - for _ in range(num_markers): - line = f.readline() - assert line.startswith("MARKER_TAG") - marker_tag = get_rhs(line).strip() - num_elems = int(get_rhs(f.readline())) - marker_elems = [ - [int(e) for e in f.readline().split()[-2:]] - for _ in range(num_elems) - ] - marker_dict[marker_tag] = marker_elems - elif line.startswith("NELEM"): - edges = [] - triangles = [] - quads = [] - num_edges = int(get_rhs(line)) - for _ in range(num_edges): - elem = [int(p) for p in f.readline().split()] - if elem[0] == SU2_SHAPE_IDS["triangle"]: - n = 3 - triangles.append(elem[1 : 1 + n]) - elif elem[0] == SU2_SHAPE_IDS["quad"]: - n = 4 - quads.append(elem[1 : 1 + n]) - else: - raise NotImplementedError - elem = elem[1 : 1 + n] - edges += [[elem[i], elem[(i + 1) % n]] for i in range(n)] - edges = np.array(edges, dtype=np.compat.long).transpose() - elems = [triangles, quads] - return nodes, edges, elems, marker_dict - - -class MeshAirfoilDataset(io.Dataset): - """Dataset for `MeshAirfoil`. - - Args: - input_keys (Tuple[str, ...]): Name of input data. - label_keys (Tuple[str, ...]): Name of label data. - data_dir (str): Directory of MeshAirfoil data. - mesh_graph_path (str): Path of mesh graph. - transpose_edges (bool, optional): Whether transpose the edges array from (2, num_edges) to (num_edges, 2) for convenient of slicing. - - Examples: - >>> import ppsci - >>> dataset = ppsci.data.dataset.MeshAirfoilDataset( - ... "input_keys": ("input",), - ... "label_keys": ("output",), - ... "data_dir": "/path/to/MeshAirfoilDataset", - ... "mesh_graph_path": "/path/to/file.su2", - ... "transpose_edges": False, - ... ) # doctest: +SKIP - """ - - # Whether support batch indexing for speeding up fetching process. - batch_index: bool = False - - use_pgl: bool = True - - def __init__( - self, - input_keys: Tuple[str, ...], - label_keys: Tuple[str, ...], - data_dir: str, - mesh_graph_path: str, - transpose_edges: bool = False, - ): - self.input_keys = input_keys - self.label_keys = label_keys - self.data_dir = data_dir - self.file_list = os.listdir(self.data_dir) - self.len = len(self.file_list) - self.mesh_graph = _get_mesh_graph(mesh_graph_path) - - with open(osp.join(osp.dirname(self.data_dir), "train_max_min.pkl"), "rb") as f: - self.normalization_factors = pickle.load(f) - - self.nodes = self.mesh_graph[0] - self.edges = self.mesh_graph[1] - if transpose_edges: - self.edges = self.edges.transpose([1, 0]) - self.elems_list = self.mesh_graph[2] - self.marker_dict = self.mesh_graph[3] - self.node_markers = np.full([self.nodes.shape[0], 1], fill_value=-1) - for i, (marker_tag, marker_elems) in enumerate(self.marker_dict.items()): - for elem in marker_elems: - self.node_markers[elem[0]] = i - self.node_markers[elem[1]] = i - - self.raw_graphs = [self.get(i) for i in range(len(self))] - - def __len__(self): - return self.len - - def __getitem__(self, idx): - return ( - { - self.input_keys[0]: self.raw_graphs[idx], - }, - { - self.label_keys[0]: self.raw_graphs[idx], - }, - None, - ) - - def get(self, idx): - with open(osp.join(self.data_dir, self.file_list[idx]), "rb") as f: - fields = pickle.load(f) - fields = self._preprocess(fields) - aoa, reynolds, mach = self._get_params_from_name(self.file_list[idx]) - # aoa = aoa - mach_or_reynolds = mach if reynolds is None else reynolds - # mach_or_reynolds = mach_or_reynolds - norm_aoa = aoa / 10 - norm_mach_or_reynolds = ( - mach_or_reynolds if reynolds is None else (mach_or_reynolds - 1.5e6) / 1.5e6 - ) - - nodes = np.concatenate( - [ - self.nodes, - np.repeat(a=norm_aoa, repeats=self.nodes.shape[0])[:, np.newaxis], - np.repeat(a=norm_mach_or_reynolds, repeats=self.nodes.shape[0])[ - :, np.newaxis - ], - self.node_markers, - ], - axis=-1, - ).astype(paddle.get_default_dtype()) - - data = pgl.Graph( - num_nodes=nodes.shape[0], - edges=self.edges, - ) - data.x = nodes - data.y = fields - data.pos = self.nodes - data.edge_index = self.edges - - sender = data.x[data.edge_index[0]] - receiver = data.x[data.edge_index[1]] - relation_pos = sender[:, 0:2] - receiver[:, 0:2] - post = np.linalg.norm(relation_pos, ord=2, axis=1, keepdims=True).astype( - paddle.get_default_dtype() - ) - data.edge_attr = post - std_epsilon = [1e-8] - a = np.mean(data.edge_attr, axis=0) - b = data.edge_attr.std(axis=0) - b = np.maximum(b, std_epsilon).astype(paddle.get_default_dtype()) - data.edge_attr = (data.edge_attr - a) / b - data.aoa = aoa - data.norm_aoa = norm_aoa - data.mach_or_reynolds = mach_or_reynolds - data.norm_mach_or_reynolds = norm_mach_or_reynolds - return data - - def _preprocess(self, tensor_list, stack_output=True): - data_max, data_min = self.normalization_factors - normalized_tensors = [] - for i in range(len(tensor_list)): - normalized = (tensor_list[i] - data_min[i]) / ( - data_max[i] - data_min[i] - ) * 2 - 1 - normalized_tensors.append(normalized) - if stack_output: - normalized_tensors = np.stack(normalized_tensors, axis=1) - return normalized_tensors - - def _get_params_from_name(self, filename): - s = filename.rsplit(".", 1)[0].split("_") - aoa = np.array(s[s.index("aoa") + 1])[np.newaxis].astype( - paddle.get_default_dtype() - ) - reynolds = s[s.index("re") + 1] - reynolds = ( - np.array(reynolds)[np.newaxis].astype(paddle.get_default_dtype()) - if reynolds != "None" - else None - ) - mach = np.array(s[s.index("mach") + 1])[np.newaxis].astype( - paddle.get_default_dtype() - ) - return aoa, reynolds, mach diff --git a/examples/smc_reac/ppsci/data/dataset/array_dataset.py b/examples/smc_reac/ppsci/data/dataset/array_dataset.py deleted file mode 100644 index 65ac0ca7c4..0000000000 --- a/examples/smc_reac/ppsci/data/dataset/array_dataset.py +++ /dev/null @@ -1,390 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import Callable -from typing import Dict -from typing import Optional - -import numpy as np -import paddle -from paddle import distributed as dist -from paddle import io -from paddle import vision - -from ppsci.utils import logger - - -def _group_array_into_ranks( - data: Optional[np.ndarray], rank: int, world_size: int -) -> Optional[np.ndarray]: - """ - Group data into different ranks. For example, if data is [1, 2, 3, 4, 5, 6, 7, 8, 9] and - world_size is 3, then the result will be rank0: [1, 4, 7], rank1: [2, 5, 8], rank2: [3, 6, 9]. - - Args: - data (Optional[np.ndarray]): Data to be grouped, can be np.ndarray or None. - rank (int): Rank number. - world_size (int): Number of workers. - - Returns: - np.ndarray: Grouped data. - """ - if data is None: - # skip grouping if data is None - return None - - # check if data can be grouped evenly into different ranks - if len(data) < world_size: - raise ValueError( - f"Length of data to be grouped{len(data)} must be larger than world_size." - ) - if len(data) % world_size != 0: - raise ValueError( - f"Length of data to be grouped{len(data)} must be divisible by world_size." - ) - - return data[rank::world_size] - - -def _group_dict_into_ranks( - data_dict: Optional[Dict[str, Optional[np.ndarray]]], rank: int, world_size: int -) -> Optional[Dict[str, Optional[np.ndarray]]]: - """ - Group data dict into different ranks for each key-value pair. - - Args: - data_dict (Dict[str, Optional[np.ndarray]]): Data to be grouped, can be Dict[str, Optional[np.ndarray]] or None. - rank (int): Rank number. - world_size (int): Number of workers. - - Returns: - Optional[Dict[str, Optional[np.ndarray]]]: Grouped data dict. - """ - - if data_dict is None: - return data_dict - - return { - k: _group_array_into_ranks(v, rank, world_size) for k, v in data_dict.items() - } - - -class NamedArrayDataset(io.Dataset): - """Class for Named Array Dataset. - - Args: - input (Dict[str, np.ndarray]): Input dict. - label (Optional[Dict[str, np.ndarray]]): Label dict. Defaults to None. - weight (Optional[Dict[str, np.ndarray]]): Weight dict. Defaults to None. - transforms (Optional[vision.Compose]): Compose object contains sample wise - transform(s). Defaults to None. - - Examples: - >>> import ppsci - >>> input = {"x": np.random.randn(100, 1)} - >>> output = {"u": np.random.randn(100, 1)} - >>> weight = {"u": np.random.randn(100, 1)} - >>> dataset = ppsci.data.dataset.NamedArrayDataset(input, output, weight) - """ - - # Whether support batch indexing for speeding up fetching process. - batch_index: bool = True - - def __init__( - self, - input: Dict[str, np.ndarray], - label: Optional[Dict[str, np.ndarray]] = None, - weight: Optional[Dict[str, np.ndarray]] = None, - transforms: Optional[vision.Compose] = None, - ): - super().__init__() - self.input = input - self.label = {} if label is None else label - self.input_keys = tuple(input.keys()) - self.label_keys = tuple(self.label.keys()) - self.weight = {} if weight is None else weight - self.transforms = transforms - self._len = len(next(iter(input.values()))) - for key in input: - if key in self.label and len(input[key]) != len(self.label[key]): - logger.warning( - f"The length of input {key}({len(input[key])}) is not equal to " - f"the length of label {key}({len(self.label[key])})." - ) - - def __getitem__(self, idx): - input_item = {key: value[idx] for key, value in self.input.items()} - label_item = {key: value[idx] for key, value in self.label.items()} - weight_item = {key: value[idx] for key, value in self.weight.items()} - - if self.transforms is not None: - input_item, label_item, weight_item = self.transforms( - input_item, label_item, weight_item - ) - - return (input_item, label_item, weight_item) - - def __len__(self): - return self._len - - -class IterableNamedArrayDataset(io.IterableDataset): - """IterableNamedArrayDataset for full-data loading. - - Args: - input (Dict[str, np.ndarray]): Input dict. - label (Optional[Dict[str, np.ndarray]]): Label dict. Defaults to None. - weight (Optional[Dict[str, np.ndarray]]): Weight dict. Defaults to None. - transforms (Optional[vision.Compose]): Compose object contains sample wise - transform(s). Defaults to None. - - Examples: - >>> import ppsci - >>> input = {"x": np.random.randn(100, 1)} - >>> label = {"u": np.random.randn(100, 1)} - >>> weight = {"u": np.random.randn(100, 1)} - >>> dataset = ppsci.data.dataset.IterableNamedArrayDataset(input, label, weight) - """ - - # Whether support batch indexing for speeding up fetching process. - batch_index: bool = False - - def __init__( - self, - input: Dict[str, np.ndarray], - label: Optional[Dict[str, np.ndarray]] = None, - weight: Optional[Dict[str, np.ndarray]] = None, - transforms: Optional[vision.Compose] = None, - ): - super().__init__() - self.input = {key: paddle.to_tensor(value) for key, value in input.items()} - self.label = ( - {key: paddle.to_tensor(value) for key, value in label.items()} - if label is not None - else {} - ) - self.input_keys = tuple(input.keys()) - self.label_keys = tuple(self.label.keys()) - self.weight = ( - { - key: paddle.to_tensor(value, paddle.get_default_dtype()) - for key, value in weight.items() - } - if weight is not None - else None - ) - self._len = len(next(iter(self.input.values()))) - self.transforms = transforms - self.world_size_ = dist.get_world_size() - self.rank_ = dist.get_rank() - - @property - def num_samples(self): - """Number of samples within current dataset.""" - return self._len - - def __iter__(self): - if callable(self.transforms): - input_, label_, weight_ = self.transforms( - self.input, self.label, self.weight - ) - else: - input_, label_, weight_ = self.input, self.label, self.weight - - if self.world_size_ > 1: - input_ = _group_dict_into_ranks(input_, self.rank_, self.world_size_) - label_ = _group_dict_into_ranks(label_, self.rank_, self.world_size_) - weight_ = _group_dict_into_ranks(weight_, self.rank_, self.world_size_) - - yield input_, label_, weight_ - - def __len__(self): - return 1 - - -class ContinuousNamedArrayDataset(io.IterableDataset): - """ContinuousNamedArrayDataset for iterable sampling. - - Args: - input (Callable): Function generate input dict. - label (Callable): Function generate label dict. - weight (Optional[Callable]): Function generate weight dict. Defaults to None. - transforms (Optional[vision.Compose]): Compose object contains sample wise - transform(s). Defaults to None. - - Examples: - >>> import ppsci - >>> import numpy as np - >>> input = lambda : {"x": np.random.randn(100, 1)} - >>> label = lambda inp: {"u": np.random.randn(100, 1)} - >>> weight = lambda inp, label: {"u": 1 - (label["u"] ** 2)} - >>> dataset = ppsci.data.dataset.ContinuousNamedArrayDataset(input, label, weight) - >>> input_batch, label_batch, weight_batch = next(iter(dataset)) - >>> print(input_batch["x"].shape) - [100, 1] - >>> print(label_batch["u"].shape) - [100, 1] - >>> print(weight_batch["u"].shape) - [100, 1] - """ - - # Whether support batch indexing for speeding up fetching process. - batch_index: bool = False - - def __init__( - self, - input: Callable, - label: Callable, - weight: Optional[Callable] = None, - transforms: Optional[vision.Compose] = None, - ): - super().__init__() - self.input_fn = input - self.input_keys = tuple(self.input_fn().keys()) - - self.label_fn = label - input_ = self.input_fn() - self.label_keys = tuple(self.label_fn(input_).keys()) - - self.weight_fn = weight - self.transforms = transforms - self.world_size_ = dist.get_world_size() - self.rank_ = dist.get_rank() - - @property - def num_samples(self): - """Number of samples within current dataset.""" - raise NotImplementedError( - "ContinuousNamedArrayDataset has no fixed number of samples." - ) - - def __iter__(self): - def to_tensor_dict(_dict): - if _dict is None: - return None - return {k: paddle.to_tensor(v) for k, v in _dict.items()} - - while True: - input_batch = self.input_fn() - label_batch = self.label_fn(input_batch) - if callable(self.weight_fn): - weight_batch = self.weight_fn(input_batch, label_batch) - else: - weight_batch = None - - if callable(self.transforms): - input_batch, label_batch, weight_batch = self.transforms( - input_batch, label_batch, weight_batch - ) - - if self.world_size_ > 1: - input_batch = _group_dict_into_ranks( - input_batch, self.rank_, self.world_size_ - ) - label_batch = _group_dict_into_ranks( - label_batch, self.rank_, self.world_size_ - ) - weight_batch = _group_dict_into_ranks( - weight_batch, self.rank_, self.world_size_ - ) - - yield to_tensor_dict(input_batch), to_tensor_dict( - label_batch - ), to_tensor_dict(weight_batch) - - def __len__(self): - return 1 - - -class ChipHeatDataset(io.Dataset): - """ChipHeatDataset for data loading of multi-branch DeepONet model. - - Args: - input (Dict[str, np.ndarray]): Input dict. - label (Optional[Dict[str, np.ndarray]]): Label dict. Defaults to None. - index (tuple[str, ...]): Key of input dict. - data_type (str): One of key of input dict. - weight (Optional[Dict[str, np.ndarray]]): Weight dict. Defaults to None. - transforms (Optional[vision.Compose]): Compose object contains sample wise - transform(s). Defaults to None. - - Examples: - >>> import ppsci - >>> input = {"x": np.random.randn(100, 1)} - >>> label = {"u": np.random.randn(100, 1)} - >>> index = ('x', 'u', 'bc', 'bc_data') - >>> data_type = 'u' - >>> weight = {"u": np.random.randn(100, 1)} - >>> dataset = ppsci.data.dataset.ChipHeatDataset(input, label, index, data_type, weight) - """ - - def __init__( - self, - input: Dict[str, np.ndarray], - label: Dict[str, np.ndarray], - index: tuple[str, ...], - data_type: str, - weight: Optional[Dict[str, float]] = None, - transforms: Optional[vision.Compose] = None, - ): - super().__init__() - self.input = input - self.label = label - self.input_keys = tuple(input.keys()) - self.label_keys = tuple(label.keys()) - self.index = index - self.data_type = data_type - self.weight = {} if weight is None else weight - self.transforms = transforms - - def __getitem__(self, idx): - quotient = idx - index_ir = dict() - for i in self.index: - index_ir[i] = 0 - - for i in index_ir: - num = len(self.input[i]) - index_ir[i] = quotient % num - quotient = quotient // num - - input_item = {} - for key in self.input: - if key == "y": - input_item[key] = self.input[key][index_ir["x"]] - elif key == "u_one": - input_item[key] = self.input[key][ - len(self.input[self.data_type]) * index_ir["x"] - + index_ir[self.data_type] - ] - else: - input_item[key] = self.input[key][index_ir[key]] - - label_item = {key: value for key, value in self.label.items()} - weight_item = {key: value for key, value in self.weight.items()} - - if self.transforms is not None: - input_item, label_item, weight_item = self.transforms( - (input_item, label_item, weight_item) - ) - - return (input_item, label_item, weight_item) - - def __len__(self): - _len = 1 - for i in self.index: - _len *= len(self.input[i]) - return _len diff --git a/examples/smc_reac/ppsci/data/dataset/atmospheric_dataset.py b/examples/smc_reac/ppsci/data/dataset/atmospheric_dataset.py deleted file mode 100644 index c0e4c9ee6d..0000000000 --- a/examples/smc_reac/ppsci/data/dataset/atmospheric_dataset.py +++ /dev/null @@ -1,1781 +0,0 @@ -# Copyright (c) 2024 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. - -from typing import List -from typing import NamedTuple -from typing import Optional -from typing import Sequence -from typing import Tuple - -import numpy as np -import paddle -import pandas as pd -import scipy -from paddle import io - -try: - import trimesh - import xarray -except ModuleNotFoundError: - pass - -# https://www.ecmwf.int/en/forecasts/dataset/ecmwf-reanalysis-v5 -PRESSURE_LEVELS_ERA5_37 = ( - 1, - 2, - 3, - 5, - 7, - 10, - 20, - 30, - 50, - 70, - 100, - 125, - 150, - 175, - 200, - 225, - 250, - 300, - 350, - 400, - 450, - 500, - 550, - 600, - 650, - 700, - 750, - 775, - 800, - 825, - 850, - 875, - 900, - 925, - 950, - 975, - 1000, -) - -# https://www.ecmwf.int/en/forecasts/datasets/set-i -PRESSURE_LEVELS_HRES_25 = ( - 1, - 2, - 3, - 5, - 7, - 10, - 20, - 30, - 50, - 70, - 100, - 150, - 200, - 250, - 300, - 400, - 500, - 600, - 700, - 800, - 850, - 900, - 925, - 950, - 1000, -) - -# https://agupubs.onlinelibrary.wiley.com/doi/full/10.1029/2020MS002203 -PRESSURE_LEVELS_WEATHERBENCH_13 = ( - 50, - 100, - 150, - 200, - 250, - 300, - 400, - 500, - 600, - 700, - 850, - 925, - 1000, -) - -PRESSURE_LEVELS = { - 13: PRESSURE_LEVELS_WEATHERBENCH_13, - 25: PRESSURE_LEVELS_HRES_25, - 37: PRESSURE_LEVELS_ERA5_37, -} - - -TARGET_SURFACE_VARS = ( - "2m_temperature", - "mean_sea_level_pressure", - "10m_v_component_of_wind", - "10m_u_component_of_wind", - "total_precipitation_6hr", -) -TARGET_SURFACE_NO_PRECIP_VARS = ( - "2m_temperature", - "mean_sea_level_pressure", - "10m_v_component_of_wind", - "10m_u_component_of_wind", -) -TARGET_ATMOSPHERIC_VARS = ( - "temperature", - "geopotential", - "u_component_of_wind", - "v_component_of_wind", - "vertical_velocity", - "specific_humidity", -) -TARGET_ATMOSPHERIC_NO_W_VARS = ( - "temperature", - "geopotential", - "u_component_of_wind", - "v_component_of_wind", - "specific_humidity", -) -EXTERNAL_FORCING_VARS = ("toa_incident_solar_radiation",) -GENERATED_FORCING_VARS = ( - "year_progress_sin", - "year_progress_cos", - "day_progress_sin", - "day_progress_cos", -) -FORCING_VARS = EXTERNAL_FORCING_VARS + GENERATED_FORCING_VARS -STATIC_VARS = ( - "geopotential_at_surface", - "land_sea_mask", -) - -TASK_input_variables = ( - TARGET_SURFACE_VARS + TARGET_ATMOSPHERIC_VARS + FORCING_VARS + STATIC_VARS -) -TASK_target_variables = TARGET_SURFACE_VARS + TARGET_ATMOSPHERIC_VARS -TASK_forcing_variables = FORCING_VARS -TASK_pressure_levels = PRESSURE_LEVELS_ERA5_37 -TASK_input_duration = ("12h",) - -TASK_13_input_variables = ( - TARGET_SURFACE_VARS + TARGET_ATMOSPHERIC_VARS + FORCING_VARS + STATIC_VARS -) -TASK_13_target_variables = TARGET_SURFACE_VARS + TARGET_ATMOSPHERIC_VARS -TASK_13_forcing_variables = FORCING_VARS -TASK_13_pressure_levels = PRESSURE_LEVELS_WEATHERBENCH_13 -TASK_13_input_duration = ("12h",) - - -TASK_13_PRECIP_OUT_input_variables = ( - TARGET_SURFACE_NO_PRECIP_VARS + TARGET_ATMOSPHERIC_VARS + FORCING_VARS + STATIC_VARS -) -TASK_13_PRECIP_OUT_target_variables = TARGET_SURFACE_VARS + TARGET_ATMOSPHERIC_VARS -TASK_13_PRECIP_OUT_forcing_variables = FORCING_VARS -TASK_13_PRECIP_OUT_pressure_levels = PRESSURE_LEVELS_WEATHERBENCH_13 -TASK_13_PRECIP_OUT_input_duration = ("12h",) - -_SEC_PER_HOUR = 3600 -_HOUR_PER_DAY = 24 -SEC_PER_DAY = _SEC_PER_HOUR * _HOUR_PER_DAY -_AVG_DAY_PER_YEAR = 365.24219 -AVG_SEC_PER_YEAR = SEC_PER_DAY * _AVG_DAY_PER_YEAR - -DAY_PROGRESS = "day_progress" -YEAR_PROGRESS = "year_progress" - - -def stacked_to_dataset( - stacked_array: "xarray.Variable", - template_dataset: "xarray.Dataset", - preserved_dims: Tuple[str, ...] = ("batch", "lat", "lon"), -) -> "xarray.Dataset": - """The inverse of dataset_to_stacked. - - Requires a template dataset to demonstrate the variables/shapes/coordinates - required. - All variables must have preserved_dims dimensions. - - Args: - stacked_array: Data in BHWC layout, encoded the same as dataset_to_stacked would if it was asked to encode `template_dataset`. - template_dataset: A template Dataset (or other mapping of DataArrays) demonstrating the shape of output required (variables, shapes, coordinates etc). - preserved_dims: dimensions from the target_template that were not folded in the predictions channels. The preserved_dims need to be a subset of the dims of all the variables of template_dataset. - - Returns: - An xarray.Dataset (or other mapping of DataArrays) with the same shape and type as template_dataset. - """ - unstack_from_channels_sizes = {} - var_names = sorted(template_dataset.keys()) - for name in var_names: - template_var = template_dataset[name] - if not all(dim in template_var.dims for dim in preserved_dims): - raise ValueError( - f"stacked_to_dataset requires all Variables to have {preserved_dims} " - f"dimensions, but found only {template_var.dims}." - ) - unstack_from_channels_sizes[name] = { - dim: size - for dim, size in template_var.sizes.items() - if dim not in preserved_dims - } - - channels = { - name: np.prod(list(unstack_sizes.values()), dtype=np.int64) - for name, unstack_sizes in unstack_from_channels_sizes.items() - } - total_expected_channels = sum(channels.values()) - found_channels = stacked_array.sizes["channels"] - if total_expected_channels != found_channels: - raise ValueError( - f"Expected {total_expected_channels} channels but found " - f"{found_channels}, when trying to convert a stacked array of shape " - f"{stacked_array.sizes} to a dataset of shape {template_dataset}." - ) - - data_vars = {} - index = 0 - for name in var_names: - template_var = template_dataset[name] - var = stacked_array.isel({"channels": slice(index, index + channels[name])}) - index += channels[name] - var = var.unstack({"channels": unstack_from_channels_sizes[name]}) - var = var.transpose(*template_var.dims) - data_vars[name] = xarray.DataArray( - data=var, - coords=template_var.coords, - # This might not always be the same as the name it's keyed under; it - # will refer to the original variable name, whereas the key might be - # some alias e.g. temperature_850 under which it should be logged: - name=template_var.name, - ) - return type(template_dataset)( - data_vars - ) # pytype:disable=not-callable,wrong-arg-count - - -def get_graph_spatial_features( - *, - node_lat: np.ndarray, - node_lon: np.ndarray, - senders: np.ndarray, - receivers: np.ndarray, - add_node_positions: bool, - add_node_latitude: bool, - add_node_longitude: bool, - add_relative_positions: bool, - relative_longitude_local_coordinates: bool, - relative_latitude_local_coordinates: bool, - sine_cosine_encoding: bool = False, - encoding_num_freqs: int = 10, - encoding_multiplicative_factor: float = 1.2, -) -> Tuple[np.ndarray, np.ndarray]: - """Computes spatial features for the nodes. - - Args: - node_lat: Latitudes in the [-90, 90] interval of shape [num_nodes] - node_lon: Longitudes in the [0, 360] interval of shape [num_nodes] - senders: Sender indices of shape [num_edges] - receivers: Receiver indices of shape [num_edges] - add_node_positions: Add unit norm absolute positions. - add_node_latitude: Add a feature for latitude (cos(90 - lat)) - Note even if this is set to False, the model may be able to infer the longitude from relative features, unless `relative_latitude_local_coordinates` is also True, or if there is any bias on the relative edge sizes for different longitudes. - add_node_longitude: Add features for longitude (cos(lon), sin(lon)). - Note even if this is set to False, the model may be able to infer the longitude from relative features, unless `relative_longitude_local_coordinates` is also True, or if there is any bias on the relative edge sizes for different longitudes. - add_relative_positions: Whether to relative positions in R3 to the edges. - relative_longitude_local_coordinates: If True, relative positions are computed in a local space where the receiver is at 0 longitude. - relative_latitude_local_coordinates: If True, relative positions are computed in a local space where the receiver is at 0 latitude. - sine_cosine_encoding: If True, we will transform the node/edge features with sine and cosine functions, similar to NERF. - encoding_num_freqs: frequency parameter - encoding_multiplicative_factor: used for calculating the frequency. - - Returns: - Arrays of shape: [num_nodes, num_features] and [num_edges, num_features]. - with node and edge features. - """ - - num_nodes = node_lat.shape[0] - num_edges = senders.shape[0] - dtype = node_lat.dtype - node_phi, node_theta = lat_lon_deg_to_spherical(node_lat, node_lon) - - # Computing some node features. - node_features = [] - if add_node_positions: - # Already in [-1, 1.] range. - node_features.extend(spherical_to_cartesian(node_phi, node_theta)) - - if add_node_latitude: - # Using the cos of theta. - # From 1. (north pole) to -1 (south pole). - node_features.append(np.cos(node_theta)) - - if add_node_longitude: - # Using the cos and sin, which is already normalized. - node_features.append(np.cos(node_phi)) - node_features.append(np.sin(node_phi)) - - if not node_features: - node_features = np.zeros([num_nodes, 0], dtype=dtype) - else: - node_features = np.stack(node_features, axis=-1) - - # Computing some edge features. - edge_features = [] - - if add_relative_positions: - - relative_position = get_relative_position_in_receiver_local_coordinates( - node_phi=node_phi, - node_theta=node_theta, - senders=senders, - receivers=receivers, - latitude_local_coordinates=relative_latitude_local_coordinates, - longitude_local_coordinates=relative_longitude_local_coordinates, - ) - - # Note this is L2 distance in 3d space, rather than geodesic distance. - relative_edge_distances = np.linalg.norm( - relative_position, axis=-1, keepdims=True - ) - - # Normalize to the maximum edge distance. Note that we expect to always - # have an edge that goes in the opposite direction of any given edge - # so the distribution of relative positions should be symmetric around - # zero. So by scaling by the maximum length, we expect all relative - # positions to fall in the [-1., 1.] interval, and all relative distances - # to fall in the [0., 1.] interval. - max_edge_distance = relative_edge_distances.max() - edge_features.append(relative_edge_distances / max_edge_distance) - edge_features.append(relative_position / max_edge_distance) - - if not edge_features: - edge_features = np.zeros([num_edges, 0], dtype=dtype) - else: - edge_features = np.concatenate(edge_features, axis=-1) - - if sine_cosine_encoding: - - def sine_cosine_transform(x: np.ndarray) -> np.ndarray: - freqs = encoding_multiplicative_factor ** np.arange(encoding_num_freqs) - phases = freqs * x[..., None] - x_sin = np.sin(phases) - x_cos = np.cos(phases) - x_cat = np.concatenate([x_sin, x_cos], axis=-1) - return x_cat.reshape([x.shape[0], -1]) - - node_features = sine_cosine_transform(node_features) - edge_features = sine_cosine_transform(edge_features) - - return node_features, edge_features - - -def lat_lon_to_leading_axes(grid_xarray: "xarray.DataArray") -> "xarray.DataArray": - """Reorders xarray so lat/lon axes come first.""" - # leading + ["lat", "lon"] + trailing - # to - # ["lat", "lon"] + leading + trailing - return grid_xarray.transpose("lat", "lon", ...) - - -def restore_leading_axes(grid_xarray: "xarray.DataArray") -> "xarray.DataArray": - """Reorders xarray so batch/time/level axes come first (if present).""" - - # ["lat", "lon"] + [(batch,) (time,) (level,)] + trailing - # to - # [(batch,) (time,) (level,)] + ["lat", "lon"] + trailing - - input_dims = list(grid_xarray.dims) - output_dims = list(input_dims) - for leading_key in ["level", "time", "batch"]: # reverse order for insert - if leading_key in input_dims: - output_dims.remove(leading_key) - output_dims.insert(0, leading_key) - return grid_xarray.transpose(*output_dims) - - -def lat_lon_deg_to_spherical( - node_lat: np.ndarray, - node_lon: np.ndarray, -) -> Tuple[np.ndarray, np.ndarray]: - phi = np.deg2rad(node_lon) - theta = np.deg2rad(90 - node_lat) - return phi, theta - - -def spherical_to_lat_lon( - phi: np.ndarray, - theta: np.ndarray, -) -> Tuple[np.ndarray, np.ndarray]: - lon = np.mod(np.rad2deg(phi), 360) - lat = 90 - np.rad2deg(theta) - return lat, lon - - -def cartesian_to_spherical( - x: np.ndarray, - y: np.ndarray, - z: np.ndarray, -) -> Tuple[np.ndarray, np.ndarray]: - phi = np.arctan2(y, x) - with np.errstate(invalid="ignore"): # circumventing b/253179568 - theta = np.arccos(z) # Assuming unit radius. - return phi, theta - - -def spherical_to_cartesian( - phi: np.ndarray, theta: np.ndarray -) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: - # Assuming unit radius. - return (np.cos(phi) * np.sin(theta), np.sin(phi) * np.sin(theta), np.cos(theta)) - - -def get_relative_position_in_receiver_local_coordinates( - node_phi: np.ndarray, - node_theta: np.ndarray, - senders: np.ndarray, - receivers: np.ndarray, - latitude_local_coordinates: bool, - longitude_local_coordinates: bool, -) -> np.ndarray: - """Returns relative position features for the edges. - - The relative positions will be computed in a rotated space for a local - coordinate system as defined by the receiver. The relative positions are - simply obtained by subtracting sender position minues receiver position in - that local coordinate system after the rotation in R^3. - - Args: - node_phi: [num_nodes] with polar angles. - node_theta: [num_nodes] with azimuthal angles. - senders: [num_edges] with indices. - receivers: [num_edges] with indices. - latitude_local_coordinates: Whether to rotate edges such that in the positions are computed such that the receiver is always at latitude 0. - longitude_local_coordinates: Whether to rotate edges such that in the positions are computed such that the receiver is always at longitude 0. - - Returns: - Array of relative positions in R3 [num_edges, 3] - """ - - node_pos = np.stack(spherical_to_cartesian(node_phi, node_theta), axis=-1) - - # No rotation in this case. - if not (latitude_local_coordinates or longitude_local_coordinates): - return node_pos[senders] - node_pos[receivers] - - # Get rotation matrices for the local space space for every node. - rotation_matrices = get_rotation_matrices_to_local_coordinates( - reference_phi=node_phi, - reference_theta=node_theta, - rotate_latitude=latitude_local_coordinates, - rotate_longitude=longitude_local_coordinates, - ) - - # Each edge will be rotated according to the rotation matrix of its receiver - # node. - edge_rotation_matrices = rotation_matrices[receivers] - - # Rotate all nodes to the rotated space of the corresponding edge. - # Note for receivers we can also do the matmul first and the gather second: - # ``` - # receiver_pos_in_rotated_space = rotate_with_matrices( - # rotation_matrices, node_pos)[receivers] - # ``` - # which is more efficient, however, we do gather first to keep it more - # symmetric with the sender computation. - receiver_pos_in_rotated_space = rotate_with_matrices( - edge_rotation_matrices, node_pos[receivers] - ) - sender_pos_in_in_rotated_space = rotate_with_matrices( - edge_rotation_matrices, node_pos[senders] - ) - # Note, here, that because the rotated space is chosen according to the - # receiver, if: - # * latitude_local_coordinates = True: latitude for the receivers will be - # 0, that is the z coordinate will always be 0. - # * longitude_local_coordinates = True: longitude for the receivers will be - # 0, that is the y coordinate will be 0. - - # Now we can just subtract. - # Note we are rotating to a local coordinate system, where the y-z axes are - # parallel to a tangent plane to the sphere, but still remain in a 3d space. - # Note that if both `latitude_local_coordinates` and - # `longitude_local_coordinates` are True, and edges are short, - # then the difference in x coordinate between sender and receiver - # should be small, so we could consider dropping the new x coordinate if - # we wanted to the tangent plane, however in doing so - # we would lose information about the curvature of the mesh, which may be - # important for very coarse meshes. - return sender_pos_in_in_rotated_space - receiver_pos_in_rotated_space - - -def get_rotation_matrices_to_local_coordinates( - reference_phi: np.ndarray, - reference_theta: np.ndarray, - rotate_latitude: bool, - rotate_longitude: bool, -) -> np.ndarray: - """Returns a rotation matrix to rotate to a point based on a reference vector. - - The rotation matrix is build such that, a vector in the - same coordinate system at the reference point that points towards the pole - before the rotation, continues to point towards the pole after the rotation. - - Args: - reference_phi: [leading_axis] Polar angles of the reference. - reference_theta: [leading_axis] Azimuthal angles of the reference. - rotate_latitude: Whether to produce a rotation matrix that would rotate R^3 vectors to zero latitude. - rotate_longitude: Whether to produce a rotation matrix that would rotate R^3 vectors to zero longitude. - - Returns: - Matrices of shape [leading_axis] such that when applied to the reference - position with `rotate_with_matrices(rotation_matrices, reference_pos)` - - * phi goes to 0. if "rotate_longitude" is True. - - * theta goes to np.pi / 2 if "rotate_latitude" is True. - - The rotation consists of: - * rotate_latitude = False, rotate_longitude = True: - Latitude preserving rotation. - * rotate_latitude = True, rotate_longitude = True: - Latitude preserving rotation, followed by longitude preserving rotation. - * rotate_latitude = True, rotate_longitude = False: - Latitude preserving rotation, followed by longitude preserving rotation, and the inverse of the latitude preserving rotation. Note this is computationally different from rotating the longitude only and is. We do it like this, so the polar geodesic curve, continues to be aligned with one of the axis after the rotation. - """ - - if rotate_longitude and rotate_latitude: - - # We first rotate around the z axis "minus the azimuthal angle", to get the - # point with zero longitude - azimuthal_rotation = -reference_phi - - # One then we will do a polar rotation (which can be done along the y - # axis now that we are at longitude 0.), "minus the polar angle plus 2pi" - # to get the point with zero latitude. - polar_rotation = -reference_theta + np.pi / 2 - - return scipy.spatial.transform.Rotation.from_euler( - "zy", np.stack([azimuthal_rotation, polar_rotation], axis=1) - ).as_matrix() - elif rotate_longitude: - # Just like the previous case, but applying only the azimuthal rotation. - azimuthal_rotation = -reference_phi - return scipy.spatial.transform.Rotation.from_euler( - "z", -reference_phi - ).as_matrix() - elif rotate_latitude: - # Just like the first case, but after doing the polar rotation, undoing - # the azimuthal rotation. - azimuthal_rotation = -reference_phi - polar_rotation = -reference_theta + np.pi / 2 - - return scipy.spatial.transform.Rotation.from_euler( - "zyz", - np.stack([azimuthal_rotation, polar_rotation, -azimuthal_rotation], axis=1), - ).as_matrix() - else: - raise ValueError("At least one of longitude and latitude should be rotated.") - - -def rotate_with_matrices( - rotation_matrices: np.ndarray, positions: np.ndarray -) -> np.ndarray: - return np.einsum("bji,bi->bj", rotation_matrices, positions) - - -def get_bipartite_graph_spatial_features( - *, - senders_node_lat: np.ndarray, - senders_node_lon: np.ndarray, - senders: np.ndarray, - receivers_node_lat: np.ndarray, - receivers_node_lon: np.ndarray, - receivers: np.ndarray, - add_node_positions: bool, - add_node_latitude: bool, - add_node_longitude: bool, - add_relative_positions: bool, - edge_normalization_factor: Optional[float] = None, - relative_longitude_local_coordinates: bool, - relative_latitude_local_coordinates: bool, -) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: - """Computes spatial features for the nodes. - - This function is almost identical to `get_graph_spatial_features`. The only - difference is that sender nodes and receiver nodes can be in different arrays. - This is necessary to enable combination with typed Graph. - - Args: - senders_node_lat: Latitudes in the [-90, 90] interval of shape [num_sender_nodes] - senders_node_lon: Longitudes in the [0, 360] interval of shape [num_sender_nodes] - senders: Sender indices of shape [num_edges], indices in [0, num_sender_nodes) - receivers_node_lat: Latitudes in the [-90, 90] interval of shape [num_receiver_nodes] - receivers_node_lon: Longitudes in the [0, 360] interval of shape [num_receiver_nodes] - receivers: Receiver indices of shape [num_edges], indices in [0, num_receiver_nodes) - add_node_positions: Add unit norm absolute positions. - add_node_latitude: Add a feature for latitude (cos(90 - lat)). - Note even ifthis is set to False, the model may be able to infer the longitude from relative features, unless `relative_latitude_local_coordinates` is also True, or if there is any bias on the relative edge sizes for different longitudes. - add_node_longitude: Add features for longitude (cos(lon), sin(lon)). - Note even if this is set to False, the model may be able to infer the longitude from relative features, unless `relative_longitude_local_coordinates` is also True, or if there is any bias on the relative edge sizes for different longitudes. - add_relative_positions: Whether to relative positions in R3 to the edges. - edge_normalization_factor: Allows explicitly controlling edge normalization. If None, defaults to max edge length. This supports using pre-trained model weights with a different graph structure to what it was trained on. - relative_longitude_local_coordinates: If True, relative positions are computed in a local space where the receiver is at 0 longitude. - relative_latitude_local_coordinates: If True, relative positions are computed in a local space where the receiver is at 0 latitude. - - Returns: - Arrays of shape: [num_nodes, num_features] and [num_edges, num_features]. with node and edge features. - """ - - num_senders = senders_node_lat.shape[0] - num_receivers = receivers_node_lat.shape[0] - num_edges = senders.shape[0] - dtype = senders_node_lat.dtype - assert receivers_node_lat.dtype == dtype - senders_node_phi, senders_node_theta = lat_lon_deg_to_spherical( - senders_node_lat, senders_node_lon - ) - receivers_node_phi, receivers_node_theta = lat_lon_deg_to_spherical( - receivers_node_lat, receivers_node_lon - ) - - # Computing some node features. - senders_node_features = [] - receivers_node_features = [] - if add_node_positions: - # Already in [-1, 1.] range. - senders_node_features.extend( - spherical_to_cartesian(senders_node_phi, senders_node_theta) - ) - receivers_node_features.extend( - spherical_to_cartesian(receivers_node_phi, receivers_node_theta) - ) - - if add_node_latitude: - # Using the cos of theta. - # From 1. (north pole) to -1 (south pole). - senders_node_features.append(np.cos(senders_node_theta)) - receivers_node_features.append(np.cos(receivers_node_theta)) - - if add_node_longitude: - # Using the cos and sin, which is already normalized. - senders_node_features.append(np.cos(senders_node_phi)) - senders_node_features.append(np.sin(senders_node_phi)) - - receivers_node_features.append(np.cos(receivers_node_phi)) - receivers_node_features.append(np.sin(receivers_node_phi)) - - if not senders_node_features: - senders_node_features = np.zeros([num_senders, 0], dtype=dtype) - receivers_node_features = np.zeros([num_receivers, 0], dtype=dtype) - else: - senders_node_features = np.stack(senders_node_features, axis=-1) - receivers_node_features = np.stack(receivers_node_features, axis=-1) - - # Computing some edge features. - edge_features = [] - - if add_relative_positions: - - relative_position = ( - get_bipartite_relative_position_in_receiver_local_coordinates( - senders_node_phi=senders_node_phi, - senders_node_theta=senders_node_theta, - receivers_node_phi=receivers_node_phi, - receivers_node_theta=receivers_node_theta, - senders=senders, - receivers=receivers, - latitude_local_coordinates=relative_latitude_local_coordinates, - longitude_local_coordinates=relative_longitude_local_coordinates, - ) - ) - - # Note this is L2 distance in 3d space, rather than geodesic distance. - relative_edge_distances = np.linalg.norm( - relative_position, axis=-1, keepdims=True - ) - - if edge_normalization_factor is None: - # Normalize to the maximum edge distance. Note that we expect to always - # have an edge that goes in the opposite direction of any given edge - # so the distribution of relative positions should be symmetric around - # zero. So by scaling by the maximum length, we expect all relative - # positions to fall in the [-1., 1.] interval, and all relative distances - # to fall in the [0., 1.] interval. - edge_normalization_factor = relative_edge_distances.max() - - edge_features.append(relative_edge_distances / edge_normalization_factor) - edge_features.append(relative_position / edge_normalization_factor) - - if not edge_features: - edge_features = np.zeros([num_edges, 0], dtype=dtype) - else: - edge_features = np.concatenate(edge_features, axis=-1) - - return senders_node_features, receivers_node_features, edge_features - - -def get_bipartite_relative_position_in_receiver_local_coordinates( - senders_node_phi: np.ndarray, - senders_node_theta: np.ndarray, - senders: np.ndarray, - receivers_node_phi: np.ndarray, - receivers_node_theta: np.ndarray, - receivers: np.ndarray, - latitude_local_coordinates: bool, - longitude_local_coordinates: bool, -) -> np.ndarray: - """Returns relative position features for the edges. - - This function is equivalent to - `get_relative_position_in_receiver_local_coordinates`, but adapted to work - with bipartite typed graphs. - - The relative positions will be computed in a rotated space for a local - coordinate system as defined by the receiver. The relative positions are - simply obtained by subtracting sender position minues receiver position in - that local coordinate system after the rotation in R^3. - - Args: - senders_node_phi: [num_sender_nodes] with polar angles. - senders_node_theta: [num_sender_nodes] with azimuthal angles. - senders: [num_edges] with indices into sender nodes. - receivers_node_phi: [num_sender_nodes] with polar angles. - receivers_node_theta: [num_sender_nodes] with azimuthal angles. - receivers: [num_edges] with indices into receiver nodes. - latitude_local_coordinates: Whether to rotate edges such that in the positions are computed such that the receiver is always at latitude 0. - longitude_local_coordinates: Whether to rotate edges such that in the positions are computed such that the receiver is always at longitude 0. - - Returns: - Array of relative positions in R3 [num_edges, 3] - """ - - senders_node_pos = np.stack( - spherical_to_cartesian(senders_node_phi, senders_node_theta), axis=-1 - ) - - receivers_node_pos = np.stack( - spherical_to_cartesian(receivers_node_phi, receivers_node_theta), axis=-1 - ) - - # No rotation in this case. - if not (latitude_local_coordinates or longitude_local_coordinates): - return senders_node_pos[senders] - receivers_node_pos[receivers] - - # Get rotation matrices for the local space space for every receiver node. - receiver_rotation_matrices = get_rotation_matrices_to_local_coordinates( - reference_phi=receivers_node_phi, - reference_theta=receivers_node_theta, - rotate_latitude=latitude_local_coordinates, - rotate_longitude=longitude_local_coordinates, - ) - - # Each edge will be rotated according to the rotation matrix of its receiver - # node. - edge_rotation_matrices = receiver_rotation_matrices[receivers] - - # Rotate all nodes to the rotated space of the corresponding edge. - # Note for receivers we can also do the matmul first and the gather second: - # ``` - # receiver_pos_in_rotated_space = rotate_with_matrices( - # rotation_matrices, node_pos)[receivers] - # ``` - # which is more efficient, however, we do gather first to keep it more - # symmetric with the sender computation. - receiver_pos_in_rotated_space = rotate_with_matrices( - edge_rotation_matrices, receivers_node_pos[receivers] - ) - sender_pos_in_in_rotated_space = rotate_with_matrices( - edge_rotation_matrices, senders_node_pos[senders] - ) - # Note, here, that because the rotated space is chosen according to the - # receiver, if: - # * latitude_local_coordinates = True: latitude for the receivers will be - # 0, that is the z coordinate will always be 0. - # * longitude_local_coordinates = True: longitude for the receivers will be - # 0, that is the y coordinate will be 0. - - # Now we can just subtract. - # Note we are rotating to a local coordinate system, where the y-z axes are - # parallel to a tangent plane to the sphere, but still remain in a 3d space. - # Note that if both `latitude_local_coordinates` and - # `longitude_local_coordinates` are True, and edges are short, - # then the difference in x coordinate between sender and receiver - # should be small, so we could consider dropping the new x coordinate if - # we wanted to the tangent plane, however in doing so - # we would lose information about the curvature of the mesh, which may be - # important for very coarse meshes. - return sender_pos_in_in_rotated_space - receiver_pos_in_rotated_space - - -class GraphGridMesh: - """Graph datatype of GraphCast. - - Args: - mesh_size (int): size of mesh. - radius_query_fraction_edge_length (float): _description_ - mesh2grid_edge_normalization_factor (float): Normalization factor of edge in Mesh2Grid GNN. - resolution (float): resolution of atmospheric data. - mesh2mesh_src_index (np.array, optional): Index of Mesh2Mesh source node. Defaults to None. - mesh2mesh_dst_index (np.array, optional): Index of Mesh2Mesh destination node. Defaults to None. - grid2mesh_src_index (np.array, optional): Index of Grid2Mesh source node. Defaults to None. - grid2mesh_dst_index (np.array, optional): Index of Grid2Mesh destination node. - mesh2grid_src_index (np.array, optional): Index of Mesh2Grid source node. Defaults to None. - mesh2grid_dst_index (np.array, optional): Index of Mesh2Grid destination node. Defaults to None. - mesh_num_nodes (int, optional): Number of mesh nodes. Defaults to None. - grid_num_nodes (int, optional): Number of grid nodes. Defaults to None. - mesh_num_edges (int, optional): Number of mesh edges. Defaults to None. - grid2mesh_num_edges (int, optional): Number of edges in Grid2Mesh GNN. Defaults to None. - mesh2grid_num_edges (int, optional): Number of edges in Mesh2Grid GNN. Defaults to None. - grid_node_feat (np.array, optional): Feature of grid nodes. Defaults to None. - mesh_node_feat (np.array, optional): Feature of mehs nodes. Defaults to None. - mesh_edge_feat (np.array, optional): Feature of mesh edges. Defaults to None. - grid2mesh_edge_feat (np.array, optional): Feature of edges in Grid2Mesh GNN. Defaults to None. - mesh2grid_edge_feat (np.array, optional): Feature of edges in Mesh2Grid GNN. Defaults to None. - """ - - def __init__( - self, - mesh_size: int, - radius_query_fraction_edge_length: float, - mesh2grid_edge_normalization_factor: float, - resolution: float, - mesh2mesh_src_index: np.array = None, - mesh2mesh_dst_index: np.array = None, - grid2mesh_src_index: np.array = None, - grid2mesh_dst_index: np.array = None, - mesh2grid_src_index: np.array = None, - mesh2grid_dst_index: np.array = None, - mesh_num_nodes: int = None, - grid_num_nodes: int = None, - mesh_num_edges: int = None, - grid2mesh_num_edges: np.array = None, - mesh2grid_num_edges: np.array = None, - grid_node_feat: np.array = None, - mesh_node_feat: np.array = None, - mesh_edge_feat: np.array = None, - grid2mesh_edge_feat: np.array = None, - mesh2grid_edge_feat: np.array = None, - ): - self.meshes = get_hierarchy_of_triangular_meshes_for_sphere(mesh_size) - - all_input_vars = [ - mesh2mesh_src_index, - mesh2mesh_dst_index, - grid2mesh_src_index, - grid2mesh_dst_index, - mesh2grid_src_index, - mesh2grid_dst_index, - mesh_num_nodes, - grid_num_nodes, - mesh_num_edges, - grid2mesh_num_edges, - mesh2grid_num_edges, - grid_node_feat, - mesh_node_feat, - mesh_edge_feat, - grid2mesh_edge_feat, - mesh2grid_edge_feat, - ] - should_init = any(var is None for var in all_input_vars) - - if should_init: - self.query_radius = ( - self._get_max_edge_distance(self.finest_mesh) - * radius_query_fraction_edge_length - ) - self._mesh2grid_edge_normalization_factor = ( - mesh2grid_edge_normalization_factor - ) - self._spatial_features_kwargs = dict( - add_node_positions=False, - add_node_latitude=True, - add_node_longitude=True, - add_relative_positions=True, - relative_longitude_local_coordinates=True, - relative_latitude_local_coordinates=True, - ) - - self.init_mesh_properties() - self._init_grid_properties( - grid_lat=np.arange(-90.0, 90.0 + resolution, resolution), - grid_lon=np.arange(0.0, 360.0, resolution), - ) - self._grid2mesh_graph_structure = self._init_grid2mesh_graph() - self._mesh_graph_structure = self._init_mesh_graph() - self._mesh2grid_graph_structure = self._init_mesh2grid_graph() - else: - self.mesh2mesh_src_index = mesh2mesh_src_index - self.mesh2mesh_dst_index = mesh2mesh_dst_index - self.grid2mesh_src_index = grid2mesh_src_index - self.grid2mesh_dst_index = grid2mesh_dst_index - self.mesh2grid_src_index = mesh2grid_src_index - self.mesh2grid_dst_index = mesh2grid_dst_index - - self.mesh_num_nodes = mesh_num_nodes - self.grid_num_nodes = grid_num_nodes - - self.mesh_num_edges = mesh_num_edges - self.grid2mesh_num_edges = grid2mesh_num_edges - self.mesh2grid_num_edges = mesh2grid_num_edges - - self.grid_node_feat = grid_node_feat - self.mesh_node_feat = mesh_node_feat - self.mesh_edge_feat = mesh_edge_feat - self.grid2mesh_edge_feat = grid2mesh_edge_feat - self.mesh2grid_edge_feat = mesh2grid_edge_feat - - def update(self, name, value): - if hasattr(self, name): - setattr(self, name, value) - else: - raise ValueError - - def tensor(self): - self.mesh2mesh_src_index = paddle.to_tensor( - self.mesh2mesh_src_index, dtype=paddle.int64 - ) - - self.mesh2mesh_dst_index = paddle.to_tensor( - self.mesh2mesh_dst_index, dtype=paddle.int64 - ) - self.grid2mesh_src_index = paddle.to_tensor( - self.grid2mesh_src_index, dtype=paddle.int64 - ) - self.grid2mesh_dst_index = paddle.to_tensor( - self.grid2mesh_dst_index, dtype=paddle.int64 - ) - self.mesh2grid_src_index = paddle.to_tensor( - self.mesh2grid_src_index, dtype=paddle.int64 - ) - self.mesh2grid_dst_index = paddle.to_tensor( - self.mesh2grid_dst_index, dtype=paddle.int64 - ) - self.grid_node_feat = paddle.to_tensor( - self.grid_node_feat, dtype=paddle.get_default_dtype() - ) - self.mesh_node_feat = paddle.to_tensor( - self.mesh_node_feat, dtype=paddle.get_default_dtype() - ) - self.mesh_edge_feat = paddle.to_tensor( - self.mesh_edge_feat, dtype=paddle.get_default_dtype() - ) - self.grid2mesh_edge_feat = paddle.to_tensor( - self.grid2mesh_edge_feat, dtype=paddle.get_default_dtype() - ) - self.mesh2grid_edge_feat = paddle.to_tensor( - self.mesh2grid_edge_feat, dtype=paddle.get_default_dtype() - ) - return self - - @property - def finest_mesh(self): - return self.meshes[-1] - - def init_mesh_properties(self): - """Inits static properties that have to do with mesh nodes.""" - self.mesh_num_nodes = self.finest_mesh.vertices.shape[0] - mesh_phi, mesh_theta = cartesian_to_spherical( - self.finest_mesh.vertices[:, 0], - self.finest_mesh.vertices[:, 1], - self.finest_mesh.vertices[:, 2], - ) - (mesh_nodes_lat, mesh_nodes_lon) = spherical_to_lat_lon( - phi=mesh_phi, - theta=mesh_theta, - ) - # Convert to f32 to ensure the lat/lon features aren't in f64. - self._mesh_nodes_lat = mesh_nodes_lat.astype(np.float32) - self._mesh_nodes_lon = mesh_nodes_lon.astype(np.float32) - - def _init_grid_properties(self, grid_lat: np.ndarray, grid_lon: np.ndarray): - """Inits static properties that have to do with grid nodes.""" - self._grid_lat = grid_lat.astype(np.float32) - self._grid_lon = grid_lon.astype(np.float32) - # Initialized the counters. - self.grid_num_nodes = grid_lat.shape[0] * grid_lon.shape[0] - - # Initialize lat and lon for the grid. - grid_nodes_lon, grid_nodes_lat = np.meshgrid(grid_lon, grid_lat) - self._grid_nodes_lon = grid_nodes_lon.reshape([-1]).astype(np.float32) - self._grid_nodes_lat = grid_nodes_lat.reshape([-1]).astype(np.float32) - - def _init_grid2mesh_graph(self): - """Build Grid2Mesh graph.""" - - # Create some edges according to distance between mesh and grid nodes. - assert self._grid_lat is not None and self._grid_lon is not None - (grid_indices, mesh_indices) = radius_query_indices( - grid_latitude=self._grid_lat, - grid_longitude=self._grid_lon, - mesh=self.finest_mesh, - radius=self.query_radius, - ) - - # Edges sending info from grid to mesh. - senders = grid_indices - receivers = mesh_indices - - # Precompute structural node and edge features according to config options. - # Structural features are those that depend on the fixed values of the - # latitude and longitudes of the nodes. - ( - senders_node_features, - _, - edge_features, - ) = get_bipartite_graph_spatial_features( - senders_node_lat=self._grid_nodes_lat, - senders_node_lon=self._grid_nodes_lon, - receivers_node_lat=self._mesh_nodes_lat, - receivers_node_lon=self._mesh_nodes_lon, - senders=senders, - receivers=receivers, - edge_normalization_factor=None, - **self._spatial_features_kwargs, - ) - - self.grid_node_feat = np.expand_dims(senders_node_features, axis=1) - - self.grid2mesh_src_index = senders - self.grid2mesh_dst_index = receivers - self.grid2mesh_edge_feat = np.expand_dims(edge_features, axis=1) - self.grid2mesh_num_edges = len(edge_features) - - def _init_mesh_graph(self): - """Build Mesh graph.""" - merged_mesh = merge_meshes(self.meshes) - # Work simply on the mesh edges. - senders, receivers = faces_to_edges(merged_mesh.faces) - # Precompute structural node and edge features according to config options. - # Structural features are those that depend on the fixed values of the - # latitude and longitudes of the nodes. - assert self._mesh_nodes_lat is not None and self._mesh_nodes_lon is not None - node_features, edge_features = get_graph_spatial_features( - node_lat=self._mesh_nodes_lat, - node_lon=self._mesh_nodes_lon, - senders=senders, - receivers=receivers, - **self._spatial_features_kwargs, - ) - - self.mesh_node_feat = np.expand_dims(node_features, axis=1) - self.mesh2mesh_src_index = senders - self.mesh2mesh_dst_index = receivers - self.mesh_edge_feat = np.expand_dims(edge_features, axis=1) - self.mesh_num_edges = len(edge_features) - - def _init_mesh2grid_graph(self): - """Build Mesh2Grid graph.""" - - # Create some edges according to how the grid nodes are contained by - # mesh triangles. - (grid_indices, mesh_indices) = in_mesh_triangle_indices( - grid_latitude=self._grid_lat, - grid_longitude=self._grid_lon, - mesh=self.finest_mesh, - ) - - # Edges sending info from mesh to grid. - senders = mesh_indices - receivers = grid_indices - - # Precompute structural node and edge features according to config options. - assert self._mesh_nodes_lat is not None and self._mesh_nodes_lon is not None - (_, _, edge_features) = get_bipartite_graph_spatial_features( - senders_node_lat=self._mesh_nodes_lat, - senders_node_lon=self._mesh_nodes_lon, - receivers_node_lat=self._grid_nodes_lat, - receivers_node_lon=self._grid_nodes_lon, - senders=senders, - receivers=receivers, - edge_normalization_factor=self._mesh2grid_edge_normalization_factor, - **self._spatial_features_kwargs, - ) - - self.mesh2grid_src_index = senders - self.mesh2grid_dst_index = receivers - self.mesh2grid_edge_feat = np.expand_dims(edge_features, axis=1) - self.mesh2grid_num_edges = len(edge_features) - - @staticmethod - def _get_max_edge_distance(mesh): - senders, receivers = faces_to_edges(mesh.faces) - edge_distances = np.linalg.norm( - mesh.vertices[senders] - mesh.vertices[receivers], axis=-1 - ) - return edge_distances.max() - - def grid_node_outputs_to_prediction( - self, - grid_node_outputs: np.ndarray, - targets_template: "xarray.Dataset", - ) -> "xarray.Dataset": - """[num_grid_nodes, batch, num_outputs] -> xarray.""" - # numpy array with shape [lat_lon_node, batch, channels] - # to xarray `DataArray` (batch, lat, lon, channels) - assert self._grid_lat is not None and self._grid_lon is not None - grid_shape = (self._grid_lat.shape[0], self._grid_lon.shape[0]) - grid_outputs_lat_lon_leading = grid_node_outputs.reshape( - grid_shape + grid_node_outputs.shape[1:] - ) - dims = ("lat", "lon", "batch", "channels") - grid_xarray_lat_lon_leading = xarray.DataArray( - data=grid_outputs_lat_lon_leading, dims=dims - ) - grid_xarray = restore_leading_axes(grid_xarray_lat_lon_leading) - - # xarray `DataArray` (batch, lat, lon, channels) - # to xarray `Dataset` (batch, one time step, lat, lon, level, multiple vars) - return stacked_to_dataset(grid_xarray.variable, targets_template) - - -class TriangularMesh(NamedTuple): - vertices: np.ndarray - faces: np.ndarray - - -def merge_meshes(mesh_list: Sequence[TriangularMesh]) -> TriangularMesh: - for i in range(len(mesh_list) - 1): - mesh_i, mesh_ip1 = mesh_list[i], mesh_list[i + 1] - num_nodes_mesh_i = mesh_i.vertices.shape[0] - assert np.allclose(mesh_i.vertices, mesh_ip1.vertices[:num_nodes_mesh_i]) - - return TriangularMesh( - vertices=mesh_list[-1].vertices, - faces=np.concatenate([mesh.faces for mesh in mesh_list], axis=0), - ) - - -def get_icosahedron(): - phi = (1 + np.sqrt(5)) / 2 - product = [[1.0, phi], [1.0, -phi], [-1.0, phi], [-1.0, -phi]] - vertices = [] - for p in product: - c1 = p[0] - c2 = p[1] - vertices.append((c1, c2, 0.0)) - vertices.append((0.0, c1, c2)) - vertices.append((c2, 0.0, c1)) - - vertices = np.array(vertices, dtype=np.float32) - vertices /= np.linalg.norm([1.0, phi]) - - faces = [ - (0, 1, 2), - (0, 6, 1), - (8, 0, 2), - (8, 4, 0), - (3, 8, 2), - (3, 2, 7), - (7, 2, 1), - (0, 4, 6), - (4, 11, 6), - (6, 11, 5), - (1, 5, 7), - (4, 10, 11), - (4, 8, 10), - (10, 8, 3), - (10, 3, 9), - (11, 10, 9), - (11, 9, 5), - (5, 9, 7), - (9, 3, 7), - (1, 6, 5), - ] - - angle_between_faces = 2 * np.arcsin(phi / np.sqrt(3)) - rotation_angle = (np.pi - angle_between_faces) / 2 - rotation = scipy.spatial.transform.Rotation.from_euler( - seq="y", angles=rotation_angle - ) - rotation_matrix = rotation.as_matrix() - vertices = np.dot(vertices, rotation_matrix) - - return TriangularMesh( - vertices=vertices.astype(np.float32), faces=np.array(faces, dtype=np.int32) - ) - - -def get_hierarchy_of_triangular_meshes_for_sphere( - splits: int, -) -> List[TriangularMesh]: - current_mesh = get_icosahedron() - output_meshes = [current_mesh] - for _ in range(splits): - current_mesh = _two_split_unit_sphere_triangle_faces(current_mesh) - output_meshes.append(current_mesh) - return output_meshes - - -def _two_split_unit_sphere_triangle_faces( - triangular_mesh: TriangularMesh, -) -> TriangularMesh: - """Splits each triangular face into 4 triangles keeping the orientation.""" - new_vertices_builder = _ChildVerticesBuilder(triangular_mesh.vertices) - - new_faces = [] - for ind1, ind2, ind3 in triangular_mesh.faces: - ind12 = new_vertices_builder.get_new_child_vertex_index((ind1, ind2)) - ind23 = new_vertices_builder.get_new_child_vertex_index((ind2, ind3)) - ind31 = new_vertices_builder.get_new_child_vertex_index((ind3, ind1)) - new_faces.extend( - [ - [ind1, ind12, ind31], # 1 - [ind12, ind2, ind23], # 2 - [ind31, ind23, ind3], # 3 - [ind12, ind23, ind31], # 4 - ] - ) - return TriangularMesh( - vertices=new_vertices_builder.get_all_vertices(), - faces=np.array(new_faces, dtype=np.int32), - ) - - -class _ChildVerticesBuilder: - """Bookkeeping of new child vertices added to an existing set of vertices.""" - - def __init__(self, parent_vertices): - self._child_vertices_index_mapping = {} - self._parent_vertices = parent_vertices - # We start with all previous vertices. - self._all_vertices_list = list(parent_vertices) - - def _get_child_vertex_key(self, parent_vertex_indices): - return tuple(sorted(parent_vertex_indices)) - - def _create_child_vertex(self, parent_vertex_indices): - """Creates a new vertex.""" - # Position for new vertex is the middle point, between the parent points, - # projected to unit sphere. - child_vertex_position = self._parent_vertices[list(parent_vertex_indices)].mean( - 0 - ) - child_vertex_position /= np.linalg.norm(child_vertex_position) - - # Add the vertex to the output list. The index for this new vertex will - # match the length of the list before adding it. - child_vertex_key = self._get_child_vertex_key(parent_vertex_indices) - self._child_vertices_index_mapping[child_vertex_key] = len( - self._all_vertices_list - ) - self._all_vertices_list.append(child_vertex_position) - - def get_new_child_vertex_index(self, parent_vertex_indices): - """Returns index for a child vertex, creating it if necessary.""" - # Get the key to see if we already have a new vertex in the middle. - child_vertex_key = self._get_child_vertex_key(parent_vertex_indices) - if child_vertex_key not in self._child_vertices_index_mapping: - self._create_child_vertex(parent_vertex_indices) - return self._child_vertices_index_mapping[child_vertex_key] - - def get_all_vertices(self): - """Returns an array with old vertices.""" - return np.array(self._all_vertices_list) - - -def faces_to_edges(faces: np.ndarray): - """Transforms polygonal faces to sender and receiver indices. - - It does so by transforming every face into N_i edges. Such if the triangular - face has indices [0, 1, 2], three edges are added 0->1, 1->2, and 2->0. - - If all faces have consistent orientation, and the surface represented by the - faces is closed, then every edge in a polygon with a certain orientation - is also part of another polygon with the opposite orientation. In this - situation, the edges returned by the method are always bidirectional. - - Args: - faces: Integer array of shape [num_faces, 3]. Contains node indices adjacent to each face. - Returns: - Tuple with sender/receiver indices, each of shape [num_edges=num_faces*3]. - """ - - assert faces.ndim == 2 - assert faces.shape[-1] == 3 - senders = np.concatenate([faces[:, 0], faces[:, 1], faces[:, 2]]) - receivers = np.concatenate([faces[:, 1], faces[:, 2], faces[:, 0]]) - return senders, receivers - - -def _grid_lat_lon_to_coordinates( - grid_latitude: np.ndarray, grid_longitude: np.ndarray -) -> np.ndarray: - """Lat [num_lat] lon [num_lon] to 3d coordinates [num_lat, num_lon, 3].""" - # Convert to spherical coordinates phi and theta defined in the grid. - # Each [num_latitude_points, num_longitude_points] - phi_grid, theta_grid = np.meshgrid( - np.deg2rad(grid_longitude), np.deg2rad(90 - grid_latitude) - ) - - # [num_latitude_points, num_longitude_points, 3] - # Note this assumes unit radius, since for now we model the earth as a - # sphere of unit radius, and keep any vertical dimension as a regular grid. - return np.stack( - [ - np.cos(phi_grid) * np.sin(theta_grid), - np.sin(phi_grid) * np.sin(theta_grid), - np.cos(theta_grid), - ], - axis=-1, - ) - - -def radius_query_indices( - *, - grid_latitude: np.ndarray, - grid_longitude: np.ndarray, - mesh: TriangularMesh, - radius: float, -) -> Tuple[np.ndarray, np.ndarray]: - """Returns mesh-grid edge indices for radius query. - - Args: - grid_latitude: Latitude values for the grid [num_lat_points] - grid_longitude: Longitude values for the grid [num_lon_points] - mesh: Mesh object. - radius: Radius of connectivity in R3. for a sphere of unit radius. - - Returns: - tuple with `grid_indices` and `mesh_indices` indicating edges between the grid and the mesh such that the distances in a straight line (not geodesic) are smaller than or equal to `radius`. - grid_indices: Indices of shape [num_edges], that index into a - [num_lat_points, num_lon_points] grid, after flattening the leading axes. - mesh_indices: Indices of shape [num_edges], that index into mesh.vertices. - """ - - # [num_grid_points=num_lat_points * num_lon_points, 3] - grid_positions = _grid_lat_lon_to_coordinates( - grid_latitude, grid_longitude - ).reshape([-1, 3]) - - # [num_mesh_points, 3] - mesh_positions = mesh.vertices - kd_tree = scipy.spatial.cKDTree(mesh_positions) - - # [num_grid_points, num_mesh_points_per_grid_point] - # Note `num_mesh_points_per_grid_point` is not constant, so this is a list - # of arrays, rather than a 2d array. - query_indices = kd_tree.query_ball_point(x=grid_positions, r=radius) - - grid_edge_indices = [] - mesh_edge_indices = [] - for grid_index, mesh_neighbors in enumerate(query_indices): - grid_edge_indices.append(np.repeat(grid_index, len(mesh_neighbors))) - mesh_edge_indices.append(mesh_neighbors) - - # [num_edges] - grid_edge_indices = np.concatenate(grid_edge_indices, axis=0).astype(int) - mesh_edge_indices = np.concatenate(mesh_edge_indices, axis=0).astype(int) - - return grid_edge_indices, mesh_edge_indices - - -def in_mesh_triangle_indices( - *, grid_latitude: np.ndarray, grid_longitude: np.ndarray, mesh: TriangularMesh -) -> Tuple[np.ndarray, np.ndarray]: - """Returns mesh-grid edge indices for grid points contained in mesh triangles. - - Args: - grid_latitude: Latitude values for the grid [num_lat_points] - grid_longitude: Longitude values for the grid [num_lon_points] - mesh: Mesh object. - - Returns: - tuple with `grid_indices` and `mesh_indices` indicating edges between the grid and the mesh vertices of the triangle that contain each grid point. The number of edges is always num_lat_points * num_lon_points * 3 - grid_indices: Indices of shape [num_edges], that index into a [num_lat_points, num_lon_points] grid, after flattening the leading axes. - mesh_indices: Indices of shape [num_edges], that index into mesh.vertices. - """ - - # [num_grid_points=num_lat_points * num_lon_points, 3] - grid_positions = _grid_lat_lon_to_coordinates( - grid_latitude, grid_longitude - ).reshape([-1, 3]) - - mesh_trimesh = trimesh.Trimesh(vertices=mesh.vertices, faces=mesh.faces) - - # [num_grid_points] with mesh face indices for each grid point. - _, _, query_face_indices = trimesh.proximity.closest_point( - mesh_trimesh, grid_positions - ) - - # [num_grid_points, 3] with mesh node indices for each grid point. - mesh_edge_indices = mesh.faces[query_face_indices] - - # [num_grid_points, 3] with grid node indices, where every row simply contains - # the row (grid_point) index. - grid_indices = np.arange(grid_positions.shape[0]) - grid_edge_indices = np.tile(grid_indices.reshape([-1, 1]), [1, 3]) - - # Flatten to get a regular list. - # [num_edges=num_grid_points*3] - mesh_edge_indices = mesh_edge_indices.reshape([-1]) - grid_edge_indices = grid_edge_indices.reshape([-1]) - - return grid_edge_indices, mesh_edge_indices - - -def get_year_progress(seconds_since_epoch: np.ndarray) -> np.ndarray: - """Computes year progress for times in seconds. - Args: - seconds_since_epoch: Times in seconds since the "epoch" (the point at which UNIX time starts). - Returns: - Year progress normalized to be in the `[0, 1)` interval for each time point. - """ - # Start with the pure integer division, and then float at the very end. - # We will try to keep as much precision as possible. - years_since_epoch = ( - seconds_since_epoch / SEC_PER_DAY / np.float64(_AVG_DAY_PER_YEAR) - ) - # Note depending on how these ops are down, we may end up with a "weak_type" - # which can cause issues in subtle ways, and hard to track here. - # In any case, casting to float32 should get rid of the weak type. - # [0, 1.) Interval. - return np.mod(years_since_epoch, 1.0).astype(np.float32) - - -def get_day_progress( - seconds_since_epoch: np.ndarray, - longitude: np.ndarray, -) -> np.ndarray: - """Computes day progress for times in seconds at each longitude. - Args: - seconds_since_epoch: 1D array of times in seconds since the 'epoch' (the point at which UNIX time starts). - longitude: 1D array of longitudes at which day progress is computed. - Returns: - 2D array of day progress values normalized to be in the [0, 1) interval for each time point at each longitude. - """ - # [0.0, 1.0) Interval. - day_progress_greenwich = np.mod(seconds_since_epoch, SEC_PER_DAY) / SEC_PER_DAY - # Offset the day progress to the longitude of each point on Earth. - longitude_offsets = np.deg2rad(longitude) / (2 * np.pi) - day_progress = np.mod( - day_progress_greenwich[..., np.newaxis] + longitude_offsets, 1.0 - ) - return day_progress.astype(np.float32) - - -def datetime_features(seconds_since_epoch, longitude_offsets): - year_progress = get_year_progress(seconds_since_epoch) - day_progress = get_day_progress(seconds_since_epoch, longitude_offsets) - year_progress_phase = year_progress * (2 * np.pi) - day_progress_phase = day_progress * (2 * np.pi) - returned_data = { - "year_progress_sin": np.sin(year_progress_phase), - "year_progress_cos": np.cos(year_progress_phase), - "day_progress_sin": np.sin(day_progress_phase), - "day_progress_cos": np.cos(day_progress_phase), - } - return returned_data - - -def add_var_into_nc_dataset( - nc_dataset, - var_name, - var_value, - var_dims=( - "batch", - "time", - ), -): - new_var = nc_dataset.createVariable(var_name, "f8", var_dims) - new_var[:] = var_value - return nc_dataset - - -def extract_input_target_times( - dataset: "xarray.Dataset", - input_duration: str, - target_lead_times: str, -): - (target_lead_times, target_duration) = _process_target_lead_times_and_get_duration( - target_lead_times - ) - time = dataset.coords["time"] - dataset = dataset.assign_coords(time=time + target_duration - time[-1]) - - targets = dataset.sel({"time": target_lead_times}) - - input_duration = pd.Timedelta(input_duration) - zero = pd.Timedelta(0) - epsilon = pd.Timedelta(1, "ns") - inputs = dataset.sel({"time": slice(-input_duration + epsilon, zero)}) - return inputs, targets - - -def _process_target_lead_times_and_get_duration(target_lead_times: str): - """Returns the minimum duration for the target lead times.""" - if isinstance(target_lead_times, slice): - if target_lead_times.start is None: - target_lead_times = slice( - pd.Timedelta(1, "ns"), target_lead_times.stop, target_lead_times.step - ) - target_duration = pd.Timedelta(target_lead_times.stop) - else: - if not isinstance(target_lead_times, (list, tuple, set)): - target_lead_times = [target_lead_times] - - target_lead_times = [pd.Timedelta(x) for x in target_lead_times] - target_lead_times.sort() - target_duration = target_lead_times[-1] - return target_lead_times, target_duration - - -def variable_to_stacked( - variable: "xarray.Variable", - sizes: "xarray.core.utils.Frozen", - preserved_dims=("batch", "lat", "lon"), -) -> "xarray.Variable": - """Converts an xarray.Variable to preserved_dims + ("channels",). - - Any dimensions other than those included in preserved_dims get stacked into a final "channels" dimension. If any of the preserved_dims are missing then they are added, with the data broadcast/tiled to match the sizes specified in `sizes`. - - Args: - variable: An xarray.Variable. - sizes: Mapping including sizes for any dimensions which are not present in `variable` but are needed for the output. This may be needed for example for a static variable with only ("lat", "lon") dims, or if you want to encode just the latitude coordinates (a variable with dims ("lat",)). - preserved_dims: dimensions of variable to not be folded in channels. - - Returns: - An xarray.Variable with dimensions preserved_dims + ("channels",). - """ - stack_to_channels_dims = [d for d in variable.dims if d not in preserved_dims] - if stack_to_channels_dims: - variable = variable.stack(channels=stack_to_channels_dims) - dims = {dim: variable.sizes.get(dim) or sizes[dim] for dim in preserved_dims} - dims["channels"] = variable.sizes.get("channels", 1) - return variable.set_dims(dims) - - -def dataset_to_stacked( - dataset: "xarray.Dataset", - sizes=None, - preserved_dims=("batch", "lat", "lon"), -) -> "xarray.DataArray": - """Converts an xarray.Dataset to a single stacked array. - - This takes each consistuent data_var, converts it into BHWC layout - using `variable_to_stacked`, then concats them all along the channels axis. - - Args: - dataset: An xarray.Dataset. - sizes: Mapping including sizes for any dimensions which are not present in the `dataset` but are needed for the output. See variable_to_stacked. - preserved_dims: dimensions from the dataset that should not be folded in the predictions channels. - - Returns: - An xarray.DataArray with dimensions preserved_dims + ("channels",). Existing coordinates for preserved_dims axes will be preserved, however there will be no coordinates for "channels". - """ - data_vars = [ - variable_to_stacked( - dataset.variables[name], sizes or dataset.sizes, preserved_dims - ) - for name in sorted(dataset.data_vars.keys()) - ] - coords = { - dim: coord for dim, coord in dataset.coords.items() if dim in preserved_dims - } - return xarray.DataArray( - data=xarray.Variable.concat(data_vars, dim="channels"), coords=coords - ) - - -class GridMeshAtmosphericDataset(io.Dataset): - """This class is used to process ERA5 re-analyze data, and is used to generate the dataset generator supported by MindSpore. This class inherits the Data class. - - Args: - input_keys (Tuple[str, ...]): Name of input data. - label_keys (Tuple[str, ...]): Name of label data. - data_path: Path of atmospheric datafile. - mean_path: Path of mean datafile. - stddev_path: Path of standard deviation datafile. - stddev_diffs_path: Path of standard deviation different datafile. - type: Type of GraphCast network. - mesh_size: Size of mesh. - mesh2grid_edge_normalization_factor: Factor of normalization of edges in Mesh2Grid GNN. - radius_query_fraction_edge_length: Length of radius query fraction edges. - resolution: Resolution of atmospheric data. - - Examples: - >>> import ppsci - >>> dataset = ppsci.data.dataset.GridMeshAtmosphericDataset( - ... "input_keys": ("input",), - ... "label_keys": ("output",), - ... "data_path": "/path/to/file.nc", - ... "mean_path": "/path/to/file.nc", - ... "stddev_path": "/path/to/file.nc", - ... "stddev_diffs_path": "/path/to/file.nc", - ... "type": "graphcast_small", - ... "mesh_size": 5, - ... "mesh2grid_edge_normalization_factor": 0.06, - ... "radius_query_fraction_edge_length": 0.6180338738074472, - ... "resolution": 1, - ... ) # doctest: +SKIP - """ - - use_graph_grid_mesh: bool = True - - def __init__( - self, - input_keys: Tuple[str, ...], - label_keys: Tuple[str, ...], - data_path: str, - mean_path: str, - stddev_path: str, - stddev_diffs_path: str, - type: str, - mesh_size: int, - mesh2grid_edge_normalization_factor: float, - radius_query_fraction_edge_length: float, - resolution: float, - ): - super().__init__() - self.input_keys = input_keys - self.label_keys = label_keys - if type == "graphcast": - self.input_variables = TASK_input_variables - self.forcing_variables = TASK_forcing_variables - self.target_variables = TASK_target_variables - self.level_variables = PRESSURE_LEVELS[37] - elif type == "graphcast_small": - self.input_variables = TASK_13_input_variables - self.forcing_variables = TASK_13_forcing_variables - self.target_variables = TASK_13_target_variables - self.level_variables = PRESSURE_LEVELS[13] - elif type == "graphcast_operational": - self.input_variables = TASK_13_PRECIP_OUT_input_variables - self.forcing_variables = TASK_13_PRECIP_OUT_forcing_variables - self.target_variables = TASK_13_PRECIP_OUT_target_variables - self.level_variables = PRESSURE_LEVELS[13] - - nc_dataset = xarray.open_dataset(data_path) - - longitude_offsets = nc_dataset.coords["lon"].data - second_since_epoch = ( - nc_dataset.coords["datetime"].data.astype("datetime64[s]").astype(np.int64) - ) - datetime_feats = datetime_features(second_since_epoch, longitude_offsets) - nc_dataset.update( - { - "year_progress_sin": xarray.Variable( - ("batch", "time"), datetime_feats["year_progress_sin"] - ), - "year_progress_cos": xarray.Variable( - ("batch", "time"), datetime_feats["year_progress_cos"] - ), - "day_progress_sin": xarray.Variable( - ("batch", "time", "lon"), datetime_feats["day_progress_sin"] - ), - "day_progress_cos": xarray.Variable( - ("batch", "time", "lon"), datetime_feats["day_progress_cos"] - ), - } - ) - - inputs, targets = extract_input_target_times( - nc_dataset, input_duration="12h", target_lead_times="6h" - ) - - stddev_data = xarray.open_dataset(stddev_path).sel( - level=list(self.level_variables) - ) - stddev_diffs_data = xarray.open_dataset(stddev_diffs_path).sel( - level=list(self.level_variables) - ) - mean_data = xarray.open_dataset(mean_path).sel(level=list(self.level_variables)) - - missing_variables = set(self.target_variables) - set(self.input_variables) - exist_variables = set(self.target_variables) - missing_variables - targets_stddev = stddev_diffs_data[list(exist_variables)] - target_mean = inputs[list(exist_variables)].isel(time=-1) - if missing_variables: - targets_stddev.update({var: stddev_data[var] for var in missing_variables}) - target_mean.update( - {var: mean_data.variables[var] for var in missing_variables} - ) - - stacked_targets_stddev = dataset_to_stacked(targets_stddev, preserved_dims=()) - stacked_targets_mean = dataset_to_stacked(target_mean) - stacked_targets_mean = stacked_targets_mean.transpose("lat", "lon", ...) - - inputs = inputs[list(self.input_variables)] - forcings = targets[list(self.forcing_variables)] - targets = targets[list(self.target_variables)] - inputs = self.normalize(inputs, stddev_data, mean_data) - forcings = self.normalize(forcings, stddev_data, mean_data) - - self.targets_template = targets - - stacked_inputs = dataset_to_stacked(inputs) - stacked_forcings = dataset_to_stacked(forcings) - stacked_targets = dataset_to_stacked(targets) - stacked_inputs = xarray.concat( - [stacked_inputs, stacked_forcings], dim="channels" - ) - - stacked_inputs = stacked_inputs.transpose("lat", "lon", ...) - stacked_targets = stacked_targets.transpose("lat", "lon", ...) - - lat_dim, lon_dim, batch_dim, feat_dim = stacked_inputs.shape - stacked_inputs = stacked_inputs.data.reshape(lat_dim * lon_dim, batch_dim, -1) - stacked_targets = stacked_targets.data.reshape(lat_dim * lon_dim, batch_dim, -1) - self.stacked_targets_stddev = stacked_targets_stddev.data - self.stacked_targets_mean = stacked_targets_mean.data.reshape( - lat_dim * lon_dim, batch_dim, -1 - ) - - self.input_data = [] - self.target_data = [] - - graph = GraphGridMesh( - mesh_size=mesh_size, - radius_query_fraction_edge_length=radius_query_fraction_edge_length, - mesh2grid_edge_normalization_factor=mesh2grid_edge_normalization_factor, - resolution=resolution, - ) - - graph.grid_node_feat = np.concatenate( - [stacked_inputs, graph.grid_node_feat], axis=-1 - ) - mesh_node_feat = np.zeros([graph.mesh_num_nodes, batch_dim, feat_dim]) - graph.mesh_node_feat = np.concatenate( - [mesh_node_feat, graph.mesh_node_feat], axis=-1 - ) - - self.input_data.append(graph) - self.target_data.append(stacked_targets) - - def __len__(self): - return len(self.input_data) - - def __getitem__(self, idx): - return ( - { - self.input_keys[0]: self.input_data[idx], - }, - { - self.label_keys[0]: self.target_data[idx], - }, - None, - ) - - def normalize(self, inputs_data, stddev_data, mean_data): - for name in list(inputs_data.keys()): - inputs_data[name] = (inputs_data[name] - mean_data[name]) / stddev_data[ - name - ] - return inputs_data - - def denormalize(self, inputs_data): - return inputs_data * self.stacked_targets_stddev + self.stacked_targets_mean diff --git a/examples/smc_reac/ppsci/data/dataset/cgcnn_dataset.py b/examples/smc_reac/ppsci/data/dataset/cgcnn_dataset.py deleted file mode 100644 index 0c04c5e319..0000000000 --- a/examples/smc_reac/ppsci/data/dataset/cgcnn_dataset.py +++ /dev/null @@ -1,312 +0,0 @@ -# Copyright (c) 2024 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. - -import csv -import functools -import json -import os -import random -import warnings -from typing import Tuple - -import numpy as np -import paddle -from paddle import io - -try: - from pymatgen.core.structure import Structure -except ModuleNotFoundError: - pass - - -def collate_pool(dataset_list): - - """ - Collate a list of data and return a batch for predicting crystal properties. - - Args: - dataset_list (list): A list of tuples for each data point containing: - - atom_fea (paddle.Tensor): Shape (n_i, atom_fea_len). - - nbr_fea (paddle.Tensor): Shape (n_i, M, nbr_fea_len). - - nbr_fea_idx (paddle.Tensor): Shape (n_i, M). - - target (paddle.Tensor): Shape (1,). - - cif_id (str or int). - - Returns: - tuple: Contains the following: - - batch_atom_fea (paddle.Tensor): Shape (N, orig_atom_fea_len). Atom features from atom type. - - batch_nbr_fea (paddle.Tensor): Shape (N, M, nbr_fea_len). Bond features of each atom's M neighbors. - - batch_nbr_fea_idx (paddle.Tensor): Shape (N, M). Indices of M neighbors of each atom. - - crystal_atom_idx (list): List of paddle.Tensor of length N0. Mapping from the crystal idx to atom idx. - - target (paddle.Tensor): Shape (N, 1). Target value for prediction. - - batch_cif_ids (list): List of CIF IDs. - - Notes: - - N = sum(n_i); N0 = sum(i) - """ - batch_atom_fea, batch_nbr_fea, batch_nbr_fea_idx = [], [], [] - crystal_atom_idx, batch_target = [], [] - batch_cif_ids = [] - base_idx = 0 - for i, item in enumerate(dataset_list): - input: Tuple[np.ndarray, np.ndarray, np.ndarray] = item[0]["i"] - label = item[1]["l"] - id = item[2]["c"] - atom_fea, nbr_fea, nbr_fea_idx = input - target = label - cif_id = id - n_i = atom_fea.shape[0] # number of atoms for this crystal - batch_atom_fea.append(atom_fea) - batch_nbr_fea.append(nbr_fea) - batch_nbr_fea_idx.append(nbr_fea_idx + base_idx) - new_idx = np.arange(n_i, dtype="int64") + int(base_idx) - crystal_atom_idx.append(new_idx) - batch_target.append(target) - batch_cif_ids.append(cif_id) - base_idx += n_i - # Debugging: print shapes of the tensors to ensure they are consistent - # print("Shapes of batch_atom_fea:", [x.shape for x in batch_atom_fea]) - # print("Shapes of batch_nbr_fea:", [x.shape for x in batch_nbr_fea]) - # print("Shapes of batch_nbr_fea_idx:", [x.shape for x in batch_nbr_fea_idx]) - # Ensure all tensors in the lists have consistent shapes before concatenation - batch_atom_fea = np.concatenate(batch_atom_fea, axis=0) - batch_nbr_fea = np.concatenate(batch_nbr_fea, axis=0) - batch_nbr_fea_idx = np.concatenate(batch_nbr_fea_idx, axis=0) - return ( - { - "i": ( - np.array(batch_atom_fea, dtype="float32"), - np.array(batch_nbr_fea, dtype="float32"), - np.array(batch_nbr_fea_idx), - [np.array(crys_idx) for crys_idx in crystal_atom_idx], - ) - }, - {"l": np.array(np.stack(batch_target, axis=0))}, - {"c": batch_cif_ids}, - ) - - -class GaussianDistance(object): - """ - Expands the distance by Gaussian basis. - - Args: - dmin (float): Minimum interatomic distance. - dmax (float): Maximum interatomic distance. - step (float): Step size for the Gaussian filter. - """ - - def __init__(self, dmin, dmax, step, var=None): - assert dmin < dmax - assert dmax - dmin > step - self.filter = np.arange(dmin, dmax + step, step) - if var is None: - var = step - self.var = var - - def expand(self, distances): - """ - Apply Gaussian distance filter to a numpy distance array. - - Args: - distance (np.array): n-dimensional distance matrix of any shape. - - Returns: - np.array: Expanded distance matrix with the last dimension of length len(self.filter). - """ - - return np.exp( - -((distances[..., np.newaxis] - self.filter) ** 2) / self.var**2 - ) - - -class AtomInitializer(object): - """ - Base class for initializing the vector representation for atoms. - - !!! Use one AtomInitializer per dataset !!! - """ - - def __init__(self, atom_types): - self.atom_types = set(atom_types) - self._embedding = {} - - def get_atom_fea(self, atom_type): - assert atom_type in self.atom_types - return self._embedding[atom_type] - - def load_state_dict(self, state_dict): - self._embedding = state_dict - self.atom_types = set(self._embedding.keys()) - self._decodedict = { - idx: atom_type for atom_type, idx in self._embedding.items() - } - - def state_dict(self): - return self._embedding - - def decode(self, idx): - if not hasattr(self, "_decodedict"): - self._decodedict = { - idx: atom_type for atom_type, idx in self._embedding.items() - } - return self._decodedict[idx] - - -class AtomCustomJSONInitializer(AtomInitializer): - """ - Initialize atom feature vectors using a JSON file, which is a Python dictionary mapping from element number to a list representing the feature vector of the element. - - Args: - elem_embedding_file (str): The path to the .json file. - """ - - def __init__(self, elem_embedding_file): - with open(elem_embedding_file) as f: - elem_embedding = json.load(f) - elem_embedding = {int(key): value for key, value in elem_embedding.items()} - atom_types = set(elem_embedding.keys()) - super(AtomCustomJSONInitializer, self).__init__(atom_types) - for key, value in elem_embedding.items(): - self._embedding[key] = np.array(value, dtype=float) - - -class CIFData(io.Dataset): - """ - The CIFData dataset is a wrapper for a dataset where the crystal structures - are stored in the form of CIF files. The dataset should have the following - directory structure: - - root_dir - ├── id_prop.csv - ├── atom_init.json - ├── id0.cif - ├── id1.cif - ├── ... - - id_prop.csv: a CSV file with two columns. The first column recodes a - unique ID for each crystal, and the second column recodes the value of - target property. - - atom_init.json: a JSON file that stores the initialization vector for each element. - - ID.cif: a CIF file that recodes the crystal structure, where ID is the - unique ID for the crystal. - - Args - root_dir (str): The path to the root directory of the dataset - max_num_nbr (int): The maximum number of neighbors while constructing the crystal graph - radius (float): The cutoff radius for searching neighbors - dmin (float): The minimum distance for constructing GaussianDistance - step (float): The step size for constructing GaussianDistance - random_seed (int): Random seed for shuffling the dataset - - - Returns - atom_fea (paddle.Tensor): Shape (n_i, atom_fea_len) - nbr_fea (paddle.Tensor): Shape (n_i, M, nbr_fea_len) - nbr_fea_idx (paddle.Tensor): Shape (n_i, M) - target (paddle.Tensor): Shape (1, ) - cif_id (str or int) - - Examples: - >>> import ppsci - >>> dataset = ppsci.data.dataset.CGCNNDataset( - ... "file_path": "/path/to/CGCNNDataset", - ... "input_keys": "i", - ... "label_keys": "l", - ... "id_keys": "c", - ... ) # doctest: +SKIP - """ - - def __init__( - self, - root_dir: str, - input_keys: Tuple[str, ...], - label_keys: Tuple[str, ...], - id_keys: Tuple[str, ...], - max_num_nbr: int = 12, - radius: int = 8, - dmin: int = 0, - step: float = 0.2, - random_seed: int = 123, - ): - super().__init__() - self.input_keys = input_keys - self.label_keys = label_keys - self.id_keys = id_keys - self.root_dir = root_dir - self.max_num_nbr, self.radius = max_num_nbr, radius - assert os.path.exists(root_dir), "root_dir does not exist!" - id_prop_file = os.path.join(self.root_dir, "id_prop.csv") - assert os.path.exists(id_prop_file), "id_prop.csv does not exist!" - with open(id_prop_file) as f: - reader = csv.reader(f) - self.id_prop_data = [row for row in reader] - random.seed(random_seed) - random.shuffle(self.id_prop_data) - atom_init_file = os.path.join(self.root_dir, "atom_init.json") - assert os.path.exists(atom_init_file), f"{atom_init_file} does not exist!" - self.ari = AtomCustomJSONInitializer(atom_init_file) - self.gdf = GaussianDistance(dmin=dmin, dmax=self.radius, step=step) - self.raw_data = [self.get(i) for i in range(len(self))] - - def __len__(self): - return len(self.id_prop_data) - - @functools.lru_cache(maxsize=None) # Cache loaded structures - def __getitem__(self, idx): - return ( - {self.input_keys[0]: self.raw_data[idx][0]}, - {self.label_keys[0]: self.raw_data[idx][1]}, - {self.id_keys[0]: self.raw_data[idx][2]}, - ) - - def get(self, idx): - cif_id, target = self.id_prop_data[idx] - crystal = Structure.from_file(os.path.join(self.root_dir, cif_id + ".cif")) - atom_fea = np.vstack( - [ - self.ari.get_atom_fea(crystal[i].specie.number) - for i in range(len(crystal)) - ] - ) - atom_fea = paddle.Tensor(atom_fea) - all_nbrs = crystal.get_all_neighbors(self.radius, include_index=True) - all_nbrs = [sorted(nbrs, key=lambda x: x[1]) for nbrs in all_nbrs] - nbr_fea_idx, nbr_fea = [], [] - for nbr in all_nbrs: - if len(nbr) < self.max_num_nbr: - warnings.warn( - f"{cif_id} not find enough neighbors to build graph. " - "If it happens frequently, consider increase " - "radius." - ) - nbr_fea_idx.append( - list(map(lambda x: x[2], nbr)) + [0] * (self.max_num_nbr - len(nbr)) - ) - nbr_fea.append( - list(map(lambda x: x[1], nbr)) - + [self.radius + 1.0] * (self.max_num_nbr - len(nbr)) - ) - else: - nbr_fea_idx.append(list(map(lambda x: x[2], nbr[: self.max_num_nbr]))) - nbr_fea.append(list(map(lambda x: x[1], nbr[: self.max_num_nbr]))) - nbr_fea_idx, nbr_fea = np.array(nbr_fea_idx), np.array(nbr_fea) - nbr_fea = self.gdf.expand(nbr_fea) - atom_fea = np.array(atom_fea) - nbr_fea = np.array(nbr_fea) - nbr_fea_idx = np.array(nbr_fea_idx, dtype="int64") - target = np.array([float(target)], dtype="float32") - return (atom_fea, nbr_fea, nbr_fea_idx), target, cif_id diff --git a/examples/smc_reac/ppsci/data/dataset/csv_dataset.py b/examples/smc_reac/ppsci/data/dataset/csv_dataset.py deleted file mode 100644 index c14bb107da..0000000000 --- a/examples/smc_reac/ppsci/data/dataset/csv_dataset.py +++ /dev/null @@ -1,287 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import Callable -from typing import Dict -from typing import Optional -from typing import Tuple -from typing import Union - -import numpy as np -import paddle -from paddle import io -from paddle import vision - -from ppsci.utils import misc -from ppsci.utils import reader - - -class CSVDataset(io.Dataset): - """Dataset class for .csv file. - - Args: - file_path (str): CSV file path. - input_keys (Tuple[str, ...]): List of input keys. - label_keys (Tuple[str, ...]): List of label keys. - alias_dict (Optional[Dict[str, str]]): Dict of alias(es) for input and label keys. - i.e. {inner_key: outer_key}. Defaults to None. - weight_dict (Optional[Dict[str, Union[Callable, float]]]): Define the weight of - each constraint variable. Defaults to None. - timestamps (Optional[Tuple[float, ...]]): The number of repetitions of the data - in the time dimension. Defaults to None. - transforms (Optional[vision.Compose]): Compose object contains sample wise - transform(s). Defaults to None. - - Examples: - >>> import ppsci - >>> dataset = ppsci.data.dataset.CSVDataset( - ... "/path/to/file.csv", - ... ("x",), - ... ("u",), - ... ) # doctest: +SKIP - """ - - # Whether support batch indexing for speeding up fetching process. - batch_index: bool = True - - def __init__( - self, - file_path: str, - input_keys: Tuple[str, ...], - label_keys: Tuple[str, ...], - alias_dict: Optional[Dict[str, str]] = None, - weight_dict: Optional[Dict[str, Union[Callable, float]]] = None, - timestamps: Optional[Tuple[float, ...]] = None, - transforms: Optional[vision.Compose] = None, - ): - super().__init__() - self.input_keys = input_keys - self.label_keys = label_keys - - # read raw data from file - raw_data = reader.load_csv_file( - file_path, - input_keys + label_keys, - alias_dict, - ) - # filter raw data by given timestamps if specified - if timestamps is not None: - if "t" in raw_data: - # filter data according to given timestamps - raw_time_array = raw_data["t"] - mask = [] - for ti in timestamps: - mask.append(np.nonzero(np.isclose(raw_time_array, ti).flatten())[0]) - raw_data = misc.convert_to_array( - raw_data, self.input_keys + self.label_keys - ) - mask = np.concatenate(mask, 0) - raw_data = raw_data[mask] - raw_data = misc.convert_to_dict( - raw_data, self.input_keys + self.label_keys - ) - else: - # repeat data according to given timestamps - raw_data = misc.convert_to_array( - raw_data, self.input_keys + self.label_keys - ) - raw_data = misc.combine_array_with_time(raw_data, timestamps) - self.input_keys = ("t",) + tuple(self.input_keys) - raw_data = misc.convert_to_dict( - raw_data, self.input_keys + self.label_keys - ) - - # fetch input data - self.input = { - key: value for key, value in raw_data.items() if key in self.input_keys - } - # fetch label data - self.label = { - key: value for key, value in raw_data.items() if key in self.label_keys - } - - # prepare weights - self.weight = ( - {key: np.ones_like(next(iter(self.label.values()))) for key in self.label} - if weight_dict is not None - else {} - ) - if weight_dict is not None: - for key, value in weight_dict.items(): - if isinstance(value, (int, float)): - self.weight[key] = np.full_like( - next(iter(self.label.values())), value - ) - elif callable(value): - func = value - self.weight[key] = func(self.input) - if isinstance(self.weight[key], (int, float)): - self.weight[key] = np.full_like( - next(iter(self.label.values())), self.weight[key] - ) - else: - raise NotImplementedError(f"type of {type(value)} is invalid yet.") - - self.transforms = transforms - self._len = len(next(iter(self.input.values()))) - - def __getitem__(self, idx): - input_item = {key: value[idx] for key, value in self.input.items()} - label_item = {key: value[idx] for key, value in self.label.items()} - weight_item = {key: value[idx] for key, value in self.weight.items()} - - if self.transforms is not None: - input_item, label_item, weight_item = self.transforms( - input_item, label_item, weight_item - ) - - return (input_item, label_item, weight_item) - - def __len__(self): - return self._len - - -class IterableCSVDataset(io.IterableDataset): - """IterableCSVDataset for full-data loading. - - Args: - file_path (str): CSV file path. - input_keys (Tuple[str, ...]): List of input keys. - label_keys (Tuple[str, ...]): List of label keys. - alias_dict (Optional[Dict[str, str]]): Dict of alias(es) for input and label keys. - Defaults to None. - weight_dict (Optional[Dict[str, Union[Callable, float]]]): Define the weight of - each constraint variable. Defaults to None. - timestamps (Optional[Tuple[float, ...]]): The number of repetitions of the data - in the time dimension. Defaults to None. - transforms (Optional[vision.Compose]): Compose object contains sample wise - transform(s). Defaults to None. - - Examples: - >>> import ppsci - >>> dataset = ppsci.data.dataset.IterableCSVDataset( - ... "/path/to/file.csv" - ... ("x",), - ... ("u",), - ... ) # doctest: +SKIP - """ - - # Whether support batch indexing for speeding up fetching process. - batch_index: bool = False - - def __init__( - self, - file_path: str, - input_keys: Tuple[str, ...], - label_keys: Tuple[str, ...], - alias_dict: Optional[Dict[str, str]] = None, - weight_dict: Optional[Dict[str, Union[Callable, float]]] = None, - timestamps: Optional[Tuple[float, ...]] = None, - transforms: Optional[vision.Compose] = None, - ): - super().__init__() - self.input_keys = input_keys - self.label_keys = label_keys - - # read raw data from file - raw_data = reader.load_csv_file( - file_path, - input_keys + label_keys, - alias_dict, - ) - # filter raw data by given timestamps if specified - if timestamps is not None: - if "t" in raw_data: - # filter data according to given timestamps - raw_time_array = raw_data["t"] - mask = [] - for ti in timestamps: - mask.append(np.nonzero(np.isclose(raw_time_array, ti).flatten())[0]) - raw_data = misc.convert_to_array( - raw_data, self.input_keys + self.label_keys - ) - mask = np.concatenate(mask, 0) - raw_data = raw_data[mask] - raw_data = misc.convert_to_dict( - raw_data, self.input_keys + self.label_keys - ) - else: - # repeat data according to given timestamps - raw_data = misc.convert_to_array( - raw_data, self.input_keys + self.label_keys - ) - raw_data = misc.combine_array_with_time(raw_data, timestamps) - self.input_keys = ("t",) + tuple(self.input_keys) - raw_data = misc.convert_to_dict( - raw_data, self.input_keys + self.label_keys - ) - - # fetch input data - self.input = { - key: value for key, value in raw_data.items() if key in self.input_keys - } - # fetch label data - self.label = { - key: value for key, value in raw_data.items() if key in self.label_keys - } - - # prepare weights - self.weight = ( - {key: np.ones_like(next(iter(self.label.values()))) for key in self.label} - if weight_dict is not None - else {} - ) - if weight_dict is not None: - for key, value in weight_dict.items(): - if isinstance(value, (int, float)): - self.weight[key] = np.full_like( - next(iter(self.label.values())), value - ) - elif callable(value): - func = value - self.weight[key] = func(self.input) - if isinstance(self.weight[key], (int, float)): - self.weight[key] = np.full_like( - next(iter(self.label.values())), self.weight[key] - ) - else: - raise NotImplementedError(f"type of {type(value)} is invalid yet.") - - self.input = {key: paddle.to_tensor(value) for key, value in self.input.items()} - self.label = {key: paddle.to_tensor(value) for key, value in self.label.items()} - self.weight = { - key: paddle.to_tensor(value) for key, value in self.weight.items() - } - - self.transforms = transforms - self._len = len(next(iter(self.input.values()))) - - @property - def num_samples(self): - """Number of samples within current dataset.""" - return self._len - - def __iter__(self): - if callable(self.transforms): - input_, label_, weight_ = self.transforms( - self.input, self.label, self.weight - ) - yield input_, label_, weight_ - else: - yield self.input, self.label, self.weight - - def __len__(self): - return 1 diff --git a/examples/smc_reac/ppsci/data/dataset/cylinder_dataset.py b/examples/smc_reac/ppsci/data/dataset/cylinder_dataset.py deleted file mode 100644 index 3a49b7d436..0000000000 --- a/examples/smc_reac/ppsci/data/dataset/cylinder_dataset.py +++ /dev/null @@ -1,215 +0,0 @@ -# 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. - -from __future__ import annotations - -import os -from os import path as osp -from typing import Tuple - -import numpy as np -import paddle -from paddle import io - -from ppsci.data.dataset import airfoil_dataset - -try: - import pgl -except ModuleNotFoundError: - pass - -SU2_SHAPE_IDS = { - "line": 3, - "triangle": 5, - "quad": 9, -} - - -class MeshCylinderDataset(io.Dataset): - """Dataset for `MeshCylinder`. - - Args: - input_keys (Tuple[str, ...]): Name of input data. - label_keys (Tuple[str, ...]): Name of label data. - data_dir (str): Directory of MeshCylinder data. - mesh_graph_path (str): Path of mesh graph. - - Examples: - >>> import ppsci - >>> dataset = ppsci.data.dataset.MeshAirfoilDataset( - ... "input_keys": ("input",), - ... "label_keys": ("output",), - ... "data_dir": "/path/to/MeshAirfoilDataset", - ... "mesh_graph_path": "/path/to/file.su2", - ... ) # doctest: +SKIP - """ - - # Whether support batch indexing for speeding up fetching process. - batch_index: bool = False - - use_pgl: bool = True - - def __init__( - self, - input_keys: Tuple[str, ...], - label_keys: Tuple[str, ...], - data_dir: str, - mesh_graph_path: str, - ): - self.input_keys = input_keys - self.label_keys = label_keys - self.data_dir = data_dir - self.file_list = os.listdir(self.data_dir) - self.len = len(self.file_list) - self.mesh_graph = airfoil_dataset._get_mesh_graph(mesh_graph_path) - - self.normalization_factors = np.array( - [[978.6001, 48.9258, 24.8404], [-692.3159, -6.9950, -24.8572]], - dtype=paddle.get_default_dtype(), - ) - - self.nodes = self.mesh_graph[0] - self.meshnodes = self.mesh_graph[0] - self.edges = self.mesh_graph[1] - self.elems_list = self.mesh_graph[2] - self.marker_dict = self.mesh_graph[3] - self.bounder = [] - self.node_markers = np.full([self.nodes.shape[0], 1], fill_value=-1) - for i, (marker_tag, marker_elems) in enumerate(self.marker_dict.items()): - for elem in marker_elems: - self.node_markers[elem[0]] = i - self.node_markers[elem[1]] = i - - self.raw_graphs = [self.get(i) for i in range(len(self))] - - def __len__(self): - return self.len - - def __getitem__(self, idx): - return ( - { - self.input_keys[0]: self.raw_graphs[idx], - }, - { - self.label_keys[0]: self.raw_graphs[idx], - }, - None, - ) - - def get(self, idx): - with open(osp.join(self.data_dir, self.file_list[idx]), "r") as f: - field = [] - pos = [] - for line in f.read().splitlines()[1:]: - lines_pos = line.split(",")[1:3] - lines_field = line.split(",")[3:] - numbers_float = list(eval(i) for i in lines_pos) - array = np.array(numbers_float, paddle.get_default_dtype()) - pos.append(array) - numbers_float = list(eval(i) for i in lines_field) - array = np.array(numbers_float, paddle.get_default_dtype()) - field.append(array) - - field = np.stack(field, axis=0) - pos = np.stack(pos, axis=0) - indexlist = [] - for i in range(self.meshnodes.shape[0]): - b = self.meshnodes[i : (i + 1)] - b = np.squeeze(b) - index = np.nonzero( - np.sum((pos == b), axis=1, dtype=paddle.get_default_dtype()) - == pos.shape[1] - ) - indexlist.append(index) - indexlist = np.stack(indexlist, axis=0) - indexlist = np.squeeze(indexlist) - fields = field[indexlist] - velocity = self._get_params_from_name(self.file_list[idx]) - - norm_aoa = velocity / 40 - # add physics parameters to graph - nodes = np.concatenate( - [ - self.nodes, - np.repeat(a=norm_aoa, repeats=self.nodes.shape[0])[:, np.newaxis], - self.node_markers, - ], - axis=-1, - ).astype(paddle.get_default_dtype()) - - data = pgl.Graph( - num_nodes=nodes.shape[0], - edges=self.edges, - ) - data.x = nodes - data.y = fields - data.pos = self.nodes - data.edge_index = self.edges - data.velocity = velocity - - sender = data.x[data.edge_index[0]] - receiver = data.x[data.edge_index[1]] - relation_pos = sender[:, 0:2] - receiver[:, 0:2] - post = np.linalg.norm(relation_pos, ord=2, axis=1, keepdims=True).astype( - paddle.get_default_dtype() - ) - data.edge_attr = post - std_epsilon = [1e-8] - a = np.mean(data.edge_attr, axis=0) - b = data.edge_attr.std(axis=0) - b = np.maximum(b, std_epsilon).astype(paddle.get_default_dtype()) - data.edge_attr = (data.edge_attr - a) / b - a = np.mean(data.y, axis=0) - b = data.y.std(axis=0) - b = np.maximum(b, std_epsilon).astype(paddle.get_default_dtype()) - data.y = (data.y - a) / b - data.norm_max = a - data.norm_min = b - - # find the face of the boundary,our cylinder dataset come from fluent solver - with open(osp.join(osp.dirname(self.data_dir), "bounder"), "r") as f: - field = [] - pos = [] - for line in f.read().splitlines()[1:]: - lines_pos = line.split(",")[1:3] - lines_field = line.split(",")[3:] - numbers_float = list(eval(i) for i in lines_pos) - array = np.array(numbers_float, paddle.get_default_dtype()) - pos.append(array) - numbers_float = list(eval(i) for i in lines_field) - array = np.array(numbers_float, paddle.get_default_dtype()) - field.append(array) - - field = np.stack(field, axis=0) - pos = np.stack(pos, axis=0) - - indexlist = [] - for i in range(pos.shape[0]): - b = pos[i : (i + 1)] - b = np.squeeze(b) - index = np.nonzero( - np.sum((self.nodes == b), axis=1, dtype=paddle.get_default_dtype()) - == self.nodes.shape[1] - ) - indexlist.append(index) - - indexlist = np.stack(indexlist, axis=0) - indexlist = np.squeeze(indexlist) - self.bounder = indexlist - return data - - def _get_params_from_name(self, filename): - s = filename.rsplit(".", 1)[0] - reynolds = np.array(s[13:])[np.newaxis].astype(paddle.get_default_dtype()) - return reynolds diff --git a/examples/smc_reac/ppsci/data/dataset/darcyflow_dataset.py b/examples/smc_reac/ppsci/data/dataset/darcyflow_dataset.py deleted file mode 100644 index 3e748eb785..0000000000 --- a/examples/smc_reac/ppsci/data/dataset/darcyflow_dataset.py +++ /dev/null @@ -1,296 +0,0 @@ -from pathlib import Path -from typing import Dict -from typing import Optional -from typing import Tuple - -import numpy as np -import paddle -from paddle import io - - -# normalization, pointwise gaussian -class UnitGaussianNormalizer: - def __init__(self, x, eps=1e-7, reduce_dim=[0], verbose=True): - super().__init__() - n_samples, *shape = x.shape - self.sample_shape = shape - self.verbose = verbose - self.reduce_dim = reduce_dim - - # x could be in shape of ntrain*n or ntrain*T*n or ntrain*n*T - self.mean = paddle.mean(x, reduce_dim, keepdim=True).squeeze(0) - self.std = paddle.std(x, reduce_dim, keepdim=True).squeeze(0) - self.eps = eps - - if verbose: - print( - f"UnitGaussianNormalizer init on {n_samples}, reducing over {reduce_dim}, samples of shape {shape}." - ) - print(f" Mean and std of shape {self.mean.shape}, eps={eps}") - - def encode(self, x): - - x -= self.mean - x /= self.std + self.eps - return x - - def decode(self, x, sample_idx=None): - if sample_idx is None: - std = self.std + self.eps # n - mean = self.mean - else: - if len(self.mean.shape) == len(sample_idx[0].shape): - std = self.std[sample_idx] + self.eps # batch*n - mean = self.mean[sample_idx] - if len(self.mean.shape) > len(sample_idx[0].shape): - std = self.std[:, sample_idx] + self.eps # T*batch*n - mean = self.mean[:, sample_idx] - - # x is in shape of batch*n or T*batch*n - x *= std - x += mean - - return x - - -def get_grid_positional_encoding( - input_tensor, grid_boundaries=[[0, 1], [0, 1]], channel_dim=1 -): - """Appends grid positional encoding to an input tensor, concatenating as additional dimensions along the channels. - - Args: - input_tensor (paddle.Tensor): The input tensor. - grid_boundaries (list, optional): The boundaries of the grid. Defaults to [[0, 1], [0, 1]]. - channel_dim (int, optional): The location of unsqueeze. Defaults to 1. - """ - - shape = list(input_tensor.shape) - if len(shape) == 2: - height, width = shape - else: - _, height, width = shape - - xt = paddle.linspace(grid_boundaries[0][0], grid_boundaries[0][1], height + 1)[:-1] - yt = paddle.linspace(grid_boundaries[1][0], grid_boundaries[1][1], width + 1)[:-1] - - grid_x, grid_y = paddle.meshgrid(xt, yt, indexing="ij") - - if len(shape) == 2: - grid_x = grid_x.unsqueeze(channel_dim) - grid_y = grid_y.unsqueeze(channel_dim) - else: - grid_x = grid_x.unsqueeze(0).unsqueeze(channel_dim) - grid_y = grid_y.unsqueeze(0).unsqueeze(channel_dim) - - return grid_x, grid_y - - -def regular_grid(spatial_dims, grid_boundaries=[[0, 1], [0, 1]]): - """ - Appends grid positional encoding to an input tensor, concatenating as additional dimensions along the channels - """ - height, width = spatial_dims - - xt = paddle.linspace(grid_boundaries[0][0], grid_boundaries[0][1], height + 1)[:-1] - yt = paddle.linspace(grid_boundaries[1][0], grid_boundaries[1][1], width + 1)[:-1] - - grid_x, grid_y = paddle.meshgrid(xt, yt, indexing="ij") - - grid_x = grid_x.tile((1, 1)) - grid_y = grid_y.tile((1, 1)) - - return grid_x, grid_y - - -class PositionalEmbedding2D: - def __init__(self, grid_boundaries=[[0, 1], [0, 1]]): - self.grid_boundaries = grid_boundaries - self._grid = None - self._res = None - - def grid(self, spatial_dims, dtype): - """Grid generates 2D grid needed for pos encoding - and caches the grid associated with MRU resolution - - Args: - spatial_dims (tuple[int,...]): Sizes of spatial resolution. - dtype (str): Dtype to encode data. - - Returns: - paddle.Tensor: Output grids to concatenate - """ - # handle case of multiple train resolutions - if self._grid is None or self._res != spatial_dims: - grid_x, grid_y = regular_grid( - spatial_dims, grid_boundaries=self.grid_boundaries - ) - - grid_x = grid_x.astype(dtype).unsqueeze(0).unsqueeze(0) - grid_y = grid_y.astype(dtype).unsqueeze(0).unsqueeze(0) - self._grid = grid_x, grid_y - self._res = spatial_dims - - return self._grid - - def __call__(self, data): - if data.ndim == 3: - data = data.unsqueeze(0) - x, y = self.grid(data.shape[-2:], data.dtype) - out = paddle.concat( - (data, x.expand([1, -1, -1, -1]), y.expand([1, -1, -1, -1])), axis=1 - ) - return out.squeeze(0) - - -class DarcyFlowDataset(io.Dataset): - """Loads a small Darcy-Flow dataset - - Training contains 1000 samples in resolution 16x16. - Testing contains 100 samples at resolution 16x16 and - 50 samples at resolution 32x32. - - Args: - input_keys (Tuple[str, ...]): Input keys, such as ("input",). - label_keys (Tuple[str, ...]): Output keys, such as ("output",). - data_dir (str): The directory to load data from. - weight_dict (Optional[Dict[str, float]], optional): Define the weight of each constraint variable. Defaults to None. - test_resolutions (List[int,...]): The resolutions to test dataset. Default is [16, 32]. - grid_boundaries (List[int,...]): The boundaries of the grid. Default is [[0,1],[0,1]]. - positional_encoding (bool): Whether to use positional encoding. Default is True - encode_input (bool): Whether to encode the input. Default is False - encode_output (bool): Whether to encode the output. Default is True - encoding (str): The type of encoding. Default is 'channel-wise'. - channel_dim (int): The location of unsqueeze. Default is 1. - where to put the channel dimension. Defaults size is batch, channel, height, width - data_split (str): Wether to use training or test dataset. Default is 'train'. - """ - - def __init__( - self, - input_keys: Tuple[str, ...], - label_keys: Tuple[str, ...], - data_dir: str, - weight_dict: Optional[Dict[str, float]] = None, - test_resolutions: Tuple[int, ...] = [32], - train_resolution: int = 32, - grid_boundaries: Tuple[Tuple[int, ...], ...] = [[0, 1], [0, 1]], - positional_encoding: bool = True, - encode_input: bool = False, - encode_output: bool = True, - encoding: str = "channel-wise", - channel_dim: int = 1, - data_split: str = "train", - ): - super().__init__() - for res in test_resolutions: - if res not in [16, 32]: - raise ValueError( - f"Only 32 and 64 are supported for test resolution, but got {test_resolutions}" - ) - - self.input_keys = input_keys - self.label_keys = label_keys - self.data_dir = data_dir - self.weight_dict = {} if weight_dict is None else weight_dict - if weight_dict is not None: - self.weight_dict = {key: 1.0 for key in self.label_keys} - self.weight_dict.update(weight_dict) - - self.test_resolutions = test_resolutions - self.train_resolution = train_resolution - self.grid_boundaries = grid_boundaries - self.positional_encoding = positional_encoding - self.encode_input = encode_input - self.encode_output = encode_output - self.encoding = encoding - self.channel_dim = channel_dim - self.data_split = data_split - - # train path - path_train = ( - Path(self.data_dir) - .joinpath(f"darcy_train_{self.train_resolution}.npy") - .as_posix() - ) - self.x_train, self.y_train = self.read_data(path_train) - # test path - path_test_1 = ( - Path(self.data_dir) - .joinpath(f"darcy_test_{self.test_resolutions[0]}.npy") - .as_posix() - ) - self.x_test_1, self.y_test_1 = self.read_data(path_test_1) - path_test_2 = ( - Path(self.data_dir) - .joinpath(f"darcy_test_{self.test_resolutions[1]}.npy") - .as_posix() - ) - self.x_test_2, self.y_test_2 = self.read_data(path_test_2) - - # input encoder - if self.encode_input: - self.input_encoder = self.encode_data(self.x_train) - self.x_train = self.input_encoder.encode(self.x_train) - self.x_test_1 = self.input_encoder.encode(self.x_test_1) - self.x_test_2 = self.input_encoder.encode(self.x_test_2) - else: - self.input_encoder = None - # output encoder - if self.encode_output: - self.output_encoder = self.encode_data(self.y_train) - self.y_train = self.output_encoder.encode(self.y_train) - else: - self.output_encoder = None - - if positional_encoding: - self.transform_x = PositionalEmbedding2D(grid_boundaries) - - def read_data(self, path): - # load with numpy - data = np.load(path, allow_pickle=True).item() - x = ( - paddle.to_tensor(data["x"]) - .unsqueeze(self.channel_dim) - .astype("float32") - .clone() - ) - y = paddle.to_tensor(data["y"]).unsqueeze(self.channel_dim).clone() - del data - return x, y - - def encode_data(self, data): - if self.encoding == "channel-wise": - reduce_dims = list(range(data.ndim)) - elif self.encoding == "pixel-wise": - reduce_dims = [0] - input_encoder = UnitGaussianNormalizer(data, reduce_dim=reduce_dims) - return input_encoder - - def __len__(self): - if self.data_split == "train": - return self.x_train.shape[0] - elif self.data_split == "test_16x16": - return self.x_test_1.shape[0] - else: - return self.x_test_2.shape[0] - - def __getitem__(self, index): - if self.data_split == "train": - x = self.x_train[index] - y = self.y_train[index] - - elif self.data_split == "test_16x16": - x = self.x_test_1[index] - y = self.y_test_1[index] - else: - x = self.x_test_2[index] - y = self.y_test_2[index] - - if self.transform_x is not None: - x = self.transform_x(x) - - input_item = {self.input_keys[0]: x} - label_item = {self.label_keys[0]: y} - weight_item = self.weight_dict - - return input_item, label_item, weight_item diff --git a/examples/smc_reac/ppsci/data/dataset/dgmr_dataset.py b/examples/smc_reac/ppsci/data/dataset/dgmr_dataset.py deleted file mode 100644 index 8490f679bc..0000000000 --- a/examples/smc_reac/ppsci/data/dataset/dgmr_dataset.py +++ /dev/null @@ -1,95 +0,0 @@ -# Copyright (c) 2024 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. - -from __future__ import annotations - -import importlib -from typing import Tuple - -import numpy as np -from numpy.random import default_rng -from paddle import io - - -class DGMRDataset(io.Dataset): - """ - Dataset class for DGMR (Deep Generative Model for Radar) model. - This open-sourced UK dataset has been mirrored to HuggingFace Datasets https://huggingface.co/datasets/openclimatefix/nimrod-uk-1km. - If the reader cannot load the dataset from Hugging Face, please manually download it and modify the dataset_path to the local path for loading. - - Args: - input_keys (Tuple[str, ...]): Input keys, such as ("input",). - label_keys (Tuple[str, ...]): Output keys, such as ("output",). - split (str, optional): The split of the dataset, "validation" or "train". Defaults to "validation". - num_input_frames (int, optional): Number of input frames. Defaults to 4. - num_target_frames (int, optional): Number of target frames. Defaults to 18. - dataset_path (str, optional): Path to the dataset. Defaults to "openclimatefix/nimrod-uk-1km". - - Examples: - >>> import ppsci - >>> dataset = ppsci.data.dataset.DGMRDataset(("input", ), ("output", )) # doctest: +SKIP - """ - - def __init__( - self, - input_keys: Tuple[str, ...], - label_keys: Tuple[str, ...], - split: str = "validation", - num_input_frames: int = 4, - num_target_frames: int = 18, - dataset_path: str = "openclimatefix/nimrod-uk-1km", - ): - super().__init__() - self.input_keys = input_keys - self.label_keys = label_keys - self.num_input_frames = num_input_frames - self.num_target_frames = num_target_frames - if not importlib.util.find_spec("datasets"): - raise ModuleNotFoundError( - "Please install datasets with `pip install datasets`" - " before exporting onnx model." - ) - import datasets - - self.reader = datasets.load_dataset( - dataset_path, "sample", split=split, streaming=True, trust_remote_code=True - ) - self.iter_reader = self.reader - - def __len__(self): - return 1000 - - def __getitem__(self, idx): - try: - row = next(self.iter_reader) - except Exception: - rng = default_rng(42) - self.iter_reader = iter( - self.reader.shuffle( - seed=rng.integers(low=0, high=100000), buffer_size=1000 - ) - ) - row = next(self.iter_reader) - radar_frames = row["radar_frames"] - input_frames = radar_frames[ - -self.num_target_frames - self.num_input_frames : -self.num_target_frames - ] - target_frames = radar_frames[-self.num_target_frames :] - input_item = { - self.input_keys[0]: np.moveaxis(input_frames, [0, 1, 2, 3], [0, 2, 3, 1]) - } - label_item = { - self.label_keys[0]: np.moveaxis(target_frames, [0, 1, 2, 3], [0, 2, 3, 1]) - } - return input_item, label_item diff --git a/examples/smc_reac/ppsci/data/dataset/drivaernet_dataset.py b/examples/smc_reac/ppsci/data/dataset/drivaernet_dataset.py deleted file mode 100644 index 23634bebb2..0000000000 --- a/examples/smc_reac/ppsci/data/dataset/drivaernet_dataset.py +++ /dev/null @@ -1,316 +0,0 @@ -# Copyright (c) 2024 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. - -""" -Created on Tue Dec 19 20:54:56 2023 - -@author: Mohamed Elrefaie, mohamed.elrefaie@mit.edu mohamed.elrefaie@tum.de - -This module is part of the research presented in the paper": -"DrivAerNet: A Parametric Car Dataset for Data-driven Aerodynamic Design and Graph-Based Drag Prediction". -""" -from __future__ import annotations - -import logging -import os -from typing import Callable -from typing import Dict -from typing import Optional -from typing import Tuple - -import numpy as np -import paddle -import pandas as pd - - -class DataAugmentation: - """ - Class encapsulating various data augmentation techniques for point clouds. - """ - - @staticmethod - def translate_pointcloud( - pointcloud: np.ndarray, - translation_range: Tuple[float, float] = (2.0 / 3.0, 3.0 / 2.0), - ) -> np.ndarray: - """ - Translates the pointcloud by a random factor within a given range. - - Args: - pointcloud: The input point cloud as a np.ndarray. - translation_range: A tuple specifying the range for translation factors. - - Returns: - Translated point cloud as a np.ndarray. - """ - xyz1 = np.random.uniform( - low=translation_range[0], high=translation_range[1], size=[3] - ) - xyz2 = np.random.uniform(low=-0.2, high=0.2, size=[3]) - translated_pointcloud = np.add(np.multiply(pointcloud, xyz1), xyz2).astype( - "float32" - ) - return paddle.to_tensor(data=translated_pointcloud, dtype="float32") - - @staticmethod - def jitter_pointcloud( - pointcloud: np.ndarray, sigma: float = 0.01, clip: float = 0.02 - ) -> np.ndarray: - """ - Adds Gaussian noise to the pointcloud. - - Args: - pointcloud: The input point cloud as a np.ndarray. - sigma: Standard deviation of the Gaussian noise. - clip: Maximum absolute value for noise. - - Returns: - Jittered point cloud as a np.ndarray. - """ - N, C = tuple(pointcloud.shape) - jittered_pointcloud = pointcloud + paddle.clip( - x=sigma * paddle.randn(shape=[N, C]), min=-clip, max=clip - ) - return jittered_pointcloud - - @staticmethod - def drop_points(pointcloud: np.ndarray, drop_rate: float = 0.1) -> np.ndarray: - """ - Randomly removes points from the point cloud based on the drop rate. - - Args: - pointcloud: The input point cloud as a np.ndarray. - drop_rate: The percentage of points to be randomly dropped. - - Returns: - The point cloud with points dropped as a np.ndarray. - """ - num_drop = int(drop_rate * pointcloud.shape[0]) - drop_indices = np.random.choice(pointcloud.shape[0], num_drop, replace=False) - keep_indices = np.setdiff1d(np.arange(pointcloud.shape[0]), drop_indices) - dropped_pointcloud = pointcloud[keep_indices, :] - return dropped_pointcloud - - -class DrivAerNetDataset(paddle.io.Dataset): - """ - Paddle Dataset class for the DrivAerNet dataset, handling loading, transforming, and augmenting 3D car models. - - This dataset is specifically designed for aerodynamic tasks, including training machine learning models - to predict aerodynamic coefficients such as drag coefficient (Cd) from 3D car models. - - Args: - input_keys (Tuple[str, ...]): Tuple specifying the keys for input features. - These keys correspond to the attributes of the dataset used as input to the model. - For example, "vertices" represents the 3D point cloud vertices of car models. - label_keys (Tuple[str, ...]): Tuple specifying the keys for ground-truth labels. - These keys correspond to the target values, such as aerodynamic coefficients like Cd. - Example: ("cd_value",) - weight_keys (Tuple[str, ...]): Tuple specifying the keys for optional sample weights. - These keys represent weighting factors that may be used to adjust loss computation - during model training. Useful for handling sample imbalance. - Example: ("weight_keys",) - subset_dir (str): Path to the directory containing subset information. - This directory typically contains files that divide the dataset into training, - validation, and test subsets using a list of model IDs. - ids_file (str): Path to the text file containing model IDs for the current subset. - Each line in the file corresponds to a unique model ID that defines which - models belong to the subset (e.g., training set or test set). - root_dir (str): Directory containing the STL files of 3D car models. - Each STL file is expected to represent a single car model and is named according - to the corresponding model ID. This is the primary data source. - csv_file (str): Path to the CSV file containing metadata for car models. - This file typically includes aerodynamic properties (e.g., drag coefficient) - and other descriptive attributes mapped to each model ID. - num_points (int): Fixed number of points to sample from each 3D model. - If a 3D model has more points than `num_points`, it will be randomly subsampled. - If it has fewer points, it will be zero-padded to reach the desired number. - transform (Optional[Callable]): An optional callable for applying data transformations. - This can include augmentations such as scaling, rotation, jittering, or other preprocessing - steps applied to the 3D point clouds before they are passed to the model. - pointcloud_exist (bool): Whether the point clouds are pre-processed and saved as `.pt` files. - If `True`, the dataset will directly load the pre-saved point clouds instead of generating them from STL files. - train_fractions (float): Fraction of the training data to use. Useful for experiments where only a portion of the data is needed. - mode (str): Mode of operation, either "train", "eval", or "test". Determines how the dataset behaves. - - Examples: - >>> import ppsci - >>> dataset = ppsci.data.dataset.DrivAerNetDataset( - ... input_keys=("vertices",), - ... label_keys=("cd_value",), - ... weight_keys=("weight_keys",), - ... subset_dir="/path/to/subset_dir", - ... ids_file="train_ids.txt", - ... root_dir="/path/to/DrivAerNetDataset", - ... csv_file="/path/to/aero_metadata.csv", - ... num_points=1024, - ... transform=None, - ... ) # doctest: +SKIP - """ - - def __init__( - self, - input_keys: Tuple[str, ...], - label_keys: Tuple[str, ...], - weight_keys: Tuple[str, ...], - subset_dir: str, - ids_file: str, - root_dir: str, - csv_file: str, - num_points: int, - transform: Optional[Callable] = None, - pointcloud_exist: bool = True, - train_fractions=1.0, - mode="eval", - ): - - super().__init__() - self.root_dir = root_dir - try: - self.data_frame = pd.read_csv(csv_file) - except Exception as e: - logging.error(f"Failed to load CSV file: {csv_file}. Error: {e}") - raise - self.input_keys = input_keys - self.label_keys = label_keys - self.weight_keys = weight_keys - self.subset_dir = subset_dir - self.ids_file = ids_file - self.transform = transform - self.num_points = num_points - self.pointcloud_exist = pointcloud_exist - self.mode = mode - self.train_fractions = train_fractions - self.augmentation = DataAugmentation() - self.cache = {} - - try: - with open(os.path.join(self.subset_dir, self.ids_file), "r") as file: - subset_ids = file.read().split() - except FileNotFoundError as e: - raise FileNotFoundError(f"Error loading subset file {self.ids_file}: {e}") - - self.subset_indices = self.data_frame[ - self.data_frame["Design"].isin(subset_ids) - ].index.tolist() - self.data_frame = self.data_frame.loc[self.subset_indices].reset_index( - drop=True - ) - - if self.mode == "train": - self.data_frame = self.data_frame.sample(frac=self.train_fractions) - else: - self.data_frame = self.data_frame - - def __len__(self) -> int: - """Returns the total number of samples in the dataset.""" - return len(self.data_frame) - - def _sample_or_pad_vertices( - self, vertices: paddle.Tensor, num_points: int - ) -> paddle.Tensor: - """ - Subsamples or pads the vertices of the model to a fixed number of points. - - Args: - vertices: The vertices of the 3D model as a paddle.Tensor. - num_points: The desired number of points for the model. - - Returns: - The vertices standardized to the specified number of points. - """ - num_vertices = vertices.shape[0] - if num_vertices > num_points: - indices = np.random.choice(num_vertices, num_points, replace=False) - vertices = vertices[indices] - elif num_vertices < num_points: - padding = paddle.zeros( - shape=(num_points - num_vertices, 3), dtype="float32" - ) - vertices = paddle.concat(x=(vertices, padding), axis=0) - return vertices - - def _load_point_cloud(self, design_id: str) -> Optional[paddle.Tensor]: - load_path = os.path.join(self.root_dir, f"{design_id}.paddle_tensor") - if os.path.exists(load_path) and os.path.getsize(load_path) > 0: - try: - vertices = paddle.load(path=str(load_path)) - num_vertices = vertices.shape[0] - - if num_vertices > self.num_points: - indices = np.random.choice( - num_vertices, self.num_points, replace=False - ) - vertices = vertices.numpy()[indices] - vertices = paddle.to_tensor(vertices) - - return vertices - except (EOFError, RuntimeError, ValueError) as e: - raise Exception( - f"Error loading point cloud from {load_path}: {e}" - ) from e - - def __getitem__( - self, idx: int, apply_augmentations: bool = True - ) -> Tuple[Dict[str, np.ndarray], Dict[str, np.ndarray], Dict[str, np.ndarray],]: - """ - Retrieves a sample and its corresponding label from the dataset, with an option to apply augmentations. - - Args: - idx (int): Index of the sample to retrieve. - apply_augmentations (bool, optional): Whether to apply data augmentations. Defaults to True. - - Tuple[Dict[str, np.ndarray], Dict[str, np.ndarray], Dict[str, np.ndarray]]: - A tuple containing three dictionaries: - - The first dictionary contains the input data (point cloud) under the key specified by `self.input_keys[0]`. - - The second dictionary contains the label (Cd value) under the key specified by `self.label_keys[0]`. - - The third dictionary contains the weight (default is 1) under the key specified by `self.weight_keys[0]`. - """ - if paddle.is_tensor(x=idx): - idx = idx.tolist() - - if idx in self.cache: - return self.cache[idx] - - row = self.data_frame.iloc[idx] - design_id = row["Design"] - cd_value = row["Average Cd"].reshape([-1]) - if self.pointcloud_exist: - try: - vertices = self._load_point_cloud(design_id) - if vertices is None: - raise ValueError( - f"Point cloud for design {design_id} is not found or corrupted." - ) - except Exception as e: - raise ValueError( - f"Failed to load point cloud for design {design_id}: {e}" - ) - if apply_augmentations: - vertices = self.augmentation.translate_pointcloud(vertices.numpy()) - vertices = self.augmentation.jitter_pointcloud(vertices) - if self.transform: - vertices = self.transform(vertices) - - self.cache[idx] = ( - {self.input_keys[0]: vertices}, - {self.label_keys[0]: cd_value}, - {self.weight_keys[0]: np.array(1, dtype=np.float32)}, - ) - - return ( - {self.input_keys[0]: vertices}, - {self.label_keys[0]: cd_value}, - {self.weight_keys[0]: np.array(1, dtype=np.float32)}, - ) diff --git a/examples/smc_reac/ppsci/data/dataset/drivaernetplusplus_dataset.py b/examples/smc_reac/ppsci/data/dataset/drivaernetplusplus_dataset.py deleted file mode 100644 index 3194ece235..0000000000 --- a/examples/smc_reac/ppsci/data/dataset/drivaernetplusplus_dataset.py +++ /dev/null @@ -1,321 +0,0 @@ -# Copyright (c) 2024 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. - -""" -@author: Mohamed Elrefaie, mohamed.elrefaie@mit.edu mohamed.elrefaie@tum.de - -This module is part of the research presented in the paper: -"DrivAerNet++: A Large-Scale Multimodal Car Dataset with Computational Fluid Dynamics Simulations and Deep Learning Benchmarks". - -The module defines two Paddle Datasets for loading and transforming 3D car models from the DrivAerNet++ dataset: -1. DrivAerNetPlusPlusDataset: Handles point cloud data, allowing loading, transforming, and augmenting 3D car models. -""" - - -from __future__ import annotations - -import logging -import os -from typing import Callable -from typing import Dict -from typing import Optional -from typing import Tuple - -import numpy as np -import paddle -import pandas as pd - - -class DataAugmentation: - """ - Class encapsulating various data augmentation techniques for point clouds. - """ - - @staticmethod - def translate_pointcloud( - pointcloud: np.ndarray, - translation_range: Tuple[float, float] = (2.0 / 3.0, 3.0 / 2.0), - ) -> np.ndarray: - """ - Translates the pointcloud by a random factor within a given range. - - Args: - pointcloud: The input point cloud as a np.ndarray. - translation_range: A tuple specifying the range for translation factors. - - Returns: - Translated point cloud as a np.ndarray. - """ - xyz1 = np.random.uniform( - low=translation_range[0], high=translation_range[1], size=[3] - ) - xyz2 = np.random.uniform(low=-0.2, high=0.2, size=[3]) - translated_pointcloud = np.add(np.multiply(pointcloud, xyz1), xyz2).astype( - "float32" - ) - return paddle.to_tensor(data=translated_pointcloud, dtype="float32") - - @staticmethod - def jitter_pointcloud( - pointcloud: np.ndarray, sigma: float = 0.01, clip: float = 0.02 - ) -> np.ndarray: - """ - Adds Gaussian noise to the pointcloud. - - Args: - pointcloud: The input point cloud as a np.ndarray. - sigma: Standard deviation of the Gaussian noise. - clip: Maximum absolute value for noise. - - Returns: - Jittered point cloud as a np.ndarray. - """ - N, C = tuple(pointcloud.shape) - jittered_pointcloud = pointcloud + paddle.clip( - x=sigma * paddle.randn(shape=[N, C]), min=-clip, max=clip - ) - return jittered_pointcloud - - @staticmethod - def drop_points(pointcloud: np.ndarray, drop_rate: float = 0.1) -> np.ndarray: - """ - Randomly removes points from the point cloud based on the drop rate. - - Args: - pointcloud: The input point cloud as a np.ndarray. - drop_rate: The percentage of points to be randomly dropped. - - Returns: - The point cloud with points dropped as a np.ndarray. - """ - num_drop = int(drop_rate * pointcloud.shape[0]) - drop_indices = np.random.choice(pointcloud.shape[0], num_drop, replace=False) - keep_indices = np.setdiff1d(np.arange(pointcloud.shape[0]), drop_indices) - dropped_pointcloud = pointcloud[keep_indices, :] - return dropped_pointcloud - - -class DrivAerNetPlusPlusDataset(paddle.io.Dataset): - """ - Paddle Dataset class for the DrivAerNet dataset, handling loading, transforming, and augmenting 3D car models. - - This dataset is designed for tasks involving aerodynamic simulations and deep learning models, - specifically for predicting aerodynamic coefficients (e.g., Cd values) from 3D car models. - - Args: - input_keys (Tuple[str, ...]): Tuple of strings specifying the input keys. - These keys correspond to the features extracted from the dataset, - typically the 3D vertices of car models. - Example: ("vertices",) - label_keys (Tuple[str, ...]): Tuple of strings specifying the label keys. - These keys correspond to the ground-truth labels, such as aerodynamic - coefficients (e.g., Cd values). - Example: ("cd_value",) - weight_keys (Tuple[str, ...]): Tuple of strings specifying the weight keys. - These keys represent optional weighting factors used during model training - to handle class imbalance or sample importance. - Example: ("weight_keys",) - subset_dir (str): Path to the directory containing subsets of the dataset. - This directory is used to divide the dataset into different subsets - (e.g., train, validation, test) based on provided IDs. - ids_file (str): Path to the file containing the list of IDs for the subset. - The file specifies which models belong to the current subset (e.g., training IDs). - root_dir (str): Root directory containing the 3D STL files of car models. - Each 3D model is expected to be stored in a file named according to its ID. - csv_file (str): Path to the CSV file containing metadata for the car models. - The CSV file includes information such as aerodynamic coefficients, - and may also map model IDs to specific attributes. - num_points (int): Number of points to sample or pad each 3D point cloud to. - If the model has more points than `num_points`, it will be subsampled. - If it has fewer points, zero-padding will be applied. - transform (Optional[Callable]): Optional transformation function applied to each sample. - This can include augmentations like scaling, rotation, or jittering. - pointcloud_exist (bool): Whether the point clouds are pre-processed and saved as `.pt` files. - If `True`, the dataset will directly load the pre-saved point clouds - instead of generating them from STL files. - - Examples: - >>> import ppsci - >>> dataset = ppsci.data.dataset.DrivAerNetPlusPlusDataset( - ... input_keys=("vertices",), - ... label_keys=("cd_value",), - ... weight_keys=("weight_keys",), - ... subset_dir="/path/to/subset_dir", - ... ids_file="train_ids.txt", - ... root_dir="/path/to/DrivAerNetPlusPlusDataset", - ... csv_file="/path/to/aero_metadata.csv", - ... num_points=1024, - ... transform=None, - ... ) # doctest: +SKIP - """ - - def __init__( - self, - input_keys: Tuple[str, ...], - label_keys: Tuple[str, ...], - weight_keys: Tuple[str, ...], - subset_dir: str, - ids_file: str, - root_dir: str, - csv_file: str, - num_points: int, - transform: Optional[Callable] = None, - pointcloud_exist: bool = True, - ): - super().__init__() - self.root_dir = root_dir - self.input_keys = input_keys - self.label_keys = label_keys - self.weight_keys = weight_keys - self.subset_dir = subset_dir - self.ids_file = ids_file - self.augmentation = DataAugmentation() - self.cache = {} - - try: - self.data_frame = pd.read_csv(csv_file) - except Exception as e: - logging.error(f"Failed to load CSV file: {csv_file}. Error: {e}") - raise - self.transform = transform - self.num_points = num_points - self.pointcloud_exist = pointcloud_exist - - try: - with open(os.path.join(self.subset_dir, self.ids_file), "r") as file: - subset_ids = file.read().split() - except FileNotFoundError as e: - raise FileNotFoundError(f"Error loading subset file {self.ids_file}: {e}") - - self.subset_indices = self.data_frame[ - self.data_frame["Design"].isin(subset_ids) - ].index.tolist() - self.data_frame = self.data_frame.loc[self.subset_indices].reset_index( - drop=True - ) - - def __len__(self) -> int: - """Returns the total number of samples in the dataset.""" - return len(self.data_frame) - - def min_max_normalize(self, data: np.ndarray) -> np.ndarray: - """ - Normalizes the data to the range [0, 1] based on min and max values. - """ - min_vals = data.min(axis=0, keepdim=True) - max_vals = data.max(axis=0, keepdim=True) - normalized_data = (data - min_vals) / (max_vals - min_vals) - return normalized_data - - def _sample_or_pad_vertices( - self, vertices: paddle.Tensor, num_points: int - ) -> paddle.Tensor: - """ - Subsamples or pads the vertices of the model to a fixed number of points. - - Args: - vertices: The vertices of the 3D model as a paddle.Tensor. - num_points: The desired number of points for the model. - - Returns: - The vertices standardized to the specified number of points. - """ - num_vertices = vertices.shape[0] - if num_vertices > num_points: - indices = np.random.choice(num_vertices, num_points, replace=False) - vertices = vertices[indices] - elif num_vertices < num_points: - padding = paddle.zeros( - shape=(num_points - num_vertices, 3), dtype="float32" - ) - vertices = paddle.concat(x=(vertices, padding), axis=0) - return vertices - - def _load_point_cloud(self, design_id: str): - load_path = os.path.join(self.root_dir, f"{design_id}.paddle_tensor") - if os.path.exists(load_path) and os.path.getsize(load_path) > 0: - try: - vertices = paddle.load(path=str(load_path)) - except (EOFError, RuntimeError, ValueError) as e: - raise Exception( - f"Error loading point cloud from {load_path}: {e}" - ) from e - num_vertices = vertices.shape[0] - - if num_vertices > self.num_points: - indices = np.random.choice(num_vertices, self.num_points, replace=False) - vertices = vertices.numpy()[indices] - - return vertices - - def __getitem__( - self, idx: int, apply_augmentations: bool = True - ) -> Tuple[Dict[str, np.ndarray], Dict[str, np.ndarray], Dict[str, np.ndarray]]: - """ - Retrieves a sample and its corresponding label from the dataset, with an option to apply augmentations. - - Args: - idx (int): Index of the sample to retrieve. - apply_augmentations (bool, optional): Whether to apply data augmentations. Defaults to True. - - Returns: - Tuple[Dict[str, np.ndarray], Dict[str, np.ndarray], Dict[str, np.ndarray]]: - A tuple containing three dictionaries: - - The first dictionary contains the input data (point cloud) under the key specified by `self.input_keys[0]`. - - The second dictionary contains the label (Cd value) under the key specified by `self.label_keys[0]`. - - The third dictionary contains the weight (default is 1) under the key specified by `self.weight_keys[0]`. - """ - if paddle.is_tensor(idx): - idx = idx.tolist() - - if idx in self.cache: - return self.cache[idx] - - row = self.data_frame.iloc[idx] - design_id = row["Design"] - cd_value = row["Average Cd"] - if self.pointcloud_exist: - try: - vertices = self._load_point_cloud(design_id) - if vertices is None: - raise ValueError( - f"Point cloud for design {design_id} is not found or corrupted." - ) - except Exception as e: - raise ValueError( - f"Failed to load point cloud for design {design_id}: {e}" - ) - - if apply_augmentations: - vertices = self.augmentation.translate_pointcloud(vertices.numpy()) - vertices = self.augmentation.jitter_pointcloud(vertices) - - if self.transform: - vertices = self.transform(vertices) - - vertices = self.min_max_normalize(vertices) - - cd_value = np.array(float(cd_value), dtype=np.float32).reshape([-1]) - - self.cache[idx] = ( - {self.input_keys[0]: vertices}, - {self.label_keys[0]: cd_value}, - {self.weight_keys[0]: np.array(1, dtype=np.float32)}, - ) - - return ( - {self.input_keys[0]: vertices}, - {self.label_keys[0]: cd_value}, - {self.weight_keys[0]: np.array(1, dtype=np.float32)}, - ) diff --git a/examples/smc_reac/ppsci/data/dataset/enso_dataset.py b/examples/smc_reac/ppsci/data/dataset/enso_dataset.py deleted file mode 100644 index 601fcec413..0000000000 --- a/examples/smc_reac/ppsci/data/dataset/enso_dataset.py +++ /dev/null @@ -1,405 +0,0 @@ -# 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. - -from __future__ import annotations - -import importlib -from pathlib import Path -from typing import Dict -from typing import Optional -from typing import Tuple - -import numpy as np -from paddle import io - -NINO_WINDOW_T = 3 # Nino index is the sliding average over sst, window size is 3 -CMIP6_SST_MAX = 10.198975563049316 -CMIP6_SST_MIN = -16.549121856689453 -CMIP5_SST_MAX = 8.991744995117188 -CMIP5_SST_MIN = -9.33076286315918 -CMIP6_NINO_MAX = 4.138188362121582 -CMIP6_NINO_MIN = -3.5832221508026123 -CMIP5_NINO_MAX = 3.8253555297851562 -CMIP5_NINO_MIN = -2.691682815551758 -SST_MAX = max(CMIP6_SST_MAX, CMIP5_SST_MAX) -SST_MIN = min(CMIP6_SST_MIN, CMIP5_SST_MIN) - - -def scale_sst(sst): - return (sst - SST_MIN) / (SST_MAX - SST_MIN) - - -def scale_back_sst(sst): - return (SST_MAX - SST_MIN) * sst + SST_MIN - - -def prepare_inputs_targets( - len_time, input_gap, input_length, pred_shift, pred_length, samples_gap -): - """Prepares the input and target indices for training. - - Args: - len_time (int): The total number of time steps in the dataset. - input_gap (int): Time gaps between two consecutive input frames. - input_length (int): The number of input frames. - pred_shift (int): The lead_time of the last target to be predicted. - pred_length (int): The number of frames to be predicted. - samples_gap (int): Stride of seq sampling. - """ - - if pred_shift < pred_length: - raise ValueError("pred_shift should be small than pred_length") - input_span = input_gap * (input_length - 1) + 1 - pred_gap = pred_shift // pred_length - input_ind = np.arange(0, input_span, input_gap) - target_ind = np.arange(0, pred_shift, pred_gap) + input_span + pred_gap - 1 - ind = np.concatenate([input_ind, target_ind]).reshape(1, input_length + pred_length) - max_n_sample = len_time - (input_span + pred_shift - 1) - ind = ind + np.arange(max_n_sample)[:, np.newaxis] @ np.ones( - (1, input_length + pred_length), dtype=int - ) - return ind[::samples_gap] - - -def fold(data, size=36, stride=12): - """Inverse of unfold/sliding window operation - only applicable to the case where the size of the sliding windows is n*stride - - Args: - data (tuple[int,...]): The input data.(N, size, *). - size (int, optional): The size of a single datum.The Defaults to 36. - stride (int, optional): The step.Defaults to 12. - - Returns: - outdata (np.array): (N_, *).N/size is the number/width of sliding blocks - """ - if size % stride != 0: - raise ValueError("size modulo stride should be zero") - times = size // stride - remain = (data.shape[0] - 1) % times - if remain > 0: - ls = list(data[::times]) + [data[-1, -(remain * stride) :]] - outdata = np.concatenate(ls, axis=0) # (36*(151//3+1)+remain*stride, *, 15) - else: - outdata = np.concatenate(data[::times], axis=0) # (36*(151/3+1), *, 15) - assert ( - outdata.shape[0] == size * ((data.shape[0] - 1) // times + 1) + remain * stride - ) - return outdata - - -def data_transform(data, num_years_per_model): - """The transform of the input data. - - Args: - data (Tuple[list,...]): The input data.Shape of (N, 36, *). - num_years_per_model (int): The number of years associated with each model.151/140. - """ - length = data.shape[0] - assert length % num_years_per_model == 0 - num_models = length // num_years_per_model - outdata = np.stack( - np.split(data, length / num_years_per_model, axis=0), axis=-1 - ) # (151, 36, *, 15) - # cmip6sst outdata.shape = (151, 36, 24, 48, 15) = (year, month, lat, lon, model) - # cmip5sst outdata.shape = (140, 36, 24, 48, 17) - # cmip6nino outdata.shape = (151, 36, 15) - # cmip5nino outdata.shape = (140, 36, 17) - outdata = fold(outdata, size=36, stride=12) - # cmip6sst outdata.shape = (1836, 24, 48, 15), 1836 == 151 * 12 + 24 - # cmip5sst outdata.shape = (1704, 24, 48, 17) - # cmip6nino outdata.shape = (1836, 15) - # cmip5nino outdata.shape = (1704, 17) - - # check output data - assert outdata.shape[-1] == num_models - assert not np.any(np.isnan(outdata)) - return outdata - - -def read_raw_data(ds_dir, out_dir=None): - """Read and process raw cmip data from CMIP_train.nc and CMIP_label.nc - - Args: - ds_dir (str): The path of the dataset. - out_dir (str): The path of output. Defaults to None. - """ - - import xarray as xr - - train_cmip = xr.open_dataset(Path(ds_dir) / "CMIP_train.nc").transpose( - "year", "month", "lat", "lon" - ) - label_cmip = xr.open_dataset(Path(ds_dir) / "CMIP_label.nc").transpose( - "year", "month" - ) - # train_cmip.sst.values.shape = (4645, 36, 24, 48) - - # select longitudes - lon = train_cmip.lon.values - lon = lon[np.logical_and(lon >= 95, lon <= 330)] - train_cmip = train_cmip.sel(lon=lon) - - cmip6sst = data_transform( - data=train_cmip.sst.values[:2265], num_years_per_model=151 - ) - cmip5sst = data_transform( - data=train_cmip.sst.values[2265:], num_years_per_model=140 - ) - cmip6nino = data_transform( - data=label_cmip.nino.values[:2265], num_years_per_model=151 - ) - cmip5nino = data_transform( - data=label_cmip.nino.values[2265:], num_years_per_model=140 - ) - - # cmip6sst.shape = (1836, 24, 48, 15) - # cmip5sst.shape = (1704, 24, 48, 17) - assert len(cmip6sst.shape) == 4 - assert len(cmip5sst.shape) == 4 - assert len(cmip6nino.shape) == 2 - assert len(cmip5nino.shape) == 2 - # store processed data for faster data access - if out_dir is not None: - ds_cmip6 = xr.Dataset( - { - "sst": (["month", "lat", "lon", "model"], cmip6sst), - "nino": (["month", "model"], cmip6nino), - }, - coords={ - "month": np.repeat( - np.arange(1, 13)[None], cmip6nino.shape[0] // 12, axis=0 - ).flatten(), - "lat": train_cmip.lat.values, - "lon": train_cmip.lon.values, - "model": np.arange(15) + 1, - }, - ) - ds_cmip6.to_netcdf(Path(out_dir) / "cmip6.nc") - ds_cmip5 = xr.Dataset( - { - "sst": (["month", "lat", "lon", "model"], cmip5sst), - "nino": (["month", "model"], cmip5nino), - }, - coords={ - "month": np.repeat( - np.arange(1, 13)[None], cmip5nino.shape[0] // 12, axis=0 - ).flatten(), - "lat": train_cmip.lat.values, - "lon": train_cmip.lon.values, - "model": np.arange(17) + 1, - }, - ) - ds_cmip5.to_netcdf(Path(out_dir) / "cmip5.nc") - train_cmip.close() - label_cmip.close() - return cmip6sst, cmip5sst, cmip6nino, cmip5nino - - -def cat_over_last_dim(data): - """Treat different models (15 from CMIP6, 17 from CMIP5) as batch_size - e.g., cmip6sst.shape = (178, 38, 24, 48, 15), converted_cmip6sst.shape = (2670, 38, 24, 48) - e.g., cmip5sst.shape = (165, 38, 24, 48, 15), converted_cmip6sst.shape = (2475, 38, 24, 48) - """ - - return np.concatenate(np.moveaxis(data, -1, 0), axis=0) - - -class ENSODataset(io.Dataset): - """The El Niño/Southern Oscillation dataset. - - Args: - input_keys (Tuple[str, ...]): Name of input keys, such as ("input",). - label_keys (Tuple[str, ...]): Name of label keys, such as ("output",). - data_dir (str): The directory of data. - weight_dict (Optional[Dict[str, Union[Callable, float]]]): Define the weight of each constraint variable. Defaults to None. - in_len (int, optional): The length of input data. Defaults to 12. - out_len (int, optional): The length of out data. Defaults to 26. - in_stride (int, optional): The stride of input data. Defaults to 1. - out_stride (int, optional): The stride of output data. Defaults to 1. - train_samples_gap (int, optional): The stride of sequence sampling during training. Defaults to 10. - e.g., samples_gap = 10, the first seq contains [0, 1, ..., T-1] frame indices, the second seq contains [10, 11, .., T+9] - eval_samples_gap (int, optional): The stride of sequence sampling during eval. Defaults to 11. - normalize_sst (bool, optional): Whether to use normalization. Defaults to True. - batch_size (int, optional): Batch size. Defaults to 1. - num_workers (int, optional): The num of workers. Defaults to 1. - training (str, optional): Training pathse. Defaults to "train". - """ - - # Whether support batch indexing for speeding up fetching process. - batch_index: bool = False - - def __init__( - self, - input_keys: Tuple[str, ...], - label_keys: Tuple[str, ...], - data_dir: str, - weight_dict: Optional[Dict[str, float]] = None, - in_len=12, - out_len=26, - in_stride=1, - out_stride=1, - train_samples_gap=10, - eval_samples_gap=11, - normalize_sst=True, - # datamodule_only - batch_size=1, - num_workers=1, - training="train", - ): - super(ENSODataset, self).__init__() - if importlib.util.find_spec("xarray") is None: - raise ModuleNotFoundError( - "To use RadarDataset, please install 'xarray' with: `pip install " - "xarray` first." - ) - self.input_keys = input_keys - self.label_keys = label_keys - self.data_dir = data_dir - self.weight_dict = {} if weight_dict is None else weight_dict - if weight_dict is not None: - self.weight_dict = {key: 1.0 for key in self.label_keys} - self.weight_dict.update(weight_dict) - - self.in_len = in_len - self.out_len = out_len - self.in_stride = in_stride - self.out_stride = out_stride - self.train_samples_gap = train_samples_gap - self.eval_samples_gap = eval_samples_gap - self.normalize_sst = normalize_sst - # datamodule_only - self.batch_size = batch_size - if num_workers != 1: - raise ValueError( - "Current implementation does not support `num_workers != 1`!" - ) - self.num_workers = num_workers - self.training = training - - # pre-data - cmip6sst, cmip5sst, cmip6nino, cmip5nino = read_raw_data(self.data_dir) - # TODO: more flexible train/val/test split - self.sst_train = [cmip6sst, cmip5sst[..., :-2]] - self.nino_train = [cmip6nino, cmip5nino[..., :-2]] - self.sst_eval = [cmip5sst[..., -2:-1]] - self.nino_eval = [cmip5nino[..., -2:-1]] - self.sst_test = [cmip5sst[..., -1:]] - self.nino_test = [cmip5nino[..., -1:]] - - self.sst, self.target_nino = self.create_data() - - def create_data( - self, - ): - if self.training == "train": - sst_cmip6 = self.sst_train[0] - nino_cmip6 = self.nino_train[0] - sst_cmip5 = self.sst_train[1] - nino_cmip5 = self.nino_train[1] - samples_gap = self.train_samples_gap - elif self.training == "eval": - sst_cmip6 = None - nino_cmip6 = None - sst_cmip5 = self.sst_eval[0] - nino_cmip5 = self.nino_eval[0] - samples_gap = self.eval_samples_gap - elif self.training == "test": - sst_cmip6 = None - nino_cmip6 = None - sst_cmip5 = self.sst_test[0] - nino_cmip5 = self.nino_test[0] - samples_gap = self.eval_samples_gap - - # cmip6 (N, *, 15) - # cmip5 (N, *, 17) - sst = [] - target_nino = [] - - nino_idx_slice = slice( - self.in_len, self.in_len + self.out_len - NINO_WINDOW_T + 1 - ) # e.g., 12:36 - if sst_cmip6 is not None: - assert len(sst_cmip6.shape) == 4 - assert len(nino_cmip6.shape) == 2 - idx_sst = prepare_inputs_targets( - len_time=sst_cmip6.shape[0], - input_length=self.in_len, - input_gap=self.in_stride, - pred_shift=self.out_len * self.out_stride, - pred_length=self.out_len, - samples_gap=samples_gap, - ) - - sst.append(cat_over_last_dim(sst_cmip6[idx_sst])) - target_nino.append( - cat_over_last_dim(nino_cmip6[idx_sst[:, nino_idx_slice]]) - ) - if sst_cmip5 is not None: - assert len(sst_cmip5.shape) == 4 - assert len(nino_cmip5.shape) == 2 - idx_sst = prepare_inputs_targets( - len_time=sst_cmip5.shape[0], - input_length=self.in_len, - input_gap=self.in_stride, - pred_shift=self.out_len * self.out_stride, - pred_length=self.out_len, - samples_gap=samples_gap, - ) - sst.append(cat_over_last_dim(sst_cmip5[idx_sst])) - target_nino.append( - cat_over_last_dim(nino_cmip5[idx_sst[:, nino_idx_slice]]) - ) - - # sst data containing both the input and target - self.sst = np.concatenate(sst, axis=0) # (N, in_len+out_len, lat, lon) - if self.normalize_sst: - self.sst = scale_sst(self.sst) - # nino data containing the target only - self.target_nino = np.concatenate( - target_nino, axis=0 - ) # (N, out_len+NINO_WINDOW_T-1) - assert self.sst.shape[0] == self.target_nino.shape[0] - assert self.sst.shape[1] == self.in_len + self.out_len - assert self.target_nino.shape[1] == self.out_len - NINO_WINDOW_T + 1 - return self.sst, self.target_nino - - def get_datashape(self): - return {"sst": self.sst.shape, "nino target": self.target_nino.shape} - - def __len__(self): - return self.sst.shape[0] - - def __getitem__(self, idx): - sst_data = self.sst[idx].astype("float32") - sst_data = sst_data[..., np.newaxis] - in_seq = sst_data[: self.in_len, ...] # ( in_len, lat, lon, 1) - target_seq = sst_data[self.in_len :, ...] # ( in_len, lat, lon, 1) - weight_item = self.weight_dict - - if self.training == "train": - input_item = {self.input_keys[0]: in_seq} - label_item = { - self.label_keys[0]: target_seq, - } - - return input_item, label_item, weight_item - else: - input_item = {self.input_keys[0]: in_seq} - label_item = { - self.label_keys[0]: target_seq, - self.label_keys[1]: self.target_nino[idx], - } - - return input_item, label_item, weight_item diff --git a/examples/smc_reac/ppsci/data/dataset/era5_dataset.py b/examples/smc_reac/ppsci/data/dataset/era5_dataset.py deleted file mode 100644 index 75cf89c754..0000000000 --- a/examples/smc_reac/ppsci/data/dataset/era5_dataset.py +++ /dev/null @@ -1,249 +0,0 @@ -# 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. - -from __future__ import annotations - -import glob -from typing import Dict -from typing import Optional -from typing import Tuple - -try: - import h5py -except ModuleNotFoundError: - pass - -import numpy as np -import paddle -from paddle import io -from paddle import vision - - -class ERA5Dataset(io.Dataset): - """Class for ERA5 dataset. - - Args: - file_path (str): Data set path. - input_keys (Tuple[str, ...]): Input keys, such as ("input",). - label_keys (Tuple[str, ...]): Output keys, such as ("output",). - precip_file_path (Optional[str]): Precipitation data set path. Defaults to None. - weight_dict (Optional[Dict[str, float]]): Weight dictionary. Defaults to None. - vars_channel (Optional[Tuple[int, ...]]): The variable channel index in ERA5 dataset. Defaults to None. - num_label_timestamps (int, optional): Number of timestamp of label. Defaults to 1. - transforms (Optional[vision.Compose]): Compose object contains sample wise - transform(s). Defaults to None. - training (bool, optional): Whether in train mode. Defaults to True. - stride (int, optional): Stride of sampling data. Defaults to 1. - - Examples: - >>> import ppsci - >>> dataset = ppsci.data.dataset.ERA5Dataset( - ... "file_path": "/path/to/ERA5Dataset", - ... "input_keys": ("input",), - ... "label_keys": ("output",), - ... ) # doctest: +SKIP - """ - - # Whether support batch indexing for speeding up fetching process. - batch_index: bool = False - - def __init__( - self, - file_path: str, - input_keys: Tuple[str, ...], - label_keys: Tuple[str, ...], - precip_file_path: Optional[str] = None, - weight_dict: Optional[Dict[str, float]] = None, - vars_channel: Optional[Tuple[int, ...]] = None, - num_label_timestamps: int = 1, - transforms: Optional[vision.Compose] = None, - training: bool = True, - stride: int = 1, - ): - super().__init__() - self.file_path = file_path - self.input_keys = input_keys - self.label_keys = label_keys - self.precip_file_path = precip_file_path - - self.weight_dict = {} if weight_dict is None else weight_dict - if weight_dict is not None: - self.weight_dict = {key: 1.0 for key in self.label_keys} - self.weight_dict.update(weight_dict) - - self.vars_channel = list(range(20)) if vars_channel is None else vars_channel - self.num_label_timestamps = num_label_timestamps - self.transforms = transforms - self.training = training - self.stride = stride - - self.files = self.read_data(file_path) - self.n_years = len(self.files) - self.num_samples_per_year = self.files[0].shape[0] - self.num_samples = self.n_years * self.num_samples_per_year - if self.precip_file_path is not None: - self.precip_files = self.read_data(precip_file_path, "tp") - - def read_data(self, path: str, var="fields"): - paths = [path] if path.endswith(".h5") else glob.glob(path + "/*.h5") - paths.sort() - files = [] - for path_ in paths: - _file = h5py.File(path_, "r") - files.append(_file[var]) - return files - - def __len__(self): - return self.num_samples // self.stride - - def __getitem__(self, global_idx): - global_idx *= self.stride - year_idx = global_idx // self.num_samples_per_year - local_idx = global_idx % self.num_samples_per_year - step = 0 if local_idx >= self.num_samples_per_year - 1 else 1 - - if self.num_label_timestamps > 1: - if local_idx >= self.num_samples_per_year - self.num_label_timestamps: - local_idx = self.num_samples_per_year - self.num_label_timestamps - 1 - - input_file = self.files[year_idx] - label_file = ( - self.precip_files[year_idx] - if self.precip_file_path is not None - else input_file - ) - if self.precip_file_path is not None and year_idx == 0 and self.training: - # first year has 2 missing samples in precip (they are first two time points) - lim = self.num_samples_per_year - 2 - local_idx = local_idx % lim - step = 0 if local_idx >= lim - 1 else 1 - input_idx = local_idx + 2 - label_idx = local_idx + step - else: - input_idx, label_idx = local_idx, local_idx + step - - input_item = {self.input_keys[0]: input_file[input_idx, self.vars_channel]} - - label_item = {} - for i in range(self.num_label_timestamps): - if self.precip_file_path is not None: - label_item[self.label_keys[i]] = np.expand_dims( - label_file[label_idx + i], 0 - ) - else: - label_item[self.label_keys[i]] = label_file[ - label_idx + i, self.vars_channel - ] - - weight_shape = [1] * len(next(iter(label_item.values())).shape) - weight_item = { - key: np.full(weight_shape, value, paddle.get_default_dtype()) - for key, value in self.weight_dict.items() - } - - if self.transforms is not None: - input_item, label_item, weight_item = self.transforms( - input_item, label_item, weight_item - ) - - return input_item, label_item, weight_item - - -class ERA5SampledDataset(io.Dataset): - """Class for ERA5 sampled dataset. - - Args: - file_path (str): Data set path. - input_keys (Tuple[str, ...]): Input keys, such as ("input",). - label_keys (Tuple[str, ...]): Output keys, such as ("output",). - weight_dict (Optional[Dict[str, float]]): Weight dictionary. Defaults to None. - transforms (Optional[vision.Compose]): Compose object contains sample wise - transform(s). Defaults to None. - - Examples: - >>> import ppsci - >>> dataset = ppsci.data.dataset.ERA5SampledDataset( - ... "file_path": "/path/to/ERA5SampledDataset", - ... "input_keys": ("input",), - ... "label_keys": ("output",), - ... ) # doctest: +SKIP - >>> # get the length of the dataset - >>> dataset_size = len(dataset) # doctest: +SKIP - >>> # get the first sample of the data - >>> first_sample = dataset[0] # doctest: +SKIP - >>> print("First sample:", first_sample) # doctest: +SKIP - """ - - def __init__( - self, - file_path: str, - input_keys: Tuple[str, ...], - label_keys: Tuple[str, ...], - weight_dict: Optional[Dict[str, float]] = None, - transforms: Optional[vision.Compose] = None, - ): - super().__init__() - self.file_path = file_path - self.input_keys = input_keys - self.label_keys = label_keys - - self.weight_dict = {} if weight_dict is None else weight_dict - if weight_dict is not None: - self.weight_dict = {key: 1.0 for key in self.label_keys} - self.weight_dict.update(weight_dict) - - self.transforms = transforms - - self.files = self.read_data(file_path) - self.num_samples = len(self.files) - - def read_data(self, path: str): - paths = glob.glob(path + "/*.h5") - paths.sort() - files = [] - for _path in paths: - _file = h5py.File(_path, "r") - files.append(_file) - return files - - def __len__(self): - return self.num_samples - - def __getitem__(self, global_idx): - _file = self.files[global_idx] - - input_item = {} - for key in _file["input_dict"]: - input_item[key] = np.asarray( - _file["input_dict"][key], paddle.get_default_dtype() - ) - - label_item = {} - for key in _file["label_dict"]: - label_item[key] = np.asarray( - _file["label_dict"][key], paddle.get_default_dtype() - ) - - weight_shape = [1] * len(next(iter(label_item.values())).shape) - weight_item = { - key: np.full(weight_shape, value, paddle.get_default_dtype()) - for key, value in self.weight_dict.items() - } - - if self.transforms is not None: - input_item, label_item, weight_item = self.transforms( - input_item, label_item, weight_item - ) - - return input_item, label_item, weight_item diff --git a/examples/smc_reac/ppsci/data/dataset/ext_moe_enso_dataset.py b/examples/smc_reac/ppsci/data/dataset/ext_moe_enso_dataset.py deleted file mode 100644 index 5286b4bfe2..0000000000 --- a/examples/smc_reac/ppsci/data/dataset/ext_moe_enso_dataset.py +++ /dev/null @@ -1,406 +0,0 @@ -# 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. - -from __future__ import annotations - -import importlib -from pathlib import Path -from typing import Dict -from typing import Optional -from typing import Tuple - -import numpy as np -from paddle import io - -NINO_WINDOW_T = 3 # Nino index is the sliding average over sst, window size is 3 -CMIP6_SST_MAX = 10.198975563049316 -CMIP6_SST_MIN = -16.549121856689453 -CMIP5_SST_MAX = 8.991744995117188 -CMIP5_SST_MIN = -9.33076286315918 -CMIP6_NINO_MAX = 4.138188362121582 -CMIP6_NINO_MIN = -3.5832221508026123 -CMIP5_NINO_MAX = 3.8253555297851562 -CMIP5_NINO_MIN = -2.691682815551758 -SST_MAX = max(CMIP6_SST_MAX, CMIP5_SST_MAX) -SST_MIN = min(CMIP6_SST_MIN, CMIP5_SST_MIN) - - -def scale_sst(sst): - return (sst - SST_MIN) / (SST_MAX - SST_MIN) - - -def scale_back_sst(sst): - return (SST_MAX - SST_MIN) * sst + SST_MIN - - -def prepare_inputs_targets( - len_time, input_gap, input_length, pred_shift, pred_length, samples_gap -): - """Prepares the input and target indices for training. - - Args: - len_time (int): The total number of time steps in the dataset. - input_gap (int): Time gaps between two consecutive input frames. - input_length (int): The number of input frames. - pred_shift (int): The lead_time of the last target to be predicted. - pred_length (int): The number of frames to be predicted. - samples_gap (int): Stride of seq sampling. - """ - - if pred_shift < pred_length: - raise ValueError("Pred_shift should be small than pred_length") - input_span = input_gap * (input_length - 1) + 1 - pred_gap = pred_shift // pred_length - input_ind = np.arange(0, input_span, input_gap) - target_ind = np.arange(0, pred_shift, pred_gap) + input_span + pred_gap - 1 - ind = np.concatenate([input_ind, target_ind]).reshape(1, input_length + pred_length) - max_n_sample = len_time - (input_span + pred_shift - 1) - ind = ind + np.arange(max_n_sample)[:, np.newaxis] @ np.ones( - (1, input_length + pred_length), dtype=int - ) - return ind[::samples_gap] - - -def fold(data, size=36, stride=12): - """inverse of unfold/sliding window operation - only applicable to the case where the size of the sliding windows is n*stride - - Args: - data (tuple[int,...]): The input data.(N, size, *). - size (int, optional): The size of a single datum.The Defaults to 36. - stride (int, optional): The step.Defaults to 12. - - Returns: - outdata (np.ndarray): (N_, *).N/size is the number/width of sliding blocks - """ - - if size % stride != 0: - raise ValueError("size modulo stride should be zero") - times = size // stride - remain = (data.shape[0] - 1) % times - if remain > 0: - ls = list(data[::times]) + [data[-1, -(remain * stride) :]] - outdata = np.concatenate(ls, axis=0) # (36*(151//3+1)+remain*stride, *, 15) - else: - outdata = np.concatenate(data[::times], axis=0) # (36*(151/3+1), *, 15) - assert ( - outdata.shape[0] == size * ((data.shape[0] - 1) // times + 1) + remain * stride - ) - return outdata - - -def data_transform(data, num_years_per_model): - """The transform of the input data. - - Args: - data (Tuple[list,...]): The input data.Shape of (N, 36, *). - num_years_per_model (int): The number of years associated with each model.151/140. - """ - - length = data.shape[0] - assert length % num_years_per_model == 0 - num_models = length // num_years_per_model - outdata = np.stack( - np.split(data, length / num_years_per_model, axis=0), axis=-1 - ) # (151, 36, *, 15) - # cmip6sst outdata.shape = (151, 36, 24, 48, 15) = (year, month, lat, lon, model) - # cmip5sst outdata.shape = (140, 36, 24, 48, 17) - # cmip6nino outdata.shape = (151, 36, 15) - # cmip5nino outdata.shape = (140, 36, 17) - outdata = fold(outdata, size=36, stride=12) - # cmip6sst outdata.shape = (1836, 24, 48, 15), 1836 == 151 * 12 + 24 - # cmip5sst outdata.shape = (1704, 24, 48, 17) - # cmip6nino outdata.shape = (1836, 15) - # cmip5nino outdata.shape = (1704, 17) - - # check output data - assert outdata.shape[-1] == num_models - assert not np.any(np.isnan(outdata)) - return outdata - - -def read_raw_data(ds_dir, out_dir=None): - """read and process raw cmip data from CMIP_train.nc and CMIP_label.nc - - Args: - ds_dir (str): The path of the dataset. - out_dir (str): The path of output. Defaults to None. - """ - import xarray as xr - - train_cmip = xr.open_dataset( - Path(ds_dir) / "CMIP_train.nc", engine="h5netcdf" - ).transpose("year", "month", "lat", "lon") - label_cmip = xr.open_dataset( - Path(ds_dir) / "CMIP_label.nc", engine="h5netcdf" - ).transpose("year", "month") - # train_cmip.sst.values.shape = (4645, 36, 24, 48) - - # select longitudes - lon = train_cmip.lon.values - lon = lon[np.logical_and(lon >= 95, lon <= 330)] - train_cmip = train_cmip.sel(lon=lon) - - cmip6sst = data_transform( - data=train_cmip.sst.values[:2265], num_years_per_model=151 - ) - cmip5sst = data_transform( - data=train_cmip.sst.values[2265:], num_years_per_model=140 - ) - cmip6nino = data_transform( - data=label_cmip.nino.values[:2265], num_years_per_model=151 - ) - cmip5nino = data_transform( - data=label_cmip.nino.values[2265:], num_years_per_model=140 - ) - - # cmip6sst.shape = (1836, 24, 48, 15) - # cmip5sst.shape = (1704, 24, 48, 17) - assert len(cmip6sst.shape) == 4 - assert len(cmip5sst.shape) == 4 - assert len(cmip6nino.shape) == 2 - assert len(cmip5nino.shape) == 2 - # store processed data for faster data access - if out_dir is not None: - ds_cmip6 = xr.Dataset( - { - "sst": (["month", "lat", "lon", "model"], cmip6sst), - "nino": (["month", "model"], cmip6nino), - }, - coords={ - "month": np.repeat( - np.arange(1, 13)[None], cmip6nino.shape[0] // 12, axis=0 - ).flatten(), - "lat": train_cmip.lat.values, - "lon": train_cmip.lon.values, - "model": np.arange(15) + 1, - }, - ) - ds_cmip6.to_netcdf(Path(out_dir) / "cmip6.nc") - ds_cmip5 = xr.Dataset( - { - "sst": (["month", "lat", "lon", "model"], cmip5sst), - "nino": (["month", "model"], cmip5nino), - }, - coords={ - "month": np.repeat( - np.arange(1, 13)[None], cmip5nino.shape[0] // 12, axis=0 - ).flatten(), - "lat": train_cmip.lat.values, - "lon": train_cmip.lon.values, - "model": np.arange(17) + 1, - }, - ) - ds_cmip5.to_netcdf(Path(out_dir) / "cmip5.nc") - train_cmip.close() - label_cmip.close() - return cmip6sst, cmip5sst, cmip6nino, cmip5nino - - -def cat_over_last_dim(data): - """treat different models (15 from CMIP6, 17 from CMIP5) as batch_size - e.g., cmip6sst.shape = (178, 38, 24, 48, 15), converted_cmip6sst.shape = (2670, 38, 24, 48) - e.g., cmip5sst.shape = (165, 38, 24, 48, 15), converted_cmip6sst.shape = (2475, 38, 24, 48) - """ - - return np.concatenate(np.moveaxis(data, -1, 0), axis=0) - - -class ExtMoEENSODataset(io.Dataset): - """The El Niño/Southern Oscillation dataset. - - Args: - input_keys (Tuple[str, ...]): Name of input keys, such as ("input",). - label_keys (Tuple[str, ...]): Name of label keys, such as ("output",). - data_dir (str): The directory of data. - weight_dict (Optional[Dict[str, Union[Callable, float]]]): Define the weight of each constraint variable. Defaults to None. - in_len (int, optional): The length of input data. Defaults to 12. - out_len (int, optional): The length of out data. Defaults to 26. - in_stride (int, optional): The stride of input data. Defaults to 1. - out_stride (int, optional): The stride of output data. Defaults to 1. - train_samples_gap (int, optional): The stride of sequence sampling during training. Defaults to 10. - e.g., samples_gap = 10, the first seq contains [0, 1, ..., T-1] frame indices, the second seq contains [10, 11, .., T+9] - eval_samples_gap (int, optional): The stride of sequence sampling during eval. Defaults to 11. - normalize_sst (bool, optional): Whether to use normalization. Defaults to True. - batch_size (int, optional): Batch size. Defaults to 1. - num_workers (int, optional): The num of workers. Defaults to 1. - training (str, optional): Training pathse. Defaults to "train". - """ - - # Whether support batch indexing for speeding up fetching process. - batch_index: bool = False - - def __init__( - self, - input_keys: Tuple[str, ...], - label_keys: Tuple[str, ...], - data_dir: str, - weight_dict: Optional[Dict[str, float]] = None, - in_len: int = 12, - out_len: int = 26, - in_stride: int = 1, - out_stride: int = 1, - train_samples_gap: int = 10, - eval_samples_gap: int = 11, - normalize_sst: bool = True, - batch_size: int = 1, - num_workers: int = 1, - training: str = "train", - ): - super(ExtMoEENSODataset, self).__init__() - if importlib.util.find_spec("xarray") is None: - raise ModuleNotFoundError( - "To use RadarDataset, please install 'xarray' with: `pip install " - "xarray` first." - ) - self.input_keys = input_keys - self.label_keys = label_keys - self.data_dir = data_dir - self.weight_dict = {} if weight_dict is None else weight_dict - if weight_dict is not None: - self.weight_dict = {key: 1.0 for key in self.label_keys} - self.weight_dict.update(weight_dict) - - self.in_len = in_len - self.out_len = out_len - self.in_stride = in_stride - self.out_stride = out_stride - self.train_samples_gap = train_samples_gap - self.eval_samples_gap = eval_samples_gap - self.normalize_sst = normalize_sst - # datamodule_only - self.batch_size = batch_size - if num_workers != 1: - raise ValueError( - "Current implementation does not support `num_workers != 1`!" - ) - self.num_workers = num_workers - self.training = training - - # pre-data - cmip6sst, cmip5sst, cmip6nino, cmip5nino = read_raw_data(self.data_dir) - # TODO: more flexible train/val/test split - self.sst_train = [cmip6sst, cmip5sst[..., :-2]] - self.nino_train = [cmip6nino, cmip5nino[..., :-2]] - self.sst_eval = [cmip5sst[..., -2:-1]] - self.nino_eval = [cmip5nino[..., -2:-1]] - self.sst_test = [cmip5sst[..., -1:]] - self.nino_test = [cmip5nino[..., -1:]] - - self.sst, self.target_nino = self.create_data() - - def create_data( - self, - ): - if self.training == "train": - sst_cmip6 = self.sst_train[0] - nino_cmip6 = self.nino_train[0] - sst_cmip5 = self.sst_train[1] - nino_cmip5 = self.nino_train[1] - samples_gap = self.train_samples_gap - elif self.training == "eval": - sst_cmip6 = None - nino_cmip6 = None - sst_cmip5 = self.sst_eval[0] - nino_cmip5 = self.nino_eval[0] - samples_gap = self.eval_samples_gap - elif self.training == "test": - sst_cmip6 = None - nino_cmip6 = None - sst_cmip5 = self.sst_test[0] - nino_cmip5 = self.nino_test[0] - samples_gap = self.eval_samples_gap - - # cmip6 (N, *, 15) - # cmip5 (N, *, 17) - sst = [] - target_nino = [] - - nino_idx_slice = slice( - self.in_len, self.in_len + self.out_len - NINO_WINDOW_T + 1 - ) # e.g., 12:36 - if sst_cmip6 is not None: - assert len(sst_cmip6.shape) == 4 - assert len(nino_cmip6.shape) == 2 - idx_sst = prepare_inputs_targets( - len_time=sst_cmip6.shape[0], - input_length=self.in_len, - input_gap=self.in_stride, - pred_shift=self.out_len * self.out_stride, - pred_length=self.out_len, - samples_gap=samples_gap, - ) - - sst.append(cat_over_last_dim(sst_cmip6[idx_sst])) - target_nino.append( - cat_over_last_dim(nino_cmip6[idx_sst[:, nino_idx_slice]]) - ) - if sst_cmip5 is not None: - assert len(sst_cmip5.shape) == 4 - assert len(nino_cmip5.shape) == 2 - idx_sst = prepare_inputs_targets( - len_time=sst_cmip5.shape[0], - input_length=self.in_len, - input_gap=self.in_stride, - pred_shift=self.out_len * self.out_stride, - pred_length=self.out_len, - samples_gap=samples_gap, - ) - sst.append(cat_over_last_dim(sst_cmip5[idx_sst])) - target_nino.append( - cat_over_last_dim(nino_cmip5[idx_sst[:, nino_idx_slice]]) - ) - - # sst data containing both the input and target - self.sst = np.concatenate(sst, axis=0) # (N, in_len+out_len, lat, lon) - if self.normalize_sst: - self.sst = scale_sst(self.sst) - # nino data containing the target only - self.target_nino = np.concatenate( - target_nino, axis=0 - ) # (N, out_len+NINO_WINDOW_T-1) - assert self.sst.shape[0] == self.target_nino.shape[0] - assert self.sst.shape[1] == self.in_len + self.out_len - assert self.target_nino.shape[1] == self.out_len - NINO_WINDOW_T + 1 - - return self.sst, self.target_nino - - def get_datashape(self): - return {"sst": self.sst.shape, "nino target": self.target_nino.shape} - - def __len__(self): - return self.sst.shape[0] - - def __getitem__(self, idx): - sst_data = self.sst[idx].astype("float32") - sst_data = sst_data[..., np.newaxis] - in_seq = sst_data[: self.in_len, ...] # ( in_len, lat, lon, 1) - target_seq = sst_data[self.in_len :, ...] # ( in_len, lat, lon, 1) - weight_item = self.weight_dict - - if self.training == "train": - input_item = {self.input_keys[0]: in_seq, "sst_target": target_seq} - label_item = { - self.label_keys[0]: target_seq, - } - - return input_item, label_item, weight_item - else: - input_item = {self.input_keys[0]: in_seq, "sst_target": target_seq} - label_item = { - self.label_keys[0]: target_seq, - self.label_keys[1]: self.target_nino[idx], - } - - return input_item, label_item, weight_item diff --git a/examples/smc_reac/ppsci/data/dataset/fwi_dataset.py b/examples/smc_reac/ppsci/data/dataset/fwi_dataset.py deleted file mode 100644 index 7eaef93187..0000000000 --- a/examples/smc_reac/ppsci/data/dataset/fwi_dataset.py +++ /dev/null @@ -1,103 +0,0 @@ -import os -from typing import Dict -from typing import Optional -from typing import Tuple - -import numpy as np -from paddle import io -from paddle import vision - - -class FWIDataset(io.Dataset): - """Datasets for full waveform inversion tasks. - For convenience, in this class, a batch refers to a npy file instead of the batch used during training. - - Args: - input_keys (Tuple[str, ...]): List of input keys. - label_keys (Tuple[str, ...]): List of label keys. - weight: Define the weight dict for loss function. - anno: Path to annotation file. - preload: Whether to load the whole dataset into memory. - sample_ratio: Downsample ratio for seismic data. - file_size: Number of samples in each npy file. - transform_data: Transformation applied to data. - transform_label: Transformation applied to label. - - Examples: - >>> import ppsci - >>> dataset = ppsci.data.dataset.FWIDataset(("input", ), ("label", ), "path/to/anno_file") # doctest: +SKIP - """ - - def __init__( - self, - input_keys: Tuple[str, ...], - label_keys: Tuple[str, ...], - anno: str, - weight: Optional[Dict[str, np.ndarray]] = None, - preload: bool = True, - sample_ratio: int = 1, - file_size: int = 500, - transform_data: Optional[vision.Compose] = None, - transform_label: Optional[vision.Compose] = None, - ): - super().__init__() - self.input_keys = input_keys - self.label_keys = label_keys - self.weight = {} if weight is None else weight - if not os.path.exists(anno): - print(f"Annotation file {anno} does not exists") - self.preload = preload - self.sample_ratio = sample_ratio - self.file_size = file_size - self.transform_data = transform_data - self.transform_label = transform_label - with open(anno, "r") as f: - self.batches = f.readlines() - if preload: - self.data_list, self.label_list = [], [] - for batch in self.batches: - data, label = self.load_every(batch) - self.data_list.append(data) - if label is not None: - self.label_list.append(label) - - def load_every(self, batch): - batch = batch.split("\t") - data_path = batch[0] if len(batch) > 1 else batch[0][:-1] - data = np.load(data_path)[:, :, :: self.sample_ratio, :] - data = data.astype("float32") - if len(batch) > 1: - label_path = batch[1][:-1] - label = np.load(label_path) - label = label.astype("float32") - else: - label = None - - return data, label - - def __getitem__(self, idx): - batch_idx, sample_idx = idx // self.file_size, idx % self.file_size - if self.preload: - data = self.data_list[batch_idx][sample_idx] - label = ( - self.label_list[batch_idx][sample_idx] - if len(self.label_list) != 0 - else None - ) - else: - data, label = self.load_every(self.batches[batch_idx]) - data = data[sample_idx] - label = label[sample_idx] if label is not None else None - if self.transform_data: - data = self.transform_data(data) - if self.transform_label and label is not None: - label = self.transform_label(label) - - input_item = {self.input_keys[0]: data} - label_item = {self.label_keys[0]: label if label is not None else np.array([])} - weight_item = self.weight - - return input_item, label_item, weight_item - - def __len__(self): - return len(self.batches) * self.file_size diff --git a/examples/smc_reac/ppsci/data/dataset/ifm_moe_dataset.py b/examples/smc_reac/ppsci/data/dataset/ifm_moe_dataset.py deleted file mode 100644 index 43cd50f419..0000000000 --- a/examples/smc_reac/ppsci/data/dataset/ifm_moe_dataset.py +++ /dev/null @@ -1,462 +0,0 @@ -# Copyright (c) 2024 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. - -from __future__ import annotations - -import copy -import os -from typing import Tuple - -import numpy as np -import paddle -import pandas as pd -from paddle import io - -tasks_dic = { - "freesolv": ["activity"], - "esol": ["activity"], - "lipop": ["activity"], - "bace": ["activity"], - "bbbp": ["activity"], - "hiv": ["activity"], - "clintox": ["FDA_APPROVED", "CT_TOX"], - "sider": [ - "SIDER1", - "SIDER2", - "SIDER3", - "SIDER4", - "SIDER5", - "SIDER6", - "SIDER7", - "SIDER8", - "SIDER9", - "SIDER10", - "SIDER11", - "SIDER12", - "SIDER13", - "SIDER14", - "SIDER15", - "SIDER16", - "SIDER17", - "SIDER18", - "SIDER19", - "SIDER20", - "SIDER21", - "SIDER22", - "SIDER23", - "SIDER24", - "SIDER25", - "SIDER26", - "SIDER27", - ], - "tox21": [ - "NR-AR", - "NR-AR-LBD", - "NR-AhR", - "NR-Aromatase", - "NR-ER", - "NR-ER-LBD", - "NR-PPAR-gamma", - "SR-ARE", - "SR-ATAD5", - "SR-HSE", - "SR-MMP", - "SR-p53", - ], - "muv": [ - "MUV-466", - "MUV-548", - "MUV-600", - "MUV-644", - "MUV-652", - "MUV-689", - "MUV-692", - "MUV-712", - "MUV-713", - "MUV-733", - "MUV-737", - "MUV-810", - "MUV-832", - "MUV-846", - "MUV-852", - "MUV-858", - "MUV-859", - ], - "toxcast": [ - "ACEA_T47D_80hr_Negative", - "ACEA_T47D_80hr_Positive", - "APR_HepG2_CellCycleArrest_24h_dn", - "APR_HepG2_CellCycleArrest_72h_dn", - "APR_HepG2_CellLoss_24h_dn", - "APR_HepG2_CellLoss_72h_dn", - "APR_HepG2_MicrotubuleCSK_72h_up", - "APR_HepG2_MitoMass_24h_dn", - "APR_HepG2_MitoMass_72h_dn", - "APR_HepG2_MitoMembPot_24h_dn", - "APR_HepG2_MitoMembPot_72h_dn", - "APR_HepG2_MitoticArrest_24h_up", - "APR_HepG2_MitoticArrest_72h_up", - "APR_HepG2_OxidativeStress_24h_up", - "APR_HepG2_OxidativeStress_72h_up", - "APR_HepG2_StressKinase_72h_up", - "APR_HepG2_p53Act_24h_up", - "APR_HepG2_p53Act_72h_up", - "ATG_AP_1_CIS_up", - "ATG_Ahr_CIS_up", - "ATG_BRE_CIS_up", - "ATG_CMV_CIS_up", - "ATG_CRE_CIS_up", - "ATG_DR4_LXR_CIS_dn", - "ATG_DR5_CIS_up", - "ATG_EGR_CIS_up", - "ATG_ERE_CIS_up", - "ATG_ERa_TRANS_up", - "ATG_E_Box_CIS_dn", - "ATG_HIF1a_CIS_up", - "ATG_HSE_CIS_up", - "ATG_IR1_CIS_dn", - "ATG_ISRE_CIS_dn", - "ATG_MRE_CIS_up", - "ATG_NRF2_ARE_CIS_up", - "ATG_Oct_MLP_CIS_up", - "ATG_PBREM_CIS_up", - "ATG_PPARg_TRANS_up", - "ATG_PPRE_CIS_up", - "ATG_PXRE_CIS_dn", - "ATG_PXRE_CIS_up", - "ATG_PXR_TRANS_up", - "ATG_Pax6_CIS_up", - "ATG_RORE_CIS_up", - "ATG_RXRb_TRANS_up", - "ATG_SREBP_CIS_up", - "ATG_Sp1_CIS_up", - "ATG_TCF_b_cat_CIS_dn", - "ATG_VDRE_CIS_up", - "ATG_Xbp1_CIS_up", - "ATG_p53_CIS_dn", - "BSK_3C_Eselectin_down", - "BSK_3C_HLADR_down", - "BSK_3C_ICAM1_down", - "BSK_3C_IL8_down", - "BSK_3C_MCP1_down", - "BSK_3C_MIG_down", - "BSK_3C_Proliferation_down", - "BSK_3C_SRB_down", - "BSK_3C_Thrombomodulin_up", - "BSK_3C_TissueFactor_down", - "BSK_3C_VCAM1_down", - "BSK_3C_Vis_down", - "BSK_3C_uPAR_down", - "BSK_4H_Eotaxin3_down", - "BSK_4H_MCP1_down", - "BSK_4H_Pselectin_down", - "BSK_4H_SRB_down", - "BSK_4H_VCAM1_down", - "BSK_4H_VEGFRII_down", - "BSK_4H_uPAR_down", - "BSK_BE3C_HLADR_down", - "BSK_BE3C_IL1a_down", - "BSK_BE3C_IP10_down", - "BSK_BE3C_MIG_down", - "BSK_BE3C_MMP1_down", - "BSK_BE3C_MMP1_up", - "BSK_BE3C_PAI1_down", - "BSK_BE3C_SRB_down", - "BSK_BE3C_TGFb1_down", - "BSK_BE3C_tPA_down", - "BSK_BE3C_uPAR_down", - "BSK_BE3C_uPA_down", - "BSK_CASM3C_HLADR_down", - "BSK_CASM3C_IL6_down", - "BSK_CASM3C_IL8_down", - "BSK_CASM3C_LDLR_down", - "BSK_CASM3C_MCP1_down", - "BSK_CASM3C_MCSF_down", - "BSK_CASM3C_MIG_down", - "BSK_CASM3C_Proliferation_down", - "BSK_CASM3C_SAA_down", - "BSK_CASM3C_SRB_down", - "BSK_CASM3C_Thrombomodulin_up", - "BSK_CASM3C_TissueFactor_down", - "BSK_CASM3C_VCAM1_down", - "BSK_CASM3C_uPAR_down", - "BSK_KF3CT_ICAM1_down", - "BSK_KF3CT_IL1a_down", - "BSK_KF3CT_IP10_down", - "BSK_KF3CT_MCP1_down", - "BSK_KF3CT_MMP9_down", - "BSK_KF3CT_SRB_down", - "BSK_KF3CT_TGFb1_down", - "BSK_KF3CT_TIMP2_down", - "BSK_KF3CT_uPA_down", - "BSK_LPS_CD40_down", - "BSK_LPS_Eselectin_down", - "BSK_LPS_IL1a_down", - "BSK_LPS_IL8_down", - "BSK_LPS_MCP1_down", - "BSK_LPS_MCSF_down", - "BSK_LPS_PGE2_down", - "BSK_LPS_SRB_down", - "BSK_LPS_TNFa_down", - "BSK_LPS_TissueFactor_down", - "BSK_LPS_VCAM1_down", - "BSK_SAg_CD38_down", - "BSK_SAg_CD40_down", - "BSK_SAg_CD69_down", - "BSK_SAg_Eselectin_down", - "BSK_SAg_IL8_down", - "BSK_SAg_MCP1_down", - "BSK_SAg_MIG_down", - "BSK_SAg_PBMCCytotoxicity_down", - "BSK_SAg_Proliferation_down", - "BSK_SAg_SRB_down", - "BSK_hDFCGF_CollagenIII_down", - "BSK_hDFCGF_IL8_down", - "BSK_hDFCGF_IP10_down", - "BSK_hDFCGF_MCSF_down", - "BSK_hDFCGF_MIG_down", - "BSK_hDFCGF_MMP1_down", - "BSK_hDFCGF_PAI1_down", - "BSK_hDFCGF_Proliferation_down", - "BSK_hDFCGF_SRB_down", - "BSK_hDFCGF_TIMP1_down", - "BSK_hDFCGF_VCAM1_down", - "CEETOX_H295R_11DCORT_dn", - "CEETOX_H295R_ANDR_dn", - "CEETOX_H295R_CORTISOL_dn", - "CEETOX_H295R_ESTRONE_dn", - "CEETOX_H295R_ESTRONE_up", - "NHEERL_ZF_144hpf_TERATOSCORE_up", - "NVS_NR_bER", - "NVS_NR_hER", - "NVS_NR_hPPARg", - "NVS_NR_hPXR", - "NVS_NR_mERa", - "OT_AR_ARSRC1_0960", - "OT_ER_ERaERb_0480", - "OT_ER_ERaERb_1440", - "OT_ER_ERbERb_0480", - "OT_ER_ERbERb_1440", - "OT_ERa_EREGFP_0120", - "OT_FXR_FXRSRC1_0480", - "OT_NURR1_NURR1RXRa_0480", - "TOX21_ARE_BLA_agonist_ratio", - "TOX21_AR_BLA_Antagonist_ratio", - "TOX21_AR_LUC_MDAKB2_Antagonist", - "TOX21_AR_LUC_MDAKB2_Antagonist2", - "TOX21_AhR_LUC_Agonist", - "TOX21_Aromatase_Inhibition", - "TOX21_ERa_BLA_Antagonist_ratio", - "TOX21_ERa_LUC_BG1_Agonist", - "TOX21_FXR_BLA_antagonist_ratio", - "TOX21_MMP_ratio_down", - "TOX21_TR_LUC_GH3_Antagonist", - "TOX21_p53_BLA_p1_ratio", - "TOX21_p53_BLA_p2_ch2", - "TOX21_p53_BLA_p2_ratio", - "TOX21_p53_BLA_p2_viability", - "TOX21_p53_BLA_p3_ratio", - "TOX21_p53_BLA_p4_ratio", - "TOX21_p53_BLA_p5_ratio", - "Tanguay_ZF_120hpf_AXIS_up", - "Tanguay_ZF_120hpf_ActivityScore", - "Tanguay_ZF_120hpf_JAW_up", - "Tanguay_ZF_120hpf_MORT_up", - "Tanguay_ZF_120hpf_PE_up", - "Tanguay_ZF_120hpf_SNOU_up", - "Tanguay_ZF_120hpf_YSE_up", - ], -} - - -def standardize(col): - return (col - np.mean(col)) / np.std(col) - - -def get_pos_weight(Ys): - Ys = paddle.to_tensor(np.nan_to_num(Ys), dtype=paddle.float32) - num_pos = paddle.sum(Ys, axis=0) - num_indices = paddle.to_tensor(len(Ys)) - return (num_indices - num_pos) / num_pos - - -class IFMMoeDataset(io.Dataset): - """Dataset for `IFMMoe`. - - Args: - input_keys (Tuple[str, ...]): Name of input data. - label_keys (Tuple[str, ...]): Name of label data. - data_dir (str): Directory of IFMMoe data. - data_label (str): IFMMoe data label in tox21/esol/freesolv/lipop... - data_mode (str): train/val/test mode data. - - Examples: - >>> import ppsci - >>> dataset = ppsci.data.dataset.IFMMoeDataset( - ... "input_keys": ("input",), - ... "label_keys": ("output",), - ... "data_dir": "/path/to/IFMMoeDataset", - ... "data_label": "tox21", - ... "data_mode": "train", - ... ) # doctest: +SKIP - """ - - # Whether support batch indexing for speeding up fetching process. - batch_index: bool = False - use_pgl: bool = False - - def __init__( - self, - input_keys: Tuple[str, ...], - label_keys: Tuple[str, ...], - data_dir: str, - data_label: str, - data_mode: str, - ): - self.input_keys = input_keys - self.label_keys = label_keys - - self.data_label = data_label - self.data_dir = data_dir - self.data_mode = data_mode - - if data_label == "esol" or data_label == "freesolv" or data_label == "lipop": - self.task_type = "reg" - self.reg = True - # metric = "rmse" - else: - self.task_type = "cla" - self.reg = False - # metric = "roc_auc" - - self.task_dict = tasks_dic - - self.Xs = None - self.Ys = None - self.mask = None - self.process_data() - - def process_data(self): - file_name = os.path.join(self.data_dir, self.data_label + "_moe_pubsubfp.csv") - # preprocess data - dataset_all = pd.read_csv(file_name) - if self.data_label == "freesolv": - dataset_all.drop(columns=["vsa_pol", "h_emd", "a_donacc"], inplace=True) - elif self.data_label == "esol": - dataset_all.drop(columns=["logS", "h_logS", "SlogP"], inplace=True) - else: - dataset_all.drop(columns=["SlogP", "h_logD", "logS"], inplace=True) - tasks = tasks_dic[self.data_label] - cols = copy.deepcopy(tasks) - cols.extend(dataset_all.columns[len(tasks) + 1 :]) - dataset = dataset_all[cols] - x_cols = dataset_all.columns[len(tasks) + 1 :] - # remove the features with na - if self.data_label != "hiv": - rm_cols1 = ( - dataset[x_cols] - .isnull() - .any()[dataset[x_cols].isnull().any() == True] # noqa: E712 - .index - ) - dataset.drop(columns=rm_cols1, inplace=True) - else: - rm_indx1 = ( - dataset[x_cols] - .isnull() - .T.any()[dataset[x_cols].isnull().T.any() == True] # noqa: E712 - .index - ) - dataset.drop(index=rm_indx1, inplace=True) - x_cols = dataset.columns.drop(tasks) - - # Removing features with low variance - # threshold = 0.05 - data_fea_var = dataset[x_cols].var() - del_fea1 = list(data_fea_var[data_fea_var <= 0.05].index) - dataset.drop(columns=del_fea1, inplace=True) - x_cols = dataset.columns.drop(tasks) - - # pair correlations - # threshold = 0.95 - data_fea_corr = dataset[x_cols].corr() - del_fea2_col = [] - del_fea2_ind = [] - length = data_fea_corr.shape[1] - for i in range(length): - for j in range(i + 1, length): - if abs(data_fea_corr.iloc[i, j]) >= 0.95: - del_fea2_col.append(data_fea_corr.columns[i]) - del_fea2_ind.append(data_fea_corr.index[j]) - dataset.drop(columns=del_fea2_ind, inplace=True) - # standardize the features - cols_ = dataset.columns[len(tasks) + 1 :] - # print('the retained features for %s is %d' % (args.task, len(cols_))) - dataset[cols_] = dataset[cols_].apply(standardize, axis=0) - - dataseta = pd.read_csv( - os.path.join( - self.data_dir, "dataset_used_for_modeling", self.data_label + ".csv" - ) - ) - data_tr = dataset[dataseta.group == "train"] - data_va = dataset[dataseta.group == "valid"] - data_te = dataset[dataseta.group == "test"] - - # training set - data_tr_y = data_tr[tasks].values.reshape(-1, len(tasks)) - data_tr_x = data_tr.iloc[:, len(tasks) :].values # 249 - # data_tr_x = data_tr.iloc[:, len(tasks):].values - # test set - data_te_y = data_te[tasks].values.reshape(-1, len(tasks)) - data_te_x = data_te.iloc[:, len(tasks) :].values - # data_te_x = data_te.iloc[:, len(tasks):].values - - # validation set - data_va_y = data_va[tasks].values.reshape(-1, len(tasks)) - data_va_x = data_va.iloc[:, len(tasks) :].values - # data_va_x = data_va.iloc[:, len(tasks):].values - - # dataloader - # train_dataset = MyDataset(data_tr_x, data_tr_y) - # validation_dataset = MyDataset(data_va_x, data_va_y) - # test_dataset = MyDataset(data_te_x, data_te_y) - if self.data_mode == "train": - Xs, Ys = data_tr_x, data_tr_y - elif self.data_mode == "val": - Xs, Ys = data_va_x, data_va_y - elif self.data_mode == "test": - Xs, Ys = data_te_x, data_te_y - if not self.reg: - self.pos_weights = get_pos_weight(dataset[tasks].values) - - self.data_tr_x = data_tr_x - self.Xs = Xs - self.Ys = np.nan_to_num(Ys) - self.mask = ~np.isnan(Ys) * 1.0 - - def __len__(self): - return len(self.Ys) - - def __getitem__(self, idx): - return ( - { - self.input_keys[0]: paddle.to_tensor(self.Xs[idx], dtype="float32"), - }, - { - self.label_keys[0]: paddle.to_tensor(self.Ys[idx], dtype="float32"), - self.label_keys[1]: paddle.to_tensor(self.mask[idx], dtype="float32"), - }, - {}, - ) diff --git a/examples/smc_reac/ppsci/data/dataset/mat_dataset.py b/examples/smc_reac/ppsci/data/dataset/mat_dataset.py deleted file mode 100644 index 609e35aeaa..0000000000 --- a/examples/smc_reac/ppsci/data/dataset/mat_dataset.py +++ /dev/null @@ -1,287 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import Callable -from typing import Dict -from typing import Optional -from typing import Tuple -from typing import Union - -import numpy as np -import paddle -from paddle import io -from paddle import vision - -from ppsci.utils import misc -from ppsci.utils import reader - - -class MatDataset(io.Dataset): - """Dataset class for .mat file. - - Args: - file_path (str): Mat file path. - input_keys (Tuple[str, ...]): List of input keys. - label_keys (Tuple[str, ...], optional): List of label keys. Defaults to (). - alias_dict (Optional[Dict[str, str]]): Dict of alias(es) for input and label keys. - i.e. {inner_key: outer_key}. Defaults to None. - weight_dict (Optional[Dict[str, Union[Callable, float]]]): Define the weight of - each constraint variable. Defaults to None. - timestamps (Optional[Tuple[float, ...]]): The number of repetitions of the data - in the time dimension. Defaults to None. - transforms (Optional[vision.Compose]): Compose object contains sample wise - transform(s). Defaults to None. - - Examples: - >>> import ppsci - >>> dataset = ppsci.data.dataset.MatDataset( - ... "/path/to/file.mat" - ... ("x",), - ... ("u",), - ... ) # doctest: +SKIP - """ - - # Whether support batch indexing for speeding up fetching process. - batch_index: bool = True - - def __init__( - self, - file_path: str, - input_keys: Tuple[str, ...], - label_keys: Tuple[str, ...] = (), - alias_dict: Optional[Dict[str, str]] = None, - weight_dict: Optional[Dict[str, Union[Callable, float]]] = None, - timestamps: Optional[Tuple[float, ...]] = None, - transforms: Optional[vision.Compose] = None, - ): - super().__init__() - self.input_keys = input_keys - self.label_keys = label_keys - - # read raw data from file - raw_data = reader.load_mat_file( - file_path, - input_keys + label_keys, - alias_dict, - ) - # filter raw data by given timestamps if specified - if timestamps is not None: - if "t" in raw_data: - # filter data according to given timestamps - raw_time_array = raw_data["t"] - mask = [] - for ti in timestamps: - mask.append(np.nonzero(np.isclose(raw_time_array, ti).flatten())[0]) - raw_data = misc.convert_to_array( - raw_data, self.input_keys + self.label_keys - ) - mask = np.concatenate(mask, 0) - raw_data = raw_data[mask] - raw_data = misc.convert_to_dict( - raw_data, self.input_keys + self.label_keys - ) - else: - # repeat data according to given timestamps - raw_data = misc.convert_to_array( - raw_data, self.input_keys + self.label_keys - ) - raw_data = misc.combine_array_with_time(raw_data, timestamps) - self.input_keys = ("t",) + tuple(self.input_keys) - raw_data = misc.convert_to_dict( - raw_data, self.input_keys + self.label_keys - ) - - # fetch input data - self.input = { - key: value for key, value in raw_data.items() if key in self.input_keys - } - # fetch label data - self.label = { - key: value for key, value in raw_data.items() if key in self.label_keys - } - - # prepare weights - self.weight = ( - {key: np.ones_like(next(iter(self.label.values()))) for key in self.label} - if weight_dict is not None - else {} - ) - if weight_dict is not None: - for key, value in weight_dict.items(): - if isinstance(value, (int, float)): - self.weight[key] = np.full_like( - next(iter(self.label.values())), value - ) - elif callable(value): - func = value - self.weight[key] = func(self.input) - if isinstance(self.weight[key], (int, float)): - self.weight[key] = np.full_like( - next(iter(self.label.values())), self.weight[key] - ) - else: - raise NotImplementedError(f"type of {type(value)} is invalid yet.") - - self.transforms = transforms - self._len = len(next(iter(self.input.values()))) - - def __getitem__(self, idx): - input_item = {key: value[idx] for key, value in self.input.items()} - label_item = {key: value[idx] for key, value in self.label.items()} - weight_item = {key: value[idx] for key, value in self.weight.items()} - - if self.transforms is not None: - input_item, label_item, weight_item = self.transforms( - input_item, label_item, weight_item - ) - - return (input_item, label_item, weight_item) - - def __len__(self): - return self._len - - -class IterableMatDataset(io.IterableDataset): - """IterableMatDataset for full-data loading. - - Args: - file_path (str): Mat file path. - input_keys (Tuple[str, ...]): List of input keys. - label_keys (Tuple[str, ...], optional): List of label keys. Defaults to (). - alias_dict (Optional[Dict[str, str]]): Dict of alias(es) for input and label keys. - i.e. {inner_key: outer_key}. Defaults to None. - weight_dict (Optional[Dict[str, Union[Callable, float]]]): Define the weight of - each constraint variable. Defaults to None. - timestamps (Optional[Tuple[float, ...]]): The number of repetitions of the data - in the time dimension. Defaults to None. - transforms (Optional[vision.Compose]): Compose object contains sample wise - transform(s). Defaults to None. - - Examples: - >>> import ppsci - >>> dataset = ppsci.data.dataset.IterableMatDataset( - ... "/path/to/file.mat" - ... ("x",), - ... ("u",), - ... ) # doctest: +SKIP - """ - - # Whether support batch indexing for speeding up fetching process. - batch_index: bool = False - - def __init__( - self, - file_path: str, - input_keys: Tuple[str, ...], - label_keys: Tuple[str, ...] = (), - alias_dict: Optional[Dict[str, str]] = None, - weight_dict: Optional[Dict[str, Union[Callable, float]]] = None, - timestamps: Optional[Tuple[float, ...]] = None, - transforms: Optional[vision.Compose] = None, - ): - super().__init__() - self.input_keys = input_keys - self.label_keys = label_keys - - # read raw data from file - raw_data = reader.load_mat_file( - file_path, - input_keys + label_keys, - alias_dict, - ) - # filter raw data by given timestamps if specified - if timestamps is not None: - if "t" in raw_data: - # filter data according to given timestamps - raw_time_array = raw_data["t"] - mask = [] - for ti in timestamps: - mask.append(np.nonzero(np.isclose(raw_time_array, ti).flatten())[0]) - raw_data = misc.convert_to_array( - raw_data, self.input_keys + self.label_keys - ) - mask = np.concatenate(mask, 0) - raw_data = raw_data[mask] - raw_data = misc.convert_to_dict( - raw_data, self.input_keys + self.label_keys - ) - else: - # repeat data according to given timestamps - raw_data = misc.convert_to_array( - raw_data, self.input_keys + self.label_keys - ) - raw_data = misc.combine_array_with_time(raw_data, timestamps) - self.input_keys = ("t",) + tuple(self.input_keys) - raw_data = misc.convert_to_dict( - raw_data, self.input_keys + self.label_keys - ) - - # fetch input data - self.input = { - key: value for key, value in raw_data.items() if key in self.input_keys - } - # fetch label data - self.label = { - key: value for key, value in raw_data.items() if key in self.label_keys - } - - # prepare weights - self.weight = ( - {key: np.ones_like(next(iter(self.label.values()))) for key in self.label} - if weight_dict is not None - else {} - ) - if weight_dict is not None: - for key, value in weight_dict.items(): - if isinstance(value, (int, float)): - self.weight[key] = np.full_like( - next(iter(self.label.values())), value - ) - elif callable(value): - func = value - self.weight[key] = func(self.input) - if isinstance(self.weight[key], (int, float)): - self.weight[key] = np.full_like( - next(iter(self.label.values())), self.weight[key] - ) - else: - raise NotImplementedError(f"type of {type(value)} is invalid yet.") - - self.input = {key: paddle.to_tensor(value) for key, value in self.input.items()} - self.label = {key: paddle.to_tensor(value) for key, value in self.label.items()} - self.weight = { - key: paddle.to_tensor(value) for key, value in self.weight.items() - } - - self.transforms = transforms - self._len = len(next(iter(self.input.values()))) - - @property - def num_samples(self): - """Number of samples within current dataset.""" - return self._len - - def __iter__(self): - if callable(self.transforms): - input_, label_, weight_ = self.transforms( - self.input, self.label, self.weight - ) - yield input_, label_, weight_ - else: - yield self.input, self.label, self.weight - - def __len__(self): - return 1 diff --git a/examples/smc_reac/ppsci/data/dataset/moflow_dataset.py b/examples/smc_reac/ppsci/data/dataset/moflow_dataset.py deleted file mode 100644 index 627365c1dc..0000000000 --- a/examples/smc_reac/ppsci/data/dataset/moflow_dataset.py +++ /dev/null @@ -1,437 +0,0 @@ -# Copyright (c) 2024 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. - -# Copyright 2020 Chengxi Zang - -from __future__ import annotations - -import os -from typing import Callable -from typing import Dict -from typing import List -from typing import Optional -from typing import Tuple - -import numpy as np -import pandas as pd -from paddle import io -from tqdm import tqdm - -from ppsci.utils import logger - -try: - from rdkit import Chem - from rdkit.Chem import rdmolops -except ModuleNotFoundError: - pass - - -class MolGraph: - """ - Args: - max_atoms (int): Max number of atoms for each molecule, if the - number of atoms is more than this value, this data is simply - ignored. - Setting negative value indicates no limit for max atoms. - out_size (int): It specifies the size of array returned by - `get_input_features`. - If the number of atoms in the molecule is less than this value, - the returned arrays is padded to have fixed size. - Setting negative value indicates do not pad returned array. - add_Hs (bool): If True, implicit Hs are added. - kekulize (bool): If True, Kekulizes the molecule. - """ - - def __init__(self, max_atoms=-1, out_size=-1, add_Hs=False, kekulize=False): - super(MolGraph, self).__init__() - self.add_Hs = add_Hs - self.kekulize = kekulize - if max_atoms >= 0 and out_size >= 0 and max_atoms > out_size: - raise ValueError( - f"max_atoms {max_atoms} must be less or equal to out_size {out_size}" - ) - self.max_atoms = max_atoms - self.out_size = out_size - - def get_input_features(self, mol): - """ - get input features - Args: - mol (Mol): mol instance - - Returns: - (tuple): (`atom`, `adj`) - - """ - self.type_check_num_atoms(mol, self.max_atoms) - atom_array = self.construct_atomic_number_array(mol, out_size=self.out_size) - adj_array = self.construct_discrete_edge_matrix(mol, out_size=self.out_size) - return atom_array, adj_array - - def prepare_smiles_and_mol(self, mol): - """Prepare `smiles` and `mol` used in following preprocessing. - This method is called before `get_input_features` is called, by parser - class. - This method may be overridden to support custom `smile`/`mol` extraction - Args: - mol (mol): mol instance - - Returns (tuple): (`smiles`, `mol`) - """ - canonical_smiles = Chem.MolToSmiles(mol, isomericSmiles=False, canonical=True) - mol = Chem.MolFromSmiles(canonical_smiles) - if self.add_Hs: - mol = Chem.AddHs(mol) - if self.kekulize: - Chem.Kekulize(mol) - return canonical_smiles, mol - - def get_label(self, mol, label_names=None): - """Extracts label information from a molecule. - This method extracts properties whose keys are - specified by ``label_names`` from a molecule ``mol`` - and returns these values as a list. - The order of the values is same as that of ``label_names``. - If the molecule does not have a - property with some label, this function fills the corresponding - index of the returned list with ``None``. - - Args: - mol (rdkit.Chem.Mol): molecule whose features to be extracted - label_names (None or iterable): list of label names. - - Returns: - list of str: label information. Its length is equal to - that of ``label_names``. If ``label_names`` is ``None``, - this function returns an empty list. - - """ - if label_names is None: - return [] - label_list = [] - for label_name in label_names: - if mol.HasProp(label_name): - label_list.append(mol.GetProp(label_name)) - else: - label_list.append(None) - return label_list - - def type_check_num_atoms(self, mol, num_max_atoms=-1): - """Check number of atoms in `mol` does not exceed `num_max_atoms` - If number of atoms in `mol` exceeds the number `num_max_atoms`, it will - raise `MolGraphError` exception. - - Args: - mol (Mol): - num_max_atoms (int): If negative value is set, not check number of - atoms. - - """ - num_atoms = mol.GetNumAtoms() - if num_max_atoms >= 0 and num_atoms > num_max_atoms: - raise MolGraphError( - f"Number of atoms in mol {num_atoms} exceeds num_max_atoms {num_max_atoms}" - ) - - def construct_atomic_number_array(self, mol, out_size=-1): - """Returns atomic numbers of atoms consisting a molecule. - - Args: - mol (rdkit.Chem.Mol): Input molecule. - out_size (int): The size of returned array. - If this option is negative, it does not take any effect. - Otherwise, it must be larger than the number of atoms - in the input molecules. In that case, the tail of - the array is padded with zeros. - - Returns: - numpy.ndarray: an array consisting of atomic numbers - of atoms in the molecule. - """ - atom_list = [a.GetAtomicNum() for a in mol.GetAtoms()] - n_atom = len(atom_list) - if out_size < 0: - return np.array(atom_list, dtype=np.int32) - elif out_size >= n_atom: - atom_array = np.zeros(out_size, dtype=np.int32) - atom_array[:n_atom] = np.array(atom_list, dtype=np.int32) - return atom_array - else: - raise ValueError( - f"`out_size` (={out_size}) must be negative or larger than or equal to " - f"the number of atoms in the input molecules (={n_atom})." - ) - - def construct_adj_matrix(self, mol, out_size=-1, self_connection=True): - """Returns the adjacent matrix of the given molecule. - - This function returns the adjacent matrix of the given molecule. - Contrary to the specification of - :func:`rdkit.Chem.rdmolops.GetAdjacencyMatrix`, - The diagonal entries of the returned matrix are all-one. - - Args: - mol (rdkit.Chem.Mol): Input molecule. - out_size (int): The size of the returned matrix. - If this option is negative, it does not take any effect. - Otherwise, it must be larger than the number of atoms - in the input molecules. In that case, the adjacent - matrix is expanded and zeros are padded to right - columns and bottom rows. - self_connection (bool): Add self connection or not. - If True, diagonal element of adjacency matrix is filled with 1. - - Returns: - adj_array (numpy.ndarray): The adjacent matrix of the input molecule. - It is 2-dimensional array with shape (atoms1, atoms2), where - atoms1 & atoms2 represent from and to of the edge respectively. - If ``out_size`` is non-negative, the returned - its size is equal to that value. Otherwise, - it is equal to the number of atoms in the the molecule. - """ - adj = rdmolops.GetAdjacencyMatrix(mol) - s0, s1 = tuple(adj.shape) - if s0 != s1: - raise ValueError( - f"The adjacent matrix of the input moleculehas an invalid shape: ({s0}, " - f"{s1}). It must be square." - ) - if self_connection: - adj = adj + np.eye(s0) - if out_size < 0: - adj_array = adj.astype(np.float32) - elif out_size >= s0: - adj_array = np.zeros((out_size, out_size), dtype=np.float32) - adj_array[:s0, :s1] = adj - else: - raise ValueError( - f"`out_size` (={out_size}) must be negative or larger than or equal to " - f"the number of atoms in the input molecules (={s0})." - ) - return adj_array - - def construct_discrete_edge_matrix(self, mol, out_size=-1): - """Returns the edge-type dependent adjacency matrix of the given molecule. - - Args: - mol (rdkit.Chem.Mol): Input molecule. - out_size (int): The size of the returned matrix. - If this option is negative, it does not take any effect. - Otherwise, it must be larger than the number of atoms - in the input molecules. In that case, the adjacent - matrix is expanded and zeros are padded to right - columns and bottom rows. - - Returns: - adj_array (numpy.ndarray): The adjacent matrix of the input molecule. - It is 3-dimensional array with shape (edge_type, atoms1, atoms2), - where edge_type represents the bond type, - atoms1 & atoms2 represent from and to of the edge respectively. - If ``out_size`` is non-negative, its size is equal to that value. - Otherwise, it is equal to the number of atoms in the the molecule. - """ - if mol is None: - raise MolGraphError("mol is None") - N = mol.GetNumAtoms() - if out_size < 0: - size = N - elif out_size >= N: - size = out_size - else: - raise ValueError( - f"out_size {out_size} is smaller than number of atoms in mol {N}" - ) - adjs = np.zeros((4, size, size), dtype=np.float32) - bond_type_to_channel = { - Chem.BondType.SINGLE: 0, - Chem.BondType.DOUBLE: 1, - Chem.BondType.TRIPLE: 2, - Chem.BondType.AROMATIC: 3, - } - for bond in mol.GetBonds(): - bond_type = bond.GetBondType() - ch = bond_type_to_channel[bond_type] - i = bond.GetBeginAtomIdx() - j = bond.GetEndAtomIdx() - adjs[ch, i, j] = 1.0 - adjs[ch, j, i] = 1.0 - return adjs - - -class MolGraphError(Exception): - pass - - -class MOlFLOWDataset(io.Dataset): - """Class for moflow qm9 and zinc250k Dataset of a tuple of datasets. - - It combines multiple datasets into one dataset. Each example is represented - by a tuple whose ``i``-th item corresponds to the i-th dataset. - And each ``i``-th dataset is expected to be an instance of numpy.ndarray. - - Args: - file_path (str): Data set path. - data_name (str): Data name, "qm9" or "zinc250k" - valid_idx (List[int, ...]): Data for validate - mode (str): "train" or "eval", output Data - input_keys (Tuple[str, ...]): Input keys, such as ("nodes","edges",). - label_keys (Tuple[str, ...]): labels (str or list or None) . - smiles_col (str): smiles column - weight_dict (Optional[Dict[str, Union[Callable, float]]]): Define the weight of each constraint variable. Defaults to None. - transform_fn: An optional function applied to an item bofre returning - """ - - # Whether support batch indexing for speeding up fetching process. - batch_index: bool = True - - def __init__( - self, - file_path: str, - data_name: str, - valid_idx: List[int, ...], - mode: str, - input_keys: Tuple[str, ...], - label_keys: Tuple[str, ...], - smiles_col: str, - weight_dict: Optional[Dict[str, float]] = None, - transform_fn: Optional[Callable] = None, - ): - super().__init__() - self.file_path = file_path - self.data_name = data_name - self.input_keys = input_keys - self.label_keys = label_keys - self.smiles_col = smiles_col - self.weight_dict = weight_dict - - if data_name == "qm9": - max_atoms = 9 - elif data_name == "zinc250k": - max_atoms = 38 - - self.molgraph = MolGraph(out_size=max_atoms, kekulize=True) - self.logger = logger - # read and deal data from file - inputs, labels = self.load_csv_file(file_path, data_name + ".csv") - train_idx = [t for t in range(len(inputs[0])) if t not in valid_idx] - self.train_idx = train_idx - # data train or test - if mode == "train": - inputs = [ - np.array(list(io.Subset(dataset=in_put, indices=train_idx))) - for in_put in inputs - ] - labels = np.array(list(io.Subset(dataset=labels, indices=train_idx))) - elif mode == "eval": - inputs = [ - np.array(list(io.Subset(dataset=in_put, indices=valid_idx))) - for in_put in inputs - ] - labels = np.array(list(io.Subset(dataset=labels, indices=valid_idx))) - - # fetch input data - self.input = {key: inputs[i] for i, key in enumerate(self.input_keys)} - # fetch label data - self.label = {"label": labels} - - self.logger.message( - f"Dataload finished. MODE {mode}, " - f"inputs {len(next(iter(self.input.values())))}, " - f"labelS {len(next(iter(self.label.values())))}" - ) - - self._length = len(next(iter(self.input.values()))) - self.transform = transform_fn - - def __getitem__(self, index: int): - input_item = {key: value[index] for key, value in self.input.items()} - label_item = {key: value[index] for key, value in self.label.items()} - - if self.transform: - input_item, label_item = self.transform_func(input_item, label_item) - - return (input_item, label_item, {}) - - def __len__(self): - return self._length - - def load_csv_file(self, path: str, name: str): - """Parse DataFrame using `MolGraph` and prepare a dataset instance - Labels are extracted from `labels` columns and input features are - extracted from smiles information in `smiles` column. - """ - file = os.path.join(path, name) - df = pd.read_csv(file, index_col=0) - all_nodes = [] - all_edges = [] - # inputs = [] - - total_count = df.shape[0] - fail_count = 0 - success_count = 0 - if isinstance(self.molgraph, MolGraph): - for smiles in tqdm(df[self.smiles_col], total=df.shape[0]): - try: - mol = Chem.MolFromSmiles(smiles) - if mol is None: - fail_count += 1 - continue - canonical_smiles, mol = self.molgraph.prepare_smiles_and_mol(mol) - nodes, edges = self.molgraph.get_input_features(mol) - - except MolGraphError as e: - fail_count += 1 - self.logger.warning(f"parse(), type: {type(e).__name__}, {e.args}") - continue - except Exception as e: - self.logger.warning(f"parse(), type: {type(e).__name__}, {e.args}") - fail_count += 1 - continue - # raw_data = misc.convert_to_dict(np.array([nodes, edges]), self.input_keys) - - all_nodes.append(nodes) - all_edges.append(edges) - # inputs.append(raw_data) - - success_count += 1 - - labels = np.array( - [*(df[label_col].values for label_col in self.label_keys)] - ).T - result = [np.array(all_nodes), np.array(all_edges)], labels - self.logger.message( - f"Preprocess finished. FAIL {fail_count}, " - f"SUCCESS {success_count}, TOTAL {total_count}" - ) - else: - raise NotImplementedError - - return result - - def transform_func(self, data_dict, label_dict): - items = [] - length = len(next(iter(data_dict.values()))) - for idx in range(length): - input_item = [value[idx] for key, value in data_dict.items()] - label_item = [value[idx] for key, value in label_dict.items()] - item = input_item + label_item - if self.transform: - item = self.transform(item) - items.append(item) - items = np.array(items, dtype=object).T - - data_dict = {key: np.stack(items[i], axis=0) for i, key in enumerate(data_dict)} - label_dict = {key: np.vstack(item[2]) for key in label_dict} - - return data_dict, label_dict diff --git a/examples/smc_reac/ppsci/data/dataset/mrms_dataset.py b/examples/smc_reac/ppsci/data/dataset/mrms_dataset.py deleted file mode 100644 index bee3337f7e..0000000000 --- a/examples/smc_reac/ppsci/data/dataset/mrms_dataset.py +++ /dev/null @@ -1,251 +0,0 @@ -# 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. - -from __future__ import annotations - -import glob -import os.path as osp -from datetime import datetime -from datetime import timedelta -from typing import Dict -from typing import List -from typing import Optional -from typing import Tuple - -try: - import h5py -except ModuleNotFoundError: - pass -import numpy as np -import paddle -from paddle import io -from paddle import vision - - -class MRMSDataset(io.Dataset): - """Class for MRMS dataset. MRMS day's data is stored in a .h5 file. Each file includes keys "date"/"time_interval"/"dataset". - - Args: - file_path (str): Dataset path. - input_keys (Tuple[str, ...]): Input keys, usually there is only one, such as ("input",). - label_keys (Tuple[str, ...]): Output keys, usually there is only one, such as ("output",). - weight_dict (Optional[Dict[str, float]]): Weight dictionary. Defaults to None. - date_period (Tuple[str,...], optional): Dates of data. Scale is [start_date, end_date] with format "%Y%m%d". Defaults to ("20230101","20230101"). - num_input_timestamps (int, optional): Number of timestamp of input. Defaults to 1. - num_label_timestamps (int, optional): Number of timestamp of label. Defaults to 1. - stride (int, optional): Stride of sampling data. Defaults to 1. - transforms (Optional[vision.Compose]): Composed transform functor(s). Defaults to None. - - Examples: - >>> import ppsci - >>> dataset = ppsci.data.dataset.MRMSDataset( - ... "file_path": "/path/to/MRMSDataset", - ... "input_keys": ("input",), - ... "label_keys": ("output",), - ... "date_period": ("20230101","20230131"), - ... "num_input_timestamps": 9, - ... "num_label_timestamps": 20, - ... "transforms": transform, - ... "stride": 1, - ... ) # doctest: +SKIP - """ - - # Whether support batch indexing for speeding up fetching process. - batch_index: bool = False - - def __init__( - self, - file_path: str, - input_keys: Tuple[str, ...], - label_keys: Tuple[str, ...], - weight_dict: Optional[Dict[str, float]] = None, - date_period: Tuple[str, ...] = ("20230101", "20230101"), - num_input_timestamps: int = 1, - num_label_timestamps: int = 1, - stride: int = 1, - transforms: Optional[vision.Compose] = None, - ): - super().__init__() - self.file_path = file_path - self.input_keys = input_keys - self.label_keys = label_keys - - self.weight_dict = {} if weight_dict is None else weight_dict - if weight_dict is not None: - self.weight_dict = {key: 1.0 for key in self.label_keys} - self.weight_dict.update(weight_dict) - - self.date_list = self._get_date_strs(date_period) - self.num_input_timestamps = num_input_timestamps - self.num_label_timestamps = num_label_timestamps - self.stride = stride - self.transforms = transforms - - self.files = self._read_data(file_path) - self.num_samples_per_day = self.files[0].shape[0] - self.num_samples = self.num_samples_per_day * len(self.date_list) - - def _get_date_strs(self, date_period: Tuple[str, ...]) -> List: - """Get a string list of all dates within given period. - - Args: - date_period (Tuple[str,...]): Dates of data. Scale is [start_date, end_date] with format "%Y%m%d". - """ - start_time = datetime.strptime(date_period[0], "%Y%m%d") - end_time = datetime.strptime(date_period[1], "%Y%m%d") - results = [] - current_time = start_time - while current_time <= end_time: - date_str = current_time.strftime("%Y%m%d") - results.append(date_str) - current_time += timedelta(days=1) - return results - - def _read_data(self, path: str): - if path.endswith(".h5"): - paths = [path] - else: - paths = [ - _path - for _path in glob.glob(osp.join(path, "*.h5")) - if _path.split(".h5")[0].split("_")[-1] in self.date_list - ] - assert len(paths) == len( - self.date_list - ), f"Data of {len(self.date_list)} days wanted but only {len(paths)} days be found" - paths.sort() - - files = [h5py.File(_path, "r")["dataset"] for _path in paths] - return files - - def __len__(self): - return ( - self.num_samples // self.stride - - self.num_input_timestamps - - self.num_label_timestamps - + 1 - ) - - def __getitem__(self, global_idx): - global_idx *= self.stride - _samples = np.empty( - ( - self.num_input_timestamps + self.num_label_timestamps, - *self.files[0].shape[1:], - ), - dtype=paddle.get_default_dtype(), - ) - for idx in range(self.num_input_timestamps + self.num_label_timestamps): - sample_idx = global_idx + idx * self.stride - day_idx = sample_idx // self.num_samples_per_day - local_idx = sample_idx % self.num_samples_per_day - _samples[idx] = self.files[day_idx][local_idx] - - input_item = {self.input_keys[0]: _samples[: self.num_input_timestamps]} - label_item = {self.label_keys[0]: _samples[self.num_input_timestamps :]} - - weight_shape = [1] * len(next(iter(label_item.values())).shape) - weight_item = { - key: np.full(weight_shape, value, paddle.get_default_dtype()) - for key, value in self.weight_dict.items() - } - - if self.transforms is not None: - input_item, label_item, weight_item = self.transforms( - input_item, label_item, weight_item - ) - - return input_item, label_item, weight_item - - -class MRMSSampledDataset(io.Dataset): - """Class for MRMS sampled dataset. MRMS one sample's data is stored in a .h5 file. Each file includes keys "date"/"time_interval"/"dataset". - The class just return data by input_item and values of label_item are empty for all label_keys. - - Args: - file_path (str): Dataset path. - input_keys (Tuple[str, ...]): Input keys, such as ("input",). - label_keys (Tuple[str, ...]): Output keys, such as ("output",). - weight_dict (Optional[Dict[str, float]]): Weight dictionary. Defaults to None. - num_total_timestamps (int, optional): Number of timestamp of input+label. Defaults to 1. - transforms (Optional[vision.Compose]): Composed transform functor(s). Defaults to None. - - Examples: - >>> import ppsci - >>> dataset = ppsci.data.dataset.MRMSSampledDataset( - ... "file_path": "/path/to/MRMSSampledDataset", - ... "input_keys": ("input",), - ... "label_keys": ("output",), - ... "num_total_timestamps": 29, - ... ) # doctest: +SKIP - >>> # get the length of the dataset - >>> dataset_size = len(dataset) # doctest: +SKIP - >>> # get the first sample of the data - >>> first_sample = dataset[0] # doctest: +SKIP - >>> print("First sample:", first_sample) # doctest: +SKIP - """ - - def __init__( - self, - file_path: str, - input_keys: Tuple[str, ...], - label_keys: Tuple[str, ...], - weight_dict: Optional[Dict[str, float]] = None, - num_total_timestamps: int = 1, - transforms: Optional[vision.Compose] = None, - ): - super().__init__() - self.file_path = file_path - self.input_keys = input_keys - self.label_keys = label_keys - - self.weight_dict = {} if weight_dict is None else weight_dict - if weight_dict is not None: - self.weight_dict = {key: 1.0 for key in self.label_keys} - self.weight_dict.update(weight_dict) - - self.num_total_timestamps = num_total_timestamps - self.transforms = transforms - - self.files = self._read_data(file_path) - self.num_samples = len(self.files) - - def _read_data(self, path: str): - paths = glob.glob(osp.join(path, "*.h5")) - paths.sort() - files = [h5py.File(_path, "r")["dataset"] for _path in paths] - return files - - def __len__(self): - return self.num_samples - self.num_total_timestamps + 1 - - def __getitem__(self, global_idx): - _samples = [] - for idx in range(global_idx, global_idx + self.num_total_timestamps): - _samples.append(np.expand_dims(self.files[idx], axis=0)) - - input_item = { - self.input_keys[0]: np.concatenate(_samples, axis=0).astype( - paddle.get_default_dtype() - ) - } - label_item = {} - weight_item = {} - - if self.transforms is not None: - input_item, label_item, weight_item = self.transforms( - input_item, label_item, weight_item - ) - - return input_item, label_item, weight_item diff --git a/examples/smc_reac/ppsci/data/dataset/npz_dataset.py b/examples/smc_reac/ppsci/data/dataset/npz_dataset.py deleted file mode 100644 index 76d737d021..0000000000 --- a/examples/smc_reac/ppsci/data/dataset/npz_dataset.py +++ /dev/null @@ -1,279 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import Callable -from typing import Dict -from typing import Optional -from typing import Tuple -from typing import Union - -import numpy as np -import paddle -from paddle import io -from paddle import vision - -from ppsci.utils import misc -from ppsci.utils import reader - - -class NPZDataset(io.Dataset): - """Dataset class for .npz file. - - Args: - file_path (str): Npz file path. - input_keys (Tuple[str, ...]): List of input keys. - label_keys (Tuple[str, ...], optional): List of label keys. Defaults to (). - alias_dict (Optional[Dict[str, str]]): Dict of alias(es) for input and label keys. - i.e. {inner_key: outer_key}. Defaults to None. - weight_dict (Optional[Dict[str, Union[Callable, float]]]): Define the weight of - each constraint variable. Defaults to None. - timestamps (Optional[Tuple[float, ...]]): The number of repetitions of the data - in the time dimension. Defaults to None. - transforms (Optional[vision.Compose]): Compose object contains sample wise - transform(s). Defaults to None. - - Examples: - >>> import ppsci - >>> dataset = ppsci.data.dataset.NPZDataset( - ... "/path/to/file.npz" - ... ("x",), - ... ("u",), - ... ) # doctest: +SKIP - """ - - # Whether support batch indexing for speeding up fetching process. - batch_index: bool = True - - def __init__( - self, - file_path: str, - input_keys: Tuple[str, ...], - label_keys: Tuple[str, ...] = (), - alias_dict: Optional[Dict[str, str]] = None, - weight_dict: Optional[Dict[str, Union[Callable, float]]] = None, - timestamps: Optional[Tuple[float, ...]] = None, - transforms: Optional[vision.Compose] = None, - ): - super().__init__() - self.input_keys = input_keys - self.label_keys = label_keys - - # read raw data from file - raw_data = reader.load_npz_file( - file_path, - input_keys + label_keys, - alias_dict, - ) - # filter raw data by given timestamps if specified - if timestamps is not None: - if "t" in raw_data: - # filter data according to given timestamps - raw_time_array = raw_data["t"] - mask = [] - for ti in timestamps: - mask.append(np.nonzero(np.isclose(raw_time_array, ti).flatten())[0]) - raw_data = misc.convert_to_array( - raw_data, self.input_keys + self.label_keys - ) - mask = np.concatenate(mask, 0) - raw_data = raw_data[mask] - raw_data = misc.convert_to_dict( - raw_data, self.input_keys + self.label_keys - ) - else: - # repeat data according to given timestamps - raw_data = misc.convert_to_array( - raw_data, self.input_keys + self.label_keys - ) - raw_data = misc.combine_array_with_time(raw_data, timestamps) - self.input_keys = ("t",) + tuple(self.input_keys) - raw_data = misc.convert_to_dict( - raw_data, self.input_keys + self.label_keys - ) - - # fetch input data - self.input = { - key: value for key, value in raw_data.items() if key in self.input_keys - } - # fetch label data - self.label = { - key: value for key, value in raw_data.items() if key in self.label_keys - } - - # prepare weights - self.weight = {} - if weight_dict is not None: - for key, value in weight_dict.items(): - if isinstance(value, (int, float)): - self.weight[key] = np.full_like( - next(iter(self.label.values())), value - ) - elif callable(value): - func = value - self.weight[key] = func(self.input) - if isinstance(self.weight[key], (int, float)): - self.weight[key] = np.full_like( - next(iter(self.label.values())), self.weight[key] - ) - else: - raise NotImplementedError(f"type of {type(value)} is invalid yet.") - - self.transforms = transforms - self._len = len(next(iter(self.input.values()))) - - def __getitem__(self, idx): - input_item = {key: value[idx] for key, value in self.input.items()} - label_item = {key: value[idx] for key, value in self.label.items()} - weight_item = {key: value[idx] for key, value in self.weight.items()} - - if self.transforms is not None: - input_item, label_item, weight_item = self.transforms( - input_item, label_item, weight_item - ) - - return (input_item, label_item, weight_item) - - def __len__(self): - return self._len - - -class IterableNPZDataset(io.IterableDataset): - """IterableNPZDataset for full-data loading. - - Args: - file_path (str): Npz file path. - input_keys (Tuple[str, ...]): List of input keys. - label_keys (Tuple[str, ...], optional): List of label keys. Defaults to (). - alias_dict (Optional[Dict[str, str]]): Dict of alias(es) for input and label keys. - i.e. {inner_key: outer_key}. Defaults to None. - weight_dict (Optional[Dict[str, Union[Callable, float]]]): Define the weight of - each constraint variable. Defaults to None. - timestamps (Optional[Tuple[float, ...]]): The number of repetitions of the data - in the time dimension. Defaults to None. - transforms (Optional[vision.Compose]): Compose object contains sample wise - transform(s). Defaults to None. - - Examples: - >>> import ppsci - >>> dataset = ppsci.data.dataset.IterableNPZDataset( - ... "/path/to/file.npz" - ... ("x",), - ... ("u",), - ... ) # doctest: +SKIP - """ - - # Whether support batch indexing for speeding up fetching process. - batch_index: bool = False - - def __init__( - self, - file_path: str, - input_keys: Tuple[str, ...], - label_keys: Tuple[str, ...] = (), - alias_dict: Optional[Dict[str, str]] = None, - weight_dict: Optional[Dict[str, Union[Callable, float]]] = None, - timestamps: Optional[Tuple[float, ...]] = None, - transforms: Optional[vision.Compose] = None, - ): - super().__init__() - self.input_keys = input_keys - self.label_keys = label_keys - - # read raw data from file - raw_data = reader.load_npz_file( - file_path, - input_keys + label_keys, - alias_dict, - ) - # filter raw data by given timestamps if specified - if timestamps is not None: - if "t" in raw_data: - # filter data according to given timestamps - raw_time_array = raw_data["t"] - mask = [] - for ti in timestamps: - mask.append(np.nonzero(np.isclose(raw_time_array, ti).flatten())[0]) - raw_data = misc.convert_to_array( - raw_data, self.input_keys + self.label_keys - ) - mask = np.concatenate(mask, 0) - raw_data = raw_data[mask] - raw_data = misc.convert_to_dict( - raw_data, self.input_keys + self.label_keys - ) - else: - # repeat data according to given timestamps - raw_data = misc.convert_to_array( - raw_data, self.input_keys + self.label_keys - ) - raw_data = misc.combine_array_with_time(raw_data, timestamps) - self.input_keys = ("t",) + tuple(self.input_keys) - raw_data = misc.convert_to_dict( - raw_data, self.input_keys + self.label_keys - ) - - # fetch input data - self.input = { - key: value for key, value in raw_data.items() if key in self.input_keys - } - # fetch label data - self.label = { - key: value for key, value in raw_data.items() if key in self.label_keys - } - - # prepare weights - self.weight = {} - if weight_dict is not None: - for key, value in weight_dict.items(): - if isinstance(value, (int, float)): - self.weight[key] = np.full_like( - next(iter(self.label.values())), value - ) - elif callable(value): - func = value - self.weight[key] = func(self.input) - if isinstance(self.weight[key], (int, float)): - self.weight[key] = np.full_like( - next(iter(self.label.values())), self.weight[key] - ) - else: - raise NotImplementedError(f"type of {type(value)} is invalid yet.") - - self.input = {key: paddle.to_tensor(value) for key, value in self.input.items()} - self.label = {key: paddle.to_tensor(value) for key, value in self.label.items()} - self.weight = { - key: paddle.to_tensor(value) for key, value in self.weight.items() - } - - self.transforms = transforms - self._len = len(next(iter(self.input.values()))) - - @property - def num_samples(self): - """Number of samples within current dataset.""" - return self._len - - def __iter__(self): - if callable(self.transforms): - input_, label_, weight_ = self.transforms( - self.input, self.label, self.weight - ) - yield input_, label_, weight_ - else: - yield self.input, self.label, self.weight - - def __len__(self): - return 1 diff --git a/examples/smc_reac/ppsci/data/dataset/pems_dataset.py b/examples/smc_reac/ppsci/data/dataset/pems_dataset.py deleted file mode 100644 index 1e3dde0a3d..0000000000 --- a/examples/smc_reac/ppsci/data/dataset/pems_dataset.py +++ /dev/null @@ -1,151 +0,0 @@ -import os -from typing import Dict -from typing import Optional -from typing import Tuple - -import numpy as np -import pandas as pd -from paddle.io import Dataset -from paddle.vision.transforms import Compose - - -class StandardScaler: - def __init__(self, mean, std): - self.mean = mean - self.std = std - - def transform(self, data): - return (data - self.mean) / self.std - - def inverse_transform(self, data): - return (data * self.std) + self.mean - - -def add_window_horizon(data, in_step=12, out_step=12): - length = len(data) - end_index = length - out_step - in_step - X = [] - Y = [] - for i in range(end_index + 1): - X.append(data[i : i + in_step]) - Y.append(data[i + in_step : i + in_step + out_step]) - return X, Y - - -def get_edge_index(file_path, bi=True, reduce="mean"): - TYPE_DICT = {0: np.int64, 1: np.int64, 2: np.float32} - df = pd.read_csv( - os.path.join(file_path, "dist.csv"), - skiprows=1, - header=None, - sep=",", - dtype=TYPE_DICT, - ) - - edge_index = df.loc[:, [0, 1]].values.T - edge_attr = df.loc[:, 2].values - - if bi: - re_edge_index = np.concatenate((edge_index[1:, :], edge_index[:1, :]), axis=0) - edge_index = np.concatenate((edge_index, re_edge_index), axis=-1) - edge_attr = np.concatenate((edge_attr, edge_attr), axis=0) - - num = np.max(edge_index) + 1 - adj = np.zeros((num, num), dtype=np.float32) - - if reduce == "sum": - adj[edge_index[0], edge_index[1]] = 1.0 - elif reduce == "mean": - adj[edge_index[0], edge_index[1]] = 1.0 - adj = adj / adj.sum(axis=-1) - else: - raise ValueError - - return edge_index, edge_attr, adj - - -class PEMSDataset(Dataset): - """Dataset class for PEMSD4 and PEMSD8 dataset. - - Args: - file_path (str): Dataset root path. - split (str): Dataset split label. - input_keys (Tuple[str, ...]): A tuple of input keys. - label_keys (Tuple[str, ...]): A tuple of label keys. - weight_dict (Optional[Dict[str, float]]): Define the weight of each constraint variable. Defaults to None. - transforms (Optional[Compose]): Compose object contains sample wise transform(s). Defaults to None. - norm_input (bool): Whether to normalize the input. Defaults to True. - norm_label (bool): Whether to normalize the output. Defaults to False. - input_len (int): The input timesteps. Defaults to 12. - label_len (int): The output timesteps. Defaults to 12. - - Examples: - >>> import ppsci - >>> dataset = ppsci.data.dataset.PEMSDataset( - ... "./Data/PEMSD4", - ... "train", - ... ("input",), - ... ("label",), - ... ) # doctest: +SKIP - """ - - def __init__( - self, - file_path: str, - split: str, - input_keys: Tuple[str, ...], - label_keys: Tuple[str, ...], - weight_dict: Optional[Dict[str, float]] = None, - transforms: Optional[Compose] = None, - norm_input: bool = True, - norm_label: bool = False, - input_len: int = 12, - label_len: int = 12, - ): - super().__init__() - - self.input_keys = input_keys - self.label_keys = label_keys - self.weight_dict = weight_dict - - self.transforms = transforms - self.norm_input = norm_input - self.norm_label = norm_label - - data = np.load(os.path.join(file_path, f"{split}.npy")).astype(np.float32) - - self.mean = np.load(os.path.join(file_path, "mean.npy")).astype(np.float32) - self.std = np.load(os.path.join(file_path, "std.npy")).astype(np.float32) - self.scaler = StandardScaler(self.mean, self.std) - - X, Y = add_window_horizon(data, input_len, label_len) - if norm_input: - X = self.scaler.transform(X) - if norm_label: - Y = self.scaler.transform(Y) - - self._len = X.shape[0] - - self.input = {input_keys[0]: X} - self.label = {label_keys[0]: Y} - - if weight_dict is not None: - self.weight_dict = {key: np.array(1.0) for key in self.label_keys} - self.weight_dict.update(weight_dict) - else: - self.weight = {} - - def __getitem__(self, idx): - input_item = {key: value[idx] for key, value in self.input.items()} - label_item = {key: value[idx] for key, value in self.label.items()} - weight_item = {key: value[idx] for key, value in self.weight.items()} - - if self.transforms is not None: - input_item, label_item, weight_item = self.transforms( - input_item, label_item, weight_item - ) - - return (input_item, label_item, weight_item) - - def __len__(self): - return self._len diff --git a/examples/smc_reac/ppsci/data/dataset/radar_dataset.py b/examples/smc_reac/ppsci/data/dataset/radar_dataset.py deleted file mode 100644 index e484558455..0000000000 --- a/examples/smc_reac/ppsci/data/dataset/radar_dataset.py +++ /dev/null @@ -1,146 +0,0 @@ -# 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. - -from __future__ import annotations - -import os -from typing import Dict -from typing import Optional -from typing import Tuple - -try: - import cv2 -except ModuleNotFoundError: - pass - -import importlib - -import numpy as np -import paddle -from paddle import io - - -class RadarDataset(io.Dataset): - """Class for Radar dataset. - - Args: - input_keys (Tuple[str, ...]): Input keys, such as ("input",). - label_keys (Tuple[str, ...]): Output keys, such as ("output",). - image_width (int): Image width. - image_height (int): Image height. - total_length (int): Total length. - dataset_path (str): Dataset path. - data_type (str): Input and output data type. Defaults to paddle.get_default_dtype(). - weight_dict (Optional[Dict[str, float]]): Weight dictionary. Defaults to None. - - Examples: - >>> import ppsci - >>> dataset = ppsci.data.dataset.RadarDataset( - ... "input_keys": ("input",), - ... "label_keys": ("output",), - ... "image_width": 512, - ... "image_height": 512, - ... "total_length": 29, - ... "dataset_path": "datasets/mrms/figure", - ... "data_type": paddle.get_default_dtype(), - ... ) # doctest: +SKIP - """ - - # Whether support batch indexing for speeding up fetching process. - batch_index: bool = False - - def __init__( - self, - input_keys: Tuple[str, ...], - label_keys: Tuple[str, ...], - image_width: int, - image_height: int, - total_length: int, - dataset_path: str, - data_type: str = paddle.get_default_dtype(), - weight_dict: Optional[Dict[str, float]] = None, - ): - super().__init__() - if importlib.util.find_spec("cv2") is None: - raise ModuleNotFoundError( - "To use RadarDataset, please install 'opencv-python' with: `pip install " - "opencv-python` first." - ) - self.input_keys = input_keys - self.label_keys = label_keys - self.img_width = image_width - self.img_height = image_height - self.length = total_length - self.dataset_path = dataset_path - self.data_type = data_type - - self.weight_dict = {} if weight_dict is None else weight_dict - if weight_dict is not None: - self.weight_dict = {key: 1.0 for key in self.label_keys} - self.weight_dict.update(weight_dict) - - self.case_list = [] - name_list = os.listdir(self.dataset_path) - name_list.sort() - for name in name_list: - case = [] - for i in range(29): - case.append( - self.dataset_path - + "/" - + name - + "/" - + name - + "-" - + str(i).zfill(2) - + ".png" - ) - self.case_list.append(case) - - def _load(self, index): - data = [] - for img_path in self.case_list[index]: - img = cv2.imread(img_path, 2) - data.append(np.expand_dims(img, axis=0)) - data = np.concatenate(data, axis=0).astype(self.data_type) / 10.0 - 3.0 - assert data.shape[1] <= 1024 and data.shape[2] <= 1024 - return data - - def __getitem__(self, index): - data = self._load(index)[-self.length :].copy() - mask = np.ones_like(data) - mask[data < 0] = 0 - data[data < 0] = 0 - data = np.clip(data, 0, 128) - vid = np.zeros( - (self.length, self.img_height, self.img_width, 2), dtype=self.data_type - ) - vid[..., 0] = data - vid[..., 1] = mask - - input_item = {self.input_keys[0]: vid} - label_item = {} - weight_item = {} - for key in self.label_keys: - label_item[key] = np.asarray([], paddle.get_default_dtype()) - if len(label_item) > 0: - weight_shape = [1] * len(next(iter(label_item.values())).shape) - weight_item = { - key: np.full(weight_shape, value, paddle.get_default_dtype()) - for key, value in self.weight_dict.items() - } - return input_item, label_item, weight_item - - def __len__(self): - return len(self.case_list) diff --git a/examples/smc_reac/ppsci/data/dataset/sevir_dataset.py b/examples/smc_reac/ppsci/data/dataset/sevir_dataset.py deleted file mode 100644 index 42ae274c2b..0000000000 --- a/examples/smc_reac/ppsci/data/dataset/sevir_dataset.py +++ /dev/null @@ -1,814 +0,0 @@ -import datetime -import os -from copy import deepcopy -from typing import Dict -from typing import Optional -from typing import Sequence -from typing import Tuple -from typing import Union - -try: - import h5py -except ModuleNotFoundError: - pass -import numpy as np -import paddle -import paddle.nn.functional as F -import pandas as pd -from paddle import io - -# SEVIR Dataset constants -SEVIR_DATA_TYPES = ["vis", "ir069", "ir107", "vil", "lght"] -SEVIR_RAW_DTYPES = { - "vis": np.int16, - "ir069": np.int16, - "ir107": np.int16, - "vil": np.uint8, - "lght": np.int16, -} -LIGHTING_FRAME_TIMES = np.arange(-120.0, 125.0, 5) * 60 -SEVIR_DATA_SHAPE = { - "lght": (48, 48), -} -PREPROCESS_SCALE_SEVIR = { - "vis": 1, # Not utilized in original paper - "ir069": 1 / 1174.68, - "ir107": 1 / 2562.43, - "vil": 1 / 47.54, - "lght": 1 / 0.60517, -} -PREPROCESS_OFFSET_SEVIR = { - "vis": 0, # Not utilized in original paper - "ir069": 3683.58, - "ir107": 1552.80, - "vil": -33.44, - "lght": -0.02990, -} -PREPROCESS_SCALE_01 = { - "vis": 1, - "ir069": 1, - "ir107": 1, - "vil": 1 / 255, # currently the only one implemented - "lght": 1, -} -PREPROCESS_OFFSET_01 = { - "vis": 0, - "ir069": 0, - "ir107": 0, - "vil": 0, # currently the only one implemented - "lght": 0, -} - - -def change_layout_np(data, in_layout="NHWT", out_layout="NHWT", ret_contiguous=False): - # first convert to 'NHWT' - if in_layout == "NHWT": - pass - elif in_layout == "NTHW": - data = np.transpose(data, axes=(0, 2, 3, 1)) - elif in_layout == "NWHT": - data = np.transpose(data, axes=(0, 2, 1, 3)) - elif in_layout == "NTCHW": - data = data[:, :, 0, :, :] - data = np.transpose(data, axes=(0, 2, 3, 1)) - elif in_layout == "NTHWC": - data = data[:, :, :, :, 0] - data = np.transpose(data, axes=(0, 2, 3, 1)) - elif in_layout == "NTWHC": - data = data[:, :, :, :, 0] - data = np.transpose(data, axes=(0, 3, 2, 1)) - elif in_layout == "TNHW": - data = np.transpose(data, axes=(1, 2, 3, 0)) - elif in_layout == "TNCHW": - data = data[:, :, 0, :, :] - data = np.transpose(data, axes=(1, 2, 3, 0)) - else: - raise NotImplementedError(f"{in_layout} is invalid.") - - if out_layout == "NHWT": - pass - elif out_layout == "NTHW": - data = np.transpose(data, axes=(0, 3, 1, 2)) - elif out_layout == "NWHT": - data = np.transpose(data, axes=(0, 2, 1, 3)) - elif out_layout == "NTCHW": - data = np.transpose(data, axes=(0, 3, 1, 2)) - data = np.expand_dims(data, axis=2) - elif out_layout == "NTHWC": - data = np.transpose(data, axes=(0, 3, 1, 2)) - data = np.expand_dims(data, axis=-1) - elif out_layout == "NTWHC": - data = np.transpose(data, axes=(0, 3, 2, 1)) - data = np.expand_dims(data, axis=-1) - elif out_layout == "TNHW": - data = np.transpose(data, axes=(3, 0, 1, 2)) - elif out_layout == "TNCHW": - data = np.transpose(data, axes=(3, 0, 1, 2)) - data = np.expand_dims(data, axis=2) - else: - raise NotImplementedError(f"{out_layout} is invalid.") - if ret_contiguous: - data = data.ascontiguousarray() - return data - - -def change_layout_paddle( - data, in_layout="NHWT", out_layout="NHWT", ret_contiguous=False -): - # first convert to 'NHWT' - if in_layout == "NHWT": - pass - elif in_layout == "NTHW": - data = data.transpose(perm=[0, 2, 3, 1]) - elif in_layout == "NTCHW": - data = data[:, :, 0, :, :] - data = data.transpose(perm=[0, 2, 3, 1]) - elif in_layout == "NTHWC": - data = data[:, :, :, :, 0] - data = data.transpose(perm=[0, 2, 3, 1]) - elif in_layout == "TNHW": - data = data.transpose(perm=[1, 2, 3, 0]) - elif in_layout == "TNCHW": - data = data[:, :, 0, :, :] - data = data.transpose(perm=[1, 2, 3, 0]) - else: - raise NotImplementedError(f"{in_layout} is invalid.") - - if out_layout == "NHWT": - pass - elif out_layout == "NTHW": - data = data.transpose(perm=[0, 3, 1, 2]) - elif out_layout == "NTCHW": - data = data.transpose(perm=[0, 3, 1, 2]) - data = paddle.unsqueeze(data, axis=2) - elif out_layout == "NTHWC": - data = data.transpose(perm=[0, 3, 1, 2]) - data = paddle.unsqueeze(data, axis=-1) - elif out_layout == "TNHW": - data = data.transpose(perm=[3, 0, 1, 2]) - elif out_layout == "TNCHW": - data = data.transpose(perm=[3, 0, 1, 2]) - data = paddle.unsqueeze(data, axis=2) - else: - raise NotImplementedError(f"{out_layout} is invalid.") - return data - - -def path_splitall(path): - allparts = [] - while 1: - parts = os.path.split(path) - if parts[0] == path: # sentinel for absolute paths - allparts.insert(0, parts[0]) - break - elif parts[1] == path: # sentinel for relative paths - allparts.insert(0, parts[1]) - break - else: - path = parts[0] - allparts.insert(0, parts[1]) - return allparts - - -class SEVIRDataset(io.Dataset): - """The Storm EVent ImagRy dataset. - - Args: - input_keys (Tuple[str, ...]): Name of input keys, such as ("input",). - label_keys (Tuple[str, ...]): Name of label keys, such as ("output",). - data_dir (str): The path of the dataset. - weight_dict (Optional[Dict[str, Union[Callable, float]]]): Define the weight of each constraint variable. Defaults to None. - data_types (Sequence[str], optional): A subset of SEVIR_DATA_TYPES. Defaults to [ "vil", ]. - seq_len (int, optional): The length of the data sequences. Should be smaller than the max length raw_seq_len. Defaults to 49. - raw_seq_len (int, optional): The length of the raw data sequences. Defaults to 49. - sample_mode (str, optional): The mode of sampling, eg.'random' or 'sequent'. Defaults to "sequent". - stride (int, optional): Useful when sample_mode == 'sequent' - stride must not be smaller than out_len to prevent data leakage in testing. Defaults to 12. - batch_size (int, optional): The batch size. Defaults to 1. - layout (str, optional): Consists of batch_size 'N', seq_len 'T', channel 'C', height 'H', width 'W' - The layout of sampled data. Raw data layout is 'NHWT'. - valid layout: 'NHWT', 'NTHW', 'NTCHW', 'TNHW', 'TNCHW'. Defaults to "NHWT". - in_len (int, optional): The length of input data. Defaults to 13. - out_len (int, optional): The length of output data. Defaults to 12. - num_shard (int, optional): Split the whole dataset into num_shard parts for distributed training. Defaults to 1. - rank (int, optional): Rank of the current process within num_shard. Defaults to 0. - split_mode (str, optional): If 'ceil', all `num_shard` dataloaders have the same length = ceil(total_len / num_shard). - Different dataloaders may have some duplicated data batches, if the total size of datasets is not divided by num_shard. - if 'floor', all `num_shard` dataloaders have the same length = floor(total_len / num_shard). - The last several data batches may be wasted, if the total size of datasets is not divided by num_shard. - if 'uneven', the last datasets has larger length when the total length is not divided by num_shard. - The uneven split leads to synchronization error in dist.all_reduce() or dist.barrier(). - See related issue: https://github.com/pytorch/pytorch/issues/33148 - Notice: this also affects the behavior of `self.use_up`. Defaults to "uneven". - start_date (datetime.datetime, optional): Start time of SEVIR samples to generate. Defaults to None. - end_date (datetime.datetime, optional): End time of SEVIR samples to generate. Defaults to None. - datetime_filter (function, optional): Mask function applied to time_utc column of catalog (return true to keep the row). - Pass function of the form lambda t : COND(t) - Example: lambda t: np.logical_and(t.dt.hour>=13,t.dt.hour<=21) # Generate only day-time events. Defaults to None. - catalog_filter (function, optional): Function or None or 'default' - Mask function applied to entire catalog dataframe (return true to keep row). - Pass function of the form lambda catalog: COND(catalog) - Example: lambda c: [s[0]=='S' for s in c.id] # Generate only the 'S' events - shuffle (bool, optional): If True, data samples are shuffled before each epoch. Defaults to False. - shuffle_seed (int, optional): Seed to use for shuffling. Defaults to 1. - output_type (np.dtype, optional): The type of generated tensors. Defaults to np.float32. - preprocess (bool, optional): If True, self.preprocess_data_dict(data_dict) is called before each sample generated. Defaults to True. - rescale_method (str, optional): The method of rescale. Defaults to "01". - downsample_dict (Dict[str, Sequence[int]], optional): Downsample_dict.keys() == data_types. - downsample_dict[key] is a Sequence of (t_factor, h_factor, w_factor),representing the downsampling factors of all dimensions. Defaults to None. - verbose (bool, optional): Verbose when opening raw data files. Defaults to False. - training (str, optional): Training pathse. Defaults to "train". - """ - - # Whether support batch indexing for speeding up fetching process. - batch_index: bool = False - - def __init__( - self, - input_keys: Tuple[str, ...], - label_keys: Tuple[str, ...], - data_dir: str, - weight_dict: Optional[Dict[str, float]] = None, - data_types: Sequence[str] = [ - "vil", - ], - seq_len: int = 49, - raw_seq_len: int = 49, - sample_mode: str = "sequent", - stride: int = 12, - batch_size: int = 1, - layout: str = "NHWT", - in_len: int = 13, - out_len: int = 12, - num_shard: int = 1, - rank: int = 0, - split_mode: str = "uneven", - start_date: datetime.datetime = None, - end_date: datetime.datetime = None, - datetime_filter=None, - catalog_filter="default", - shuffle: bool = False, - shuffle_seed: int = 1, - output_type=np.float32, - preprocess: bool = True, - rescale_method: str = "01", - downsample_dict: Dict[str, Sequence[int]] = None, - verbose: bool = False, - training="train", - ): - super(SEVIRDataset, self).__init__() - self.input_keys = input_keys - self.label_keys = label_keys - self.data_dir = data_dir - self.weight_dict = {} if weight_dict is None else weight_dict - if weight_dict is not None: - self.weight_dict = {key: 1.0 for key in self.label_keys} - self.weight_dict.update(weight_dict) - - # sevir - SEVIR_ROOT_DIR = os.path.join(self.data_dir, "sevir") - sevir_catalog = os.path.join(SEVIR_ROOT_DIR, "CATALOG.csv") - sevir_data_dir = os.path.join(SEVIR_ROOT_DIR, "data") - # sevir-lr - # SEVIR_ROOT_DIR = os.path.join(self.data_dir, "sevir_lr") - # SEVIR_CATALOG = os.path.join(SEVIR_ROOT_DIR, "CATALOG.csv") - # SEVIR_DATA_DIR = os.path.join(SEVIR_ROOT_DIR, "data") - - if data_types is None: - data_types = SEVIR_DATA_TYPES - else: - assert set(data_types).issubset(SEVIR_DATA_TYPES) - - # configs which should not be modified - self._dtypes = SEVIR_RAW_DTYPES - self.lght_frame_times = LIGHTING_FRAME_TIMES - self.data_shape = SEVIR_DATA_SHAPE - - self.raw_seq_len = raw_seq_len - self.seq_len = seq_len - - if seq_len > raw_seq_len: - raise ValueError("seq_len must be small than raw_seq_len") - - if sample_mode not in ["random", "sequent"]: - raise ValueError("sample_mode must be 'random' or 'sequent'.") - - self.sample_mode = sample_mode - self.stride = stride - self.batch_size = batch_size - valid_layout = ("NHWT", "NTHW", "NTCHW", "NTHWC", "TNHW", "TNCHW") - if layout not in valid_layout: - raise ValueError( - f"Invalid layout = {layout}! Must be one of {valid_layout}." - ) - self.layout = layout - self.in_len = in_len - self.out_len = out_len - - self.num_shard = num_shard - self.rank = rank - valid_split_mode = ("ceil", "floor", "uneven") - if split_mode not in valid_split_mode: - raise ValueError( - f"Invalid split_mode: {split_mode}! Must be one of {valid_split_mode}." - ) - self.split_mode = split_mode - self._samples = None - self._hdf_files = {} - self.data_types = data_types - if isinstance(sevir_catalog, str): - self.catalog = pd.read_csv( - sevir_catalog, parse_dates=["time_utc"], low_memory=False - ) - else: - self.catalog = sevir_catalog - self.sevir_data_dir = sevir_data_dir - self.datetime_filter = datetime_filter - self.catalog_filter = catalog_filter - self.start_date = start_date - self.end_date = end_date - # train val test split - self.start_date = ( - datetime.datetime(*start_date) if start_date is not None else None - ) - self.end_date = datetime.datetime(*end_date) if end_date is not None else None - - self.shuffle = shuffle - self.shuffle_seed = int(shuffle_seed) - self.output_type = output_type - self.preprocess = preprocess - self.downsample_dict = downsample_dict - self.rescale_method = rescale_method - self.verbose = verbose - - if self.start_date is not None: - self.catalog = self.catalog[self.catalog.time_utc > self.start_date] - if self.end_date is not None: - self.catalog = self.catalog[self.catalog.time_utc <= self.end_date] - if self.datetime_filter: - self.catalog = self.catalog[self.datetime_filter(self.catalog.time_utc)] - - if self.catalog_filter is not None: - if self.catalog_filter == "default": - self.catalog_filter = lambda c: c.pct_missing == 0 - self.catalog = self.catalog[self.catalog_filter(self.catalog)] - - self._compute_samples() - self._open_files(verbose=self.verbose) - - def _compute_samples(self): - """ - Computes the list of samples in catalog to be used. This sets self._samples - """ - # locate all events containing colocated data_types - imgt = self.data_types - imgts = set(imgt) - filtcat = self.catalog[ - np.logical_or.reduce([self.catalog.img_type == i for i in imgt]) - ] - # remove rows missing one or more requested img_types - filtcat = filtcat.groupby("id").filter( - lambda x: imgts.issubset(set(x["img_type"])) - ) - # If there are repeated IDs, remove them (this is a bug in SEVIR) - # TODO: is it necessary to keep one of them instead of deleting them all - filtcat = filtcat.groupby("id").filter(lambda x: x.shape[0] == len(imgt)) - self._samples = filtcat.groupby("id").apply( - lambda df: self._df_to_series(df, imgt) - ) - if self.shuffle: - self.shuffle_samples() - - def shuffle_samples(self): - self._samples = self._samples.sample(frac=1, random_state=self.shuffle_seed) - - def _df_to_series(self, df, imgt): - d = {} - df = df.set_index("img_type") - for i in imgt: - s = df.loc[i] - idx = s.file_index if i != "lght" else s.id - d.update({f"{i}_filename": [s.file_name], f"{i}_index": [idx]}) - - return pd.DataFrame(d) - - def _open_files(self, verbose=True): - """ - Opens HDF files - """ - imgt = self.data_types - hdf_filenames = [] - for t in imgt: - hdf_filenames += list(np.unique(self._samples[f"{t}_filename"].values)) - self._hdf_files = {} - for f in hdf_filenames: - if verbose: - print("Opening HDF5 file for reading", f) - self._hdf_files[f] = h5py.File(self.sevir_data_dir + "/" + f, "r") - - def close(self): - """ - Closes all open file handles - """ - for f in self._hdf_files: - self._hdf_files[f].close() - self._hdf_files = {} - - @property - def num_seq_per_event(self): - return 1 + (self.raw_seq_len - self.seq_len) // self.stride - - @property - def total_num_seq(self): - """ - The total number of sequences within each shard. - Notice that it is not the product of `self.num_seq_per_event` and `self.total_num_event`. - """ - return int(self.num_seq_per_event * self.num_event) - - @property - def total_num_event(self): - """ - The total number of events in the whole dataset, before split into different shards. - """ - return int(self._samples.shape[0]) - - @property - def start_event_idx(self): - """ - The event idx used in certain rank should satisfy event_idx >= start_event_idx - """ - return self.total_num_event // self.num_shard * self.rank - - @property - def end_event_idx(self): - """ - The event idx used in certain rank should satisfy event_idx < end_event_idx - - """ - if self.split_mode == "ceil": - _last_start_event_idx = ( - self.total_num_event // self.num_shard * (self.num_shard - 1) - ) - _num_event = self.total_num_event - _last_start_event_idx - return self.start_event_idx + _num_event - elif self.split_mode == "floor": - return self.total_num_event // self.num_shard * (self.rank + 1) - else: # self.split_mode == 'uneven': - if self.rank == self.num_shard - 1: # the last process - return self.total_num_event - else: - return self.total_num_event // self.num_shard * (self.rank + 1) - - @property - def num_event(self): - """ - The number of events split into each rank - """ - return self.end_event_idx - self.start_event_idx - - def __len__(self): - """ - Used only when self.sample_mode == 'sequent' - """ - return self.total_num_seq // self.batch_size - - def _read_data(self, row, data): - """ - Iteratively read data into data dict. Finally data[imgt] gets shape (batch_size, height, width, raw_seq_len). - - Args: - row (Dict,optional): A series with fields IMGTYPE_filename, IMGTYPE_index, IMGTYPE_time_index. - data (Dict,optional): , data[imgt] is a data tensor with shape = (tmp_batch_size, height, width, raw_seq_len). - - Returns: - data (np.array): Updated data. Updated shape = (tmp_batch_size + 1, height, width, raw_seq_len). - """ - - imgtyps = np.unique([x.split("_")[0] for x in list(row.keys())]) - for t in imgtyps: - fname = row[f"{t}_filename"] - idx = row[f"{t}_index"] - t_slice = slice(0, None) - # Need to bin lght counts into grid - if t == "lght": - lght_data = self._hdf_files[fname][idx][:] - data_i = self._lght_to_grid(lght_data, t_slice) - else: - data_i = self._hdf_files[fname][t][idx : idx + 1, :, :, t_slice] - data[t] = ( - np.concatenate((data[t], data_i), axis=0) if (t in data) else data_i - ) - return data - - def _lght_to_grid(self, data, t_slice=slice(0, None)): - """ - Converts Nx5 lightning data matrix into a 2D grid of pixel counts - """ - # out_size = (48,48,len(self.lght_frame_times)-1) if isinstance(t_slice,(slice,)) else (48,48) - out_size = ( - (*self.data_shape["lght"], len(self.lght_frame_times)) - if t_slice.stop is None - else (*self.data_shape["lght"], 1) - ) - if data.shape[0] == 0: - return np.zeros((1,) + out_size, dtype=np.float32) - - # filter out points outside the grid - x, y = data[:, 3], data[:, 4] - m = np.logical_and.reduce([x >= 0, x < out_size[0], y >= 0, y < out_size[1]]) - data = data[m, :] - if data.shape[0] == 0: - return np.zeros((1,) + out_size, dtype=np.float32) - - # Filter/separate times - t = data[:, 0] - if t_slice.stop is not None: # select only one time bin - if t_slice.stop > 0: - if t_slice.stop < len(self.lght_frame_times): - tm = np.logical_and( - t >= self.lght_frame_times[t_slice.stop - 1], - t < self.lght_frame_times[t_slice.stop], - ) - else: - tm = t >= self.lght_frame_times[-1] - else: # special case: frame 0 uses lght from frame 1 - tm = np.logical_and( - t >= self.lght_frame_times[0], t < self.lght_frame_times[1] - ) - # tm=np.logical_and( (t>=FRAME_TIMES[t_slice],t self.end_event_idx: - pad_size = event_idx_slice_end - self.end_event_idx - event_idx_slice_end = self.end_event_idx - pd_batch = self._samples.iloc[event_idx:event_idx_slice_end] - data = {} - for index, row in pd_batch.iterrows(): - data = self._read_data(row, data) - if pad_size > 0: - event_batch = [] - for t in self.data_types: - pad_shape = [ - pad_size, - ] + list(data[t].shape[1:]) - data_pad = np.concatenate( - ( - data[t].astype(self.output_type), - np.zeros(pad_shape, dtype=self.output_type), - ), - axis=0, - ) - event_batch.append(data_pad) - else: - event_batch = [data[t].astype(self.output_type) for t in self.data_types] - return event_batch - - def __iter__(self): - return self - - @staticmethod - def preprocess_data_dict( - data_dict, data_types=None, layout="NHWT", rescale="01" - ) -> Dict[str, Union[np.ndarray, paddle.Tensor]]: - """The preprocess of data dict. - - Args: - data_dict (Dict[str, Union[np.ndarray, paddle.Tensor]]): The dict of data. - data_types (Sequence[str]) : The data types that we want to rescale. This mainly excludes "mask" from preprocessing. - layout (str) : consists of batch_size 'N', seq_len 'T', channel 'C', height 'H', width 'W'. - rescale (str): - 'sevir': use the offsets and scale factors in original implementation. - '01': scale all values to range 0 to 1, currently only supports 'vil'. - - Returns: - data_dict (Dict[str, Union[np.ndarray, paddle.Tensor]]): preprocessed data. - """ - - if rescale == "sevir": - scale_dict = PREPROCESS_SCALE_SEVIR - offset_dict = PREPROCESS_OFFSET_SEVIR - elif rescale == "01": - scale_dict = PREPROCESS_SCALE_01 - offset_dict = PREPROCESS_OFFSET_01 - else: - raise ValueError(f"Invalid rescale option: {rescale}.") - if data_types is None: - data_types = data_dict.keys() - for key, data in data_dict.items(): - if key in data_types: - if isinstance(data, np.ndarray): - data = scale_dict[key] * ( - data.astype(np.float32) + offset_dict[key] - ) - data = change_layout_np( - data=data, in_layout="NHWT", out_layout=layout - ) - elif isinstance(data, paddle.Tensor): - data = scale_dict[key] * (data.astype("float32") + offset_dict[key]) - data = change_layout_paddle( - data=data, in_layout="NHWT", out_layout=layout - ) - data_dict[key] = data - return data_dict - - @staticmethod - def process_data_dict_back(data_dict, data_types=None, rescale="01"): - if rescale == "sevir": - scale_dict = PREPROCESS_SCALE_SEVIR - offset_dict = PREPROCESS_OFFSET_SEVIR - elif rescale == "01": - scale_dict = PREPROCESS_SCALE_01 - offset_dict = PREPROCESS_OFFSET_01 - else: - raise ValueError(f"Invalid rescale option: {rescale}.") - if data_types is None: - data_types = data_dict.keys() - for key in data_types: - data = data_dict[key] - data = data.astype("float32") / scale_dict[key] - offset_dict[key] - data_dict[key] = data - return data_dict - - @staticmethod - def data_dict_to_tensor(data_dict, data_types=None): - """ - Convert each element in data_dict to paddle.Tensor (copy without grad). - """ - ret_dict = {} - if data_types is None: - data_types = data_dict.keys() - for key, data in data_dict.items(): - if key in data_types: - if isinstance(data, paddle.Tensor): - ret_dict[key] = data.detach().clone() - elif isinstance(data, np.ndarray): - ret_dict[key] = paddle.to_tensor(data) - else: - raise ValueError( - f"Invalid data type: {type(data)}. Should be paddle.Tensor or np.ndarray" - ) - else: # key == "mask" - ret_dict[key] = data - return ret_dict - - @staticmethod - def downsample_data_dict( - data_dict, data_types=None, factors_dict=None, layout="NHWT" - ) -> Dict[str, paddle.Tensor]: - """The downsample of data. - - Args: - data_dict (Dict[str, Union[np.array, paddle.Tensor]]): The dict of data. - factors_dict (Optional[Dict[str, Sequence[int]]]):each element `factors` is - a Sequence of int, representing (t_factor, h_factor, w_factor). - - Returns: - downsampled_data_dict (Dict[str, paddle.Tensor]): Modify on a deep copy of - data_dict instead of directly modifying the original data_dict. - """ - - if factors_dict is None: - factors_dict = {} - if data_types is None: - data_types = data_dict.keys() - downsampled_data_dict = SEVIRDataset.data_dict_to_tensor( - data_dict=data_dict, data_types=data_types - ) # make a copy - for key, data in data_dict.items(): - factors = factors_dict.get(key, None) - if factors is not None: - downsampled_data_dict[key] = change_layout_paddle( - data=downsampled_data_dict[key], in_layout=layout, out_layout="NTHW" - ) - # downsample t dimension - t_slice = [ - slice(None, None), - ] * 4 - t_slice[1] = slice(None, None, factors[0]) - downsampled_data_dict[key] = downsampled_data_dict[key][tuple(t_slice)] - # downsample spatial dimensions - downsampled_data_dict[key] = F.avg_pool2d( - input=downsampled_data_dict[key], - kernel_size=(factors[1], factors[2]), - ) - - downsampled_data_dict[key] = change_layout_paddle( - data=downsampled_data_dict[key], in_layout="NTHW", out_layout=layout - ) - - return downsampled_data_dict - - def layout_to_in_out_slice( - self, - ): - t_axis = self.layout.find("T") - num_axes = len(self.layout) - in_slice = [ - slice(None, None), - ] * num_axes - out_slice = deepcopy(in_slice) - in_slice[t_axis] = slice(None, self.in_len) - if self.out_len is None: - out_slice[t_axis] = slice(self.in_len, None) - else: - out_slice[t_axis] = slice(self.in_len, self.in_len + self.out_len) - return in_slice, out_slice - - def __getitem__(self, index): - event_idx = (index * self.batch_size) // self.num_seq_per_event - seq_idx = (index * self.batch_size) % self.num_seq_per_event - num_sampled = 0 - sampled_idx_list = [] # list of (event_idx, seq_idx) records - while num_sampled < self.batch_size: - sampled_idx_list.append({"event_idx": event_idx, "seq_idx": seq_idx}) - seq_idx += 1 - if seq_idx >= self.num_seq_per_event: - event_idx += 1 - seq_idx = 0 - num_sampled += 1 - - start_event_idx = sampled_idx_list[0]["event_idx"] - event_batch_size = sampled_idx_list[-1]["event_idx"] - start_event_idx + 1 - - event_batch = self._load_event_batch( - event_idx=start_event_idx, event_batch_size=event_batch_size - ) - ret_dict = {} - for sampled_idx in sampled_idx_list: - batch_slice = [ - sampled_idx["event_idx"] - start_event_idx, - ] # use [] to keepdim - seq_slice = slice( - sampled_idx["seq_idx"] * self.stride, - sampled_idx["seq_idx"] * self.stride + self.seq_len, - ) - for imgt_idx, imgt in enumerate(self.data_types): - sampled_seq = event_batch[imgt_idx][batch_slice, :, :, seq_slice] - if imgt in ret_dict: - ret_dict[imgt] = np.concatenate( - (ret_dict[imgt], sampled_seq), axis=0 - ) - else: - ret_dict.update({imgt: sampled_seq}) - - ret_dict = self.data_dict_to_tensor( - data_dict=ret_dict, data_types=self.data_types - ) - if self.preprocess: - ret_dict = self.preprocess_data_dict( - data_dict=ret_dict, - data_types=self.data_types, - layout=self.layout, - rescale=self.rescale_method, - ) - - if self.downsample_dict is not None: - ret_dict = self.downsample_data_dict( - data_dict=ret_dict, - data_types=self.data_types, - factors_dict=self.downsample_dict, - layout=self.layout, - ) - in_slice, out_slice = self.layout_to_in_out_slice() - data_seq = ret_dict["vil"] - if isinstance(data_seq, paddle.Tensor): - data_seq = data_seq.numpy() - x = data_seq[in_slice[0], in_slice[1], in_slice[2], in_slice[3], in_slice[4]] - y = data_seq[ - out_slice[0], out_slice[1], out_slice[2], out_slice[3], out_slice[4] - ] - - weight_item = self.weight_dict - input_item = {self.input_keys[0]: x} - label_item = { - self.label_keys[0]: y, - } - - return input_item, label_item, weight_item diff --git a/examples/smc_reac/ppsci/data/dataset/spherical_swe_dataset.py b/examples/smc_reac/ppsci/data/dataset/spherical_swe_dataset.py deleted file mode 100644 index 68e29e7883..0000000000 --- a/examples/smc_reac/ppsci/data/dataset/spherical_swe_dataset.py +++ /dev/null @@ -1,104 +0,0 @@ -from pathlib import Path -from typing import Dict -from typing import Optional -from typing import Tuple - -import numpy as np -from paddle import io - - -class SphericalSWEDataset(io.Dataset): - """Loads a Spherical Shallow Water equations dataset - - Training contains 200 samples in resolution 32x64. - Testing contains 50 samples at resolution 32x64 and 50 samples at resolution 64x128. - - Args: - input_keys (Tuple[str, ...]): Input keys, such as ("input",). - label_keys (Tuple[str, ...]): Output keys, such as ("output",). - data_dir (str): The directory to load data from. - weight_dict (Optional[Dict[str, float]], optional): Define the weight of each constraint variable. - Defaults to None. - test_resolutions (Tuple[str, ...], optional): The resolutions to test dataset. Defaults to ["34x64", "64x128"]. - train_resolution (str, optional): The resolutions to train dataset. Defaults to "34x64". - data_split (str, optional): Specify the dataset split, either 'train' , 'test_32x64',or 'test_64x128'. - Defaults to "train". - """ - - def __init__( - self, - input_keys: Tuple[str, ...], - label_keys: Tuple[str, ...], - data_dir: str, - weight_dict: Optional[Dict[str, float]] = None, - test_resolutions: Tuple[str, ...] = ["34x64", "64x128"], - train_resolution: str = "34x64", - data_split: str = "train", - ): - super().__init__() - self.input_keys = input_keys - self.label_keys = label_keys - self.data_dir = data_dir - self.weight_dict = {} if weight_dict is None else weight_dict - if weight_dict is not None: - self.weight_dict = {key: 1.0 for key in self.label_keys} - self.weight_dict.update(weight_dict) - - self.test_resolutions = test_resolutions - self.train_resolution = train_resolution - self.data_split = data_split - - # train path - path_train = ( - Path(self.data_dir) - .joinpath(f"train_SWE_{self.train_resolution}.npy") - .as_posix() - ) - self.x_train, self.y_train = self.read_data(path_train) - # test path - path_test_1 = ( - Path(self.data_dir) - .joinpath(f"test_SWE_{self.test_resolutions[0]}.npy") - .as_posix() - ) - self.x_test_1, self.y_test_1 = self.read_data(path_test_1) - path_test_2 = ( - Path(self.data_dir) - .joinpath(f"test_SWE_{self.test_resolutions[1]}.npy") - .as_posix() - ) - self.x_test_2, self.y_test_2 = self.read_data(path_test_2) - - def read_data(self, path): - # load with numpy - data = np.load(path, allow_pickle=True).item() - x = data["x"].astype("float32") - y = data["y"].astype("float32") - del data - return x, y - - def __len__(self): - if self.data_split == "train": - return self.x_train.shape[0] - elif self.data_split == "test_32x64": - return self.x_test_1.shape[0] - else: - return self.x_test_2.shape[0] - - def __getitem__(self, index): - if self.data_split == "train": - x = self.x_train[index] - y = self.y_train[index] - - elif self.data_split == "test_32x64": - x = self.x_test_1[index] - y = self.y_test_1[index] - else: - x = self.x_test_2[index] - y = self.y_test_2[index] - - input_item = {self.input_keys[0]: x} - label_item = {self.label_keys[0]: y} - weight_item = self.weight_dict - - return input_item, label_item, weight_item diff --git a/examples/smc_reac/ppsci/data/dataset/trphysx_dataset.py b/examples/smc_reac/ppsci/data/dataset/trphysx_dataset.py deleted file mode 100644 index 3160951530..0000000000 --- a/examples/smc_reac/ppsci/data/dataset/trphysx_dataset.py +++ /dev/null @@ -1,326 +0,0 @@ -# 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. - -""" -Code below is heavily based on [transformer-physx](https://github.com/zabaras/transformer-physx) -""" - -from __future__ import annotations - -import os -from typing import Dict -from typing import Optional -from typing import Tuple - -try: - import h5py -except ModuleNotFoundError: - pass -import numpy as np -import paddle -from paddle import io - -from ppsci.arch import base - - -class LorenzDataset(io.Dataset): - """Dataset for training Lorenz model. - - Args: - file_path (str): Data set path. - input_keys (Tuple[str, ...]): Input keys, such as ("states",). - label_keys (Tuple[str, ...]): Output keys, such as ("pred_states", "recover_states"). - block_size (int): Data block size. - stride (int): Data stride. - ndata (Optional[int]): Number of data series to use. Defaults to None. - weight_dict (Optional[Dict[str, float]]): Weight dictionary. Defaults to None. - embedding_model (Optional[base.Arch]): Embedding model. Defaults to None. - - Examples: - >>> import ppsci - >>> dataset = ppsci.data.dataset.LorenzDataset( - ... "file_path": "/path/to/LorenzDataset", - ... "input_keys": ("x",), - ... "label_keys": ("v",), - ... "block_size": 32, - ... "stride": 16, - ... ) # doctest: +SKIP - """ - - # Whether support batch indexing for speeding up fetching process. - batch_index: bool = False - - def __init__( - self, - file_path: str, - input_keys: Tuple[str, ...], - label_keys: Tuple[str, ...], - block_size: int, - stride: int, - ndata: Optional[int] = None, - weight_dict: Optional[Dict[str, float]] = None, - embedding_model: Optional[base.Arch] = None, - ): - super().__init__() - if not os.path.exists(file_path): - raise FileNotFoundError( - f"file_path({file_path}) not exists. Please download dataset first. " - "Training: https://paddle-org.bj.bcebos.com/paddlescience/datasets/transformer_physx/lorenz_training_rk.hdf5. " - "Valid: https://paddle-org.bj.bcebos.com/paddlescience/datasets/transformer_physx/lorenz_valid_rk.hdf5." - ) - - self.file_path = file_path - self.input_keys = input_keys - self.label_keys = label_keys - - self.block_size = block_size - self.stride = stride - self.ndata = ndata - self.weight_dict = {key: 1.0 for key in self.label_keys} - if weight_dict is not None: - self.weight_dict.update(weight_dict) - - self.data = self.read_data(file_path, block_size, stride) - self.embedding_model = embedding_model - if embedding_model is None: - self.embedding_data = None - else: - embedding_model.eval() - with paddle.no_grad(): - data_tensor = paddle.to_tensor(self.data) - embedding_data_tensor = embedding_model.encoder(data_tensor) - self.embedding_data = embedding_data_tensor.numpy() - - def read_data(self, file_path: str, block_size: int, stride: int): - data = [] - with h5py.File(file_path, "r") as f: - data_num = 0 - for key in f.keys(): - data_series = np.asarray(f[key], dtype=paddle.get_default_dtype()) - for i in range(0, data_series.shape[0] - block_size + 1, stride): - data.append(data_series[i : i + block_size]) - data_num += 1 - if self.ndata is not None and data_num >= self.ndata: - break - return np.asarray(data) - - def __len__(self): - return len(self.data) - - def __getitem__(self, idx): - # when embedding data is None - if self.embedding_data is None: - data_item = self.data[idx] - input_item = {self.input_keys[0]: data_item} - label_item = { - self.label_keys[0]: data_item[1:, :], - self.label_keys[1]: data_item, - } - else: - data_item = self.embedding_data[idx] - input_item = {self.input_keys[0]: data_item[:-1, :]} - label_item = {self.label_keys[0]: data_item[1:, :]} - if len(self.label_keys) == 2: - label_item[self.label_keys[1]] = self.data[idx][1:, :] - - weight_shape = [1] * len(data_item.shape) - weight_item = { - key: np.full(weight_shape, value, paddle.get_default_dtype()) - for key, value in self.weight_dict.items() - } - return (input_item, label_item, weight_item) - - -class RosslerDataset(LorenzDataset): - """Dataset for training Rossler model. - - Args: - file_path (str): Data set path. - input_keys (Tuple[str, ...]): Input keys, such as ("states",). - label_keys (Tuple[str, ...]): Output keys, such as ("pred_states", "recover_states"). - block_size (int): Data block size. - stride (int): Data stride. - ndata (Optional[int]): Number of data series to use. Defaults to None. - weight_dict (Optional[Dict[str, float]]): Weight dictionary. Defaults to None. - embedding_model (Optional[base.Arch]): Embedding model. Defaults to None. - - Examples: - >>> import ppsci - >>> dataset = ppsci.data.dataset.RosslerDataset( - ... "file_path": "/path/to/RosslerDataset", - ... "input_keys": ("x",), - ... "label_keys": ("v",), - ... "block_size": 32, - ... "stride": 16, - ... ) # doctest: +SKIP - """ - - # Whether support batch indexing for speeding up fetching process. - batch_index: bool = False - - def __init__( - self, - file_path: str, - input_keys: Tuple[str, ...], - label_keys: Tuple[str, ...], - block_size: int, - stride: int, - ndata: Optional[int] = None, - weight_dict: Optional[Dict[str, float]] = None, - embedding_model: Optional[base.Arch] = None, - ): - if not os.path.exists(file_path): - raise FileNotFoundError( - f"file_path({file_path}) not exists. Please download dataset first. " - "Training: https://paddle-org.bj.bcebos.com/paddlescience/datasets/transformer_physx/rossler_training.hdf5. " - "Valid: https://paddle-org.bj.bcebos.com/paddlescience/datasets/transformer_physx/rossler_valid.hdf5." - ) - super().__init__( - file_path, - input_keys, - label_keys, - block_size, - stride, - ndata, - weight_dict, - embedding_model, - ) - - -class CylinderDataset(io.Dataset): - """Dataset for training Cylinder model. - - Args: - file_path (str): Data set path. - input_keys (Tuple[str, ...]): Input keys, such as ("states","visc"). - label_keys (Tuple[str, ...]): Output keys, such as ("pred_states", "recover_states"). - block_size (int): Data block size. - stride (int): Data stride. - ndata (Optional[int]): Number of data series to use. Defaults to None. - weight_dict (Optional[Dict[str, float]]): Weight dictionary. Defaults to None. - embedding_model (Optional[base.Arch]): Embedding model. Defaults to None. - embedding_batch_size (int, optional): The batch size of embedding model. Defaults to 64. - - Examples: - >>> import ppsci - >>> dataset = ppsci.data.dataset.CylinderDataset( - ... "file_path": "/path/to/CylinderDataset", - ... "input_keys": ("x",), - ... "label_keys": ("v",), - ... "block_size": 32, - ... "stride": 16, - ... ) # doctest: +SKIP - """ - - # Whether support batch indexing for speeding up fetching process. - batch_index: bool = False - - def __init__( - self, - file_path: str, - input_keys: Tuple[str, ...], - label_keys: Tuple[str, ...], - block_size: int, - stride: int, - ndata: Optional[int] = None, - weight_dict: Optional[Dict[str, float]] = None, - embedding_model: Optional[base.Arch] = None, - embedding_batch_size: int = 64, - ): - if not os.path.exists(file_path): - raise FileNotFoundError( - f"file_path({file_path}) not exists. Please download dataset first. " - "Training: https://paddle-org.bj.bcebos.com/paddlescience/datasets/transformer_physx/cylinder_training.hdf5. " - "Valid: https://paddle-org.bj.bcebos.com/paddlescience/datasets/transformer_physx/cylinder_valid.hdf5." - ) - super().__init__() - self.file_path = file_path - self.input_keys = input_keys - self.label_keys = label_keys - - self.block_size = block_size - self.stride = stride - self.ndata = ndata - self.weight_dict = {key: 1.0 for key in self.label_keys} - if weight_dict is not None: - self.weight_dict.update(weight_dict) - - self.data, self.visc = self.read_data(file_path, block_size, stride) - self.embedding_model = embedding_model - if embedding_model is None: - self.embedding_data = None - else: - embedding_model.eval() - with paddle.no_grad(): - data_tensor = paddle.to_tensor(self.data) - visc_tensor = paddle.to_tensor(self.visc) - embedding_data = [] - for i in range(0, len(data_tensor), embedding_batch_size): - start, end = i, min(i + embedding_batch_size, len(data_tensor)) - embedding_data_batch = embedding_model.encoder( - data_tensor[start:end], visc_tensor[start:end] - ) - embedding_data.append(embedding_data_batch.numpy()) - self.embedding_data = np.concatenate(embedding_data) - - def read_data(self, file_path: str, block_size: int, stride: int): - data = [] - visc = [] - with h5py.File(file_path, "r") as f: - data_num = 0 - for key in f.keys(): - visc0 = 2.0 / float(key) - ux = np.asarray(f[key + "/ux"], dtype=paddle.get_default_dtype()) - uy = np.asarray(f[key + "/uy"], dtype=paddle.get_default_dtype()) - p = np.asarray(f[key + "/p"], dtype=paddle.get_default_dtype()) - data_series = np.stack([ux, uy, p], axis=1) - - for i in range(0, data_series.shape[0] - block_size + 1, stride): - data.append(data_series[i : i + block_size]) - visc.append([visc0]) - - data_num += 1 - if self.ndata is not None and data_num >= self.ndata: - break - - data = np.asarray(data) - visc = np.asarray(visc, dtype=paddle.get_default_dtype()) - return data, visc - - def __len__(self): - return len(self.data) - - def __getitem__(self, i): - if self.embedding_data is None: - data_item = self.data[i] - input_item = { - self.input_keys[0]: data_item, - self.input_keys[1]: self.visc[i], - } - label_item = { - self.label_keys[0]: data_item[1:], - self.label_keys[1]: data_item, - } - else: - data_item = self.embedding_data[i] - input_item = {self.input_keys[0]: data_item[:-1, :]} - label_item = {self.label_keys[0]: data_item[1:, :]} - if len(self.label_keys) == 2: - label_item[self.label_keys[1]] = data_item[1:, :] - weight_shape = [1] * len(data_item.shape) - weight_item = { - key: np.full(weight_shape, value, paddle.get_default_dtype()) - for key, value in self.weight_dict.items() - } - return (input_item, label_item, weight_item) diff --git a/examples/smc_reac/ppsci/data/dataset/vtu_dataset.py b/examples/smc_reac/ppsci/data/dataset/vtu_dataset.py deleted file mode 100644 index fb0c9201b7..0000000000 --- a/examples/smc_reac/ppsci/data/dataset/vtu_dataset.py +++ /dev/null @@ -1,106 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import Dict -from typing import Optional -from typing import Tuple - -import numpy as np -from paddle import io -from paddle import vision - -from ppsci.utils import reader - - -class VtuDataset(io.Dataset): - """Dataset class for .vtu file. - - Args: - file_path (str): *.vtu file path. - input_keys (Optional[Tuple[str, ...]]): Tuple of input keys. Defaults to None. - label_keys (Optional[Tuple[str, ...]]): Tuple of label keys. Defaults to None. - time_step (Optional[int]): Time step with unit second. Defaults to None. - time_index (Optional[Tuple[int, ...]]): Time index tuple in increasing order. - labels (Optional[Dict[str, float]]): Temporary variable for [load_vtk_with_time_file]. - transforms (vision.Compose, optional): Compose object contains sample wise. - transform(s). - - Examples: - >>> from ppsci.data.dataset import VtuDataset - - >>> dataset = VtuDataset(file_path='example.vtu') # doctest: +SKIP - - >>> # get the length of the dataset - >>> dataset_size = len(dataset) # doctest: +SKIP - >>> # get the first sample of the data - >>> first_sample = dataset[0] # doctest: +SKIP - >>> print("First sample:", first_sample) # doctest: +SKIP - """ - - # Whether support batch indexing for speeding up fetching process. - batch_index: bool = True - - def __init__( - self, - file_path: str, - input_keys: Optional[Tuple[str, ...]] = None, - label_keys: Optional[Tuple[str, ...]] = None, - time_step: Optional[int] = None, - time_index: Optional[Tuple[int, ...]] = None, - labels: Optional[Dict[str, float]] = None, - transforms: Optional[vision.Compose] = None, - ): - super().__init__() - - # load data from file - if time_step is not None and time_index is not None: - _input, _label = reader.load_vtk_file( - file_path, time_step, time_index, input_keys, label_keys - ) - _label = {key: _label[key] for key in label_keys} - elif time_step is None and time_index is None: - _input = reader.load_vtk_with_time_file(file_path) - _label = {} - for key, value in labels.items(): - if isinstance(value, (int, float)): - _label[key] = np.full_like( - next(iter(_input.values())), value, "float32" - ) - else: - _label[key] = value - else: - raise ValueError( - "Error, read vtu with time_step and time_index, or neither" - ) - - # transform - _input = transforms(_input) - _label = transforms(_label) - - self.input = _input - self.label = _label - self.input_keys = input_keys - self.label_keys = label_keys - self.transforms = transforms - self.num_samples = len(next(iter(self.input.values()))) - - def __getitem__(self, idx): - input_item = {key: value[idx] for key, value in self.input.items()} - label_item = {key: value[idx] for key, value in self.label.items()} - return (input_item, label_item, {}) - - def __len__(self): - return self.num_samples diff --git a/examples/smc_reac/ppsci/data/process/__init__.py b/examples/smc_reac/ppsci/data/process/__init__.py deleted file mode 100644 index f46c8dd9cf..0000000000 --- a/examples/smc_reac/ppsci/data/process/__init__.py +++ /dev/null @@ -1,21 +0,0 @@ -# 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. - -from ppsci.data.process import batch_transform -from ppsci.data.process import transform - -__all__ = [ - "batch_transform", - "transform", -] diff --git a/examples/smc_reac/ppsci/data/process/batch_transform/__init__.py b/examples/smc_reac/ppsci/data/process/batch_transform/__init__.py deleted file mode 100644 index 9e98f39264..0000000000 --- a/examples/smc_reac/ppsci/data/process/batch_transform/__init__.py +++ /dev/null @@ -1,135 +0,0 @@ -# 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. - -import copy -import numbers -from collections.abc import Mapping -from collections.abc import Sequence -from typing import Any -from typing import Callable -from typing import List -from typing import Optional - -import numpy as np -import paddle - -from ppsci.data.process import transform -from ppsci.data.process.batch_transform.preprocess import FunctionalBatchTransform - -try: - import pgl -except ModuleNotFoundError: - pass - - -__all__ = [ - "build_batch_transforms", - "default_collate_fn", - "FunctionalBatchTransform", -] - - -def default_collate_fn(batch: List[Any]) -> Any: - """Default_collate_fn for paddle dataloader. - - NOTE: This `default_collate_fn` is different from official `default_collate_fn` - which specially adapt case where sample is `None` and `pgl.Graph`. - - ref: https://github.com/PaddlePaddle/Paddle/blob/develop/python/paddle/io/dataloader/collate.py#L25 - - Args: - batch (List[Any]): Batch of samples to be collated. - - Returns: - Any: Collated batch data. - """ - sample = batch[0] - if sample is None: - return None - elif isinstance(sample, np.ndarray): - batch = np.stack(batch, axis=0) - return batch - elif isinstance(sample, (paddle.Tensor, paddle.framework.core.eager.Tensor)): - return paddle.stack(batch, axis=0) - elif isinstance(sample, numbers.Number): - batch = np.array(batch) - return batch - elif isinstance(sample, (str, bytes)): - return batch - elif isinstance(sample, Mapping): - return {key: default_collate_fn([d[key] for d in batch]) for key in sample} - elif isinstance(sample, Sequence): - sample_fields_num = len(sample) - if not all(len(sample) == sample_fields_num for sample in iter(batch)): - raise RuntimeError("Fields number not same among samples in a batch") - return [default_collate_fn(fields) for fields in zip(*batch)] - elif str(type(sample)) == "": - # use str(type()) instead of isinstance() in case of pgl is not installed. - graph = pgl.Graph(num_nodes=sample.num_nodes, edges=sample.edges) - graph.x = np.concatenate([g.x for g in batch]) - graph.y = np.concatenate([g.y for g in batch]) - graph.edge_index = np.concatenate([g.edge_index for g in batch], axis=1) - - graph.edge_attr = np.concatenate([g.edge_attr for g in batch]) - graph.pos = np.concatenate([g.pos for g in batch]) - if hasattr(sample, "aoa"): - graph.aoa = np.concatenate([g.aoa for g in batch]) - if hasattr(sample, "mach_or_reynolds"): - graph.mach_or_reynolds = np.concatenate([g.mach_or_reynolds for g in batch]) - graph.tensor() - graph.shape = [len(batch)] - return graph - elif ( - str(type(sample)) - == "" - ): - graph = sample - graph.tensor() - graph.shape = [1] - return graph - raise TypeError( - "batch data can only contains: paddle.Tensor, numpy.ndarray, " - f"dict, list, number, None, pgl.Graph, GraphGridMesh, but got {type(sample)}" - ) - - -def build_transforms(cfg): - if not cfg: - return transform.Compose([]) - cfg = copy.deepcopy(cfg) - - transform_list = [] - for _item in cfg: - transform_cls = next(iter(_item.keys())) - transform_cfg = _item[transform_cls] - transform_obj = eval(transform_cls)(**transform_cfg) - transform_list.append(transform_obj) - - return transform.Compose(transform_list) - - -def build_batch_transforms(cfg, collate_fn: Optional[Callable]): - cfg = copy.deepcopy(cfg) - batch_transforms: Callable[[List[Any]], List[Any]] = build_transforms(cfg) - if collate_fn is None: - collate_fn = default_collate_fn - - def collate_fn_batch_transforms(batch: List[Any]): - # apply batch transform on separate samples - batch = batch_transforms(batch) - - # then collate separate samples into batched data - return collate_fn(batch) - - return collate_fn_batch_transforms diff --git a/examples/smc_reac/ppsci/data/process/batch_transform/preprocess.py b/examples/smc_reac/ppsci/data/process/batch_transform/preprocess.py deleted file mode 100644 index 62ca5d3be1..0000000000 --- a/examples/smc_reac/ppsci/data/process/batch_transform/preprocess.py +++ /dev/null @@ -1,74 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import Any -from typing import Callable -from typing import Dict -from typing import List -from typing import Optional -from typing import Tuple - -import numpy as np - - -class FunctionalBatchTransform: - """Functional data transform class, which allows to use custom data transform function from given transform_func for special cases. - - Args: - transform_func (Callable): Function of batch data transform. - - Examples: - >>> import ppsci - >>> from typing import Tuple, Dict, Optional - >>> def batch_transform_func( - ... data_list: List[ - ... Tuple[Dict[str, np.ndarray], Dict[str, np.ndarray], Optional[Dict[str, np.ndarray]]] - ... ], - ... ) -> List[Tuple[Dict[str, np.ndarray], Dict[str, np.ndarray], Optional[Dict[str, np.ndarray]]]]: - ... input_dicts, label_dicts, weight_dicts = zip(*data_list) - ... - ... for input_dict in input_dicts: - ... for key in input_dict: - ... input_dict[key] = input_dict[key] * 2 - ... - ... for label_dict in label_dicts: - ... for key in label_dict: - ... label_dict[key] = label_dict[key] + 1.0 - ... - ... return list(zip(input_dicts, label_dicts, weight_dicts)) - ... - >>> # Create a FunctionalBatchTransform object with the batch_transform_func function - >>> transform = ppsci.data.batch_transform.FunctionalBatchTransform(batch_transform_func) - >>> # Define some sample data, labels, and weights - >>> data = [({'x': 1}, {'y': 2}, None), ({'x': 11}, {'y': 22}, None)] - >>> transformed_data = transform(data) - >>> for tuple in transformed_data: - ... print(tuple) - ({'x': 2}, {'y': 3.0}, None) - ({'x': 22}, {'y': 23.0}, None) - """ - - def __init__( - self, - transform_func: Callable[[List[Any]], List[Any]], - ): - self.transform_func = transform_func - - def __call__( - self, - data_list: List[Tuple[Optional[Dict[str, np.ndarray]], ...]], - ) -> List[Tuple[Optional[Dict[str, np.ndarray]], ...]]: - return self.transform_func(data_list) diff --git a/examples/smc_reac/ppsci/data/process/transform/__init__.py b/examples/smc_reac/ppsci/data/process/transform/__init__.py deleted file mode 100644 index f5a4baa287..0000000000 --- a/examples/smc_reac/ppsci/data/process/transform/__init__.py +++ /dev/null @@ -1,72 +0,0 @@ -# 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. - -import copy -import traceback -from typing import Any -from typing import Tuple - -from paddle import vision - -from ppsci.data.process.transform.preprocess import CropData -from ppsci.data.process.transform.preprocess import FunctionalTransform -from ppsci.data.process.transform.preprocess import Log1p -from ppsci.data.process.transform.preprocess import Normalize -from ppsci.data.process.transform.preprocess import Scale -from ppsci.data.process.transform.preprocess import SqueezeData -from ppsci.data.process.transform.preprocess import Translate - -__all__ = [ - "CropData", - "FunctionalTransform", - "Log1p", - "Normalize", - "Scale", - "SqueezeData", - "Translate", - "build_transforms", -] - - -class Compose(vision.Compose): - """Custom Compose for multiple items in given data.""" - - def __call__(self, *data: Tuple[Any, ...]): - for f in self.transforms: - try: - # NOTE: This is different from vision.Compose to allow receive multiple data items - data = f(*data) - except Exception as e: - stack_info = traceback.format_exc() - print( - f"fail to perform transform [{f}] with error: " - f"{e} and stack:\n{str(stack_info)}" - ) - raise e - return data - - -def build_transforms(cfg): - if not cfg: - return Compose([]) - cfg = copy.deepcopy(cfg) - - transform_list = [] - for _item in cfg: - transform_cls = next(iter(_item.keys())) - transform_cfg = _item[transform_cls] - transform = eval(transform_cls)(**transform_cfg) - transform_list.append(transform) - - return Compose(transform_list) diff --git a/examples/smc_reac/ppsci/data/process/transform/preprocess.py b/examples/smc_reac/ppsci/data/process/transform/preprocess.py deleted file mode 100644 index 48f0fa1222..0000000000 --- a/examples/smc_reac/ppsci/data/process/transform/preprocess.py +++ /dev/null @@ -1,331 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import Callable -from typing import Dict -from typing import Tuple -from typing import Union - -import numpy as np - - -class Translate: - """Translate class. - - Args: - offset (Dict[str, float]): Shift the input data according to the variable name - and coefficient specified in offset. - - Examples: - >>> import ppsci - >>> import numpy as np - - >>> input_dict = {"x": np.array([5.0, 10.0]), "y": np.array([20.0, 40.0])} - >>> label_dict = {"x": np.array([1.0, 2.0]), "y": np.array([3.0, 4.0])} - >>> weight_dict = {"x": np.array([10.0, 20.0]), "y": np.array([30.0, 40.0])} - - >>> translate = ppsci.data.transform.Translate({"x": 1.0, "y": -1.0}) - >>> translated_input_dict, translated_label_dict, translated_weight_dict = translate(input_dict, label_dict, weight_dict) - - >>> print(translated_input_dict) - {'x': array([ 6., 11.]), 'y': array([19., 39.])} - >>> print(translated_label_dict) - {'x': array([1., 2.]), 'y': array([3., 4.])} - >>> print(translated_weight_dict) - {'x': array([10., 20.]), 'y': array([30., 40.])} - """ - - def __init__(self, offset: Dict[str, float]): - self.offset = offset - - def __call__(self, input_dict, label_dict, weight_dict): - input_dict_copy = {**input_dict} - for key in self.offset: - if key in input_dict: - input_dict_copy[key] += self.offset[key] - return input_dict_copy, label_dict, weight_dict - - -class Scale: - """Scale class for data transformation. - - Args: - scale (Dict[str, float]): Scale the input data according to the variable name - and coefficient specified in scale. - - Examples: - >>> import ppsci - >>> translate = ppsci.data.transform.Scale({"x": 1.5, "y": 2.0}) - >>> input_dict = {"x": 10, "y": 20} - >>> label_dict = {"x": 100, "y": 200} - >>> weight_dict = {"x": 1000, "y": 2000} - >>> input_dict_scaled, label_dict_scaled, weight_dict_scaled = translate(input_dict, label_dict, weight_dict) - >>> print(input_dict_scaled) - {'x': 15.0, 'y': 40.0} - >>> print(label_dict_scaled) - {'x': 100, 'y': 200} - >>> print(weight_dict_scaled) - {'x': 1000, 'y': 2000} - """ - - def __init__(self, scale: Dict[str, float]): - self.scale = scale - - def __call__(self, input_dict, label_dict, weight_dict): - input_dict_copy = {**input_dict} - for key in self.scale: - if key in input_dict: - input_dict_copy[key] *= self.scale[key] - return input_dict_copy, label_dict, weight_dict - - -class Normalize: - """Normalize data class. - - NOTE: This transform will modify the input data dict inplace. - - Args: - mean (Union[np.ndarray, Tuple[float, ...]]): Mean of training dataset. - std (Union[np.ndarray, Tuple[float, ...]]): Standard Deviation of training dataset. - apply_keys (Tuple[str, ...], optional): Which data is the normalization method applied to. Defaults to ("input", "label"). - - Examples: - >>> import ppsci - >>> normalize = ppsci.data.transform.Normalize((0.0, 0.0, 0.0), (1.0, 1.0, 1.0)) - >>> input_item = {"data": np.array([1.0, 2.0, 3.0])} - >>> label_item = {"data": np.array([4.0, 5.0, 6.0])} - >>> weight_item = np.array([0.1, 0.2, 0.3]) - >>> normalized_item = normalize(input_item, label_item, weight_item) - >>> print(normalized_item) - ({'data': array([1., 2., 3.])}, {'data': array([4., 5., 6.])}, array([0.1, 0.2, 0.3])) - """ - - def __init__( - self, - mean: Union[np.ndarray, Tuple[float, ...]], - std: Union[np.ndarray, Tuple[float, ...]], - apply_keys: Tuple[str, ...] = ("input", "label"), - ): - if len(apply_keys) == 0 or len(set(apply_keys) | {"input", "label"}) > 2: - raise ValueError( - f"apply_keys should be a non empty subset of ('input', 'label'), but got {apply_keys}" - ) - self.mean = mean - self.std = std - self.apply_keys = apply_keys - - def __call__(self, input_item, label_item, weight_item): - if "input" in self.apply_keys: - for key, value in input_item.items(): - input_item[key] = (value - self.mean) / self.std - if "label" in self.apply_keys: - for key, value in label_item.items(): - label_item[key] = (value - self.mean) / self.std - return input_item, label_item, weight_item - - -class Log1p: - """Calculates the natural logarithm of one plus the data, element-wise. - - NOTE: This transform will modify the input data dict inplace. - - Args: - scale (float, optional): Scale data. Defaults to 1.0. - apply_keys (Tuple[str, ...], optional): Which data is the log1p method applied to. Defaults to ("input", "label"). - - Examples: - >>> import ppsci - >>> log1p = ppsci.data.transform.Log1p(1e-5) - >>> input_item = {"data": np.array([1.0, 2.0, 3.0])} - >>> label_item = {"data": np.array([4.0, 5.0, 6.0])} - >>> weight_item = np.array([0.1, 0.2, 0.3]) - >>> input_item_transformed, label_item_transformed, weight_item_transformed = log1p(input_item, label_item, weight_item) - >>> print(input_item_transformed) - {'data': array([11.51293546, 12.20607765, 12.61154109])} - >>> print(label_item_transformed) - {'data': array([12.89922233, 13.12236538, 13.3046866 ])} - >>> print(weight_item_transformed) - [0.1 0.2 0.3] - """ - - def __init__( - self, - scale: float = 1.0, - apply_keys: Tuple[str, ...] = ("input", "label"), - ): - if len(apply_keys) == 0 or len(set(apply_keys) | {"input", "label"}) > 2: - raise ValueError( - f"apply_keys should be a non empty subset of ('input', 'label'), but got {apply_keys}" - ) - self.scale = scale - self.apply_keys = apply_keys - - def __call__(self, input_item, label_item, weight_item): - if "input" in self.apply_keys: - for key, value in input_item.items(): - input_item[key] = np.log1p(value / self.scale) - if "label" in self.apply_keys: - for key, value in label_item.items(): - label_item[key] = np.log1p(value / self.scale) - return input_item, label_item, weight_item - - -class CropData: - """Crop data class. - - This class is used to crop data based on a specified bounding box. - - NOTE: This transform will modify the input data dict inplace. - - Args: - xmin (Tuple[int, ...]): Bottom left corner point, [x0, y0]. - xmax (Tuple[int, ...]): Top right corner point, [x1, y1]. - apply_keys (Tuple[str, ...], optional): Which data is the crop method applied to. Defaults to ("input", "label"). - - Examples: - >>> import ppsci - >>> import numpy as np - >>> crop_data = ppsci.data.transform.CropData((0, 0), (256, 512)) - >>> input_item = {"input": np.zeros((3, 720, 1440))} - >>> label_item = {"label": np.zeros((3, 720, 1440))} - >>> weight_item = {"weight": np.ones((3, 720, 1440))} - >>> input_item, label_item, weight_item = crop_data(input_item, label_item, weight_item) - >>> print(input_item["input"].shape) - (3, 256, 512) - >>> print(label_item["label"].shape) - (3, 256, 512) - """ - - def __init__( - self, - xmin: Tuple[int, ...], - xmax: Tuple[int, ...], - apply_keys: Tuple[str, ...] = ("input", "label"), - ): - if len(apply_keys) == 0 or len(set(apply_keys) | {"input", "label"}) > 2: - raise ValueError( - f"apply_keys should be a non empty subset of ('input', 'label'), but got {apply_keys}" - ) - self.xmin = xmin - self.xmax = xmax - self.apply_keys = apply_keys - - def __call__(self, input_item, label_item, weight_item): - if "input" in self.apply_keys: - for key, value in input_item.items(): - input_item[key] = value[ - :, self.xmin[0] : self.xmax[0], self.xmin[1] : self.xmax[1] - ] - if "label" in self.apply_keys: - for key, value in label_item.items(): - label_item[key] = value[ - :, self.xmin[0] : self.xmax[0], self.xmin[1] : self.xmax[1] - ] - return input_item, label_item, weight_item - - -class SqueezeData: - """Squeeze data class. - - NOTE: This transform will modify the input data dict inplace. - - Args: - apply_keys (Tuple[str, ...], optional): Which data is the squeeze method applied to. Defaults to ("input", "label"). - - Examples: - >>> import ppsci - >>> import numpy as np - >>> squeeze_data = ppsci.data.transform.SqueezeData() - >>> input_data = {"input": np.random.rand(10, 224, 224)} - >>> label_data = {"label": np.random.rand(10, 224, 224)} - >>> weight_data = {"weight": np.random.rand(10, 224, 224)} - >>> input_data_squeezed, label_data_squeezed, weight_data_squeezed = squeeze_data(input_data, label_data, weight_data) - """ - - def __init__(self, apply_keys: Tuple[str, ...] = ("input", "label")): - if len(apply_keys) == 0 or len(set(apply_keys) | {"input", "label"}) > 2: - raise ValueError( - f"apply_keys should be a non empty subset of ('input', 'label'), but got {apply_keys}" - ) - self.apply_keys = apply_keys - - def __call__(self, input_item, label_item, weight_item): - if "input" in self.apply_keys: - for key, value in input_item.items(): - if value.ndim == 4: - B, C, H, W = value.shape - input_item[key] = value.reshape((B * C, H, W)) - if value.ndim != 3: - raise ValueError( - f"Only support squeeze data to ndim=3 now, but got ndim={value.ndim}" - ) - if "label" in self.apply_keys: - for key, value in label_item.items(): - if value.ndim == 4: - B, C, H, W = value.shape - label_item[key] = value.reshape((B * C, H, W)) - if value.ndim != 3: - raise ValueError( - f"Only support squeeze data to ndim=3 now, but got ndim={value.ndim}" - ) - return input_item, label_item, weight_item - - -class FunctionalTransform: - """Functional data transform class, which allows to use custom data transform function from given transform_func for special cases. - - Args: - transform_func (Callable): Function of data transform. - - Examples: - >>> # This is the transform_func function. It takes three dictionaries as input: data_dict, label_dict, and weight_dict. - >>> # The function will perform some transformations on the data in data_dict, convert all labels in label_dict to uppercase, - >>> # and modify the weights in weight_dict by dividing each weight by 10. - >>> # Finally, it returns the transformed data, labels, and weights as a tuple. - >>> import ppsci - >>> def transform_func(data_dict, label_dict, weight_dict): - ... for key in data_dict: - ... data_dict[key] = data_dict[key] * 2 - ... for key in label_dict: - ... label_dict[key] = label_dict[key] + 1.0 - ... for key in weight_dict: - ... weight_dict[key] = weight_dict[key] / 10 - ... return data_dict, label_dict, weight_dict - >>> transform = ppsci.data.transform.FunctionalTransform(transform_func) - >>> # Define some sample data, labels, and weights - >>> data = {'feature1': np.array([1, 2, 3]), 'feature2': np.array([4, 5, 6])} - >>> label = {'class': 0.0, 'instance': 0.1} - >>> weight = {'weight1': 0.5, 'weight2': 0.5} - >>> # Apply the transform function to the data, labels, and weights using the FunctionalTransform instance - >>> transformed_data = transform(data, label, weight) - >>> print(transformed_data) - ({'feature1': array([2, 4, 6]), 'feature2': array([ 8, 10, 12])}, {'class': 1.0, 'instance': 1.1}, {'weight1': 0.05, 'weight2': 0.05}) - """ - - def __init__( - self, - transform_func: Callable, - ): - self.transform_func = transform_func - - def __call__( - self, *data: Tuple[Dict[str, np.ndarray], ...] - ) -> Tuple[Dict[str, np.ndarray], ...]: - data_dict, label_dict, weight_dict = data - data_dict_copy = {**data_dict} - label_dict_copy = {**label_dict} - weight_dict_copy = {**weight_dict} if weight_dict is not None else {} - return self.transform_func(data_dict_copy, label_dict_copy, weight_dict_copy) diff --git a/examples/smc_reac/ppsci/equation/__init__.py b/examples/smc_reac/ppsci/equation/__init__.py deleted file mode 100644 index bcffef4060..0000000000 --- a/examples/smc_reac/ppsci/equation/__init__.py +++ /dev/null @@ -1,76 +0,0 @@ -# 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. - -import copy - -from ppsci.equation.fpde import FractionalPoisson -from ppsci.equation.ide import Volterra -from ppsci.equation.pde import DETACH_FUNC_NAME -from ppsci.equation.pde import NLSMB -from ppsci.equation.pde import PDE -from ppsci.equation.pde import AllenCahn -from ppsci.equation.pde import Biharmonic -from ppsci.equation.pde import HeatExchanger -from ppsci.equation.pde import Helmholtz -from ppsci.equation.pde import Laplace -from ppsci.equation.pde import LinearElasticity -from ppsci.equation.pde import NavierStokes -from ppsci.equation.pde import NormalDotVec -from ppsci.equation.pde import Poisson -from ppsci.equation.pde import Vibration -from ppsci.utils import logger -from ppsci.utils import misc - -__all__ = [ - "PDE", - "DETACH_FUNC_NAME", - "AllenCahn", - "Biharmonic", - "HeatExchanger", - "Helmholtz", - "Laplace", - "LinearElasticity", - "NavierStokes", - "NormalDotVec", - "Poisson", - "Vibration", - "Volterra", - "NLSMB", - "FractionalPoisson", - "build_equation", -] - - -def build_equation(cfg): - """Build equation(s) - - Args: - cfg (List[DictConfig]): Equation(s) config list. - - Returns: - Dict[str, Equation]: Equation(s) in dict. - """ - if cfg is None: - return None - cfg = copy.deepcopy(cfg) - eq_dict = misc.PrettyOrderedDict() - for _item in cfg: - eq_cls = next(iter(_item.keys())) - eq_cfg = _item[eq_cls] - eq_name = eq_cfg.pop("name", eq_cls) - eq_dict[eq_name] = eval(eq_cls)(**eq_cfg) - - logger.debug(str(eq_dict[eq_name])) - - return eq_dict diff --git a/examples/smc_reac/ppsci/equation/fpde/__init__.py b/examples/smc_reac/ppsci/equation/fpde/__init__.py deleted file mode 100644 index 3e74ec56c7..0000000000 --- a/examples/smc_reac/ppsci/equation/fpde/__init__.py +++ /dev/null @@ -1,19 +0,0 @@ -# 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. - -from ppsci.equation.fpde.fractional_poisson import FractionalPoisson - -__all__ = [ - "FractionalPoisson", -] diff --git a/examples/smc_reac/ppsci/equation/fpde/fractional_poisson.py b/examples/smc_reac/ppsci/equation/fpde/fractional_poisson.py deleted file mode 100644 index 01b6fc929f..0000000000 --- a/examples/smc_reac/ppsci/equation/fpde/fractional_poisson.py +++ /dev/null @@ -1,196 +0,0 @@ -# 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. - -from __future__ import annotations - -import math -from typing import Tuple - -import numpy as np -import paddle -from paddle import sparse -from scipy import special - -from ppsci import geometry -from ppsci.equation.pde import PDE -from ppsci.utils import misc - - -class FractionalPoisson(PDE): - """(TODO)Docstring of this class will be refined in the future. - - Args: - alpha (float): Alpha. - geom (geometry.Geometry): Computation geometry. - resolution (Tuple[int, ...]): Resolution. - - Examples: - >>> import ppsci - >>> geom_disk = ppsci.geometry.Disk([0, 0], 1) - >>> ALPHA = 0.5 - >>> fpde = ppsci.equation.FractionalPoisson(ALPHA, geom_disk, [8, 100]) - """ - - dtype = paddle.get_default_dtype() - - def __init__( - self, alpha: float, geom: geometry.Geometry, resolution: Tuple[int, ...] - ): - super().__init__() - self.alpha = alpha - self.geom = geom - self.resolution = resolution - self._w_init = self._init_weights() - - def compute_fpde_func(out): - x = paddle.concat((out["x"], out["y"]), axis=1) - y = out["u"] - indices, values, shape = self.int_mat - int_mat = sparse.sparse_coo_tensor( - [[p[0] for p in indices], [p[1] for p in indices]], - values, - shape, - stop_gradient=False, - ) - lhs = sparse.matmul(int_mat, y) - lhs = lhs[:, 0] - lhs *= ( - special.gamma((1 - self.alpha) / 2) - * special.gamma((2 + self.alpha) / 2) - / (2 * np.pi**1.5) - ) - x = x[: paddle.numel(lhs)] - rhs = ( - 2**self.alpha - * special.gamma(2 + self.alpha / 2) - * special.gamma(1 + self.alpha / 2) - * (1 - (1 + self.alpha / 2) * paddle.sum(x**2, axis=1)) - ) - res = lhs - rhs - return res - - self.add_equation("fpde", compute_fpde_func) - - def _init_weights(self): - n = self._dynamic_dist2npts(self.geom.diam) + 1 - w = [1.0] - for j in range(1, n): - w.append(w[-1] * (j - 1 - self.alpha) / j) - return np.array(w, dtype=self.dtype) - - def get_x(self, x_f): - if hasattr(self, "train_x"): - return self.train_x - - self.x0 = x_f - if np.any(self.geom.on_boundary(self.x0)): - raise ValueError("x0 contains boundary points.") - - if self.geom.ndim == 1: - dirns, dirn_w = [-1, 1], [1, 1] - elif self.geom.ndim == 2: - gauss_x, gauss_w = np.polynomial.legendre.leggauss(self.resolution[0]) - gauss_x, gauss_w = gauss_x.astype(self.dtype), gauss_w.astype(self.dtype) - thetas = np.pi * gauss_x + np.pi - dirns = np.vstack((np.cos(thetas), np.sin(thetas))).T - dirn_w = np.pi * gauss_w - elif self.geom.ndim == 3: - gauss_x, gauss_w = np.polynomial.legendre.leggauss(max(self.resolution[:2])) - gauss_x, gauss_w = gauss_x.astype(self.dtype), gauss_w.astype(self.dtype) - thetas = (np.pi * gauss_x[: self.resolution[0]] + np.pi) / 2 - phis = np.pi * gauss_x[: self.resolution[1]] + np.pi - dirns, dirn_w = [], [] - for i in range(self.resolution[0]): - for j in range(self.resolution[1]): - dirns.append( - [ - np.sin(thetas[i]) * np.cos(phis[j]), - np.sin(thetas[i]) * np.sin(phis[j]), - np.cos(thetas[i]), - ] - ) - dirn_w.append(gauss_w[i] * gauss_w[j] * np.sin(thetas[i])) - dirn_w = np.pi**2 / 2 * np.array(dirn_w) - - x, self.w = [], [] - for x0i in self.x0: - xi = list( - map( - lambda dirn: self.background_points( - x0i, dirn, self._dynamic_dist2npts, 0 - ), - dirns, - ) - ) - wi = list( - map( - lambda i: dirn_w[i] - * np.linalg.norm(xi[i][1] - xi[i][0]) ** (-self.alpha) - * self.get_weight(len(xi[i]) - 1), - range(len(dirns)), - ) - ) - # first order - # xi, wi = zip(self.modify_first_order(xij, wij) for xij, wij in zip(xi, wi)) - xi, wi = zip(*map(self.modify_first_order, xi, wi)) - # second order - # xi, wi = zip(*map(self.modify_second_order, xi, wi)) - # third order - # xi, wi = zip(*map(self.modify_third_order, xi, wi)) - x.append(np.vstack(xi)) - self.w.append(np.hstack(wi)) - self.x = np.vstack([self.x0] + x) - self.int_mat = self._get_int_matrix(self.x0) - self.train_x = misc.convert_to_dict(self.x, ("x", "y")) - return self.train_x - - def get_weight(self, n): - return self._w_init[: n + 1] - - def background_points(self, x, dirn, dist2npt, shift): - dirn = dirn / np.linalg.norm(dirn) - dx = self.distance2boundary_unitdirn(x, -dirn) - n = max(dist2npt(dx), 1) - h = dx / n - pts = x - np.arange(-shift, n - shift + 1, dtype=self.dtype)[:, None] * h * dirn - return pts - - def distance2boundary_unitdirn(self, x, dirn): - # https://en.wikipedia.org/wiki/Line%E2%80%93sphere_intersection - xc = x - self.geom.center - xc = xc - ad = np.dot(xc, dirn) - return ( - -ad + (ad**2 - np.sum(xc * xc, axis=-1) + self.geom.radius**2) ** 0.5 - ).astype(self.dtype) - - def modify_first_order(self, x, w): - x = np.vstack(([2 * x[0] - x[1]], x[:-1])) - if not self.geom.is_inside(x[0:1])[0]: - return x[1:], w[1:] - return x, w - - def _dynamic_dist2npts(self, dx): - return int(math.ceil(self.resolution[-1] * dx)) - - def _get_int_matrix(self, x: np.ndarray) -> np.ndarray: - dense_shape = (x.shape[0], self.x.shape[0]) - indices, values = [], [] - beg = x.shape[0] - for i in range(x.shape[0]): - for _ in range(self.w[i].shape[0]): - indices.append([i, beg]) - beg += 1 - values = np.hstack((values, self.w[i])) - return indices, values.astype(self.dtype), dense_shape diff --git a/examples/smc_reac/ppsci/equation/ide/__init__.py b/examples/smc_reac/ppsci/equation/ide/__init__.py deleted file mode 100644 index 4d4cab56cb..0000000000 --- a/examples/smc_reac/ppsci/equation/ide/__init__.py +++ /dev/null @@ -1,19 +0,0 @@ -# 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. - -from ppsci.equation.ide.volterra import Volterra - -__all__ = [ - "Volterra", -] diff --git a/examples/smc_reac/ppsci/equation/ide/volterra.py b/examples/smc_reac/ppsci/equation/ide/volterra.py deleted file mode 100644 index 77fb6f3173..0000000000 --- a/examples/smc_reac/ppsci/equation/ide/volterra.py +++ /dev/null @@ -1,127 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import Callable - -import numpy as np -import paddle - -from ppsci.equation.pde import PDE - - -class Volterra(PDE): - r"""A second kind of volterra integral equation with Gaussian quadrature algorithm. - - $$ - x(t) - f(t)=\int_a^t K(t, s) x(s) d s - $$ - - [Volterra integral equation](https://en.wikipedia.org/wiki/Volterra_integral_equation) - - [Gaussian quadrature](https://en.wikipedia.org/wiki/Gaussian_quadrature#Change_of_interval) - - Args: - bound (float): Lower bound `a` for Volterra integral equation. - num_points (int): Sampled points in integral interval. - quad_deg (int): Number of quadrature. - kernel_func (Callable): Kernel func `K(t,s)`. - func (Callable): `x(t) - f(t)` in Volterra integral equation. - - Examples: - >>> import ppsci - >>> import numpy as np - >>> vol_eq = ppsci.equation.Volterra( - ... 0, 12, 20, lambda t, s: np.exp(s - t), lambda out: out["u"], - ... ) - """ - - dtype = paddle.get_default_dtype() - - def __init__( - self, - bound: float, - num_points: int, - quad_deg: int, - kernel_func: Callable, - func: Callable, - ): - super().__init__() - self.bound = bound - self.num_points = num_points - self.quad_deg = quad_deg - self.kernel_func = kernel_func - self.func = func - - self.quad_x, self.quad_w = np.polynomial.legendre.leggauss(quad_deg) - self.quad_x = self.quad_x.astype(Volterra.dtype).reshape([-1, 1]) # [Q, 1] - self.quad_x = paddle.to_tensor(self.quad_x) # [Q, 1] - - self.quad_w = self.quad_w.astype(Volterra.dtype) # [Q, ] - - def compute_volterra_func(out): - x, u = out["x"], out["u"] - lhs = self.func(out) - - int_mat = paddle.to_tensor(self._get_int_matrix(x), stop_gradient=False) - rhs = paddle.mm(int_mat, u) # (N, 1) - - volterra = lhs[: len(rhs)] - rhs - return volterra - - self.add_equation("volterra", compute_volterra_func) - - def get_quad_points(self, t: paddle.Tensor) -> paddle.Tensor: - """Scale and transform quad_x from [-1, 1] to range [a, b]. - - reference: https://en.wikipedia.org/wiki/Gaussian_quadrature#Change_of_interval - - Args: - t (paddle.Tensor): Tensor array of upper bounds 't' for integral. - - Returns: - paddle.Tensor: Transformed points in desired range with shape of [N, Q]. - """ - a, b = self.bound, t - return ((b - a) / 2) @ self.quad_x.T + (b + a) / 2 - - def _get_quad_weights(self, t: float) -> np.ndarray: - """Scale weights to range according to given t and lower bound of integral. - - reference: https://en.wikipedia.org/wiki/Gaussian_quadrature#Change_of_interval - - Args: - t (float): Array of upper bound 't' for integral. - - Returns: - np.ndarray: Transformed weights in desired range with shape of [Q, ]. - """ - a, b = self.bound, t - return (b - a) / 2 * self.quad_w - - def _get_int_matrix(self, x: np.ndarray) -> np.ndarray: - int_mat = np.zeros( - (self.num_points, self.num_points + (self.num_points * self.quad_deg)), - dtype=Volterra.dtype, - ) - for i in range(self.num_points): - xi = float(x[i]) - beg = self.num_points + self.quad_deg * i - end = self.num_points + self.quad_deg * (i + 1) - K = np.ravel( - self.kernel_func(np.full((self.quad_deg, 1), xi), x[beg:end].numpy()) - ) - int_mat[i, beg:end] = self._get_quad_weights(xi) * K - return int_mat diff --git a/examples/smc_reac/ppsci/equation/pde/__init__.py b/examples/smc_reac/ppsci/equation/pde/__init__.py deleted file mode 100644 index 0dbcea2a8f..0000000000 --- a/examples/smc_reac/ppsci/equation/pde/__init__.py +++ /dev/null @@ -1,43 +0,0 @@ -# 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. - -from ppsci.equation.pde.allen_cahn import AllenCahn -from ppsci.equation.pde.base import DETACH_FUNC_NAME -from ppsci.equation.pde.base import PDE -from ppsci.equation.pde.biharmonic import Biharmonic -from ppsci.equation.pde.heat_exchanger import HeatExchanger -from ppsci.equation.pde.helmholtz import Helmholtz -from ppsci.equation.pde.laplace import Laplace -from ppsci.equation.pde.linear_elasticity import LinearElasticity -from ppsci.equation.pde.navier_stokes import NavierStokes -from ppsci.equation.pde.nls_m_b import NLSMB -from ppsci.equation.pde.normal_dot_vec import NormalDotVec -from ppsci.equation.pde.poisson import Poisson -from ppsci.equation.pde.viv import Vibration - -__all__ = [ - "PDE", - "DETACH_FUNC_NAME", - "AllenCahn", - "Biharmonic", - "HeatExchanger", - "Helmholtz", - "Laplace", - "LinearElasticity", - "NavierStokes", - "NLSMB", - "NormalDotVec", - "Poisson", - "Vibration", -] diff --git a/examples/smc_reac/ppsci/equation/pde/allen_cahn.py b/examples/smc_reac/ppsci/equation/pde/allen_cahn.py deleted file mode 100644 index 44e0ec899f..0000000000 --- a/examples/smc_reac/ppsci/equation/pde/allen_cahn.py +++ /dev/null @@ -1,64 +0,0 @@ -# Copyright (c) 2024 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. - -from __future__ import annotations - -from typing import Optional -from typing import Tuple - -from ppsci.autodiff import jacobian -from ppsci.equation.pde import base - - -class AllenCahn(base.PDE): - r"""Class for Allen-Cahn equation. - - $$ - \dfrac{\partial u}{\partial t} - \epsilon^2 \Delta u + 5u^3 - 5u = 0 - $$ - - Args: - eps (float): Represents the characteristicscale of interfacial width, - influencing the thickness and dynamics of phase boundaries. - detach_keys (Optional[Tuple[str, ...]]): Keys used for detach during computing. - Defaults to None. - - Examples: - >>> import ppsci - >>> pde = ppsci.equation.AllenCahn(eps=0.01) - """ - - def __init__( - self, - eps: float, - detach_keys: Optional[Tuple[str, ...]] = None, - ): - super().__init__() - self.detach_keys = detach_keys - self.eps = eps - # t, x = self.create_symbols("t x") - # invars = (t, x, ) - # u = self.create_function("u", invars) - # allen_cahn = u.diff(t) + 5 * u**3 - 5 * u - 0.0001 * u.diff(x, 2) - - # TODO: Pow(u,3) seems cause slightly larger L2 error than multiply(u*u*u) - def allen_cahn(out): - t, x = out["t"], out["x"] - u = out["u"] - u__t, u__x = jacobian(u, [t, x]) - u__x__x = jacobian(u__x, x) - - return u__t - (self.eps**2) * u__x__x + 5 * u * u * u - 5 * u - - self.add_equation("allen_cahn", allen_cahn) diff --git a/examples/smc_reac/ppsci/equation/pde/base.py b/examples/smc_reac/ppsci/equation/pde/base.py deleted file mode 100644 index 41f54b3861..0000000000 --- a/examples/smc_reac/ppsci/equation/pde/base.py +++ /dev/null @@ -1,243 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import Callable -from typing import Dict -from typing import List -from typing import Optional -from typing import Tuple -from typing import Union - -import paddle -import sympy as sp -from paddle import nn - -DETACH_FUNC_NAME = "detach" - - -class PDE: - """Base class for Partial Differential Equation.""" - - def __init__(self): - super().__init__() - self.equations: Dict[str, Union[Callable, sp.Basic]] = {} - # for PDE which has learnable parameter(s) - self.learnable_parameters = nn.ParameterList() - - self.detach_keys: Optional[Tuple[str, ...]] = None - - @staticmethod - def create_symbols( - symbol_str: str, - ) -> Union[sp.Symbol, Tuple[sp.Symbol, ...]]: - """Create symbolic variables. - - Args: - symbol_str (str): String contains symbols, such as "x", "x y z". - - Returns: - Union[sympy.Symbol, Tuple[sympy.Symbol, ...]]: Created symbol(s). - - Examples: - >>> import ppsci - >>> pde = ppsci.equation.PDE() - >>> symbol_x = pde.create_symbols('x') - >>> symbols_xyz = pde.create_symbols('x y z') - >>> print(symbol_x) - x - >>> print(symbols_xyz) - (x, y, z) - """ - return sp.symbols(symbol_str) - - def create_function(self, name: str, invars: Tuple[sp.Symbol, ...]) -> sp.Function: - """Create named function depending on given invars. - - Args: - name (str): Function name. such as "u", "v", and "f". - invars (Tuple[sympy.Symbol, ...]): List of independent variable of function. - - Returns: - sympy.Function: Named sympy function. - - Examples: - >>> import ppsci - >>> pde = ppsci.equation.PDE() - >>> x, y, z = pde.create_symbols('x y z') - >>> u = pde.create_function('u', (x, y)) - >>> f = pde.create_function('f', (x, y, z)) - >>> print(u) - u(x, y) - >>> print(f) - f(x, y, z) - """ - expr = sp.Function(name)(*invars) - - return expr - - def _apply_detach(self): - """ - Wrap detached sub_expr into detach(sub_expr) to prevent gradient back-propagation, only for those items specified in self.detach_keys. - - NOTE: This function is expected to be called after self.equations is ready in PDE.__init__. - - Examples: - >>> import ppsci - >>> ns = ppsci.equation.NavierStokes(1.0, 1.0, 2, False) - >>> print(ns) - NavierStokes - continuity: Derivative(u(x, y), x) + Derivative(v(x, y), y) - momentum_x: u(x, y)*Derivative(u(x, y), x) + v(x, y)*Derivative(u(x, y), y) + 1.0*Derivative(p(x, y), x) - 1.0*Derivative(u(x, y), (x, 2)) - 1.0*Derivative(u(x, y), (y, 2)) - momentum_y: u(x, y)*Derivative(v(x, y), x) + v(x, y)*Derivative(v(x, y), y) + 1.0*Derivative(p(x, y), y) - 1.0*Derivative(v(x, y), (x, 2)) - 1.0*Derivative(v(x, y), (y, 2)) - >>> detach_keys = ("u", "v__y") - >>> ns = ppsci.equation.NavierStokes(1.0, 1.0, 2, False, detach_keys=detach_keys) - >>> print(ns) - NavierStokes - continuity: detach(Derivative(v(x, y), y)) + Derivative(u(x, y), x) - momentum_x: detach(u(x, y))*Derivative(u(x, y), x) + v(x, y)*Derivative(u(x, y), y) + 1.0*Derivative(p(x, y), x) - 1.0*Derivative(u(x, y), (x, 2)) - 1.0*Derivative(u(x, y), (y, 2)) - momentum_y: detach(u(x, y))*Derivative(v(x, y), x) + detach(Derivative(v(x, y), y))*v(x, y) + 1.0*Derivative(p(x, y), y) - 1.0*Derivative(v(x, y), (x, 2)) - 1.0*Derivative(v(x, y), (y, 2)) - """ - if self.detach_keys is None: - return - - from copy import deepcopy - - from sympy.core.traversal import postorder_traversal - - from ppsci.utils.symbolic import _cvt_to_key - - for name, expr in self.equations.items(): - if not isinstance(expr, sp.Basic): - continue - # only process sympy expression - expr_ = deepcopy(expr) - for item in postorder_traversal(expr): - if _cvt_to_key(item) in self.detach_keys: - # inplace all related sub_expr into detach(sub_expr) - expr_ = expr_.replace(item, sp.Function(DETACH_FUNC_NAME)(item)) - - # remove all detach wrapper for more-than-once wrapped items to prevent duplicated wrapping - expr_ = expr_.replace( - sp.Function(DETACH_FUNC_NAME)( - sp.Function(DETACH_FUNC_NAME)(item) - ), - sp.Function(DETACH_FUNC_NAME)(item), - ) - - # remove unccessary detach wrapping for the first arg of Derivative - for item_ in list(postorder_traversal(expr_)): - if isinstance(item_, sp.Derivative): - if item_.args[0].name == DETACH_FUNC_NAME: - expr_ = expr_.replace( - item_, - sp.Derivative( - item_.args[0].args[0], *item_.args[1:] - ), - ) - - self.equations[name] = expr_ - - def add_equation(self, name: str, equation: Callable): - """Add an equation. - - Args: - name (str): Name of equation - equation (Callable): Computation function for equation. - - Examples: - >>> import ppsci - >>> import sympy - >>> pde = ppsci.equation.PDE() - >>> x, y = pde.create_symbols('x y') - >>> u = x**2 + y**2 - >>> equation = sympy.diff(u, x) + sympy.diff(u, y) - >>> pde.add_equation('linear_pde', equation) - >>> print(pde) - PDE - linear_pde: 2*x + 2*y - """ - self.equations.update({name: equation}) - - def parameters(self) -> List[paddle.Tensor]: - """Return learnable parameters contained in PDE. - - Returns: - List[Tensor]: A list of learnable parameters. - - Examples: - >>> import ppsci - >>> pde = ppsci.equation.Vibration(2, -4, 0) - >>> print(pde.parameters()) - [Parameter containing: - Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=False, - -4.), Parameter containing: - Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=False, - 0.)] - """ - return self.learnable_parameters.parameters() - - def state_dict(self) -> Dict[str, paddle.Tensor]: - """Return named learnable parameters in dict. - - Returns: - Dict[str, Tensor]: A dict of states(str) and learnable parameters(Tensor). - - Examples: - >>> import ppsci - >>> pde = ppsci.equation.Vibration(2, -4, 0) - >>> print(pde.state_dict()) - OrderedDict([('0', Parameter containing: - Tensor(shape=[], dtype=float64, place=Place(gpu:0), stop_gradient=False, - -4.)), ('1', Parameter containing: - Tensor(shape=[], dtype=float64, place=Place(gpu:0), stop_gradient=False, - 0.))]) - """ - return self.learnable_parameters.state_dict() - - def set_state_dict( - self, state_dict: Dict[str, paddle.Tensor] - ) -> Tuple[List[str], List[str]]: - """Set state dict from dict. - - Args: - state_dict (Dict[str, paddle.Tensor]): The state dict to be set. - - Returns: - Tuple[List[str], List[str]]: List of missing_keys and unexpected_keys. - Expected to be two empty tuples mostly. - - Examples: - >>> import paddle - >>> import ppsci - >>> paddle.set_default_dtype("float64") - >>> pde = ppsci.equation.Vibration(2, -4, 0) - >>> state = pde.state_dict() - >>> state['0'] = paddle.to_tensor(-3.1) - >>> pde.set_state_dict(state) - ([], []) - >>> print(state) - OrderedDict([('0', Tensor(shape=[], dtype=float64, place=Place(gpu:0), stop_gradient=True, - -3.10000000)), ('1', Parameter containing: - Tensor(shape=[], dtype=float64, place=Place(gpu:0), stop_gradient=False, - 0.))]) - """ - return self.learnable_parameters.set_state_dict(state_dict) - - def __str__(self): - return "\n".join( - [self.__class__.__name__] - + [f" {name}: {eq}" for name, eq in self.equations.items()] - ) diff --git a/examples/smc_reac/ppsci/equation/pde/biharmonic.py b/examples/smc_reac/ppsci/equation/pde/biharmonic.py deleted file mode 100644 index 933888ac60..0000000000 --- a/examples/smc_reac/ppsci/equation/pde/biharmonic.py +++ /dev/null @@ -1,74 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import Optional -from typing import Tuple -from typing import Union - -import sympy - -from ppsci.equation.pde import base - - -class Biharmonic(base.PDE): - r"""Class for biharmonic equation with supporting special load. - - $$ - \nabla^4 \varphi = \dfrac{q}{D} - $$ - - Args: - dim (int): Dimension of equation. - q (Union[float, str, sympy.Basic]): Load. - D (Union[float, str]): Rigidity. - detach_keys (Optional[Tuple[str, ...]]): Keys used for detach during computing. - Defaults to None. - - Examples: - >>> import ppsci - >>> pde = ppsci.equation.Biharmonic(2, -1.0, 1.0) - """ - - def __init__( - self, - dim: int, - q: Union[float, str, sympy.Basic], - D: Union[float, str], - detach_keys: Optional[Tuple[str, ...]] = None, - ): - super().__init__() - self.detach_keys = detach_keys - - invars = self.create_symbols("x y z")[:dim] - u = self.create_function("u", invars) - - if isinstance(q, str): - q = self.create_function("q", invars) - if isinstance(D, str): - D = self.create_function("D", invars) - - self.dim = dim - self.q = q - self.D = D - - biharmonic = -self.q / self.D - for invar_i in invars: - for invar_j in invars: - biharmonic += u.diff(invar_i, 2).diff(invar_j, 2) - - self.add_equation("biharmonic", biharmonic) - - self._apply_detach() diff --git a/examples/smc_reac/ppsci/equation/pde/heat_exchanger.py b/examples/smc_reac/ppsci/equation/pde/heat_exchanger.py deleted file mode 100644 index c2e0107ff3..0000000000 --- a/examples/smc_reac/ppsci/equation/pde/heat_exchanger.py +++ /dev/null @@ -1,94 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import Union - -from ppsci.equation.pde import base - - -class HeatExchanger(base.PDE): - r"""Class for heat exchanger equation. - - $$ - \begin{aligned} - & L\left(\frac{q_m c_p}{v}\right)_{\mathrm{c}} \frac{\partial T_{\mathrm{c}}}{\partial \tau}-L\left(q_m c_p\right)_{\mathrm{c}} \frac{\partial T_{\mathrm{c}}}{\partial x}=\left(\eta_{\mathrm{o}} \alpha A\right)_{\mathrm{c}}\left(T_{\mathrm{w}}-T_{\mathrm{c}}\right), \\ - & L\left(\frac{q_m c_p}{v}\right)_{\mathrm{h}} \frac{\partial T_{\mathrm{h}}}{\partial \tau}+L\left(q_m c_p\right)_{\mathrm{h}} \frac{\partial T_{\mathrm{h}}}{\partial x}=\left(\eta_{\mathrm{o}} \alpha A\right)_{\mathrm{h}}\left(T_{\mathrm{w}}-T_{\mathrm{h}}\right), \\ - & \left(M c_p\right)_{\mathrm{w}} \frac{\partial T_{\mathrm{w}}}{\partial \tau}=\left(\eta_{\mathrm{o}} \alpha A\right)_{\mathrm{h}}\left(T_{\mathrm{h}}-T_{\mathrm{w}}\right)+\left(\eta_{\mathrm{o}} \alpha A\right)_{\mathrm{c}}\left(T_{\mathrm{c}}-T_{\mathrm{w}}\right). - \end{aligned} - $$ - - where: - - - $T$ is temperature, - - $q_m$ is mass flow rate, - - $c_p$ represents specific heat capacity, - - $v$ denotes flow velocity, - - $L$ stands for flow length, - - $\eta_{\mathrm{o}}$ signifies fin surface efficiency, - - $\alpha$ stands for heat transfer coefficient, - - $A$ indicates heat transfer area, - - $M$ represents the mass of the heat transfer structure, - - $\tau$ correspond to time, - - $x$ correspond flow direction, - - Subscripts $\mathrm{h}$, $\mathrm{c}$, and $\mathrm{w}$ denote the hot fluid side, cold fluid side, and heat transfer wall, respectively. - - Args: - alpha_h: $\frac{(\eta_o\alpha A)_h}{L(c_p)_h}$ - alpha_c: $\frac{(\eta_o\alpha A)_c}{L(c_p)_c}$ - v_h: $v_h$ - v_c: $v_c$ - w_h: $\frac{(\eta_o\alpha A)_h}{M(c_p)_w}$ - w_c: $\frac{(\eta_o\alpha A)_c}{M(c_p)_w}$ - - Examples: - >>> import ppsci - >>> pde = ppsci.equation.HeatExchanger(1.0,1.0,1.0,1.0,1.0,1.0) - """ - - def __init__( - self, - alpha_h: Union[float, str], - alpha_c: Union[float, str], - v_h: Union[float, str], - v_c: Union[float, str], - w_h: Union[float, str], - w_c: Union[float, str], - ): - super().__init__() - x, t, qm_h, qm_c = self.create_symbols("x t qm_h qm_c") - - T_h = self.create_function("T_h", (x, t, qm_h)) - T_c = self.create_function("T_c", (x, t, qm_c)) - T_w = self.create_function("T_w", (x, t)) - - T_h_x = T_h.diff(x) - T_h_t = T_h.diff(t) - T_c_x = T_c.diff(x) - T_c_t = T_c.diff(t) - T_w_t = T_w.diff(t) - - beta_h = (alpha_h * v_h) / qm_h - beta_c = (alpha_c * v_c) / qm_c - - heat_boundary = T_h_t + v_h * T_h_x - beta_h * (T_w - T_h) - cold_boundary = T_c_t - v_c * T_c_x - beta_c * (T_w - T_c) - wall = T_w_t - w_h * (T_h - T_w) - w_c * (T_c - T_w) - - self.add_equation("heat_boundary", heat_boundary) - self.add_equation("cold_boundary", cold_boundary) - self.add_equation("wall", wall) - - self._apply_detach() diff --git a/examples/smc_reac/ppsci/equation/pde/helmholtz.py b/examples/smc_reac/ppsci/equation/pde/helmholtz.py deleted file mode 100644 index e71fdbe983..0000000000 --- a/examples/smc_reac/ppsci/equation/pde/helmholtz.py +++ /dev/null @@ -1,119 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import Callable -from typing import Dict -from typing import Optional -from typing import Tuple - -import paddle - -from ppsci.equation.pde import base - - -def hvp_revrev(f: Callable, primals: Tuple[paddle.Tensor, ...]) -> paddle.Tensor: - """Compute the Hessian vector product of f with respect to primals using - double backward trick in reverse mode AD. - - Args: - f (Callable): Function to compute HVP. - primals (Tuple[paddle.Tensor, ...]): Input tensors. - - Returns: - paddle.Tensor: Hessian vector product of f with respect to primals. - """ - # TODO: Merge this option into ppsci.autodiff.ad - g = lambda primals: paddle.incubate.autograd.jvp(f, primals)[1] - tangents_out = paddle.incubate.autograd.jvp(g, primals)[1] - return tangents_out[0] - - -class Helmholtz(base.PDE): - r"""Class for helmholtz equation. - - $$ - \nabla^2 u + k^2 u = f - $$ - - $$ - \text{where } f \text{ is the source term}. - $$ - - Args: - dim (int): Dimension of equation. - k (float): The wave number, which is a parameter that affects the frequency of the solution. - detach_keys (Optional[Tuple[str, ...]]): Keys used for detach during computing. - Defaults to None. - - Examples: - >>> import ppsci - >>> model = ppsci.arch.MLP(("x", "y"), ("u",), 2, 32) - >>> pde = ppsci.equation.Helmholtz(2, -1.0, model) - """ - - def __init__( - self, - dim: int, - k: float, - model: paddle.nn.Layer, - detach_keys: Optional[Tuple[str, ...]] = None, - ): - super().__init__() - self.dim = dim - self.k = k - self.detach_keys = detach_keys - - invars = self.create_symbols("x y z")[:dim] - - # TODO: This is a hack, should be simplified in the future - self.model = model - - def helmholtz(data_dict: Dict[str, paddle.Tensor]) -> paddle.Tensor: - xs = tuple(data_dict[invar.name] for invar in invars) - - # TODO: Hard code here, for hvp_revrev requires tuple input(s) but not dict - if self.dim == 1: - u__x__x = hvp_revrev(lambda x_: self.model.forward_tensor(x_), (xs[0],)) - out = (self.k**2) * data_dict["u"] + u__x__x - elif self.dim == 2: - u__x__x = hvp_revrev( - lambda x_: self.model.forward_tensor(x_, xs[1]), (xs[0],) - ) - u__y__y = hvp_revrev( - lambda y_: self.model.forward_tensor(xs[0], y_), (xs[1],) - ) - out = (self.k**2) * data_dict["u"] + u__x__x + u__y__y - elif self.dim == 3: - u__x__x = hvp_revrev( - lambda x_: self.model.forward_tensor(x_, xs[1], xs[2]), (xs[0],) - ) - u__y__y = hvp_revrev( - lambda y_: self.model.forward_tensor(xs[0], y_, xs[2]), (xs[1],) - ) - u__z__z = hvp_revrev( - lambda z_: self.model.forward_tensor(xs[0], xs[1], z_), (xs[2],) - ) - out = (self.k**2) * data_dict["u"] + u__x__x + u__y__y + u__z__z - else: - raise NotImplementedError( - f"dim should be less or equal to 3, but got {self.dim}." - ) - - return out - - self.add_equation("helmholtz", helmholtz) - - self._apply_detach() diff --git a/examples/smc_reac/ppsci/equation/pde/laplace.py b/examples/smc_reac/ppsci/equation/pde/laplace.py deleted file mode 100644 index b99d7c8d9a..0000000000 --- a/examples/smc_reac/ppsci/equation/pde/laplace.py +++ /dev/null @@ -1,55 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import Optional -from typing import Tuple - -from ppsci.equation.pde import base - - -class Laplace(base.PDE): - r"""Class for laplace equation. - - $$ - \nabla^2 \varphi = 0 - $$ - - Args: - dim (int): Dimension of equation. - detach_keys (Optional[Tuple[str, ...]]): Keys used for detach during computing. - Defaults to None. - - Examples: - >>> import ppsci - >>> pde = ppsci.equation.Laplace(2) - """ - - def __init__(self, dim: int, detach_keys: Optional[Tuple[str, ...]] = None): - super().__init__() - self.detach_keys = detach_keys - - invars = self.create_symbols("x y z")[:dim] - u = self.create_function("u", invars) - - self.dim = dim - - laplace = 0 - for invar in invars: - laplace += u.diff(invar, 2) - - self.add_equation("laplace", laplace) - - self._apply_detach() diff --git a/examples/smc_reac/ppsci/equation/pde/linear_elasticity.py b/examples/smc_reac/ppsci/equation/pde/linear_elasticity.py deleted file mode 100644 index 289d924899..0000000000 --- a/examples/smc_reac/ppsci/equation/pde/linear_elasticity.py +++ /dev/null @@ -1,184 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import Optional -from typing import Tuple -from typing import Union - -import sympy as sp - -from ppsci.equation.pde import base - - -class LinearElasticity(base.PDE): - r"""Linear elasticity equations. - - Use either (E, nu) or (lambda_, mu) to define the material properties. - - $$ - \begin{cases} - stress\_disp_{xx} = \lambda(\dfrac{\partial u}{\partial x} + \dfrac{\partial v}{\partial y} + \dfrac{\partial w}{\partial z}) + 2\mu \dfrac{\partial u}{\partial x} - \sigma_{xx} \\ - stress\_disp_{yy} = \lambda(\dfrac{\partial u}{\partial x} + \dfrac{\partial v}{\partial y} + \dfrac{\partial w}{\partial z}) + 2\mu \dfrac{\partial v}{\partial y} - \sigma_{yy} \\ - stress\_disp_{zz} = \lambda(\dfrac{\partial u}{\partial x} + \dfrac{\partial v}{\partial y} + \dfrac{\partial w}{\partial z}) + 2\mu \dfrac{\partial w}{\partial z} - \sigma_{zz} \\ - stress\_disp_{xy} = \mu(\dfrac{\partial u}{\partial y} + \dfrac{\partial v}{\partial x}) - \sigma_{xy} \\ - stress\_disp_{xz} = \mu(\dfrac{\partial u}{\partial z} + \dfrac{\partial w}{\partial x}) - \sigma_{xz} \\ - stress\_disp_{yz} = \mu(\dfrac{\partial v}{\partial z} + \dfrac{\partial w}{\partial y}) - \sigma_{yz} \\ - equilibrium_{x} = \rho \dfrac{\partial^2 u}{\partial t^2} - (\dfrac{\partial \sigma_{xx}}{\partial x} + \dfrac{\partial \sigma_{xy}}{\partial y} + \dfrac{\partial \sigma_{xz}}{\partial z}) \\ - equilibrium_{y} = \rho \dfrac{\partial^2 u}{\partial t^2} - (\dfrac{\partial \sigma_{xy}}{\partial x} + \dfrac{\partial \sigma_{yy}}{\partial y} + \dfrac{\partial \sigma_{yz}}{\partial z}) \\ - equilibrium_{z} = \rho \dfrac{\partial^2 u}{\partial t^2} - (\dfrac{\partial \sigma_{xz}}{\partial x} + \dfrac{\partial \sigma_{yz}}{\partial y} + \dfrac{\partial \sigma_{zz}}{\partial z}) \\ - \end{cases} - $$ - - Args: - E (Optional[Union[float, str]]): The Young's modulus. Defaults to None. - nu (Optional[Union[float, str]]): The Poisson's ratio. Defaults to None. - lambda_ (Optional[Union[float, str]]): Lamé's first parameter. Defaults to None. - mu (Optional[Union[float, str]]): Lamé's second parameter (shear modulus). Defaults to None. - rho (Union[float, str], optional): Mass density. Defaults to 1. - dim (int, optional): Dimension of the linear elasticity (2 or 3). Defaults to 3. - time (bool, optional): Whether contains time data. Defaults to False. - detach_keys (Optional[Tuple[str, ...]]): Keys used for detach during computing. - Defaults to None. - - Examples: - >>> import ppsci - >>> pde = ppsci.equation.LinearElasticity( - ... E=None, nu=None, lambda_=1e4, mu=100, dim=3 - ... ) - """ - - def __init__( - self, - E: Optional[Union[float, str]] = None, - nu: Optional[Union[float, str]] = None, - lambda_: Optional[Union[float, str]] = None, - mu: Optional[Union[float, str]] = None, - rho: Union[float, str] = 1, - dim: int = 3, - time: bool = False, - detach_keys: Optional[Tuple[str, ...]] = None, - ): - super().__init__() - self.detach_keys = detach_keys - self.dim = dim - self.time = time - - t, x, y, z = self.create_symbols("t x y z") - normal_x, normal_y, normal_z = self.create_symbols("normal_x normal_y normal_z") - invars = (x, y) - if time: - invars = (t,) + invars - if self.dim == 3: - invars += (z,) - - u = self.create_function("u", invars) - v = self.create_function("v", invars) - w = self.create_function("w", invars) if dim == 3 else sp.Number(0) - - sigma_xx = self.create_function("sigma_xx", invars) - sigma_yy = self.create_function("sigma_yy", invars) - sigma_xy = self.create_function("sigma_xy", invars) - sigma_zz = ( - self.create_function("sigma_zz", invars) if dim == 3 else sp.Number(0) - ) - sigma_xz = ( - self.create_function("sigma_xz", invars) if dim == 3 else sp.Number(0) - ) - sigma_yz = ( - self.create_function("sigma_yz", invars) if dim == 3 else sp.Number(0) - ) - - # compute lambda and mu - if lambda_ is None: - if isinstance(nu, str): - nu = self.create_function(nu, invars) - if isinstance(E, str): - E = self.create_function(E, invars) - lambda_ = nu * E / ((1 + nu) * (1 - 2 * nu)) - mu = E / (2 * (1 + nu)) - else: - if isinstance(lambda_, str): - lambda_ = self.create_function(lambda_, invars) - if isinstance(mu, str): - mu = self.create_function(mu, invars) - - if isinstance(rho, str): - rho = self.create_function(rho, invars) - - self.E = E - self.nu = nu - self.lambda_ = lambda_ - self.mu = mu - self.rho = rho - - # compute stress equations - stress_disp_xx = ( - lambda_ * (u.diff(x) + v.diff(y) + w.diff(z)) - + 2 * mu * u.diff(x) - - sigma_xx - ) - stress_disp_yy = ( - lambda_ * (u.diff(x) + v.diff(y) + w.diff(z)) - + 2 * mu * v.diff(y) - - sigma_yy - ) - stress_disp_zz = ( - lambda_ * (u.diff(x) + v.diff(y) + w.diff(z)) - + 2 * mu * w.diff(z) - - sigma_zz - ) - stress_disp_xy = mu * (u.diff(y) + v.diff(x)) - sigma_xy - stress_disp_xz = mu * (u.diff(z) + w.diff(x)) - sigma_xz - stress_disp_yz = mu * (v.diff(z) + w.diff(y)) - sigma_yz - - # compute equilibrium equations - equilibrium_x = rho * ((u.diff(t)).diff(t)) - ( - sigma_xx.diff(x) + sigma_xy.diff(y) + sigma_xz.diff(z) - ) - equilibrium_y = rho * ((v.diff(t)).diff(t)) - ( - sigma_xy.diff(x) + sigma_yy.diff(y) + sigma_yz.diff(z) - ) - equilibrium_z = rho * ((w.diff(t)).diff(t)) - ( - sigma_xz.diff(x) + sigma_yz.diff(y) + sigma_zz.diff(z) - ) - - # compute traction equations - traction_x = normal_x * sigma_xx + normal_y * sigma_xy + normal_z * sigma_xz - traction_y = normal_x * sigma_xy + normal_y * sigma_yy + normal_z * sigma_yz - traction_z = normal_x * sigma_xz + normal_y * sigma_yz + normal_z * sigma_zz - - # add stress equations - self.add_equation("stress_disp_xx", stress_disp_xx) - self.add_equation("stress_disp_yy", stress_disp_yy) - self.add_equation("stress_disp_xy", stress_disp_xy) - if self.dim == 3: - self.add_equation("stress_disp_zz", stress_disp_zz) - self.add_equation("stress_disp_xz", stress_disp_xz) - self.add_equation("stress_disp_yz", stress_disp_yz) - - # add equilibrium equations - self.add_equation("equilibrium_x", equilibrium_x) - self.add_equation("equilibrium_y", equilibrium_y) - if self.dim == 3: - self.add_equation("equilibrium_z", equilibrium_z) - - # add traction equations - self.add_equation("traction_x", traction_x) - self.add_equation("traction_y", traction_y) - if self.dim == 3: - self.add_equation("traction_z", traction_z) - - self._apply_detach() diff --git a/examples/smc_reac/ppsci/equation/pde/navier_stokes.py b/examples/smc_reac/ppsci/equation/pde/navier_stokes.py deleted file mode 100644 index c0d3d193a2..0000000000 --- a/examples/smc_reac/ppsci/equation/pde/navier_stokes.py +++ /dev/null @@ -1,151 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import Optional -from typing import Tuple -from typing import Union - -import sympy as sp -from sympy.parsing import sympy_parser as sp_parser - -from ppsci.equation.pde import base - - -class NavierStokes(base.PDE): - r"""Class for navier-stokes equation. - - $$ - \begin{cases} - \dfrac{\partial u}{\partial x} + \dfrac{\partial v}{\partial y} + \dfrac{\partial w}{\partial z} = 0 \\ - \dfrac{\partial u}{\partial t} + u\dfrac{\partial u}{\partial x} + v\dfrac{\partial u}{\partial y} + w\dfrac{\partial u}{\partial z} = - - \dfrac{1}{\rho}\dfrac{\partial p}{\partial x} - + \nu( - \dfrac{\partial ^2 u}{\partial x ^2} - + \dfrac{\partial ^2 u}{\partial y ^2} - + \dfrac{\partial ^2 u}{\partial z ^2} - ) \\ - \dfrac{\partial v}{\partial t} + u\dfrac{\partial v}{\partial x} + v\dfrac{\partial v}{\partial y} + w\dfrac{\partial v}{\partial z} = - - \dfrac{1}{\rho}\dfrac{\partial p}{\partial y} - + \nu( - \dfrac{\partial ^2 v}{\partial x ^2} - + \dfrac{\partial ^2 v}{\partial y ^2} - + \dfrac{\partial ^2 v}{\partial z ^2} - ) \\ - \dfrac{\partial w}{\partial t} + u\dfrac{\partial w}{\partial x} + v\dfrac{\partial w}{\partial y} + w\dfrac{\partial w}{\partial z} = - - \dfrac{1}{\rho}\dfrac{\partial p}{\partial z} - + \nu( - \dfrac{\partial ^2 w}{\partial x ^2} - + \dfrac{\partial ^2 w}{\partial y ^2} - + \dfrac{\partial ^2 w}{\partial z ^2} - ) \\ - \end{cases} - $$ - - Args: - nu (Union[float, str]): Dynamic viscosity. - rho (Union[float, str]): Density. - dim (int): Dimension of equation. - time (bool): Whether the equation is time-dependent. - detach_keys (Optional[Tuple[str, ...]]): Keys used for detach during computing. - Defaults to None. - - Examples: - >>> import ppsci - >>> pde = ppsci.equation.NavierStokes(0.1, 1.0, 3, False) - """ - - def __init__( - self, - nu: Union[float, str], - rho: Union[float, str], - dim: int, - time: bool, - detach_keys: Optional[Tuple[str, ...]] = None, - ): - super().__init__() - self.detach_keys = detach_keys - self.dim = dim - self.time = time - - t, x, y, z = self.create_symbols("t x y z") - invars = (x, y) - if time: - invars = (t,) + invars - if dim == 3: - invars += (z,) - - if isinstance(nu, str): - nu = sp_parser.parse_expr(nu) - if isinstance(nu, sp.Symbol): - invars += (nu,) - - if isinstance(rho, str): - rho = sp_parser.parse_expr(rho) - if isinstance(rho, sp.Symbol): - invars += (rho,) - - self.nu = nu - self.rho = rho - - u = self.create_function("u", invars) - v = self.create_function("v", invars) - w = self.create_function("w", invars) if dim == 3 else sp.Number(0) - p = self.create_function("p", invars) - - continuity = u.diff(x) + v.diff(y) + w.diff(z) - momentum_x = ( - u.diff(t) - + u * u.diff(x) - + v * u.diff(y) - + w * u.diff(z) - - ( - (nu * u.diff(x)).diff(x) - + (nu * u.diff(y)).diff(y) - + (nu * u.diff(z)).diff(z) - ) - + 1 / rho * p.diff(x) - ) - momentum_y = ( - v.diff(t) - + u * v.diff(x) - + v * v.diff(y) - + w * v.diff(z) - - ( - (nu * v.diff(x)).diff(x) - + (nu * v.diff(y)).diff(y) - + (nu * v.diff(z)).diff(z) - ) - + 1 / rho * p.diff(y) - ) - momentum_z = ( - w.diff(t) - + u * w.diff(x) - + v * w.diff(y) - + w * w.diff(z) - - ( - (nu * w.diff(x)).diff(x) - + (nu * w.diff(y)).diff(y) - + (nu * w.diff(z)).diff(z) - ) - + 1 / rho * p.diff(z) - ) - self.add_equation("continuity", continuity) - self.add_equation("momentum_x", momentum_x) - self.add_equation("momentum_y", momentum_y) - if self.dim == 3: - self.add_equation("momentum_z", momentum_z) - - self._apply_detach() diff --git a/examples/smc_reac/ppsci/equation/pde/nls_m_b.py b/examples/smc_reac/ppsci/equation/pde/nls_m_b.py deleted file mode 100644 index 3db2984268..0000000000 --- a/examples/smc_reac/ppsci/equation/pde/nls_m_b.py +++ /dev/null @@ -1,101 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import Optional -from typing import Tuple -from typing import Union - -from ppsci.equation.pde import base - - -class NLSMB(base.PDE): - r"""Class for nonlinear Schrodinger-Maxwell-Bloch equation. - - $$ - \begin{cases} - \dfrac{\partial E}{\partial x} = i \alpha_1 \dfrac{\partial^2 E}{\partial t ^2} - i \alpha_2 |E|^2 E+2 p \\ - \dfrac{\partial p}{\partial t} = 2 i \omega_0 p+2 E \eta \\ - \dfrac{\partial \eta}{\partial t} = -(E p^* + E^* p) - \end{cases} - $$ - - Args: - alpha_1 (Union[float, str]): Group velocity dispersion. - alpha_2 (Union[float, str]): Kerr nonlinearity. - omega_0 (Union[float, str]): The offset of resonance frequency. - time (bool): Whether the equation is time-dependent. - detach_keys (Optional[Tuple[str, ...]]): Keys used for detach during computing. - Defaults to None. - - Examples: - >>> import ppsci - >>> pde = ppsci.equation.NLSMB(0.5, -1.0, 0.5, True) - """ - - def __init__( - self, - alpha_1: Union[float, str], - alpha_2: Union[float, str], - omega_0: Union[float, str], - time: bool, - detach_keys: Optional[Tuple[str, ...]] = None, - ): - super().__init__() - self.detach_keys = detach_keys - self.time = time - - t, x = self.create_symbols("t x") - invars = (x,) - if time: - invars = (t,) + invars - - self.alpha_1 = alpha_1 - self.alpha_2 = alpha_2 - self.omega_0 = omega_0 - - Eu = self.create_function("Eu", invars) - Ev = self.create_function("Ev", invars) - pu = self.create_function("pu", invars) - pv = self.create_function("pv", invars) - eta = self.create_function("eta", invars) - - pu_t = pu.diff(t) - pv_t = pv.diff(t) - eta_t = eta.diff(t) - - Eu_x = Eu.diff(x) - Ev_x = Ev.diff(x) - - Eu_tt = Eu.diff(t).diff(t) - Ev_tt = Ev.diff(t).diff(t) - - Schrodinger_1 = ( - alpha_1 * Eu_tt - alpha_2 * Eu * (Eu**2 + Ev**2) + 2 * pv - Ev_x - ) - Schrodinger_2 = ( - alpha_1 * Ev_tt - alpha_2 * Ev * (Eu**2 + Ev**2) - 2 * pu + Eu_x - ) - Maxwell_1 = 2 * Ev * eta - pv_t + 2 * pu * omega_0 - Maxwell_2 = -2 * Eu * eta + pu_t + 2 * pv * omega_0 - Bloch = 2 * pv * Ev + 2 * pu * Eu + eta_t - - self.add_equation("Schrodinger_1", Schrodinger_1) - self.add_equation("Schrodinger_2", Schrodinger_2) - self.add_equation("Maxwell_1", Maxwell_1) - self.add_equation("Maxwell_2", Maxwell_2) - self.add_equation("Bloch", Bloch) - - self._apply_detach() diff --git a/examples/smc_reac/ppsci/equation/pde/normal_dot_vec.py b/examples/smc_reac/ppsci/equation/pde/normal_dot_vec.py deleted file mode 100644 index a6f3942eeb..0000000000 --- a/examples/smc_reac/ppsci/equation/pde/normal_dot_vec.py +++ /dev/null @@ -1,59 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import Optional -from typing import Tuple - -from ppsci.equation.pde import base - - -class NormalDotVec(base.PDE): - r"""Normal Dot Vector. - - $$ - \mathbf{n} \cdot \mathbf{v} = 0 - $$ - - Args: - vec_keys (Tuple[str, ...]): Keys for vectors, such as ("u", "v", "w") for - velocity vector. - detach_keys (Optional[Tuple[str, ...]]): Keys used for detach during computing. - Defaults to None. - - Examples: - >>> import ppsci - >>> pde = ppsci.equation.NormalDotVec(("u", "v", "w")) - """ - - def __init__( - self, vec_keys: Tuple[str, ...], detach_keys: Optional[Tuple[str, ...]] = None - ): - super().__init__() - self.detach_keys = detach_keys - if not vec_keys: - raise ValueError(f"len(vec_keys)({len(vec_keys)}) should be larger than 0.") - - self.vec_keys = vec_keys - vec_vars = self.create_symbols(" ".join(vec_keys)) - normals = self.create_symbols("normal_x normal_y normal_z") - - normal_dot_vec = 0 - for (normal, vec) in zip(normals, vec_vars): - normal_dot_vec += normal * vec - - self.add_equation("normal_dot_vec", normal_dot_vec) - - self._apply_detach() diff --git a/examples/smc_reac/ppsci/equation/pde/poisson.py b/examples/smc_reac/ppsci/equation/pde/poisson.py deleted file mode 100644 index 4f9551a23a..0000000000 --- a/examples/smc_reac/ppsci/equation/pde/poisson.py +++ /dev/null @@ -1,53 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import Optional -from typing import Tuple - -from ppsci.equation.pde import base - - -class Poisson(base.PDE): - r"""Class for poisson equation. - - $$ - \nabla^2 \varphi = C - $$ - - Args: - dim (int): Dimension of equation. - detach_keys (Optional[Tuple[str, ...]]): Keys used for detach during computing. - Defaults to None. - - Examples: - >>> import ppsci - >>> pde = ppsci.equation.Poisson(2) - """ - - def __init__(self, dim: int, detach_keys: Optional[Tuple[str, ...]] = None): - super().__init__() - self.detach_keys = detach_keys - invars = self.create_symbols("x y z")[:dim] - p = self.create_function("p", invars) - self.dim = dim - - poisson = 0 - for invar in invars: - poisson += p.diff(invar, 2) - - self.add_equation("poisson", poisson) - - self._apply_detach() diff --git a/examples/smc_reac/ppsci/equation/pde/viv.py b/examples/smc_reac/ppsci/equation/pde/viv.py deleted file mode 100644 index c3d85895f1..0000000000 --- a/examples/smc_reac/ppsci/equation/pde/viv.py +++ /dev/null @@ -1,64 +0,0 @@ -# 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. - -from __future__ import annotations - -import paddle -import sympy as sp -from paddle.nn import initializer - -from ppsci.equation.pde import base - - -class Vibration(base.PDE): - r"""Vortex induced vibration equation. - - $$ - \rho \dfrac{\partial^2 \eta}{\partial t^2} + e^{k1} \dfrac{\partial \eta}{\partial t} + e^{k2} \eta = f - $$ - - Args: - rho (float): Generalized mass. - k1 (float): Learnable parameter for modal damping. - k2 (float): Learnable parameter for generalized stiffness. - - Examples: - >>> import ppsci - >>> pde = ppsci.equation.Vibration(1.0, 4.0, -1.0) - """ - - def __init__(self, rho: float, k1: float, k2: float): - super().__init__() - self.rho = rho - self.k1 = paddle.create_parameter( - shape=[], - dtype=paddle.get_default_dtype(), - default_initializer=initializer.Constant(k1), - ) - self.k2 = paddle.create_parameter( - shape=[], - dtype=paddle.get_default_dtype(), - default_initializer=initializer.Constant(k2), - ) - self.learnable_parameters.append(self.k1) - self.learnable_parameters.append(self.k2) - - t_f = self.create_symbols("t_f") - eta = self.create_function("eta", (t_f,)) - k1 = self.create_symbols(self.k1.name) - k2 = self.create_symbols(self.k2.name) - f = self.rho * eta.diff(t_f, 2) + sp.exp(k1) * eta.diff(t_f) + sp.exp(k2) * eta - self.add_equation("f", f) - - self._apply_detach() diff --git a/examples/smc_reac/ppsci/experimental/__init__.py b/examples/smc_reac/ppsci/experimental/__init__.py deleted file mode 100644 index 842f19428a..0000000000 --- a/examples/smc_reac/ppsci/experimental/__init__.py +++ /dev/null @@ -1,37 +0,0 @@ -# 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. - -""" -This module is for experimental API -""" - -from ppsci.experimental.math_module import bessel_i0 -from ppsci.experimental.math_module import bessel_i0e -from ppsci.experimental.math_module import bessel_i1 -from ppsci.experimental.math_module import bessel_i1e -from ppsci.experimental.math_module import fractional_diff -from ppsci.experimental.math_module import gaussian_integrate -from ppsci.experimental.math_module import montecarlo_integrate -from ppsci.experimental.math_module import trapezoid_integrate - -__all__ = [ - "bessel_i0", - "bessel_i0e", - "bessel_i1", - "bessel_i1e", - "fractional_diff", - "gaussian_integrate", - "trapezoid_integrate", - "montecarlo_integrate", -] diff --git a/examples/smc_reac/ppsci/experimental/math_module.py b/examples/smc_reac/ppsci/experimental/math_module.py deleted file mode 100644 index fc255c5671..0000000000 --- a/examples/smc_reac/ppsci/experimental/math_module.py +++ /dev/null @@ -1,646 +0,0 @@ -# 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. - -from __future__ import annotations - -import functools -from typing import Any -from typing import Callable -from typing import List -from typing import Optional -from typing import Tuple -from typing import Union - -import numpy as np -import paddle -from typing_extensions import Literal - - -def bessel_i0(x: paddle.Tensor) -> paddle.Tensor: - """Zero-order modified Bézier curve functions of the first kind. - - Args: - x (paddle.Tensor): Input data of the formula. - - Examples: - >>> import paddle - >>> import ppsci - >>> res = ppsci.experimental.bessel_i0(paddle.to_tensor([0, 1, 2, 3, 4], dtype="float32")) - """ - return paddle.i0(x) - - -def bessel_i0e(x: paddle.Tensor) -> paddle.Tensor: - """Exponentially scaled zero-order modified Bézier curve functions of the first kind. - - Args: - x (paddle.Tensor): Input data of the formula. - - Examples: - >>> import paddle - >>> import ppsci - >>> res = ppsci.experimental.bessel_i0e(paddle.to_tensor([0, 1, 2, 3, 4], dtype="float32")) - """ - return paddle.i0e(x) - - -def bessel_i1(x: paddle.Tensor) -> paddle.Tensor: - """First-order modified Bézier curve functions of the first kind. - - Args: - x (paddle.Tensor): Input data of the formula. - - Examples: - >>> import paddle - >>> import ppsci - >>> res = ppsci.experimental.bessel_i1(paddle.to_tensor([0, 1, 2, 3, 4], dtype="float32")) - """ - return paddle.i1(x) - - -def bessel_i1e(x: paddle.Tensor) -> paddle.Tensor: - """Exponentially scaled first-order modified Bézier curve functions of the first kind. - - Args: - x (paddle.Tensor): Input data of the formula. - - Examples: - >>> import paddle - >>> import ppsci - >>> res = ppsci.experimental.bessel_i1e(paddle.to_tensor([0, 1, 2, 3, 4], dtype="float32")) - """ - return paddle.i1e(x) - - -def expand_func_values_and_squeeze_integral(f: Callable): - """This decorator ensures that the trailing dimension of integrands is indeed the integrand dimension. - This is pertinent in the 1d case when the sampled values are often of shape `(N,)`. Then, to maintain backward - consistency, we squeeze the result in the 1d case so it does not have any trailing dimensions. - - Args: - f (Callable): The wrapped function. - """ - - @functools.wraps(f) - def wrap(*args, **kwargs): - # i.e we only have one dimension, or the second dimension (that of the integrand) is 1 - is_1d = len(args[0].shape) == 1 or ( - len(args[0].shape) == 2 and args[0].shape[1] == 1 - ) - if is_1d: - return paddle.squeeze( - f(paddle.unsqueeze(args[0], axis=1), *args[1:], **kwargs) - ) - return f(*args, **kwargs) - - return wrap - - -def gaussian_integrate( - fn: Callable[[Any], paddle.Tensor], - dim: int, - N: int, - integration_domains: List[List[float]], - dtype: Literal["float32", "float64"] = "float64", -) -> paddle.Tensor: - """Integrate given function using gaussian quadrature. - - Args: - fn (Callable[[Any], paddle.Tensor]): Function to be integrated. - dim (int): Dimensionality of the integrand. - N (int): Number of dicretization points. - integration_domains (List[List[float]]): Integration domains. - dtype (Literal["float32", "float64"], optional): Dtype used during computation. Defaults to "float64". - - Returns: - paddle.Tensor: Integral result. - - Examples: - >>> import numpy as np - >>> import paddle - >>> import ppsci.experimental - >>> func = lambda x: paddle.sin(x) - >>> dim = 1 - >>> N = 500 - >>> integration_domains = [[0, np.pi]] - >>> result = ppsci.experimental.gaussian_integrate(func, dim, N, integration_domains) - >>> np.testing.assert_allclose(float(result), 2.0, 1e-6) - >>> print(float(result)) - 1.9999999999999576 - """ - - def _compatible_meshgrid(*args: paddle.Tensor, **kwargs: paddle.Tensor): - # TODO(HydrogenSulfate): paddle.meshgrid do not support single Tensor, - # which will be fixed in paddle framework. - if len(args) == 1: - return args - else: - return paddle.meshgrid(*args, **kwargs) - - def _roots(N: int) -> np.ndarray: - return np.polynomial.legendre.leggauss(N)[0] - - def _calculate_grid( - N: int, - integration_domains: paddle.Tensor, - ) -> Tuple[paddle.Tensor, paddle.Tensor, int]: - """Calculate grid points, widths and N per dim - - Args: - N (int): Number of points. - integration_domain (paddle.Tensor): Integration domain. - - Returns: - Tuple[paddle.Tensor, paddle.Tensor, int]: Grid points, grid widths and - Number of grid slices per dimension. - """ - # Create grid and assemble evaluation points - grid_1d = [] - _dim = integration_domains.shape[0] - n_per_dim = int(N ** (1.0 / _dim) + 1e-8) - - # Determine for each dimension grid points and mesh width - def _resize_roots( - integration_domain: Tuple[float, float], roots: np.ndarray - ): # scale from [-1,1] to [a,b] - a = integration_domain[0] - b = integration_domain[1] - return ((b - a) / 2) * roots + ((a + b) / 2) - - for dim in range(_dim): - grid_1d.append(_resize_roots(integration_domains[dim], _roots(n_per_dim))) - h = paddle.stack([grid_1d[dim][1] - grid_1d[dim][0] for dim in range(_dim)]) - - # Get grid points - points = _compatible_meshgrid(*grid_1d) - points = paddle.stack([mg.reshape([-1]) for mg in points], axis=1) - - return points, h, n_per_dim - - def _evaluate_integrand(fn, points, weights=None, fn_args=None) -> paddle.Tensor: - """Evaluate the integrand function at the passed points. - - Args: - fn (function): Integrand function. - points (paddle.Tensor): Integration points. - weights (paddle.Tensor, optional): Integration weights. Defaults to None. - fn_args (list or tuple, optional): Any arguments required by the function. Defaults to None. - - Returns: - paddle.Tensor: Integral result. - """ - if fn_args is None: - fn_args = () - - result = fn(points, *fn_args) - if not str(result.dtype).endswith(dtype): - result = result.astype(dtype) - - if result.shape[0] != points.shape[0]: - raise ValueError( - f"The passed function was given {points.shape[0]} points but only returned {result.shape[0]} value(s)." - f"Please ensure that your function is vectorized, i.e. can be called with multiple evaluation points at once. It should return a tensor " - f"where first dimension matches length of passed elements. " - ) - - if weights is not None: - if ( - len(result.shape) > 1 - ): # if the the integrand is multi-dimensional, we need to reshape/repeat weights so they can be broadcast in the *= - integrand_shape = result.shape[1:] - weights = paddle.repeat_interleave( - paddle.unsqueeze(weights, axis=1), np.prod(integrand_shape) - ).reshape((weights.shape[0], *(integrand_shape))) - result *= weights - - return result - - def _weights(N, dim): - """Return the weights, broadcast across the dimensions, generated from the polynomial of choice. - - Args: - N (int): Number of nodes. - dim (int): Number of dimensions. - - Returns: - paddle.Tensor: Integration weights. - """ - weights = paddle.to_tensor(np.polynomial.legendre.leggauss(N)[1], dtype=dtype) - return paddle.prod( - paddle.stack(_compatible_meshgrid(*([weights] * dim)), axis=0), - axis=0, - ).reshape([-1]) - - def _apply_composite_rule(cur_dim_areas, dim, hs, domain): - """Apply "composite" rule for gaussian integrals - - cur_dim_areas will contain the areas per dimension - """ - # We collapse dimension by dimension - for cur_dim in range(dim): - cur_dim_areas = ( - 0.5 - * (domain[cur_dim][1] - domain[cur_dim][0]) - * paddle.sum( - cur_dim_areas, axis=len(cur_dim_areas.shape) - 1, dtype=dtype - ) - ) - return cur_dim_areas - - @expand_func_values_and_squeeze_integral - def _calculate_result( - function_values: paddle.Tensor, - dim: int, - n_per_dim: int, - hs: paddle.Tensor, - integration_domains: paddle.Tensor, - ) -> paddle.Tensor: - """Apply the "composite rule" to calculate a result from the evaluated integrand. - - Args: - function_values (paddle.Tensor): Output of the integrand. - dim (int): Dimensionality. - n_per_dim (int): Number of grid slices per dimension. - hs (paddle.Tensor): Distances between grid slices for each dimension. - - Returns: - paddle.Tensor: Quadrature result. - """ - # Reshape the output to be [integrand_dim,N,N,...] points instead of [integrand_dim,dim*N] points - integrand_shape = function_values.shape[1:] - dim_shape = [n_per_dim] * dim - new_shape = [*integrand_shape, *dim_shape] - - perm = list(range(len(function_values.shape))) - if len(perm) >= 2: - perm.append(perm.pop(0)) - reshaped_function_values = paddle.transpose(function_values, perm) - reshaped_function_values = reshaped_function_values.reshape(new_shape) - - assert new_shape == list( - reshaped_function_values.shape - ), f"reshaping produced shape {reshaped_function_values.shape}, expected shape was {new_shape}" - - result = _apply_composite_rule( - reshaped_function_values, dim, hs, integration_domains - ) - return result - - assert dtype in [ - "float32", - "float64", - ], f"dtype must be either 'float32' or 'float64', but got {dtype}" - - neg = False - for i, (a, b) in enumerate(integration_domains): - if a > b: - neg = not neg - integration_domains[i] = [b, a] - - integration_domains = paddle.to_tensor( - integration_domains, - dtype=dtype, - ) - - if integration_domains.shape[0] != dim: - raise ValueError( - f"The number of integration domain({integration_domains.shape[0]}) " - f"must be equal to the given 'dim'({dim})." - ) - if integration_domains.shape[1] != 2: - raise ValueError( - f"integration_domain should be in format of [[a_1, b_1], [a_2, b_2], ..., " - f"[a_dim, b_dim]], but got each range of integration is {integration_domains[0]}" - ) - grid_points, hs, n_per_dim = _calculate_grid(N, integration_domains) - - function_values = _evaluate_integrand( - fn, grid_points, weights=_weights(n_per_dim, dim) - ) - - result = _calculate_result(function_values, dim, n_per_dim, hs, integration_domains) - return result if (not neg) else -result - - -def fractional_diff( - func: Callable, alpha: float, a: float, t: float, h: float, dtype="float64" -) -> paddle.Tensor: - r"""Compute fractional derivative of given function at point t with fractional order - alpha using [Caputo derivative of fractional](https://en.wikipedia.org/wiki/Fractional_calculus#Caputo_fractional_derivative). - - $$ - D_t^\alpha f(t)=\frac{1}{\Gamma(n-\alpha)} \int_0^t \frac{f^{(n)}(s)}{(t-s)^{\alpha+1-n}} d s . - $$ - - $$ - s.t. 0 \lt \alpha \lt 1 . - $$ - - Args: - func (Callable): Function to compute the fractional derivative of. - alpha (float): Fractional order. - t (float): Point to compute the fractional derivative at. - a (float): Start point of the fractional integral. - h (float): Step size for finite difference. - dtype (str, optional): Data dtype during computation. Defaults to "float64". - - Returns: - paddle.Tensor: Fractional derivative result of the function at t. - - Examples: - >>> from ppsci.experimental import fractional_diff - >>> import numpy as np - >>> # define f(x) = x^2 - >>> def f(x): - ... return x * x - >>> # compute 0.5-order fractional derivative of f(x) at t=1.0 with step size h=1e-6 - >>> res = fractional_diff(f, alpha=0.5, a=0, t=1.0, h=1e-6, dtype="float64") - >>> np.testing.assert_allclose(float(res), 1.503547, 1e-6) - """ - - if not (0 < alpha < 1): - raise NotImplementedError( - f"Given alpha should be in range (0, 1), but got {alpha}" - ) - - def _finite_derivative( - func: Callable, x: paddle.Tensor, dx: float - ) -> paddle.Tensor: - """Compute the finite difference of a function at x using centered difference. - - Args: - func (Callable): Function to compute the finite difference of. - x (paddle.Tensor): Point to compute the finite difference at. - dx (float): Delta to use for the finite difference. - - Returns: - paddle.Tensor: First-order Finite difference of the function at x. - """ - return (func(x + dx) - func(x - dx)) / (2 * dx) - - def int_func(s): - return _finite_derivative(func, s, dx=h) / (t - s) ** (alpha) - - result = ( - 1.0 / paddle.exp(paddle.lgamma(paddle.to_tensor(1.0 - alpha, dtype=dtype))) - ) * gaussian_integrate( - int_func, dim=1, N=2**10 + 1, integration_domains=[[a, t]], dtype=dtype - ) - return result - - -def trapezoid_integrate( - y: paddle.Tensor, - x: paddle.Tensor = None, - dx: float = None, - axis: int = -1, - mode: Literal["sum", "cumsum"] = "sum", -) -> paddle.Tensor: - """ - Integrate along the given axis using the composite trapezoidal rule. Use the sum method. - - Args: - y (paddle.Tensor): Input to be integrated. - x (paddle.Tensor, optional): The sample points corresponding to the input samples. its shape should be - (1) input.shape; (2) the input.shape[axis] if axis is not default. Defaults to None. - dx (float, optional): The sample points are assumed to be evenly spaced and it is the spacing between sample points. - If 'x' and 'dx' are both default, 'dx' is set to 1 by default. Defaults to None. - axis (int, optional): The axis along which to integrate. Defaults to -1. - mode (Literal["sum", "cumsum"], optional): Which type cumulative sum function used. Defaults to "sum". - - Returns: - paddle.Tensor: Integral result. If dim of input is N, return is N-1 dim. - - Examples: - >>> import paddle - >>> import ppsci - >>> y = paddle.to_tensor([[0, 1, 2], [3, 4, 5]], dtype="float32") - >>> res = ppsci.experimental.trapezoid_integrate(y) - >>> print(res) - Tensor(shape=[2], dtype=float32, place=Place(gpu:0), stop_gradient=True, - [2., 8.]) - >>> res = ppsci.experimental.trapezoid_integrate(y, mode="cumsum") - >>> print(res) - Tensor(shape=[2, 2], dtype=float32, place=Place(gpu:0), stop_gradient=True, - [[0.50000000, 2. ], - [3.50000000, 8. ]]) - >>> res = ppsci.experimental.trapezoid_integrate( - ... y, x=paddle.to_tensor([[0, 1, 2], [3, 4, 5]], dtype="float32") - ... ) - >>> print(res) - Tensor(shape=[2], dtype=float32, place=Place(gpu:0), stop_gradient=True, - [2., 8.]) - >>> res = ppsci.experimental.trapezoid_integrate( - ... y, x=paddle.to_tensor([0, 1], dtype="float32"), axis=0 - ... ) - >>> print(res) - Tensor(shape=[3], dtype=float32, place=Place(gpu:0), stop_gradient=True, - [1.50000000, 2.50000000, 3.50000000]) - >>> res = ppsci.experimental.trapezoid_integrate( - ... y, x=paddle.to_tensor([0, 1, 2], dtype="float32"), axis=1 - ... ) - >>> print(res) - Tensor(shape=[2], dtype=float32, place=Place(gpu:0), stop_gradient=True, - [2., 8.]) - >>> res = ppsci.experimental.trapezoid_integrate(y, dx=2) - >>> print(res) - Tensor(shape=[2], dtype=float32, place=Place(gpu:0), stop_gradient=True, - [4. , 16.]) - """ - if mode == "sum": - return paddle.trapezoid(y, x, dx, axis) - elif mode == "cumsum": - return paddle.cumulative_trapezoid(y, x, dx, axis) - else: - raise ValueError(f'mode should be "sum" or "cumsum", but got {mode}') - - -def montecarlo_integrate( - fn: Callable, - dim: int, - N: int = 1000, - integration_domain: Union[List[List[float]], paddle.Tensor] = None, - seed: int = None, -) -> paddle.Tensor: - """Integrates the passed function on the passed domain using vanilla Monte - Carlo Integration. - - Args: - fn (Callable): The function to integrate over. - dim (int): Dimensionality of the function's domain over which to - integrate. - N (Optional[int]): Number of sample points to use for the integration. - Defaults to 1000. - integration_domain (Union[List[List[float]], paddle.Tensor]): Integration - domain, e.g. [[-1,1],[0,1]]. Defaults to [-1,1]^dim. - seed (Optional[int]): Random number generation seed to the sampling - point creation, only set if provided. Defaults to None. - - Raises: - ValueError: If len(integration_domain) != dim - - Returns: - paddle.Tensor: Integral result. - - Examples: - >>> import paddle - >>> import ppsci - - >>> _ = paddle.seed(1024) - >>> # The function we want to integrate, in this example - >>> # f(x0,x1) = sin(x0) + e^x1 for x0=[0,1] and x1=[-1,1] - >>> # Note that the function needs to support multiple evaluations at once (first - >>> # dimension of x here) - >>> # Expected result here is ~3.2698 - >>> def some_function(x): - ... return paddle.sin(x[:, 0]) + paddle.exp(x[:, 1]) - - >>> # Compute the function integral by sampling 10000 points over domain - >>> integral_value = ppsci.experimental.montecarlo_integrate( - ... some_function, - ... dim=2, - ... N=10000, - ... integration_domain=[[0, 1], [-1, 1]], - ... ) - - >>> print(integral_value) - Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, - 3.25152588) - """ - - @expand_func_values_and_squeeze_integral - def calculate_result(function_values, integration_domain): - """Calculate an integral result from the function evaluations - - Args: - function_values (paddle.Tensor): Output of the integrand - integration_domain (paddle.Tensor): Integration domain - - Returns: - Quadrature result - """ - scales = integration_domain[:, 1] - integration_domain[:, 0] - volume = paddle.prod(scales) - - # Integral = V / N * sum(func values) - N = function_values.shape[0] - integral = volume * paddle.sum(function_values, axis=0) / N - return integral - - def calculate_sample_points( - N: int, integration_domain: paddle.Tensor, seed: Optional[int] = None - ): - """Calculate random points for the integrand evaluation. - - Args: - N (int): Number of points - integration_domain (paddle.Tensor): Integration domain. - seed (int, optional): Random number generation seed for the sampling point creation, only set if provided. Defaults to None. - Returns: - Sample points. - """ - dim = integration_domain.shape[0] - domain_starts = integration_domain[:, 0] - domain_sizes = integration_domain[:, 1] - domain_starts - # Scale and translate random numbers via broadcasting - return ( - paddle.uniform( - shape=[N, dim], - dtype=domain_sizes.dtype, - min=0.0, - max=1.0, - seed=seed or 0, - ) - * domain_sizes - + domain_starts - ) - - if dim is not None: - if dim < 1: - raise ValueError("Dimension needs to be 1 or larger.") - if N is not None: - if N < 1 or type(N) is not int: - raise ValueError("N has to be a positive integer.") - - integration_domain = _setup_integration_domain(dim, integration_domain) - sample_points = calculate_sample_points(N, integration_domain, seed) - function_values, _ = _evaluate_integrand(fn, sample_points) - return calculate_result(function_values, integration_domain) - - -def _setup_integration_domain( - dim: int, integration_domain: Union[List[List[float]], paddle.Tensor] -) -> paddle.Tensor: - """Sets up the integration domain if unspecified by the user. - Args: - dim (int): Dimensionality of the integration domain. - integration_domain (List or Tensor): Integration domain, e.g. [[-1,1],[0,1]]. Defaults to [-1,1]^dim. - - Returns: - Integration domain. - """ - # If no integration_domain is specified, create [-1,1]^d bounds - if integration_domain is None: - integration_domain = [[-1.0, 1.0]] * dim - - integration_domain = [[float(b) for b in bounds] for bounds in integration_domain] - - integration_domain = paddle.to_tensor(integration_domain) - - if tuple(integration_domain.shape) != (dim, 2): - raise ValueError( - "The integration domain has an unexpected shape. " - f"Expected {(dim, 2)}, got {integration_domain.shape}" - ) - return integration_domain - - -def _evaluate_integrand(fn, points, weights=None, args=None): - """Evaluate the integrand function at the passed points. - - Args: - fn (Callable): Integrand function. - points (paddle.Tensor): Integration points. - weights (Optional[paddle.Tensor]): Integration weights. Defaults to None. - args (Optional[List, Tuple]): Any arguments required by the function. Defaults to None. - - Returns: - padlde.Tensor: Integrand function output. - int: Number of evaluated points. - """ - num_points = points.shape[0] - - if args is None: - args = () - - result = fn(points, *args) - num_results = result.shape[0] - if num_results != num_points: - raise ValueError( - f"The passed function was given {num_points} points but only returned {num_results} value(s)." - f"Please ensure that your function is vectorized, i.e. can be called with multiple evaluation points at once. It should return a tensor " - f"where first dimension matches length of passed elements. " - ) - - if weights is not None: - if ( - len(result.shape) > 1 - ): # if the the integrand is multi-dimensional, we need to reshape/repeat weights so they can be broadcast in the *= - integrand_shape = paddle.to_tensor(result.shape[1:]) - weights = paddle.tile( - paddle.unsqueeze(weights, axis=1), paddle.prod(integrand_shape) - ).reshape((weights.shape[0], *(integrand_shape))) - result *= weights - - return result, num_points diff --git a/examples/smc_reac/ppsci/externals/__init__.py b/examples/smc_reac/ppsci/externals/__init__.py deleted file mode 100644 index 67e62f29f3..0000000000 --- a/examples/smc_reac/ppsci/externals/__init__.py +++ /dev/null @@ -1,20 +0,0 @@ -"""External Development Packages""" - -import importlib.util - -EXTERNAL_PACKAGES_LIST = [ - "deepali", - "neuraloperator", - "open3d", - "paddle_harmonics", - "paddle_scatter", - "paddle_sparse", - "tensorly", - "warp", -] - -__all__ = [] -for package_name in EXTERNAL_PACKAGES_LIST: - if importlib.util.find_spec(package_name): - globals()[package_name] = __import__(package_name) - __all__.append(package_name) diff --git a/examples/smc_reac/ppsci/geometry/__init__.py b/examples/smc_reac/ppsci/geometry/__init__.py deleted file mode 100644 index 30b4ad0859..0000000000 --- a/examples/smc_reac/ppsci/geometry/__init__.py +++ /dev/null @@ -1,83 +0,0 @@ -# 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. - -import copy - -from ppsci.geometry.geometry import Geometry -from ppsci.geometry.geometry_1d import Interval -from ppsci.geometry.geometry_2d import Disk -from ppsci.geometry.geometry_2d import Polygon -from ppsci.geometry.geometry_2d import Rectangle -from ppsci.geometry.geometry_2d import Triangle -from ppsci.geometry.geometry_3d import Cuboid -from ppsci.geometry.geometry_3d import Sphere -from ppsci.geometry.geometry_nd import Hypercube -from ppsci.geometry.geometry_nd import Hypersphere -from ppsci.geometry.mesh import Mesh -from ppsci.geometry.mesh import SDFMesh -from ppsci.geometry.pointcloud import PointCloud -from ppsci.geometry.timedomain import TimeDomain -from ppsci.geometry.timedomain import TimeXGeometry -from ppsci.utils import logger -from ppsci.utils import misc - -__all__ = [ - "build_geometry", - "Cuboid", - "Disk", - "Geometry", - "Hypercube", - "Hypersphere", - "Interval", - "Mesh", - "SDFMesh", - "Polygon", - "Rectangle", - "Sphere", - "TimeDomain", - "TimeXGeometry", - "Triangle", - "PointCloud", -] - - -def build_geometry(cfg): - """Build geometry(ies) - - Args: - cfg (List[DictConfig]): Geometry config list. - - Returns: - Dict[str, Geometry]: Geometry(ies) in dict. - """ - if cfg is None: - return None - cfg = copy.deepcopy(cfg) - - geom_dict = misc.PrettyOrderedDict() - for _item in cfg: - geom_cls = next(iter(_item.keys())) - geom_cfg = _item[geom_cls] - geom_name = geom_cfg.pop("name", geom_cls) - if geom_cls == "TimeXGeometry": - time_cfg = geom_cfg.pop("TimeDomain") - geom_cls = next(iter(geom_cfg.keys())) - geom_dict[geom_name] = TimeXGeometry( - TimeDomain(**time_cfg), eval(geom_cls)(**geom_cfg[geom_cls]) - ) - else: - geom_dict[geom_name] = eval(geom_cls)(**geom_cfg) - - logger.debug(str(geom_dict[geom_name])) - return geom_dict diff --git a/examples/smc_reac/ppsci/geometry/csg.py b/examples/smc_reac/ppsci/geometry/csg.py deleted file mode 100644 index 87534bedd6..0000000000 --- a/examples/smc_reac/ppsci/geometry/csg.py +++ /dev/null @@ -1,337 +0,0 @@ -# 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. - -""" -Code below is heavily based on [https://github.com/lululxvi/deepxde](https://github.com/lululxvi/deepxde) -""" - -from __future__ import annotations - -import numpy as np -import paddle - -from ppsci.geometry import geometry - - -class CSGUnion(geometry.Geometry): - """Construct an object by CSG Union(except for Mesh).""" - - def __init__(self, geom1, geom2): - if geom1.ndim != geom2.ndim: - raise ValueError( - f"{geom1}.ndim({geom1.ndim}) should be equal to " - f"{geom2}.ndim({geom1.ndim})" - ) - super().__init__( - geom1.ndim, - ( - np.minimum(geom1.bbox[0], geom2.bbox[0]), - np.maximum(geom1.bbox[1], geom2.bbox[1]), - ), - geom1.diam + geom2.diam, - ) - self.geom1 = geom1 - self.geom2 = geom2 - - def is_inside(self, x): - return np.logical_or(self.geom1.is_inside(x), self.geom2.is_inside(x)) - - def on_boundary(self, x): - return np.logical_or( - np.logical_and(self.geom1.on_boundary(x), ~self.geom2.is_inside(x)), - np.logical_and(self.geom2.on_boundary(x), ~self.geom1.is_inside(x)), - ) - - def boundary_normal(self, x): - return np.logical_and(self.geom1.on_boundary(x), ~self.geom2.is_inside(x))[ - :, np.newaxis - ] * self.geom1.boundary_normal(x) + np.logical_and( - self.geom2.on_boundary(x), ~self.geom1.is_inside(x) - )[ - :, np.newaxis - ] * self.geom2.boundary_normal( - x - ) - - def random_points(self, n, random="pseudo"): - x = np.empty(shape=(n, self.ndim), dtype=paddle.get_default_dtype()) - _size = 0 - while _size < n: - points = ( - np.random.rand(n, self.ndim) * (self.bbox[1] - self.bbox[0]) - + self.bbox[0] - ) - points = points[self.is_inside(points)] - - if len(points) > n - _size: - points = points[: n - _size] - x[_size : _size + len(points)] = points - _size += len(points) - return x - - def random_boundary_points(self, n, random="pseudo"): - x = np.empty(shape=(n, self.ndim), dtype=paddle.get_default_dtype()) - _size = 0 - while _size < n: - geom1_boundary_points = self.geom1.random_boundary_points(n, random=random) - geom1_boundary_points = geom1_boundary_points[ - ~self.geom2.is_inside(geom1_boundary_points) - ] - - geom2_boundary_points = self.geom2.random_boundary_points(n, random=random) - geom2_boundary_points = geom2_boundary_points[ - ~self.geom1.is_inside(geom2_boundary_points) - ] - - points = np.concatenate((geom1_boundary_points, geom2_boundary_points)) - points = np.random.permutation(points) - - if len(points) > n - _size: - points = points[: n - _size] - x[_size : _size + len(points)] = points - _size += len(points) - return x - - def periodic_point(self, x, component): - x = np.copy(x) - on_boundary_geom1 = np.logical_and( - self.geom1.on_boundary(x), ~self.geom2.is_inside(x) - ) - x[on_boundary_geom1] = self.geom1.periodic_point(x, component)[ - on_boundary_geom1 - ] - on_boundary_geom2 = np.logical_and( - self.geom2.on_boundary(x), ~self.geom1.is_inside(x) - ) - x[on_boundary_geom2] = self.geom2.periodic_point(x, component)[ - on_boundary_geom2 - ] - return x - - def sdf_func(self, points: np.ndarray) -> np.ndarray: - """Compute signed distance field of CSG union of two geometries. - ref: https://iquilezles.org/articles/distfunctions/ - - Args: - points (np.ndarray): The coordinate points used to calculate the SDF - value, the shape is [N, D]. - - Returns: - np.ndarray: SDF values of input points without squared, the shape is [N, 1]. - """ - sdf1 = self.geom1.sdf_func(points) - sdf2 = self.geom2.sdf_func(points) - return np.minimum(sdf1, sdf2) - - -class CSGDifference(geometry.Geometry): - """Construct an object by CSG Difference.""" - - def __init__(self, geom1, geom2): - if geom1.ndim != geom2.ndim: - raise ValueError( - f"{geom1}.ndim({geom1.ndim}) should be equal to " - f"{geom2}.ndim({geom1.ndim})." - ) - super().__init__(geom1.ndim, geom1.bbox, geom1.diam) - self.geom1 = geom1 - self.geom2 = geom2 - - def is_inside(self, x): - return np.logical_and(self.geom1.is_inside(x), ~self.geom2.is_inside(x)) - - def on_boundary(self, x): - return np.logical_or( - np.logical_and(self.geom1.on_boundary(x), ~self.geom2.is_inside(x)), - np.logical_and(self.geom1.is_inside(x), self.geom2.on_boundary(x)), - ) - - def boundary_normal(self, x): - return np.logical_and(self.geom1.on_boundary(x), ~self.geom2.is_inside(x))[ - :, np.newaxis - ] * self.geom1.boundary_normal(x) + np.logical_and( - self.geom1.is_inside(x), self.geom2.on_boundary(x) - )[ - :, np.newaxis - ] * -self.geom2.boundary_normal( - x - ) - - def random_points(self, n, random="pseudo"): - x = np.empty(shape=(n, self.ndim), dtype=paddle.get_default_dtype()) - _size = 0 - while _size < n: - tmp = self.geom1.random_points(n, random=random) - tmp = tmp[~self.geom2.is_inside(tmp)] - - if len(tmp) > n - _size: - tmp = tmp[: n - _size] - x[_size : _size + len(tmp)] = tmp - _size += len(tmp) - return x - - def random_boundary_points(self, n, random="pseudo"): - x = np.empty(shape=(n, self.ndim), dtype=paddle.get_default_dtype()) - _size = 0 - while _size < n: - geom1_boundary_points = self.geom1.random_boundary_points(n, random=random) - geom1_boundary_points = geom1_boundary_points[ - ~self.geom2.is_inside(geom1_boundary_points) - ] - - geom2_boundary_points = self.geom2.random_boundary_points(n, random=random) - geom2_boundary_points = geom2_boundary_points[ - self.geom1.is_inside(geom2_boundary_points) - ] - - points = np.concatenate((geom1_boundary_points, geom2_boundary_points)) - points = np.random.permutation(points) - - if len(points) > n - _size: - points = points[: n - _size] - x[_size : _size + len(points)] = points - _size += len(points) - return x - - def periodic_point(self, x, component): - x = np.copy(x) - on_boundary_geom1 = np.logical_and( - self.geom1.on_boundary(x), ~self.geom2.is_inside(x) - ) - x[on_boundary_geom1] = self.geom1.periodic_point(x, component)[ - on_boundary_geom1 - ] - return x - - def sdf_func(self, points: np.ndarray) -> np.ndarray: - """Compute signed distance field of CSG difference of two geometries. - - Args: - points (np.ndarray): The coordinate points used to calculate the SDF - value, the shape is [N, D]. - - Returns: - np.ndarray: SDF values of input points without squared, the shape is [N, 1]. - """ - sdf1 = self.geom1.sdf_func(points) - sdf2 = self.geom2.sdf_func(points) - return np.maximum(sdf1, -sdf2) - - -class CSGIntersection(geometry.Geometry): - """Construct an object by CSG Intersection.""" - - def __init__(self, geom1, geom2): - if geom1.ndim != geom2.ndim: - raise ValueError( - f"{geom1}.ndim({geom1.ndim}) should be equal to " - f"{geom2}.ndim({geom1.ndim})" - ) - super().__init__( - geom1.ndim, - ( - np.maximum(geom1.bbox[0], geom2.bbox[0]), - np.minimum(geom1.bbox[1], geom2.bbox[1]), - ), - min(geom1.diam, geom2.diam), - ) - self.geom1 = geom1 - self.geom2 = geom2 - - def is_inside(self, x): - return np.logical_and(self.geom1.is_inside(x), self.geom2.is_inside(x)) - - def on_boundary(self, x): - return np.logical_or( - np.logical_and(self.geom1.on_boundary(x), self.geom2.is_inside(x)), - np.logical_and(self.geom1.is_inside(x), self.geom2.on_boundary(x)), - ) - - def boundary_normal(self, x): - return np.logical_and(self.geom1.on_boundary(x), self.geom2.is_inside(x))[ - :, np.newaxis - ] * self.geom1.boundary_normal(x) + np.logical_and( - self.geom1.is_inside(x), self.geom2.on_boundary(x) - )[ - :, np.newaxis - ] * self.geom2.boundary_normal( - x - ) - - def random_points(self, n, random="pseudo"): - x = np.empty(shape=(n, self.ndim), dtype=paddle.get_default_dtype()) - _size = 0 - while _size < n: - points = self.geom1.random_points(n, random=random) - points = points[self.geom2.is_inside(points)] - - if len(points) > n - _size: - points = points[: n - _size] - x[_size : _size + len(points)] = points - _size += len(points) - return x - - def random_boundary_points(self, n, random="pseudo"): - x = np.empty(shape=(n, self.ndim), dtype=paddle.get_default_dtype()) - _size = 0 - while _size < n: - geom1_boundary_points = self.geom1.random_boundary_points(n, random=random) - geom1_boundary_points = geom1_boundary_points[ - self.geom2.is_inside(geom1_boundary_points) - ] - - geom2_boundary_points = self.geom2.random_boundary_points(n, random=random) - geom2_boundary_points = geom2_boundary_points[ - self.geom1.is_inside(geom2_boundary_points) - ] - - points = np.concatenate((geom1_boundary_points, geom2_boundary_points)) - points = np.random.permutation(points) - - if len(points) > n - _size: - points = points[: n - _size] - x[_size : _size + len(points)] = points - _size += len(points) - return x - - def periodic_point(self, x, component): - x = np.copy(x) - on_boundary_geom1 = np.logical_and( - self.geom1.on_boundary(x), self.geom2.is_inside(x) - ) - x[on_boundary_geom1] = self.geom1.periodic_point(x, component)[ - on_boundary_geom1 - ] - on_boundary_geom2 = np.logical_and( - self.geom2.on_boundary(x), self.geom1.is_inside(x) - ) - x[on_boundary_geom2] = self.geom2.periodic_point(x, component)[ - on_boundary_geom2 - ] - return x - - def sdf_func(self, points: np.ndarray) -> np.ndarray: - """Compute signed distance field of CSG intersection of two geometries. - ref: https://iquilezles.org/articles/distfunctions/ - - Args: - points (np.ndarray): The coordinate points used to calculate the SDF - value the shape is [N, D]. - - Returns: - np.ndarray: SDF values of input points without squared, the shape is [N, 1]. - """ - sdf1 = self.geom1.sdf_func(points) - sdf2 = self.geom2.sdf_func(points) - return np.maximum(sdf1, sdf2) diff --git a/examples/smc_reac/ppsci/geometry/geometry.py b/examples/smc_reac/ppsci/geometry/geometry.py deleted file mode 100644 index 5bda675414..0000000000 --- a/examples/smc_reac/ppsci/geometry/geometry.py +++ /dev/null @@ -1,696 +0,0 @@ -# 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. - -""" -Code below is heavily based on [https://github.com/lululxvi/deepxde](https://github.com/lululxvi/deepxde) -""" -from __future__ import annotations - -import abc -from typing import Callable -from typing import Dict -from typing import Optional -from typing import Tuple - -import numpy as np -import paddle -from typing_extensions import Literal - -from ppsci.utils import logger -from ppsci.utils import misc - - -class Geometry: - """Base class for geometry. - - Args: - ndim (int): Number of geometry dimension. - bbox (Tuple[np.ndarray, np.ndarray]): Bounding box of upper and lower. - diam (float): Diameter of geometry. - """ - - def __init__(self, ndim: int, bbox: Tuple[np.ndarray, np.ndarray], diam: float): - self.ndim = ndim - self.bbox = bbox - self.diam = min(diam, np.linalg.norm(bbox[1] - bbox[0])) - - @property - def dim_keys(self): - return ("x", "y", "z")[: self.ndim] - - @abc.abstractmethod - def is_inside(self, x: np.ndarray) -> np.ndarray: - """Returns a boolean array where x is inside the geometry. - - Args: - x (np.ndarray): Points to check if inside the geometry. The shape is [N, D], - where D is the number of dimension of geometry. - - Returns: - np.ndarray: Boolean array where x is inside the geometry. The shape is [N]. - - Examples: - >>> import numpy as np - >>> import ppsci - >>> interval = ppsci.geometry.Interval(0, 1) - >>> x = np.array([[0], [0.5], [1.5]]) - >>> interval.is_inside(x) - array([ True, True, False]) - >>> rectangle = ppsci.geometry.Rectangle((0, 0), (1, 1)) - >>> x = np.array([[0.0, 0.0], [0.5, 0.5], [1.5, 1.5]]) - >>> rectangle.is_inside(x) - array([ True, True, False]) - >>> cuboid = ppsci.geometry.Cuboid((0, 0, 0), (1, 1, 1)) - >>> x = np.array([[0, 0, 0], [0.5, 0.5, 0.5], [1.5, 1.5, 1.5]]) - >>> cuboid.is_inside(x) - array([ True, True, False]) - """ - - @abc.abstractmethod - def on_boundary(self, x: np.ndarray) -> np.ndarray: - """Returns a boolean array where x is on geometry boundary. - - Args: - x (np.ndarray): Points to check if on the geometry boundary. The shape is [N, D], - where D is the number of dimension of geometry. - - Returns: - np.ndarray: Boolean array where x is on the geometry boundary. The shape is [N]. - - Examples: - >>> import numpy as np - >>> import ppsci - >>> interval = ppsci.geometry.Interval(0, 1) - >>> x = np.array([[0], [0.5], [1.5]]) - >>> interval.on_boundary(x) - array([ True, False, False]) - >>> rectangle = ppsci.geometry.Rectangle((0, 0), (1, 1)) - >>> x = np.array([[0, 0], [0.5, 0.5], [1, 1.5]]) - >>> rectangle.on_boundary(x) - array([ True, False, False]) - >>> cuboid = ppsci.geometry.Cuboid((0, 0, 0), (1, 1, 1)) - >>> x = np.array([[0, 0, 0], [0.5, 0.5, 0.5], [1, 1, 1.5]]) - >>> cuboid.on_boundary(x) - array([ True, False, False]) - """ - - def boundary_normal(self, x): - """Compute the unit normal at x.""" - raise NotImplementedError(f"{self}.boundary_normal is not implemented") - - def uniform_points(self, n: int, boundary: bool = True) -> np.ndarray: - """Compute the equi-spaced points in the geometry. - - Args: - n (int): Number of points. - boundary (bool): Include boundary points. Defaults to True. - - Returns: - np.ndarray: Random points in the geometry. The shape is [N, D]. - """ - logger.warning( - f"{self}.uniform_points not implemented. " f"Use random_points instead." - ) - return self.random_points(n) - - def sample_interior( - self, - n: int, - random: Literal["pseudo", "Halton", "LHS"] = "pseudo", - criteria: Optional[Callable[..., np.ndarray]] = None, - evenly: bool = False, - compute_sdf_derivatives: bool = False, - ) -> Dict[str, np.ndarray]: - """Sample random points in the geometry and return those meet criteria. - - Args: - n (int): Number of points. - random (Literal["pseudo", "Halton", "LHS"]): Random method. Defaults to "pseudo". - pseudo: Pseudo random. - Halton: Halton sequence. - LHS: Latin Hypercube Sampling. - criteria (Optional[Callable[..., np.ndarray]]): Criteria function. Given - coords from different dimension and return a boolean array with shape [n,]. - Defaults to None. - evenly (bool): Evenly sample points. Defaults to False. - compute_sdf_derivatives (bool): Compute SDF derivatives. Defaults to False. - - Returns: - Dict[str, np.ndarray]: Random points in the geometry. The shape is [N, D]. - their signed distance function. The shape is [N, 1]. - their derivatives of SDF(optional). The shape is [N, D]. - - Examples: - >>> import numpy as np - >>> import ppsci - >>> np.random.seed(42) - >>> interval = ppsci.geometry.Interval(0, 1) - >>> interval.sample_interior(2) - {'x': array([[0.37454012], - [0.9507143 ]], dtype=float32), 'sdf': array([[0.37454012], - [0.04928571]], dtype=float32)} - >>> rectangle = ppsci.geometry.Rectangle((0, 0), (1, 1)) - >>> rectangle.sample_interior(2, "pseudo", None, False, True) - {'x': array([[0.7319939 ], - [0.15601864]], dtype=float32), 'y': array([[0.5986585 ], - [0.15599452]], dtype=float32), 'sdf': array([[0.2680061 ], - [0.15599453]], dtype=float32), 'sdf__x': array([[-1.0001659 ], - [ 0.25868416]], dtype=float32), 'sdf__y': array([[-0. ], - [ 0.74118376]], dtype=float32)} - >>> cuboid = ppsci.geometry.Cuboid((0, 0, 0), (1, 1, 1)) - >>> cuboid.sample_interior(2, "pseudo", None, True, True) - {'x': array([[0.], - [0.]], dtype=float32), 'y': array([[0.], - [0.]], dtype=float32), 'z': array([[0.], - [1.]], dtype=float32), 'sdf': array([[0.], - [0.]], dtype=float32), 'sdf__x': array([[0.50008297], - [0.50008297]], dtype=float32), 'sdf__y': array([[0.50008297], - [0.50008297]], dtype=float32), 'sdf__z': array([[ 0.50008297], - [-0.49948692]], dtype=float32)} - """ - x = np.empty(shape=(n, self.ndim), dtype=paddle.get_default_dtype()) - _size, _ntry, _nsuc = 0, 0, 0 - while _size < n: - if evenly: - points = self.uniform_points(n) - else: - if misc.typename(self) == "TimeXGeometry": - points = self.random_points(n, random, criteria) - else: - points = self.random_points(n, random) - - if criteria is not None: - criteria_mask = criteria(*np.split(points, self.ndim, axis=1)).flatten() - points = points[criteria_mask] - - if len(points) > n - _size: - points = points[: n - _size] - x[_size : _size + len(points)] = points - - _size += len(points) - _ntry += 1 - if len(points) > 0: - _nsuc += 1 - - if _ntry >= 1000 and _nsuc == 0: - raise ValueError( - "Sample interior points failed, " - "please check correctness of geometry and given criteria." - ) - - # if sdf_func added, return x_dict and sdf_dict, else, only return the x_dict - if hasattr(self, "sdf_func"): - sdf = -self.sdf_func(x) - sdf_dict = misc.convert_to_dict(sdf, ("sdf",)) - sdf_derives_dict = {} - if compute_sdf_derivatives: - sdf_derives = -self.sdf_derivatives(x) - sdf_derives_dict = misc.convert_to_dict( - sdf_derives, tuple(f"sdf__{key}" for key in self.dim_keys) - ) - else: - sdf_dict = {} - sdf_derives_dict = {} - x_dict = misc.convert_to_dict(x, self.dim_keys) - - return {**x_dict, **sdf_dict, **sdf_derives_dict} - - def sample_boundary( - self, - n: int, - random: Literal["pseudo", "Halton", "LHS"] = "pseudo", - criteria: Optional[Callable[..., np.ndarray]] = None, - evenly: bool = False, - ) -> Dict[str, np.ndarray]: - """Compute the random points in the geometry and return those meet criteria. - - Args: - n (int): Number of points. - random (Literal["pseudo", "Halton", "LHS"]): Random method. Defaults to "pseudo". - pseudo: Pseudo random. - Halton: Halton sequence. - LHS: Latin Hypercube Sampling. - criteria (Optional[Callable[..., np.ndarray]]): Criteria function. Given - coords from different dimension and return a boolean array with shape [n,]. - Defaults to None. - evenly (bool): Evenly sample points. Defaults to False. - - Returns: - Dict[str, np.ndarray]: Random points in the geometry. The shape is [N, D]. - their normal vectors. The shape is [N, D]. - their area. The shape is [N, 1].(only if the geometry is a mesh) - - Examples: - >>> import numpy as np - >>> import ppsci - >>> np.random.seed(42) - >>> interval = ppsci.geometry.Interval(0, 1) - >>> interval.sample_boundary(2) - {'x': array([[0.], - [1.]], dtype=float32), 'normal_x': array([[-1.], - [ 1.]], dtype=float32)} - >>> rectangle = ppsci.geometry.Rectangle((0, 0), (1, 1)) - >>> rectangle.sample_boundary(2) - {'x': array([[1.], - [0.]], dtype=float32), 'y': array([[0.49816048], - [0.19714284]], dtype=float32), 'normal_x': array([[ 1.], - [-1.]], dtype=float32), 'normal_y': array([[0.], - [0.]], dtype=float32)} - >>> cuboid = ppsci.geometry.Cuboid((0, 0, 0), (1, 1, 1)) - >>> cuboid.sample_boundary(2) - {'x': array([[0.83244264], - [0.18182497]], dtype=float32), 'y': array([[0.21233912], - [0.1834045 ]], dtype=float32), 'z': array([[0.], - [1.]], dtype=float32), 'normal_x': array([[0.], - [0.]], dtype=float32), 'normal_y': array([[0.], - [0.]], dtype=float32), 'normal_z': array([[-1.], - [ 1.]], dtype=float32)} - """ - x = np.empty(shape=(n, self.ndim), dtype=paddle.get_default_dtype()) - _size, _ntry, _nsuc = 0, 0, 0 - while _size < n: - if evenly: - if ( - misc.typename(self) == "TimeXGeometry" - and misc.typename(self.geometry) == "Mesh" - ): - points, normal, area = self.uniform_boundary_points(n) - else: - points = self.uniform_boundary_points(n) - else: - if ( - misc.typename(self) == "TimeXGeometry" - and misc.typename(self.geometry) == "Mesh" - ): - points, normal, area = self.random_boundary_points(n, random) - else: - if misc.typename(self) == "TimeXGeometry": - points = self.random_boundary_points(n, random, criteria) - else: - points = self.random_boundary_points(n, random) - - if criteria is not None: - criteria_mask = criteria(*np.split(points, self.ndim, axis=1)).flatten() - points = points[criteria_mask] - - if len(points) > n - _size: - points = points[: n - _size] - x[_size : _size + len(points)] = points - - _size += len(points) - _ntry += 1 - if len(points) > 0: - _nsuc += 1 - - if _ntry >= 10000 and _nsuc == 0: - raise ValueError( - "Sample boundary points failed, " - "please check correctness of geometry and given criteria." - ) - - if not ( - misc.typename(self) == "TimeXGeometry" - and misc.typename(self.geometry) == "Mesh" - ): - normal = self.boundary_normal(x) - - normal_dict = misc.convert_to_dict( - normal[:, 1:] if "t" in self.dim_keys else normal, - [f"normal_{key}" for key in self.dim_keys if key != "t"], - ) - x_dict = misc.convert_to_dict(x, self.dim_keys) - if ( - misc.typename(self) == "TimeXGeometry" - and misc.typename(self.geometry) == "Mesh" - ): - area_dict = misc.convert_to_dict(area[:, 1:], ["area"]) - return {**x_dict, **normal_dict, **area_dict} - - return {**x_dict, **normal_dict} - - @abc.abstractmethod - def random_points( - self, n: int, random: Literal["pseudo", "Halton", "LHS"] = "pseudo" - ) -> np.ndarray: - """Compute the random points in the geometry. - - Args: - n (int): Number of points. - random (Literal["pseudo", "Halton", "LHS"]): Random method. Defaults to "pseudo". - pseudo: Pseudo random. - Halton: Halton sequence. - LHS: Latin Hypercube Sampling. - - Returns: - np.ndarray: Random points in the geometry. The shape is [N, D]. - - Examples: - >>> import numpy as np - >>> import ppsci - >>> np.random.seed(42) - >>> interval = ppsci.geometry.Interval(0, 1) - >>> interval.random_points(2) - array([[0.37454012], - [0.9507143 ]], dtype=float32) - >>> rectangle = ppsci.geometry.Rectangle((0, 0), (1, 1)) - >>> rectangle.random_points(2) - array([[0.7319939 , 0.5986585 ], - [0.15601864, 0.15599452]], dtype=float32) - >>> cuboid = ppsci.geometry.Cuboid((0, 0, 0), (1, 1, 1)) - >>> cuboid.random_points(2) - array([[0.05808361, 0.8661761 , 0.601115 ], - [0.7080726 , 0.02058449, 0.96990985]], dtype=float32) - """ - - def uniform_boundary_points(self, n: int) -> np.ndarray: - """Compute the equi-spaced points on the boundary(not implemented). - - Args: - n (int): Number of points. - - Returns: - np.ndarray: Random points on the boundary. The shape is [N, D]. - """ - logger.warning( - f"{self}.uniform_boundary_points not implemented. " - f"Use random_boundary_points instead." - ) - return self.random_boundary_points(n) - - @abc.abstractmethod - def random_boundary_points( - self, n: int, random: Literal["pseudo", "Halton", "LHS"] = "pseudo" - ) -> np.ndarray: - """Compute the random points on the boundary. - - Args: - n (int): Number of points. - random (Literal["pseudo", "Halton", "LHS"]): Random method. Defaults to "pseudo". - pseudo: Pseudo random. - Halton: Halton sequence. - LHS: Latin Hypercube Sampling. - - Returns: - np.ndarray: Random points on the boundary. The shape is [N, D]. - - Examples: - >>> import numpy as np - >>> import ppsci - >>> np.random.seed(42) - >>> interval = ppsci.geometry.Interval(0, 1) - >>> interval.random_boundary_points(2) - array([[0.], - [1.]], dtype=float32) - >>> rectangle = ppsci.geometry.Rectangle((0, 0), (1, 1)) - >>> rectangle.random_boundary_points(2) - array([[1. , 0.49816048], - [0. , 0.19714284]], dtype=float32) - >>> cuboid = ppsci.geometry.Cuboid((0, 0, 0), (1, 1, 1)) - >>> cuboid.random_boundary_points(2) - array([[0.83244264, 0.21233912, 0. ], - [0.18182497, 0.1834045 , 1. ]], dtype=float32) - """ - - def periodic_point(self, x: np.ndarray, component: int): - """Compute the periodic image of x(not implemented).""" - raise NotImplementedError(f"{self}.periodic_point to be implemented") - - def sdf_derivatives(self, x: np.ndarray, epsilon: float = 1e-4) -> np.ndarray: - """Compute derivatives of SDF function. - - Args: - x (np.ndarray): Points for computing SDF derivatives using central - difference. The shape is [N, D], D is the number of dimension of - geometry. - epsilon (float): Derivative step. Defaults to 1e-4. - - Returns: - np.ndarray: Derivatives of corresponding SDF function. - The shape is [N, D]. D is the number of dimension of geometry. - - Examples: - >>> import numpy as np - >>> import ppsci - >>> interval = ppsci.geometry.Interval(0, 1) - >>> x = np.array([[0], [0.5], [1.5]]) - >>> interval.sdf_derivatives(x) - array([[-1.], - [ 0.], - [ 1.]]) - >>> rectangle = ppsci.geometry.Rectangle((0, 0), (1, 1)) - >>> x = np.array([[0.0, 0.0], [0.5, 0.5], [1.5, 1.5]]) - >>> rectangle.sdf_derivatives(x) - array([[-0.5 , -0.5 ], - [ 0. , 0. ], - [ 0.70710678, 0.70710678]]) - >>> cuboid = ppsci.geometry.Cuboid((0, 0, 0), (1, 1, 1)) - >>> x = np.array([[0, 0, 0], [0.5, 0.5, 0.5], [1, 1, 1]]) - >>> cuboid.sdf_derivatives(x) - array([[-0.5, -0.5, -0.5], - [ 0. , 0. , 0. ], - [ 0.5, 0.5, 0.5]]) - """ - if not hasattr(self, "sdf_func"): - raise NotImplementedError( - f"{misc.typename(self)}.sdf_func should be implemented " - "when using 'sdf_derivatives'." - ) - # Only compute sdf derivatives for those already implement `sdf_func` method. - sdf_derives = np.empty_like(x) - for i in range(self.ndim): - h = np.zeros_like(x) - h[:, i] += epsilon / 2 - derives_at_i = (self.sdf_func(x + h) - self.sdf_func(x - h)) / epsilon - sdf_derives[:, i : i + 1] = derives_at_i - return sdf_derives - - def union(self, other: "Geometry") -> "Geometry": - """CSG Union. - - Args: - other (Geometry): The other geometry. - - Returns: - Geometry: The union of two geometries. - - Examples: - >>> import numpy as np - >>> import ppsci - >>> interval1 = ppsci.geometry.Interval(0, 1) - >>> interval2 = ppsci.geometry.Interval(0.5, 1.5) - >>> union = interval1.union(interval2) - >>> union.bbox - (array([[0.]]), array([[1.5]])) - >>> rectangle1 = ppsci.geometry.Rectangle((0, 0), (2, 3)) - >>> rectangle2 = ppsci.geometry.Rectangle((0, 0), (3, 2)) - >>> union = rectangle1.union(rectangle2) - >>> union.bbox - (array([0., 0.], dtype=float32), array([3., 3.], dtype=float32)) - >>> cuboid1 = ppsci.geometry.Cuboid((0, 0, 0), (1, 2, 2)) - >>> cuboid2 = ppsci.geometry.Cuboid((0, 0, 0), (2, 1, 1)) - >>> union = cuboid1 | cuboid2 - >>> union.bbox - (array([0., 0., 0.], dtype=float32), array([2., 2., 2.], dtype=float32)) - """ - from ppsci.geometry import csg - - return csg.CSGUnion(self, other) - - def __or__(self, other: "Geometry") -> "Geometry": - """CSG Union. - - Args: - other (Geometry): The other geometry. - - Returns: - Geometry: The union of two geometries. - - Examples: - >>> import numpy as np - >>> import ppsci - >>> interval1 = ppsci.geometry.Interval(0, 1) - >>> interval2 = ppsci.geometry.Interval(0.5, 1.5) - >>> union = interval1.__or__(interval2) - >>> union.bbox - (array([[0.]]), array([[1.5]])) - >>> rectangle1 = ppsci.geometry.Rectangle((0, 0), (2, 3)) - >>> rectangle2 = ppsci.geometry.Rectangle((0, 0), (3, 2)) - >>> union = rectangle1.__or__(rectangle2) - >>> union.bbox - (array([0., 0.], dtype=float32), array([3., 3.], dtype=float32)) - >>> cuboid1 = ppsci.geometry.Cuboid((0, 0, 0), (1, 2, 2)) - >>> cuboid2 = ppsci.geometry.Cuboid((0, 0, 0), (2, 1, 1)) - >>> union = cuboid1 | cuboid2 - >>> union.bbox - (array([0., 0., 0.], dtype=float32), array([2., 2., 2.], dtype=float32)) - """ - from ppsci.geometry import csg - - return csg.CSGUnion(self, other) - - def difference(self, other: "Geometry") -> "Geometry": - """CSG Difference. - - Args: - other (Geometry): The other geometry. - - Returns: - Geometry: The difference of two geometries. - - Examples: - >>> import numpy as np - >>> import ppsci - >>> interval1 = ppsci.geometry.Interval(0.0, 2.0) - >>> interval2 = ppsci.geometry.Interval(1.0, 3.0) - >>> difference = interval1.difference(interval2) - >>> difference.bbox - (array([[0.]]), array([[2.]])) - >>> rectangle1 = ppsci.geometry.Rectangle((0.0, 0.0), (2.0, 3.0)) - >>> rectangle2 = ppsci.geometry.Rectangle((1.0, 1.0), (2.0, 2.0)) - >>> difference = rectangle1.difference(rectangle2) - >>> difference.bbox - (array([0., 0.], dtype=float32), array([2., 3.], dtype=float32)) - >>> cuboid1 = ppsci.geometry.Cuboid((0, 0, 0), (1, 2, 2)) - >>> cuboid2 = ppsci.geometry.Cuboid((0, 0, 0), (2, 1, 1)) - >>> difference = cuboid1 - cuboid2 - >>> difference.bbox - (array([0., 0., 0.], dtype=float32), array([1., 2., 2.], dtype=float32)) - """ - from ppsci.geometry import csg - - return csg.CSGDifference(self, other) - - def __sub__(self, other: "Geometry") -> "Geometry": - """CSG Difference. - - Args: - other (Geometry): The other geometry. - - Returns: - Geometry: The difference of two geometries. - - Examples: - >>> import numpy as np - >>> import ppsci - >>> interval1 = ppsci.geometry.Interval(0.0, 2.0) - >>> interval2 = ppsci.geometry.Interval(1.0, 3.0) - >>> difference = interval1.__sub__(interval2) - >>> difference.bbox - (array([[0.]]), array([[2.]])) - >>> rectangle1 = ppsci.geometry.Rectangle((0.0, 0.0), (2.0, 3.0)) - >>> rectangle2 = ppsci.geometry.Rectangle((1.0, 1.0), (2.0, 2.0)) - >>> difference = rectangle1.__sub__(rectangle2) - >>> difference.bbox - (array([0., 0.], dtype=float32), array([2., 3.], dtype=float32)) - >>> cuboid1 = ppsci.geometry.Cuboid((0, 0, 0), (1, 2, 2)) - >>> cuboid2 = ppsci.geometry.Cuboid((0, 0, 0), (2, 1, 1)) - >>> difference = cuboid1 - cuboid2 - >>> difference.bbox - (array([0., 0., 0.], dtype=float32), array([1., 2., 2.], dtype=float32)) - """ - from ppsci.geometry import csg - - return csg.CSGDifference(self, other) - - def intersection(self, other: "Geometry") -> "Geometry": - """CSG Intersection. - - Args: - other (Geometry): The other geometry. - - Returns: - Geometry: The intersection of two geometries. - - Examples: - >>> import numpy as np - >>> import ppsci - >>> interval1 = ppsci.geometry.Interval(0.0, 1.0) - >>> interval2 = ppsci.geometry.Interval(0.5, 1.5) - >>> intersection = interval1.intersection(interval2) - >>> intersection.bbox - (array([[0.5]]), array([[1.]])) - >>> rectangle1 = ppsci.geometry.Rectangle((0.0, 0.0), (2.0, 3.0)) - >>> rectangle2 = ppsci.geometry.Rectangle((0.0, 0.0), (3.0, 2.0)) - >>> intersection = rectangle1.intersection(rectangle2) - >>> intersection.bbox - (array([0., 0.], dtype=float32), array([2., 2.], dtype=float32)) - >>> cuboid1 = ppsci.geometry.Cuboid((0, 0, 0), (1, 2, 2)) - >>> cuboid2 = ppsci.geometry.Cuboid((0, 0, 0), (2, 1, 1)) - >>> intersection = cuboid1 & cuboid2 - >>> intersection.bbox - (array([0., 0., 0.], dtype=float32), array([1., 1., 1.], dtype=float32)) - """ - from ppsci.geometry import csg - - return csg.CSGIntersection(self, other) - - def __and__(self, other: "Geometry") -> "Geometry": - """CSG Intersection. - - Args: - other (Geometry): The other geometry. - - Returns: - Geometry: The intersection of two geometries. - - Examples: - >>> import numpy as np - >>> import ppsci - >>> interval1 = ppsci.geometry.Interval(0.0, 1.0) - >>> interval2 = ppsci.geometry.Interval(0.5, 1.5) - >>> intersection = interval1.__and__(interval2) - >>> intersection.bbox - (array([[0.5]]), array([[1.]])) - >>> rectangle1 = ppsci.geometry.Rectangle((0.0, 0.0), (2.0, 3.0)) - >>> rectangle2 = ppsci.geometry.Rectangle((0.0, 0.0), (3.0, 2.0)) - >>> intersection = rectangle1.__and__(rectangle2) - >>> intersection.bbox - (array([0., 0.], dtype=float32), array([2., 2.], dtype=float32)) - >>> cuboid1 = ppsci.geometry.Cuboid((0, 0, 0), (1, 2, 2)) - >>> cuboid2 = ppsci.geometry.Cuboid((0, 0, 0), (2, 1, 1)) - >>> intersection = cuboid1 & cuboid2 - >>> intersection.bbox - (array([0., 0., 0.], dtype=float32), array([1., 1., 1.], dtype=float32)) - """ - from ppsci.geometry import csg - - return csg.CSGIntersection(self, other) - - def __str__(self) -> str: - """Return the name of class. - - Returns: - str: Meta information of geometry. - - Examples: - >>> import ppsci - >>> interval = ppsci.geometry.Interval(0, 1) - >>> interval.__str__() - "Interval, ndim = 1, bbox = (array([[0]]), array([[1]])), diam = 1, dim_keys = ('x',)" - >>> rectangle = ppsci.geometry.Rectangle((0, 0), (1, 1)) - >>> rectangle.__str__() - "Rectangle, ndim = 2, bbox = (array([0., 0.], dtype=float32), array([1., 1.], dtype=float32)), diam = 1.4142135381698608, dim_keys = ('x', 'y')" - >>> cuboid = ppsci.geometry.Cuboid((0, 0, 0), (1, 1, 1)) - >>> cuboid.__str__() - "Cuboid, ndim = 3, bbox = (array([0., 0., 0.], dtype=float32), array([1., 1., 1.], dtype=float32)), diam = 1.7320507764816284, dim_keys = ('x', 'y', 'z')" - """ - return ", ".join( - [ - self.__class__.__name__, - f"ndim = {self.ndim}", - f"bbox = {self.bbox}", - f"diam = {self.diam}", - f"dim_keys = {self.dim_keys}", - ] - ) diff --git a/examples/smc_reac/ppsci/geometry/geometry_1d.py b/examples/smc_reac/ppsci/geometry/geometry_1d.py deleted file mode 100644 index d5de01fe56..0000000000 --- a/examples/smc_reac/ppsci/geometry/geometry_1d.py +++ /dev/null @@ -1,119 +0,0 @@ -# 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. - -""" -Code below is heavily based on [https://github.com/lululxvi/deepxde](https://github.com/lululxvi/deepxde) -""" - -from __future__ import annotations - -import numpy as np -import paddle - -from ppsci.geometry import geometry -from ppsci.geometry.sampler import sample -from ppsci.utils import misc - - -class Interval(geometry.Geometry): - """Class for interval. - - Args: - l (float): Left position of interval. - r (float): Right position of interval. - - Examples: - >>> import ppsci - >>> geom = ppsci.geometry.Interval(-1, 1) - """ - - def __init__(self, l: float, r: float): - super().__init__(1, (np.array([[l]]), np.array([[r]])), r - l) - self.l = l - self.r = r - - def is_inside(self, x: np.ndarray): - return ((self.l <= x) & (x <= self.r)).flatten() - - def on_boundary(self, x: np.ndarray): - return (np.isclose(x, self.l) | np.isclose(x, self.r)).flatten() - - def boundary_normal(self, x: np.ndarray): - return -np.isclose(x, self.l).astype(paddle.get_default_dtype()) + np.isclose( - x, self.r - ).astype(paddle.get_default_dtype()) - - def uniform_points(self, n: int, boundary: bool = True): - if boundary: - return np.linspace( - self.l, self.r, n, dtype=paddle.get_default_dtype() - ).reshape([-1, 1]) - return np.linspace( - self.l, self.r, n + 1, endpoint=False, dtype=paddle.get_default_dtype() - )[1:].reshape([-1, 1]) - - def random_points(self, n: int, random: str = "pseudo"): - x = sample(n, 1, random) - return (self.l + x * self.diam).astype(paddle.get_default_dtype()) - - def uniform_boundary_points(self, n: int): - if n == 1: - return np.array([[self.l]], dtype=paddle.get_default_dtype()) - xl = np.full([n // 2, 1], self.l, dtype=paddle.get_default_dtype()) - xr = np.full([n - n // 2, 1], self.r, dtype=paddle.get_default_dtype()) - return np.concatenate((xl, xr), axis=0) - - def random_boundary_points(self, n: int, random: str = "pseudo"): - if n == 2: - return np.array([[self.l], [self.r]], dtype=paddle.get_default_dtype()) - return ( - np.random.choice([self.l, self.r], n) - .reshape([-1, 1]) - .astype(paddle.get_default_dtype()) - ) - - def periodic_point(self, x: np.ndarray, component: int = 0): - x_array = misc.convert_to_array(x, self.dim_keys) - periodic_x = x_array - periodic_x[np.isclose(x_array, self.l)] = self.r - periodic_x[np.isclose(x_array, self.r)] = self.l - periodic_x_normal = self.boundary_normal(periodic_x) - - periodic_x = misc.convert_to_dict(periodic_x, self.dim_keys) - periodic_x_normal = misc.convert_to_dict( - periodic_x_normal, [f"normal_{k}" for k in self.dim_keys] - ) - return {**periodic_x, **periodic_x_normal} - - def sdf_func(self, points: np.ndarray) -> np.ndarray: - """Compute signed distance field - - Args: - points (np.ndarray): The coordinate points used to calculate the SDF value, - the shape is [N, 1] - - Returns: - np.ndarray: SDF values of input points without squared, the shape is [N, 1]. - - NOTE: This function usually returns ndarray with negative values, because - according to the definition of SDF, the SDF value of the coordinate point inside - the object(interior points) is negative, the outside is positive, and the edge - is 0. Therefore, when used for weighting, a negative sign is often added before - the result of this function. - """ - if points.shape[1] != self.ndim: - raise ValueError( - f"Shape of given points should be [*, {self.ndim}], but got {points.shape}" - ) - return -((self.r - self.l) / 2 - np.abs(points - (self.l + self.r) / 2)) diff --git a/examples/smc_reac/ppsci/geometry/geometry_2d.py b/examples/smc_reac/ppsci/geometry/geometry_2d.py deleted file mode 100644 index 2df6293b27..0000000000 --- a/examples/smc_reac/ppsci/geometry/geometry_2d.py +++ /dev/null @@ -1,706 +0,0 @@ -# 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. - -""" -Code below is heavily based on [https://github.com/lululxvi/deepxde](https://github.com/lululxvi/deepxde) -""" - -from __future__ import annotations - -from typing import Tuple - -import numpy as np -import paddle -from scipy import spatial - -from ppsci.geometry import geometry -from ppsci.geometry import geometry_nd -from ppsci.geometry import sampler - - -class Disk(geometry.Geometry): - """Class for disk geometry - - Args: - center (Tuple[float, float]): Center point of disk [x0, y0]. - radius (float): Radius of disk. - - Examples: - >>> import ppsci - >>> geom = ppsci.geometry.Disk((0.0, 0.0), 1.0) - """ - - def __init__(self, center: Tuple[float, float], radius: float): - self.center = np.array(center, dtype=paddle.get_default_dtype()) - self.radius = radius - super().__init__(2, (self.center - radius, self.center + radius), 2 * radius) - - def is_inside(self, x): - return np.linalg.norm(x - self.center, axis=1) <= self.radius - - def on_boundary(self, x): - return np.isclose(np.linalg.norm(x - self.center, axis=1), self.radius) - - def boundary_normal(self, x): - ox = x - self.center - ox_len = np.linalg.norm(ox, axis=1, keepdims=True) - ox = (ox / ox_len) * np.isclose(ox_len, self.radius).astype( - paddle.get_default_dtype() - ) - return ox - - def random_points(self, n, random="pseudo"): - # http://mathworld.wolfram.com/DiskPointPicking.html - rng = sampler.sample(n, 2, random) - r, theta = rng[:, 0], 2 * np.pi * rng[:, 1] - x = np.sqrt(r) * np.cos(theta) - y = np.sqrt(r) * np.sin(theta) - return self.radius * np.stack((x, y), axis=1) + self.center - - def uniform_boundary_points(self, n): - theta = np.linspace( - 0, 2 * np.pi, num=n, endpoint=False, dtype=paddle.get_default_dtype() - ) - X = np.stack((np.cos(theta), np.sin(theta)), axis=1) - return self.radius * X + self.center - - def random_boundary_points(self, n, random="pseudo"): - theta = 2 * np.pi * sampler.sample(n, 1, random) - X = np.concatenate((np.cos(theta), np.sin(theta)), axis=1) - return self.radius * X + self.center - - def sdf_func(self, points: np.ndarray) -> np.ndarray: - """Compute signed distance field. - - Args: - points (np.ndarray): The coordinate points used to calculate the SDF value, - the shape is [N, 2] - - Returns: - np.ndarray: SDF values of input points without squared, the shape is [N, 1]. - - NOTE: This function usually returns ndarray with negative values, because - according to the definition of SDF, the SDF value of the coordinate point inside - the object(interior points) is negative, the outside is positive, and the edge - is 0. Therefore, when used for weighting, a negative sign is often added before - the result of this function. - """ - if points.shape[1] != self.ndim: - raise ValueError( - f"Shape of given points should be [*, {self.ndim}], but got {points.shape}" - ) - sdf = self.radius - np.linalg.norm(points - self.center, axis=1) - sdf = -sdf[..., np.newaxis] - return sdf - - -class Rectangle(geometry_nd.Hypercube): - """Class for rectangle geometry - - Args: - xmin (Tuple[float, float]): Bottom left corner point, [x0, y0]. - xmax (Tuple[float, float]): Top right corner point, [x1, y1]. - - Examples: - >>> import ppsci - >>> geom = ppsci.geometry.Rectangle((0.0, 0.0), (1.0, 1.0)) - """ - - def __init__(self, xmin, xmax): - super().__init__(xmin, xmax) - self.perimeter = 2 * np.sum(self.xmax - self.xmin) - self.area = np.prod(self.xmax - self.xmin) - - def uniform_boundary_points(self, n): - nx, ny = np.ceil(n / self.perimeter * (self.xmax - self.xmin)).astype(int) - bottom = np.hstack( - ( - np.linspace( - self.xmin[0], - self.xmax[0], - nx, - endpoint=False, - dtype=paddle.get_default_dtype(), - ).reshape([nx, 1]), - np.full([nx, 1], self.xmin[1], dtype=paddle.get_default_dtype()), - ) - ) - right = np.hstack( - ( - np.full([ny, 1], self.xmax[0], dtype=paddle.get_default_dtype()), - np.linspace( - self.xmin[1], - self.xmax[1], - ny, - endpoint=False, - dtype=paddle.get_default_dtype(), - ).reshape([ny, 1]), - ) - ) - top = np.hstack( - ( - np.linspace( - self.xmin[0], self.xmax[0], nx + 1, dtype=paddle.get_default_dtype() - )[1:].reshape([nx, 1]), - np.full([nx, 1], self.xmax[1], dtype=paddle.get_default_dtype()), - ) - ) - left = np.hstack( - ( - np.full([ny, 1], self.xmin[0], dtype=paddle.get_default_dtype()), - np.linspace( - self.xmin[1], self.xmax[1], ny + 1, dtype=paddle.get_default_dtype() - )[1:].reshape([ny, 1]), - ) - ) - x = np.vstack((bottom, right, top, left)) - if len(x) > n: - x = x[0:n] - return x - - def random_boundary_points(self, n, random="pseudo"): - l1 = self.xmax[0] - self.xmin[0] - l2 = l1 + self.xmax[1] - self.xmin[1] - l3 = l2 + l1 - u = np.ravel(sampler.sample(n + 10, 1, random)) - # Remove the possible points very close to the corners - u = u[~np.isclose(u, l1 / self.perimeter)] - u = u[~np.isclose(u, l3 / self.perimeter)] - u = u[0:n] - - u *= self.perimeter - x = [] - for l in u: - if l < l1: - x.append([self.xmin[0] + l, self.xmin[1]]) - elif l < l2: - x.append([self.xmax[0], self.xmin[1] + (l - l1)]) - elif l < l3: - x.append([self.xmax[0] - (l - l2), self.xmax[1]]) - else: - x.append([self.xmin[0], self.xmax[1] - (l - l3)]) - return np.vstack(x) - - @staticmethod - def is_valid(vertices): - """Check if the geometry is a Rectangle.""" - return ( - len(vertices) == 4 - and np.isclose(np.prod(vertices[1] - vertices[0]), 0) - and np.isclose(np.prod(vertices[2] - vertices[1]), 0) - and np.isclose(np.prod(vertices[3] - vertices[2]), 0) - and np.isclose(np.prod(vertices[0] - vertices[3]), 0) - ) - - def sdf_func(self, points: np.ndarray) -> np.ndarray: - """Compute signed distance field. - - Args: - points (np.ndarray): The coordinate points used to calculate the SDF value, - the shape of the array is [N, 2]. - - Returns: - np.ndarray: SDF values of input points without squared, the shape is [N, 1]. - - NOTE: This function usually returns ndarray with negative values, because - according to the definition of SDF, the SDF value of the coordinate point inside - the object(interior points) is negative, the outside is positive, and the edge - is 0. Therefore, when used for weighting, a negative sign is often added before - the result of this function. - """ - if points.shape[1] != self.ndim: - raise ValueError( - f"Shape of given points should be [*, {self.ndim}], but got {points.shape}" - ) - center = (self.xmin + self.xmax) / 2 - dist_to_boundary = ( - np.abs(points - center) - np.array([self.xmax - self.xmin]) / 2 - ) - return ( - np.linalg.norm(np.maximum(dist_to_boundary, 0), axis=1) - + np.minimum(np.max(dist_to_boundary, axis=1), 0) - ).reshape(-1, 1) - - -class Triangle(geometry.Geometry): - """Class for Triangle - - The order of vertices can be in a clockwise or counterclockwise direction. The - vertices will be re-ordered in counterclockwise (right hand rule). - - Args: - x1 (Tuple[float, float]): First point of Triangle [x0, y0]. - x2 (Tuple[float, float]): Second point of Triangle [x1, y1]. - x3 (Tuple[float, float]): Third point of Triangle [x2, y2]. - - Examples: - >>> import ppsci - >>> geom = ppsci.geometry.Triangle((0, 0), (1, 0), (0, 1)) - """ - - def __init__(self, x1, x2, x3): - self.area = polygon_signed_area([x1, x2, x3]) - # Clockwise - if self.area < 0: - self.area = -self.area - x2, x3 = x3, x2 - - self.x1 = np.array(x1, dtype=paddle.get_default_dtype()) - self.x2 = np.array(x2, dtype=paddle.get_default_dtype()) - self.x3 = np.array(x3, dtype=paddle.get_default_dtype()) - - self.v12 = self.x2 - self.x1 - self.v23 = self.x3 - self.x2 - self.v31 = self.x1 - self.x3 - self.l12 = np.linalg.norm(self.v12) - self.l23 = np.linalg.norm(self.v23) - self.l31 = np.linalg.norm(self.v31) - self.n12 = self.v12 / self.l12 - self.n23 = self.v23 / self.l23 - self.n31 = self.v31 / self.l31 - self.n12_normal = clockwise_rotation_90(self.n12) - self.n23_normal = clockwise_rotation_90(self.n23) - self.n31_normal = clockwise_rotation_90(self.n31) - self.perimeter = self.l12 + self.l23 + self.l31 - - super().__init__( - 2, - (np.minimum(x1, np.minimum(x2, x3)), np.maximum(x1, np.maximum(x2, x3))), - self.l12 - * self.l23 - * self.l31 - / ( - self.perimeter - * (self.l12 + self.l23 - self.l31) - * (self.l23 + self.l31 - self.l12) - * (self.l31 + self.l12 - self.l23) - ) - ** 0.5, - ) - - def is_inside(self, x): - # https://stackoverflow.com/a/2049593/12679294 - _sign = np.stack( - [ - np.cross(self.v12, x - self.x1), - np.cross(self.v23, x - self.x2), - np.cross(self.v31, x - self.x3), - ], - axis=1, - ) - return ~(np.any(_sign > 0, axis=-1) & np.any(_sign < 0, axis=-1)) - - def on_boundary(self, x): - l1 = np.linalg.norm(x - self.x1, axis=-1) - l2 = np.linalg.norm(x - self.x2, axis=-1) - l3 = np.linalg.norm(x - self.x3, axis=-1) - return np.any( - np.isclose( - [l1 + l2 - self.l12, l2 + l3 - self.l23, l3 + l1 - self.l31], - 0, - atol=1e-6, - ), - axis=0, - ) - - def boundary_normal(self, x): - l1 = np.linalg.norm(x - self.x1, axis=-1, keepdims=True) - l2 = np.linalg.norm(x - self.x2, axis=-1, keepdims=True) - l3 = np.linalg.norm(x - self.x3, axis=-1, keepdims=True) - on12 = np.isclose(l1 + l2, self.l12) - on23 = np.isclose(l2 + l3, self.l23) - on31 = np.isclose(l3 + l1, self.l31) - # Check points on the vertexes - if np.any(np.count_nonzero(np.hstack([on12, on23, on31]), axis=-1) > 1): - raise ValueError( - "{}.boundary_normal do not accept points on the vertexes.".format( - self.__class__.__name__ - ) - ) - return self.n12_normal * on12 + self.n23_normal * on23 + self.n31_normal * on31 - - def random_points(self, n, random="pseudo"): - # There are two methods for triangle point picking. - # Method 1 (used here): - # - https://math.stackexchange.com/questions/18686/uniform-random-point-in-triangle - # Method 2: - # - http://mathworld.wolfram.com/TrianglePointPicking.html - # - https://hbfs.wordpress.com/2010/10/05/random-points-in-a-triangle-generating-random-sequences-ii/ - # - https://stackoverflow.com/questions/19654251/random-point-inside-triangle-inside-java - sqrt_r1 = np.sqrt(np.random.rand(n, 1)) - r2 = np.random.rand(n, 1) - return ( - (1 - sqrt_r1) * self.x1 - + sqrt_r1 * (1 - r2) * self.x2 - + r2 * sqrt_r1 * self.x3 - ) - - def uniform_boundary_points(self, n): - density = n / self.perimeter - x12 = ( - np.linspace( - 0, - 1, - num=int(np.ceil(density * self.l12)), - endpoint=False, - dtype=paddle.get_default_dtype(), - )[:, None] - * self.v12 - + self.x1 - ) - x23 = ( - np.linspace( - 0, - 1, - num=int(np.ceil(density * self.l23)), - endpoint=False, - dtype=paddle.get_default_dtype(), - )[:, None] - * self.v23 - + self.x2 - ) - x31 = ( - np.linspace( - 0, - 1, - num=int(np.ceil(density * self.l31)), - endpoint=False, - dtype=paddle.get_default_dtype(), - )[:, None] - * self.v31 - + self.x3 - ) - x = np.vstack((x12, x23, x31)) - if len(x) > n: - x = x[0:n] - return x - - def random_boundary_points(self, n, random="pseudo"): - u = np.ravel(sampler.sample(n + 2, 1, random)) - # Remove the possible points very close to the corners - u = u[np.logical_not(np.isclose(u, self.l12 / self.perimeter))] - u = u[np.logical_not(np.isclose(u, (self.l12 + self.l23) / self.perimeter))] - u = u[:n] - - u *= self.perimeter - x = [] - for l in u: - if l < self.l12: - x.append(l * self.n12 + self.x1) - elif l < self.l12 + self.l23: - x.append((l - self.l12) * self.n23 + self.x2) - else: - x.append((l - self.l12 - self.l23) * self.n31 + self.x3) - return np.vstack(x) - - def sdf_func(self, points: np.ndarray) -> np.ndarray: - """Compute signed distance field. - - Args: - points (np.ndarray): The coordinate points used to calculate the SDF value, - the shape of the array is [N, 2]. - - Returns: - np.ndarray: SDF values of input points without squared, the shape is [N, 1]. - - NOTE: This function usually returns ndarray with negative values, because - according to the definition of SDF, the SDF value of the coordinate point inside - the object(interior points) is negative, the outside is positive, and the edge - is 0. Therefore, when used for weighting, a negative sign is often added before - the result of this function. - """ - if points.shape[1] != self.ndim: - raise ValueError( - f"Shape of given points should be [*, {self.ndim}], but got {points.shape}" - ) - v1p = points - self.x1 # v1p: vector from x1 to points - v2p = points - self.x2 - v3p = points - self.x3 - # vv12_p: vertical vector of points to v12(If the vertical point is in the extension of v12, - # the vector will be the vector from x1 to points) - vv12_p = ( - self.v12 - * np.clip(np.dot(v1p, self.v12.reshape(2, -1)) / self.l12**2, 0, 1) - - v1p - ) - vv23_p = ( - self.v23 - * np.clip(np.dot(v2p, self.v23.reshape(2, -1)) / self.l23**2, 0, 1) - - v2p - ) - vv31_p = ( - self.v31 - * np.clip(np.dot(v3p, self.v31.reshape(2, -1)) / self.l31**2, 0, 1) - - v3p - ) - is_inside = self.is_inside(points).reshape(-1, 1) * 2 - 1 - len_vv12_p = np.linalg.norm(vv12_p, axis=1, keepdims=True) - len_vv23_p = np.linalg.norm(vv23_p, axis=1, keepdims=True) - len_vv31_p = np.linalg.norm(vv31_p, axis=1, keepdims=True) - mini_dist = np.minimum(np.minimum(len_vv12_p, len_vv23_p), len_vv31_p) - return is_inside * mini_dist - - -class Polygon(geometry.Geometry): - """Class for simple polygon. - - Args: - vertices (Tuple[Tuple[float, float], ...]): The order of vertices can be in a - clockwise or counter-clockwise direction. The vertices will be re-ordered in - counterclockwise (right hand rule). - - Examples: - >>> import ppsci - >>> geom = ppsci.geometry.Polygon(((0, 0), (1, 0), (2, 1), (2, 2), (0, 2))) - """ - - def __init__(self, vertices): - self.vertices = np.array(vertices, dtype=paddle.get_default_dtype()) - if len(vertices) == 3: - raise ValueError("The polygon is a triangle. Use Triangle instead.") - if Rectangle.is_valid(self.vertices): - raise ValueError("The polygon is a rectangle. Use Rectangle instead.") - - self.area = polygon_signed_area(self.vertices) - # Clockwise - if self.area < 0: - self.area = -self.area - self.vertices = np.flipud(self.vertices) - - self.diagonals = spatial.distance.squareform( - spatial.distance.pdist(self.vertices) - ) - super().__init__( - 2, - (np.amin(self.vertices, axis=0), np.amax(self.vertices, axis=0)), - np.max(self.diagonals), - ) - self.nvertices = len(self.vertices) - self.perimeter = np.sum( - [self.diagonals[i, i + 1] for i in range(-1, self.nvertices - 1)] - ) - self.bbox = np.array( - [np.min(self.vertices, axis=0), np.max(self.vertices, axis=0)], - dtype=paddle.get_default_dtype(), - ) - - self.segments = self.vertices[1:] - self.vertices[:-1] - self.segments = np.vstack((self.vertices[0] - self.vertices[-1], self.segments)) - self.normal = clockwise_rotation_90(self.segments.T).T - self.normal = self.normal / np.linalg.norm(self.normal, axis=1).reshape(-1, 1) - - def is_inside(self, x): - def wn_PnPoly(P, V): - """Winding number algorithm. - - https://en.wikipedia.org/wiki/Point_in_polygon - http://geomalgorithms.com/a03-_inclusion.html - - Args: - P: A point. - V: Vertex points of a polygon. - - Returns: - wn: Winding number (=0 only if P is outside polygon). - """ - wn = np.zeros(len(P)) # Winding number counter - - # Repeat the first vertex at end - # Loop through all edges of the polygon - for i in range(-1, self.nvertices - 1): # Edge from V[i] to V[i+1] - tmp = np.all( - np.hstack( - [ - V[i, 1] <= P[:, 1:2], # Start y <= P[1] - V[i + 1, 1] > P[:, 1:2], # An upward crossing - is_left(V[i], V[i + 1], P) > 0, # P left of edge - ] - ), - axis=-1, - ) - wn[tmp] += 1 # Have a valid up intersect - tmp = np.all( - np.hstack( - [ - V[i, 1] > P[:, 1:2], # Start y > P[1] - V[i + 1, 1] <= P[:, 1:2], # A downward crossing - is_left(V[i], V[i + 1], P) < 0, # P right of edge - ] - ), - axis=-1, - ) - wn[tmp] -= 1 # Have a valid down intersect - return wn - - return wn_PnPoly(x, self.vertices) != 0 - - def on_boundary(self, x): - _on = np.zeros(shape=len(x), dtype=np.int) - for i in range(-1, self.nvertices - 1): - l1 = np.linalg.norm(self.vertices[i] - x, axis=-1) - l2 = np.linalg.norm(self.vertices[i + 1] - x, axis=-1) - _on[np.isclose(l1 + l2, self.diagonals[i, i + 1])] += 1 - return _on > 0 - - def random_points(self, n, random="pseudo"): - x = np.empty((0, 2), dtype=paddle.get_default_dtype()) - vbbox = self.bbox[1] - self.bbox[0] - while len(x) < n: - x_new = sampler.sample(n, 2, "pseudo") * vbbox + self.bbox[0] - x = np.vstack((x, x_new[self.is_inside(x_new)])) - return x[:n] - - def uniform_boundary_points(self, n): - density = n / self.perimeter - x = [] - for i in range(-1, self.nvertices - 1): - x.append( - np.linspace( - 0, - 1, - num=int(np.ceil(density * self.diagonals[i, i + 1])), - endpoint=False, - dtype=paddle.get_default_dtype(), - )[:, None] - * (self.vertices[i + 1] - self.vertices[i]) - + self.vertices[i] - ) - x = np.vstack(x) - if len(x) > n: - x = x[0:n] - return x - - def random_boundary_points(self, n, random="pseudo"): - u = np.ravel(sampler.sample(n + self.nvertices, 1, random)) - # Remove the possible points very close to the corners - l = 0 - for i in range(0, self.nvertices - 1): - l += self.diagonals[i, i + 1] - u = u[np.logical_not(np.isclose(u, l / self.perimeter))] - u = u[:n] - u *= self.perimeter - u.sort() - - x = [] - i = -1 - l0 = 0 - l1 = l0 + self.diagonals[i, i + 1] - v = (self.vertices[i + 1] - self.vertices[i]) / self.diagonals[i, i + 1] - for l in u: - if l > l1: - i += 1 - l0, l1 = l1, l1 + self.diagonals[i, i + 1] - v = (self.vertices[i + 1] - self.vertices[i]) / self.diagonals[i, i + 1] - x.append((l - l0) * v + self.vertices[i]) - return np.vstack(x) - - def sdf_func(self, points: np.ndarray) -> np.ndarray: - """Compute signed distance field. - - Args: - points (np.ndarray): The coordinate points used to calculate the SDF value, - the shape is [N, 2] - Returns: - np.ndarray: SDF values of input points without squared, the shape is [N, 1]. - - NOTE: This function usually returns ndarray with negative values, because - according to the definition of SDF, the SDF value of the coordinate point inside - the object(interior points) is negative, the outside is positive, and the edge - is 0. Therefore, when used for weighting, a negative sign is often added before - the result of this function. - """ - if points.shape[1] != self.ndim: - raise ValueError( - f"Shape of given points should be [*, {self.ndim}], but got {points.shape}" - ) - sdf_value = np.empty((points.shape[0], 1), dtype=paddle.get_default_dtype()) - for n in range(points.shape[0]): - distance = np.dot( - points[n] - self.vertices[0], points[n] - self.vertices[0] - ) - inside_tag = 1.0 - for i in range(self.vertices.shape[0]): - j = (self.vertices.shape[0] - 1) if i == 0 else (i - 1) - # Calculate the shortest distance from point P to each edge. - vector_ij = self.vertices[j] - self.vertices[i] - vector_in = points[n] - self.vertices[i] - distance_vector = vector_in - vector_ij * np.clip( - np.dot(vector_in, vector_ij) / np.dot(vector_ij, vector_ij), - 0.0, - 1.0, - ) - distance = np.minimum( - distance, np.dot(distance_vector, distance_vector) - ) - # Calculate the inside and outside using the Odd-even rule - odd_even_rule_number = np.array( - [ - points[n][1] >= self.vertices[i][1], - points[n][1] < self.vertices[j][1], - vector_ij[0] * vector_in[1] > vector_ij[1] * vector_in[0], - ] - ) - if odd_even_rule_number.all() or np.all(~odd_even_rule_number): - inside_tag *= -1.0 - sdf_value[n] = inside_tag * np.sqrt(distance) - return -sdf_value - - -def polygon_signed_area(vertices): - """The (signed) area of a simple polygon. - - If the vertices are in the counterclockwise direction, then the area is positive; if - they are in the clockwise direction, the area is negative. - - Shoelace formula: https://en.wikipedia.org/wiki/Shoelace_formula - - Args: - vertices (np.ndarray): Polygon vertices with shape of [N, 2]. - - Returns: - float: The (signed) area of a simple polygon. - """ - x, y = zip(*vertices) - x = np.array(list(x) + [x[0]], dtype=paddle.get_default_dtype()) - y = np.array(list(y) + [y[0]], dtype=paddle.get_default_dtype()) - return 0.5 * (np.sum(x[:-1] * y[1:]) - np.sum(x[1:] * y[:-1])) - - -def clockwise_rotation_90(v): - """Rotate a vector of 90 degrees clockwise about the origin. - - Args: - v (np.ndarray): Vector with shape of [2, N]. - - Returns: - np.ndarray: Rotated vector with shape of [2, N]. - """ - return np.array([v[1], -v[0]], dtype=paddle.get_default_dtype()) - - -def is_left(P0, P1, P2): - """Test if a point is Left|On|Right of an infinite line. - - See: the January 2001 Algorithm "Area of 2D and 3D Triangles and Polygons". - - Args: - P0 (np.ndarray): One point in the line. - P1 (np.ndarray): One point in the line. - P2 (np.ndarray): A array of point to be tested with shape of [N, 2]. - - Returns: - np.ndarray: >0 if P2 left of the line through P0 and P1, =0 if P2 on the line, <0 if P2 - right of the line. - """ - return np.cross(P1 - P0, P2 - P0, axis=-1).reshape((-1, 1)) diff --git a/examples/smc_reac/ppsci/geometry/geometry_3d.py b/examples/smc_reac/ppsci/geometry/geometry_3d.py deleted file mode 100644 index 8af958b1b9..0000000000 --- a/examples/smc_reac/ppsci/geometry/geometry_3d.py +++ /dev/null @@ -1,203 +0,0 @@ -# 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. - -""" -Code below is heavily based on [https://github.com/lululxvi/deepxde](https://github.com/lululxvi/deepxde) -""" - -from __future__ import annotations - -import itertools -from typing import Tuple - -import numpy as np -import paddle - -from ppsci.geometry import geometry_2d -from ppsci.geometry import geometry_nd - - -class Cuboid(geometry_nd.Hypercube): - """Class for Cuboid - - Args: - xmin (Tuple[float, float, float]): Bottom left corner point [x0, y0, z0]. - xmax (Tuple[float, float, float]): Top right corner point [x1, y1, z1]. - - Examples: - >>> import ppsci - >>> geom = ppsci.geometry.Cuboid((0, 0, 0), (1, 1, 1)) - """ - - def __init__( - self, xmin: Tuple[float, float, float], xmax: Tuple[float, float, float] - ): - super().__init__(xmin, xmax) - dx = self.xmax - self.xmin - self.area = 2 * np.sum(dx * np.roll(dx, 2)) - - def random_boundary_points(self, n, random="pseudo"): - pts = [] - density = n / self.area - rect = geometry_2d.Rectangle(self.xmin[:-1], self.xmax[:-1]) - for z in [self.xmin[-1], self.xmax[-1]]: - u = rect.random_points(int(np.ceil(density * rect.area)), random=random) - pts.append( - np.hstack( - (u, np.full((len(u), 1), z, dtype=paddle.get_default_dtype())) - ) - ) - rect = geometry_2d.Rectangle(self.xmin[::2], self.xmax[::2]) - for y in [self.xmin[1], self.xmax[1]]: - u = rect.random_points(int(np.ceil(density * rect.area)), random=random) - pts.append( - np.hstack( - ( - u[:, 0:1], - np.full((len(u), 1), y, dtype=paddle.get_default_dtype()), - u[:, 1:], - ) - ) - ) - rect = geometry_2d.Rectangle(self.xmin[1:], self.xmax[1:]) - for x in [self.xmin[0], self.xmax[0]]: - u = rect.random_points(int(np.ceil(density * rect.area)), random=random) - pts.append( - np.hstack( - (np.full((len(u), 1), x, dtype=paddle.get_default_dtype()), u) - ) - ) - pts = np.vstack(pts) - if len(pts) > n: - return pts[np.random.choice(len(pts), size=n, replace=False)] - return pts - - def uniform_boundary_points(self, n): - h = (self.area / n) ** 0.5 - nx, ny, nz = np.ceil((self.xmax - self.xmin) / h).astype(int) + 1 - x = np.linspace( - self.xmin[0], self.xmax[0], num=nx, dtype=paddle.get_default_dtype() - ) - y = np.linspace( - self.xmin[1], self.xmax[1], num=ny, dtype=paddle.get_default_dtype() - ) - z = np.linspace( - self.xmin[2], self.xmax[2], num=nz, dtype=paddle.get_default_dtype() - ) - - pts = [] - for v in [self.xmin[-1], self.xmax[-1]]: - u = list(itertools.product(x, y)) - pts.append( - np.hstack( - (u, np.full((len(u), 1), v, dtype=paddle.get_default_dtype())) - ) - ) - if nz > 2: - for v in [self.xmin[1], self.xmax[1]]: - u = np.array( - list(itertools.product(x, z[1:-1])), - dtype=paddle.get_default_dtype(), - ) - pts.append( - np.hstack( - ( - u[:, 0:1], - np.full((len(u), 1), v, dtype=paddle.get_default_dtype()), - u[:, 1:], - ) - ) - ) - if ny > 2 and nz > 2: - for v in [self.xmin[0], self.xmax[0]]: - u = list(itertools.product(y[1:-1], z[1:-1])) - pts.append( - np.hstack( - (np.full((len(u), 1), v, dtype=paddle.get_default_dtype()), u) - ) - ) - pts = np.vstack(pts) - if len(pts) > n: - return pts[np.random.choice(len(pts), size=n, replace=False)] - return pts - - def sdf_func(self, points: np.ndarray) -> np.ndarray: - """Compute signed distance field. - - Args: - points (np.ndarray): The coordinate points used to calculate the SDF value, - the shape is [N, 3] - - Returns: - np.ndarray: SDF values of input points without squared, the shape is [N, 1]. - - NOTE: This function usually returns ndarray with negative values, because - according to the definition of SDF, the SDF value of the coordinate point inside - the object(interior points) is negative, the outside is positive, and the edge - is 0. Therefore, when used for weighting, a negative sign is often added before - the result of this function. - """ - if points.shape[1] != self.ndim: - raise ValueError( - f"Shape of given points should be [*, {self.ndim}], but got {points.shape}" - ) - sdf = ( - ((self.xmax - self.xmin) / 2 - abs(points - (self.xmin + self.xmax) / 2)) - ).min(axis=1) - sdf = -sdf[..., np.newaxis] - return sdf - - -class Sphere(geometry_nd.Hypersphere): - """Class for Sphere - - Args: - center (Tuple[float, float, float]): Center of the sphere [x0, y0, z0]. - radius (float): Radius of the sphere. - """ - - def __init__(self, center, radius): - super().__init__(center, radius) - - def uniform_boundary_points(self, n: int): - nl = np.arange(1, n + 1).astype(paddle.get_default_dtype()) - g = (np.sqrt(5) - 1) / 2 - z = (2 * nl - 1) / n - 1 - x = np.sqrt(1 - z**2) * np.cos(2 * np.pi * nl * g) - y = np.sqrt(1 - z**2) * np.sin(2 * np.pi * nl * g) - return np.stack((x, y, z), axis=-1) - - def sdf_func(self, points: np.ndarray) -> np.ndarray: - """Compute signed distance field. - - Args: - points (np.ndarray): The coordinate points used to calculate the SDF value, - the shape is [N, 3] - - Returns: - np.ndarray: SDF values of input points without squared, the shape is [N, 1]. - - NOTE: This function usually returns ndarray with negative values, because - according to the definition of SDF, the SDF value of the coordinate point inside - the object(interior points) is negative, the outside is positive, and the edge - is 0. Therefore, when used for weighting, a negative sign is often added before - the result of this function. - """ - if points.shape[1] != self.ndim: - raise ValueError( - f"Shape of given points should be [*, {self.ndim}], but got {points.shape}" - ) - sdf = self.radius - (((points - self.center) ** 2).sum(axis=1)) ** 0.5 - sdf = -sdf[..., np.newaxis] - return sdf diff --git a/examples/smc_reac/ppsci/geometry/geometry_nd.py b/examples/smc_reac/ppsci/geometry/geometry_nd.py deleted file mode 100644 index 84a8a3edbc..0000000000 --- a/examples/smc_reac/ppsci/geometry/geometry_nd.py +++ /dev/null @@ -1,196 +0,0 @@ -# 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. - -""" -Code below is heavily based on [https://github.com/lululxvi/deepxde](https://github.com/lululxvi/deepxde) -""" -from __future__ import annotations - -import itertools -from typing import Tuple - -import numpy as np -import paddle -from scipy import stats -from sklearn import preprocessing - -from ppsci.geometry import geometry -from ppsci.geometry import sampler -from ppsci.utils import misc - - -class Hypercube(geometry.Geometry): - """Multi-dimensional hyper cube. - - Args: - xmin (Tuple[float, ...]): Lower corner point. - xmax (Tuple[float, ...]): Upper corner point. - - Examples: - >>> import ppsci - >>> geom = ppsci.geometry.Hypercube((0, 0, 0, 0), (1, 1, 1, 1)) - """ - - def __init__(self, xmin: Tuple[float, ...], xmax: Tuple[float, ...]): - if len(xmin) != len(xmax): - raise ValueError("Dimensions of xmin and xmax do not match.") - - self.xmin = np.array(xmin, dtype=paddle.get_default_dtype()) - self.xmax = np.array(xmax, dtype=paddle.get_default_dtype()) - if np.any(self.xmin >= self.xmax): - raise ValueError("xmin >= xmax") - - self.side_length = self.xmax - self.xmin - super().__init__( - len(xmin), (self.xmin, self.xmax), np.linalg.norm(self.side_length) - ) - self.volume = np.prod(self.side_length, dtype=paddle.get_default_dtype()) - - def is_inside(self, x): - return np.logical_and( - np.all(x >= self.xmin, axis=-1), np.all(x <= self.xmax, axis=-1) - ) - - def on_boundary(self, x): - _on_boundary = np.logical_or( - np.any(np.isclose(x, self.xmin), axis=-1), - np.any(np.isclose(x, self.xmax), axis=-1), - ) - return np.logical_and(self.is_inside(x), _on_boundary) - - def boundary_normal(self, x): - _n = -np.isclose(x, self.xmin).astype(paddle.get_default_dtype()) + np.isclose( - x, self.xmax - ) - # For vertices, the normal is averaged for all directions - idx = np.count_nonzero(_n, axis=-1) > 1 - if np.any(idx): - l = np.linalg.norm(_n[idx], axis=-1, keepdims=True) - _n[idx] /= l - return _n - - def uniform_points(self, n, boundary=True): - dx = (self.volume / n) ** (1 / self.ndim) - xi = [] - for i in range(self.ndim): - ni = int(np.ceil(self.side_length[i] / dx)) - if boundary: - xi.append( - np.linspace( - self.xmin[i], - self.xmax[i], - num=ni, - dtype=paddle.get_default_dtype(), - ) - ) - else: - xi.append( - np.linspace( - self.xmin[i], - self.xmax[i], - num=ni + 1, - endpoint=False, - dtype=paddle.get_default_dtype(), - )[1:] - ) - x = np.array(list(itertools.product(*xi)), dtype=paddle.get_default_dtype()) - if len(x) > n: - x = x[0:n] - return x - - def random_points(self, n, random="pseudo"): - x = sampler.sample(n, self.ndim, random) - # print(f"Hypercube's range: {self.__class__.__name__}", self.xmin, self.xmax) - return (self.xmax - self.xmin) * x + self.xmin - - def random_boundary_points(self, n, random="pseudo"): - x = sampler.sample(n, self.ndim, random) - # Randomly pick a dimension - rand_dim = np.random.randint(self.ndim, size=n) - # Replace value of the randomly picked dimension with the nearest boundary value (0 or 1) - x[np.arange(n), rand_dim] = np.round(x[np.arange(n), rand_dim]) - return (self.xmax - self.xmin) * x + self.xmin - - def periodic_point(self, x, component): - y = misc.convert_to_array(x, self.dim_keys) - _on_xmin = np.isclose(y[:, component], self.xmin[component]) - _on_xmax = np.isclose(y[:, component], self.xmax[component]) - y[:, component][_on_xmin] = self.xmax[component] - y[:, component][_on_xmax] = self.xmin[component] - y_normal = self.boundary_normal(y) - - y = misc.convert_to_dict(y, self.dim_keys) - y_normal = misc.convert_to_dict( - y_normal, [f"normal_{k}" for k in self.dim_keys] - ) - return {**y, **y_normal} - - -class Hypersphere(geometry.Geometry): - """Multi-dimensional hyper sphere. - - Args: - center (Tuple[float, ...]): Center point coordinate. - radius (Tuple[float, ...]): Radius along each dimension. - - Examples: - >>> import ppsci - >>> geom = ppsci.geometry.Hypersphere((0, 0, 0, 0), 1.0) - """ - - def __init__(self, center, radius): - self.center = np.array(center, dtype=paddle.get_default_dtype()) - self.radius = radius - super().__init__( - len(center), (self.center - radius, self.center + radius), 2 * radius - ) - - self._r2 = radius**2 - - def is_inside(self, x): - return np.linalg.norm(x - self.center, axis=-1) <= self.radius - - def on_boundary(self, x): - return np.isclose(np.linalg.norm(x - self.center, axis=-1), self.radius) - - def boundary_normal(self, x): - _n = x - self.center - l = np.linalg.norm(_n, axis=-1, keepdims=True) - _n = _n / l * np.isclose(l, self.radius) - return _n - - def random_points(self, n, random="pseudo"): - # https://math.stackexchange.com/questions/87230/picking-random-points-in-the-volume-of-sphere-with-uniform-probability - if random == "pseudo": - U = np.random.rand(n, 1).astype(paddle.get_default_dtype()) - X = np.random.normal(size=(n, self.ndim)).astype(paddle.get_default_dtype()) - else: - rng = sampler.sample(n, self.ndim + 1, random) - U, X = rng[:, 0:1], rng[:, 1:] # Error if X = [0, 0, ...] - X = stats.norm.ppf(X).astype(paddle.get_default_dtype()) - X = preprocessing.normalize(X) - X = U ** (1 / self.ndim) * X - return self.radius * X + self.center - - def random_boundary_points(self, n, random="pseudo"): - # http://mathworld.wolfram.com/HyperspherePointPicking.html - if random == "pseudo": - X = np.random.normal(size=(n, self.ndim)).astype(paddle.get_default_dtype()) - else: - U = sampler.sample( - n, self.ndim, random - ) # Error for [0, 0, ...] or [0.5, 0.5, ...] - X = stats.norm.ppf(U).astype(paddle.get_default_dtype()) - X = preprocessing.normalize(X) - return self.radius * X + self.center diff --git a/examples/smc_reac/ppsci/geometry/inflation.py b/examples/smc_reac/ppsci/geometry/inflation.py deleted file mode 100644 index 198a7cd223..0000000000 --- a/examples/smc_reac/ppsci/geometry/inflation.py +++ /dev/null @@ -1,192 +0,0 @@ -# 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. - -from __future__ import annotations - -import numpy as np -import paddle - -from ppsci.utils import checker - -if not checker.dynamic_import_to_globals(["pymesh", "open3d"]): - raise ModuleNotFoundError - -__all__ = [ - "pymesh_inflation", -] - - -def open3d_inflation( - mesh: open3d.geometry.TriangleMesh, distance: float, direction: int = 1 -) -> open3d.geometry.TriangleMesh: - """Inflate mesh geometry. - - Args: - mesh (open3d.geometry.TriangleMesh): Open3D mesh object. - distance (float): Distance along exterior normal to inflate. - direction (int): 1 for exterior normal, -1 for interior normal. Defaults to 1. - - Returns: - open3d.geometry.TriangleMesh: Inflated mesh. - """ - mesh.remove_duplicated_vertices() - mesh.remove_degenerate_triangles() - mesh.remove_duplicated_triangles() - mesh.remove_unreferenced_vertices() - triangles = np.asarray(mesh.triangles) - points = np.asarray(mesh.vertices) - - remove_ids = [] - for i, point in enumerate(points): - boolean_index = np.argwhere(triangles == i)[:, 0] - if len(boolean_index) < 3: - remove_ids.append(i) - mesh.remove_vertices_by_index(remove_ids) - - points = np.asarray(mesh.vertices, dtype=paddle.get_default_dtype()) - mesh.compute_triangle_normals() - normals = np.asarray(mesh.triangle_normals, dtype=paddle.get_default_dtype()) - mesh.orient_triangles() - triangles = np.asarray(mesh.triangles, dtype=paddle.get_default_dtype()) - new_points = [] - for i, point in enumerate(points): - boolean_index = np.argwhere(triangles == i)[:, 0] - normal = normals[boolean_index] * direction - d = np.ones(len(normal), dtype=paddle.get_default_dtype()) * distance - - new_point = np.linalg.lstsq(normal, d, rcond=None)[0].squeeze() - new_point = point + new_point - if np.linalg.norm(new_point - point) > distance * 2: - # TODO : Find a better way to solve the bad inflation - new_point = point + distance * normal.mean(axis=0) - - new_points.append(new_point) - - new_points = np.array(new_points, dtype=paddle.get_default_dtype()) - new_mesh = open3d.geometry.TriangleMesh( - open3d.utility.Vector3dVector(new_points), - open3d.utility.Vector3iVector(triangles), - ) - - new_mesh.remove_duplicated_vertices() - new_mesh.remove_degenerate_triangles() - new_mesh.remove_duplicated_triangles() - new_mesh.remove_unreferenced_vertices() - new_mesh.compute_triangle_normals() - return new_mesh - - -def pymesh_inflation(mesh: pymesh.Mesh, distance: float) -> pymesh.Mesh: - """Inflate mesh by distance. - - Args: - mesh (pymesh.Mesh): PyMesh object. - distance (float): Inflation distance. - - Returns: - pymesh.Mesh: Inflated mesh. - """ - vertices = np.array(mesh.vertices, dtype=paddle.get_default_dtype()) - faces = np.array(mesh.faces) - open3d_mesh = open3d.geometry.TriangleMesh( - open3d.utility.Vector3dVector(vertices), open3d.utility.Vector3iVector(faces) - ) - inflated_open3d_mesh = open3d_inflation( - open3d_mesh, abs(distance), 1.0 if distance >= 0.0 else -1.0 - ) - vertices = np.array(inflated_open3d_mesh.vertices, dtype=paddle.get_default_dtype()) - faces = np.array(inflated_open3d_mesh.triangles) - inflated_pymesh = pymesh.form_mesh(vertices, faces) - return inflated_pymesh - - -def offset(mesh, distance) -> open3d.geometry.TriangleMesh: - """Offset the 2D mesh - - Args: - mesh (open3d.geometry.TriangleMesh): The mesh to be offset. - distance (float): The distance to offset. - - Returns: - open3d.geometry.TriangleMesh: Result mesh. - """ - # check if the mesh is 2D - mesh.compute_triangle_normals() - normals = np.asarray(mesh.triangle_normals, dtype=paddle.get_default_dtype()) - if not np.allclose(normals[:, :-1], 0): - raise ValueError("The mesh is not 2D") - - mesh.remove_duplicated_vertices() - mesh.remove_degenerate_triangles() - mesh.remove_duplicated_triangles() - mesh.remove_unreferenced_vertices() - triangles = np.asarray(mesh.triangles, dtype=paddle.get_default_dtype()) - - edges = np.vstack( - [triangles[:, [0, 1]], triangles[:, [1, 2]], triangles[:, [2, 0]]] - ) - edges = set(map(tuple, edges)) - edges = np.array(list(edges)) - - vertices = np.asarray(mesh.vertices, dtype=paddle.get_default_dtype())[:, :-1] - edges_in_triangle = np.array( - [ - np.intersect1d( - np.argwhere(triangles == edge[0])[:, 0], - np.argwhere(triangles == edge[1])[:, 0], - ) - for edge in edges - ], - dtype=object, - ) - surface_edges = edges[[len(i) == 1 for i in edges_in_triangle]] - edges_in_triangle = [i for i in edges_in_triangle if len(i) == 1] - - edges_normals = [] - for edge, triangle in zip(surface_edges, edges_in_triangle): - triangle = triangles[triangle].squeeze() - other_point = vertices[np.setdiff1d(triangle, edge)].squeeze() - edge = vertices[edge] - u = (other_point[0] - edge[0][0]) * (edge[0][0] - edge[1][0]) + ( - other_point[1] - edge[0][1] - ) * (edge[0][1] - edge[1][1]) - u = u / np.sum((edge[0] - edge[1]) ** 2) - edge_normal = edge[0] + u * (edge[0] - edge[1]) - edge_normal = edge_normal - other_point - edges_normals.append(edge_normal) - - edges_normals = np.array(edges_normals, dtype=paddle.get_default_dtype()) - edges_normals = edges_normals / np.linalg.norm(edges_normals, axis=1)[:, None] - - new_mesh = open3d.geometry.TriangleMesh() - new_vertices = [] - for point in set(surface_edges.reshape(-1)): - index = np.argwhere(surface_edges == point)[:, 0] - normal = edges_normals[index] - d = np.ones(len(index), dtype=paddle.get_default_dtype()) * distance - new_point = np.linalg.lstsq(normal, d, rcond=None)[0] - new_point = vertices[point] + new_point - new_vertices.append(new_point) - - new_vertices = np.hstack( - ( - np.array(new_vertices, dtype=paddle.get_default_dtype()), - np.zeros((len(new_vertices), 1), dtype=paddle.get_default_dtype()), - ) - ) - new_mesh.vertices = open3d.utility.Vector3dVector(new_vertices) - new_mesh.triangles = open3d.utility.Vector3iVector(triangles) - new_mesh.compute_triangle_normals() - new_mesh.compute_vertex_normals() - return new_mesh diff --git a/examples/smc_reac/ppsci/geometry/mesh.py b/examples/smc_reac/ppsci/geometry/mesh.py deleted file mode 100644 index 8a363ca36b..0000000000 --- a/examples/smc_reac/ppsci/geometry/mesh.py +++ /dev/null @@ -1,1392 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import TYPE_CHECKING -from typing import Callable -from typing import Dict -from typing import Optional -from typing import Tuple -from typing import Union - -import numpy as np -import paddle - -try: - from stl import mesh as np_mesh_module -except ModuleNotFoundError: - pass -except ImportError: - pass - -from typing_extensions import Literal - -from ppsci.geometry import geometry -from ppsci.geometry import geometry_3d -from ppsci.geometry import sampler -from ppsci.geometry import sdf as sdf_module -from ppsci.utils import checker -from ppsci.utils import misc - -if TYPE_CHECKING: - import pymesh - - -class Mesh(geometry.Geometry): - """Class for mesh geometry. - - Args: - mesh (Union[str, Mesh]): Mesh file path or mesh object, such as "/path/to/mesh.stl". - - Examples: - >>> import ppsci - >>> geom = ppsci.geometry.Mesh("/path/to/mesh.stl") # doctest: +SKIP - """ - - def __init__(self, mesh: Union["pymesh.Mesh", str]): - # check if pymesh is installed when using Mesh Class - if not checker.dynamic_import_to_globals(["pymesh"]): - raise ImportError( - "Could not import pymesh python package." - "Please install it as https://pymesh.readthedocs.io/en/latest/installation.html." - ) - import pymesh - - if isinstance(mesh, str): - self.py_mesh = pymesh.meshio.load_mesh(mesh) - elif isinstance(mesh, pymesh.Mesh): - self.py_mesh = mesh - else: - raise ValueError("arg `mesh` should be path string or `pymesh.Mesh`") - - self.init_mesh() - - @classmethod - def from_pymesh(cls, mesh: "pymesh.Mesh") -> "Mesh": - """Instantiate Mesh object with given PyMesh object. - - Args: - mesh (pymesh.Mesh): PyMesh object. - - Returns: - Mesh: Instantiated ppsci.geometry.Mesh object. - - Examples: - >>> import ppsci - >>> import pymesh # doctest: +SKIP - >>> import numpy as np # doctest: +SKIP - >>> box = pymesh.generate_box_mesh(np.array([0, 0, 0]), np.array([1, 1, 1])) # doctest: +SKIP - >>> mesh = ppsci.geometry.Mesh.from_pymesh(box) # doctest: +SKIP - >>> print(mesh.vertices) # doctest: +SKIP - [[0. 0. 0.] - [1. 0. 0.] - [1. 1. 0.] - [0. 1. 0.] - [0. 0. 1.] - [1. 0. 1.] - [1. 1. 1.] - [0. 1. 1.]] - """ - # check if pymesh is installed when using Mesh Class - if not checker.dynamic_import_to_globals(["pymesh"]): - raise ImportError( - "Could not import pymesh python package." - "Please install it as https://pymesh.readthedocs.io/en/latest/installation.html." - ) - import pymesh - - if isinstance(mesh, pymesh.Mesh): - return cls(mesh) - else: - raise ValueError( - f"arg `mesh` should be type of `pymesh.Mesh`, but got {type(mesh)}" - ) - - def init_mesh(self): - """Initialize necessary variables for mesh""" - if "face_normal" not in self.py_mesh.get_attribute_names(): - self.py_mesh.add_attribute("face_normal") - self.face_normal = self.py_mesh.get_attribute("face_normal").reshape([-1, 3]) - - if not checker.dynamic_import_to_globals(["open3d"]): - raise ImportError( - "Could not import open3d python package. " - "Please install it with `pip install open3d`." - ) - import open3d - - self.open3d_mesh = open3d.geometry.TriangleMesh( - open3d.utility.Vector3dVector(np.array(self.py_mesh.vertices)), - open3d.utility.Vector3iVector(np.array(self.py_mesh.faces)), - ) - self.open3d_mesh.compute_vertex_normals() - - self.vertices = self.py_mesh.vertices - self.faces = self.py_mesh.faces - self.vectors = self.vertices[self.faces] - super().__init__( - self.vertices.shape[-1], - (np.amin(self.vertices, axis=0), np.amax(self.vertices, axis=0)), - np.inf, - ) - self.v0 = self.vectors[:, 0] - self.v1 = self.vectors[:, 1] - self.v2 = self.vectors[:, 2] - self.num_vertices = self.py_mesh.num_vertices - self.num_faces = self.py_mesh.num_faces - - if not checker.dynamic_import_to_globals(["pysdf"]): - raise ImportError( - "Could not import pysdf python package. " - "Please install open3d with `pip install pysdf`." - ) - import pysdf - - self.pysdf = pysdf.SDF(self.vertices, self.faces) - self.bounds = ( - ((np.min(self.vectors[:, :, 0])), np.max(self.vectors[:, :, 0])), - ((np.min(self.vectors[:, :, 1])), np.max(self.vectors[:, :, 1])), - ((np.min(self.vectors[:, :, 2])), np.max(self.vectors[:, :, 2])), - ) - - def sdf_func(self, points: np.ndarray) -> np.ndarray: - """Compute signed distance field. - - Args: - points (np.ndarray): The coordinate points used to calculate the SDF value, - the shape is [N, 3] - - Returns: - np.ndarray: SDF values of input points without squared, the shape is [N, 1]. - - NOTE: This function usually returns ndarray with negative values, because - according to the definition of SDF, the SDF value of the coordinate point inside - the object(interior points) is negative, the outside is positive, and the edge - is 0. Therefore, when used for weighting, a negative sign is often added before - the result of this function. - """ - if not checker.dynamic_import_to_globals(["pymesh"]): - raise ImportError( - "Could not import pymesh python package." - "Please install it as https://pymesh.readthedocs.io/en/latest/installation.html." - ) - import pymesh - - sdf, _, _, _ = pymesh.signed_distance_to_mesh(self.py_mesh, points) - sdf = sdf[..., np.newaxis].astype(paddle.get_default_dtype()) - return sdf - - def is_inside(self, x): - # NOTE: point on boundary is included - return self.pysdf.contains(x) - - def on_boundary(self, x): - return np.isclose(self.sdf_func(x), 0.0).ravel() - - def translate(self, translation: np.ndarray, relative: bool = True) -> "Mesh": - """Translate by given offsets. - - NOTE: This API generate a completely new Mesh object with translated geometry, - without modifying original Mesh object inplace. - - Args: - translation (np.ndarray): Translation offsets, numpy array of shape (3,): - [offset_x, offset_y, offset_z]. - relative (bool, optional): Whether translate relatively. Defaults to True. - - Returns: - Mesh: Translated Mesh object. - - Examples: - >>> import ppsci - >>> import pymesh # doctest: +SKIP - >>> import numpy as np - >>> box = pymesh.generate_box_mesh(np.array([0, 0, 0]), np.array([1, 1, 1])) # doctest: +SKIP - >>> mesh = ppsci.geometry.Mesh(box) # doctest: +SKIP - >>> print(mesh.vertices) # doctest: +SKIP - [[0. 0. 0.] - [1. 0. 0.] - [1. 1. 0.] - [0. 1. 0.] - [0. 0. 1.] - [1. 0. 1.] - [1. 1. 1.] - [0. 1. 1.]] - >>> print(mesh.translate((-0.5, 0, 0.5), False).vertices) # the center is moved to the translation vector. # doctest: +SKIP - [[-1. -0.5 0. ] - [ 0. -0.5 0. ] - [ 0. 0.5 0. ] - [-1. 0.5 0. ] - [-1. -0.5 1. ] - [ 0. -0.5 1. ] - [ 0. 0.5 1. ] - [-1. 0.5 1. ]] - >>> print(mesh.translate((-0.5, 0, 0.5), True).vertices) # the translation vector is directly added to the geometry coordinates # doctest: +SKIP - [[-0.5 0. 0.5] - [ 0.5 0. 0.5] - [ 0.5 1. 0.5] - [-0.5 1. 0.5] - [-0.5 0. 1.5] - [ 0.5 0. 1.5] - [ 0.5 1. 1.5] - [-0.5 1. 1.5]] - """ - vertices = np.array(self.vertices, dtype=paddle.get_default_dtype()) - faces = np.array(self.faces) - - if not checker.dynamic_import_to_globals(("open3d", "pymesh")): - raise ImportError( - "Could not import open3d and pymesh python package. " - "Please install open3d with `pip install open3d` and " - "pymesh as https://paddlescience-docs.readthedocs.io/zh/latest/zh/install_setup/#__tabbed_4_1" - ) - import open3d # isort:skip - import pymesh # isort:skip - - open3d_mesh = open3d.geometry.TriangleMesh( - open3d.utility.Vector3dVector(vertices), - open3d.utility.Vector3iVector(faces), - ) - open3d_mesh = open3d_mesh.translate(translation, relative) - translated_mesh = pymesh.form_mesh( - np.asarray(open3d_mesh.vertices, dtype=paddle.get_default_dtype()), faces - ) - # Generate a new Mesh object using class method - return Mesh.from_pymesh(translated_mesh) - - def scale( - self, scale: float, center: Tuple[float, float, float] = (0, 0, 0) - ) -> "Mesh": - """Scale by given scale coefficient and center coordinate. - - NOTE: This API generate a completely new Mesh object with scaled geometry, - without modifying original Mesh object inplace. - - Args: - scale (float): Scale coefficient. - center (Tuple[float,float,float], optional): Center coordinate, [x, y, z]. - Defaults to (0, 0, 0). - - Returns: - Mesh: Scaled Mesh object. - - Examples: - >>> import ppsci - >>> import pymesh # doctest: +SKIP - >>> import numpy as np - >>> box = pymesh.generate_box_mesh(np.array([0, 0, 0]), np.array([1, 1, 1])) # doctest: +SKIP - >>> mesh = ppsci.geometry.Mesh(box) # doctest: +SKIP - >>> print(mesh.vertices) # doctest: +SKIP - [[0. 0. 0.] - [1. 0. 0.] - [1. 1. 0.] - [0. 1. 0.] - [0. 0. 1.] - [1. 0. 1.] - [1. 1. 1.] - [0. 1. 1.]] - >>> mesh = mesh.scale(2, (0.25, 0.5, 0.75)) # doctest: +SKIP - >>> print(mesh.vertices) # doctest: +SKIP - [[-0.25 -0.5 -0.75] - [ 1.75 -0.5 -0.75] - [ 1.75 1.5 -0.75] - [-0.25 1.5 -0.75] - [-0.25 -0.5 1.25] - [ 1.75 -0.5 1.25] - [ 1.75 1.5 1.25] - [-0.25 1.5 1.25]] - """ - vertices = np.array(self.vertices, dtype=paddle.get_default_dtype()) - faces = np.array(self.faces, dtype=paddle.get_default_dtype()) - - if not checker.dynamic_import_to_globals(("open3d", "pymesh")): - raise ImportError( - "Could not import open3d and pymesh python package. " - "Please install open3d with `pip install open3d` and " - "pymesh as https://pymesh.readthedocs.io/en/latest/installation.html." - ) - import open3d # isort:skip - import pymesh # isort:skip - - open3d_mesh = open3d.geometry.TriangleMesh( - open3d.utility.Vector3dVector(vertices), - open3d.utility.Vector3iVector(faces), - ) - open3d_mesh = open3d_mesh.scale(scale, center) - scaled_pymesh = pymesh.form_mesh( - np.asarray(open3d_mesh.vertices, dtype=paddle.get_default_dtype()), faces - ) - # Generate a new Mesh object using class method - return Mesh.from_pymesh(scaled_pymesh) - - def uniform_boundary_points(self, n: int): - """Compute the equi-spaced points on the boundary.""" - return self.pysdf.sample_surface(n) - - def inflated_random_points(self, n, distance, random="pseudo", criteria=None): - if not isinstance(n, (tuple, list)): - n = [n] - if not isinstance(distance, (tuple, list)): - distance = [distance] - if len(n) != len(distance): - raise ValueError( - f"len(n)({len(n)}) should be equal to len(distance)({len(distance)})" - ) - - from ppsci.geometry import inflation - - all_points = [] - all_areas = [] - for _n, _dist in zip(n, distance): - inflated_mesh = Mesh(inflation.pymesh_inflation(self.py_mesh, _dist)) - points, areas = inflated_mesh.random_points(_n, random, criteria) - all_points.append(points) - all_areas.append(areas) - - all_points = np.concatenate(all_points, axis=0) - all_areas = np.concatenate(all_areas, axis=0) - return all_points, all_areas - - def _approximate_area( - self, - random: Literal["pseudo"] = "pseudo", - criteria: Optional[Callable] = None, - n_appr: int = 10000, - ) -> float: - """Approximate area with given `criteria` and `n_appr` points by Monte Carlo - algorithm. - - Args: - random (str, optional): Random method. Defaults to "pseudo". - criteria (Optional[Callable]): Criteria function. Defaults to None. - n_appr (int): Number of points for approximating area. Defaults to 10000. - - Returns: - float: Approximation area with given criteria. - """ - triangle_areas = area_of_triangles(self.v0, self.v1, self.v2) - triangle_probabilities = triangle_areas / np.linalg.norm(triangle_areas, ord=1) - triangle_index = np.arange(triangle_probabilities.shape[0]) - npoint_per_triangle = np.random.choice( - triangle_index, n_appr, p=triangle_probabilities - ) - npoint_per_triangle, _ = np.histogram( - npoint_per_triangle, - np.arange(triangle_probabilities.shape[0] + 1) - 0.5, - ) - - appr_areas = [] - if criteria is not None: - aux_points = [] - - for i, npoint in enumerate(npoint_per_triangle): - if npoint == 0: - continue - # sample points for computing criteria mask if criteria is given - if criteria is not None: - points_at_triangle_i = sample_in_triangle( - self.v0[i], self.v1[i], self.v2[i], npoint, random - ) - aux_points.append(points_at_triangle_i) - - appr_areas.append( - np.full( - (npoint, 1), triangle_areas[i] / npoint, paddle.get_default_dtype() - ) - ) - appr_areas = np.concatenate(appr_areas, axis=0) # [n_appr, 1] - - # set invalid area to 0 by computing criteria mask with auxiliary points - if criteria is not None: - aux_points = np.concatenate(aux_points, axis=0) # [n_appr, 3] - criteria_mask = criteria(*np.split(aux_points, self.ndim, 1)) - appr_areas *= criteria_mask - return appr_areas.sum() - - def random_boundary_points(self, n, random="pseudo"): - triangle_area = area_of_triangles(self.v0, self.v1, self.v2) - triangle_prob = triangle_area / np.linalg.norm(triangle_area, ord=1) - npoint_per_triangle = np.random.choice( - np.arange(len(triangle_prob)), n, p=triangle_prob - ) - npoint_per_triangle, _ = np.histogram( - npoint_per_triangle, np.arange(len(triangle_prob) + 1) - 0.5 - ) - - points = [] - normal = [] - areas = [] - for i, npoint in enumerate(npoint_per_triangle): - if npoint == 0: - continue - points_at_triangle_i = sample_in_triangle( - self.v0[i], self.v1[i], self.v2[i], npoint, random - ) - normal_at_triangle_i = np.tile(self.face_normal[i], (npoint, 1)).astype( - paddle.get_default_dtype() - ) - areas_at_triangle_i = np.full( - (npoint, 1), - triangle_area[i] / npoint, - dtype=paddle.get_default_dtype(), - ) - - points.append(points_at_triangle_i) - normal.append(normal_at_triangle_i) - areas.append(areas_at_triangle_i) - - points = np.concatenate(points, axis=0) - normal = np.concatenate(normal, axis=0) - areas = np.concatenate(areas, axis=0) - - return points, normal, areas - - def sample_boundary( - self, - n: int, - random: Literal["pseudo"] = "pseudo", - criteria: Optional[Callable[..., np.ndarray]] = None, - evenly: bool = False, - inflation_dist: Union[float, Tuple[float, ...]] = None, - ) -> Dict[str, np.ndarray]: - # TODO(sensen): Support for time-dependent points(repeat data in time) - if inflation_dist is not None: - if not isinstance(n, (tuple, list)): - n = [n] - if not isinstance(inflation_dist, (tuple, list)): - inflation_dist = [inflation_dist] - if len(n) != len(inflation_dist): - raise ValueError( - f"len(n)({len(n)}) should be equal to len(inflation_dist)({len(inflation_dist)})" - ) - - from ppsci.geometry import inflation - - inflated_data_dict = {} - for _n, _dist in zip(n, inflation_dist): - # 1. manually inflate mesh at first - inflated_mesh = Mesh(inflation.pymesh_inflation(self.py_mesh, _dist)) - # 2. compute all data by sample_boundary with `inflation_dist=None` - data_dict = inflated_mesh.sample_boundary( - _n, - random, - criteria, - evenly, - inflation_dist=None, - ) - for key, value in data_dict.items(): - if key not in inflated_data_dict: - inflated_data_dict[key] = value - else: - inflated_data_dict[key] = np.concatenate( - (inflated_data_dict[key], value), axis=0 - ) - return inflated_data_dict - else: - if evenly: - raise ValueError( - "Can't sample evenly on mesh now, please set evenly=False." - ) - _size, _ntry, _nsuc = 0, 0, 0 - all_points = [] - all_normal = [] - while _size < n: - points, normal, _ = self.random_boundary_points(n, random) - if criteria is not None: - criteria_mask = criteria( - *np.split(points, self.ndim, axis=1) - ).ravel() - points = points[criteria_mask] - normal = normal[criteria_mask] - - if len(points) > n - _size: - points = points[: n - _size] - normal = normal[: n - _size] - - all_points.append(points) - all_normal.append(normal) - - _size += len(points) - _ntry += 1 - if len(points) > 0: - _nsuc += 1 - - if _ntry >= 1000 and _nsuc == 0: - raise ValueError( - "Sample boundary points failed, " - "please check correctness of geometry and given criteria." - ) - - all_points = np.concatenate(all_points, axis=0) - all_normal = np.concatenate(all_normal, axis=0) - appr_area = self._approximate_area(random, criteria) - all_areas = np.full((n, 1), appr_area / n, paddle.get_default_dtype()) - - x_dict = misc.convert_to_dict(all_points, self.dim_keys) - normal_dict = misc.convert_to_dict( - all_normal, [f"normal_{key}" for key in self.dim_keys if key != "t"] - ) - area_dict = misc.convert_to_dict(all_areas, ["area"]) - return {**x_dict, **normal_dict, **area_dict} - - def random_points(self, n, random="pseudo", criteria=None): - _size = 0 - all_points = [] - cuboid = geometry_3d.Cuboid( - [bound[0] for bound in self.bounds], - [bound[1] for bound in self.bounds], - ) - _nsample, _nvalid = 0, 0 - while _size < n: - random_points = cuboid.random_points(n, random) - valid_mask = self.is_inside(random_points) - - if criteria: - valid_mask &= criteria( - *np.split(random_points, self.ndim, axis=1) - ).ravel() - valid_points = random_points[valid_mask] - _nvalid += len(valid_points) - - if len(valid_points) > n - _size: - valid_points = valid_points[: n - _size] - - all_points.append(valid_points) - _size += len(valid_points) - _nsample += n - - all_points = np.concatenate(all_points, axis=0) - cuboid_volume = np.prod([b[1] - b[0] for b in self.bounds]) - all_areas = np.full( - (n, 1), cuboid_volume * (_nvalid / _nsample) / n, paddle.get_default_dtype() - ) - return all_points, all_areas - - def sample_interior( - self, - n: int, - random: Literal["pseudo"] = "pseudo", - criteria: Optional[Callable[..., np.ndarray]] = None, - evenly: bool = False, - compute_sdf_derivatives: bool = False, - ): - """Sample random points in the geometry and return those meet criteria.""" - if evenly: - # TODO(sensen): Implement uniform sample for mesh interior. - raise NotImplementedError( - "uniformly sample for interior in mesh is not support yet, " - "you may need to set evenly=False in config dict of constraint" - ) - points, areas = self.random_points(n, random, criteria) - - x_dict = misc.convert_to_dict(points, self.dim_keys) - area_dict = misc.convert_to_dict(areas, ("area",)) - - # NOTE: add negative to the sdf values because weight should be positive. - sdf = -self.sdf_func(points) - sdf_dict = misc.convert_to_dict(sdf, ("sdf",)) - - sdf_derives_dict = {} - if compute_sdf_derivatives: - sdf_derives = -self.sdf_derivatives(points) - sdf_derives_dict = misc.convert_to_dict( - sdf_derives, tuple(f"sdf__{key}" for key in self.dim_keys) - ) - - return {**x_dict, **area_dict, **sdf_dict, **sdf_derives_dict} - - def union(self, other: "Mesh"): - if not checker.dynamic_import_to_globals(["pymesh"]): - raise ImportError( - "Could not import pymesh python package. " - "Please install it as https://pymesh.readthedocs.io/en/latest/installation.html." - ) - import pymesh - - csg = pymesh.CSGTree( - {"union": [{"mesh": self.py_mesh}, {"mesh": other.py_mesh}]} - ) - return Mesh(csg.mesh) - - def __or__(self, other: "Mesh"): - return self.union(other) - - def __add__(self, other: "Mesh"): - return self.union(other) - - def difference(self, other: "Mesh"): - if not checker.dynamic_import_to_globals(["pymesh"]): - raise ImportError( - "Could not import pymesh python package. " - "Please install it as https://pymesh.readthedocs.io/en/latest/installation.html." - ) - import pymesh - - csg = pymesh.CSGTree( - {"difference": [{"mesh": self.py_mesh}, {"mesh": other.py_mesh}]} - ) - return Mesh(csg.mesh) - - def __sub__(self, other: "Mesh"): - return self.difference(other) - - def intersection(self, other: "Mesh"): - if not checker.dynamic_import_to_globals(["pymesh"]): - raise ImportError( - "Could not import pymesh python package. " - "Please install it as https://pymesh.readthedocs.io/en/latest/installation.html." - ) - import pymesh - - csg = pymesh.CSGTree( - {"intersection": [{"mesh": self.py_mesh}, {"mesh": other.py_mesh}]} - ) - return Mesh(csg.mesh) - - def __and__(self, other: "Mesh"): - return self.intersection(other) - - def __str__(self) -> str: - """Return the name of class""" - return ", ".join( - [ - self.__class__.__name__, - f"num_vertices = {self.num_vertices}", - f"num_faces = {self.num_faces}", - f"bounds = {self.bounds}", - f"dim_keys = {self.dim_keys}", - ] - ) - - -class SDFMesh(geometry.Geometry): - """Class for SDF geometry, a kind of implicit surface mesh. - - Args: - vectors (np.ndarray): Vectors of triangles of mesh with shape [M, 3, 3]. - normals (np.ndarray): Unit normals of each triangle face with shape [M, 3]. - sdf_func (Callable[[np.ndarray, bool], np.ndarray]): Signed distance function - of the triangle mesh. - - Examples: - >>> import ppsci - >>> geom = ppsci.geometry.SDFMesh.from_stl("/path/to/mesh.stl") # doctest: +SKIP - """ - - eps = 1e-6 - - def __init__( - self, - vectors: np.ndarray, - normals: np.ndarray, - sdf_func: Callable[[np.ndarray, bool], np.ndarray], - ): - if vectors.shape[1:] != (3, 3): - raise ValueError( - f"The shape of `vectors` must be [M, 3, 3], but got {vectors.shape}" - ) - if normals.shape[1] != 3: - raise ValueError( - f"The shape of `normals` must be [M, 3], but got {normals.shape}" - ) - self.vectors = vectors - self.face_normal = normals - self.sdf_func = sdf_func # overwrite sdf_func - self.bounds = ( - ((np.min(self.vectors[:, :, 0])), np.max(self.vectors[:, :, 0])), - ((np.min(self.vectors[:, :, 1])), np.max(self.vectors[:, :, 1])), - ((np.min(self.vectors[:, :, 2])), np.max(self.vectors[:, :, 2])), - ) - self.ndim = 3 - super().__init__( - self.vectors.shape[-1], - (np.amin(self.vectors, axis=(0, 1)), np.amax(self.vectors, axis=(0, 1))), - np.inf, - ) - - @property - def v0(self) -> np.ndarray: - return self.vectors[:, 0] - - @property - def v1(self) -> np.ndarray: - return self.vectors[:, 1] - - @property - def v2(self) -> np.ndarray: - return self.vectors[:, 2] - - @classmethod - def from_stl(cls, mesh_file: str) -> "SDFMesh": - """Instantiate SDFMesh from given mesh file. - - Args: - mesh_file (str): Path to triangle mesh file. - - Returns: - SDFMesh: Instantiated ppsci.geometry.SDFMesh object. - - Examples: - >>> import ppsci - >>> import pymesh # doctest: +SKIP - >>> import numpy as np # doctest: +SKIP - >>> box = pymesh.generate_box_mesh(np.array([0, 0, 0]), np.array([1, 1, 1])) # doctest: +SKIP - >>> pymesh.save_mesh("box.stl", box) # doctest: +SKIP - >>> mesh = ppsci.geometry.SDFMesh.from_stl("box.stl") # doctest: +SKIP - >>> print(sdfmesh.vectors.shape) # doctest: +SKIP - (12, 3, 3) - """ - # check if pymesh is installed when using Mesh Class - if not checker.dynamic_import_to_globals(["stl"]): - raise ImportError( - "Could not import stl python package. " - "Please install numpy-stl with: pip install 'numpy-stl>=2.16,<2.17'" - ) - - np_mesh_obj = np_mesh_module.Mesh.from_file(mesh_file) - return cls( - np_mesh_obj.vectors, - np_mesh_obj.get_unit_normals(), - make_sdf(np_mesh_obj.vectors), - ) - - def sdf_func( - self, points: np.ndarray, compute_sdf_derivatives: bool = False - ) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]: - """Compute signed distance field. - - Args: - points (np.ndarray): The coordinate points used to calculate the SDF value, - the shape is [N, 3] - compute_sdf_derivatives (bool): Whether to compute SDF derivatives. - Defaults to False. - - Returns: - Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]: - If compute_sdf_derivatives is True, then return both SDF values([N, 1]) - and their derivatives([N, 3]); otherwise only return SDF values([N, 1]). - - NOTE: This function usually returns ndarray with negative values, because - according to the definition of SDF, the SDF value of the coordinate point inside - the object(interior points) is negative, the outside is positive, and the edge - is 0. Therefore, when used for weighting, a negative sign is often added before - the result of this function. - """ - # normalize triangles - x_min, y_min, z_min = np.min(points, axis=0) - x_max, y_max, z_max = np.max(points, axis=0) - max_dis = max(max((x_max - x_min), (y_max - y_min)), (z_max - z_min)) - store_triangles = np.array(self.vectors, dtype=np.float64) - store_triangles[:, :, 0] -= x_min - store_triangles[:, :, 1] -= y_min - store_triangles[:, :, 2] -= z_min - store_triangles *= 1 / max_dis - store_triangles = store_triangles.reshape([-1, 3]) - - # normalize query points - points = points.copy() - points[:, 0] -= x_min - points[:, 1] -= y_min - points[:, 2] -= z_min - points *= 1 / max_dis - points = points.astype(np.float64).ravel() - - # compute sdf values for query points - sdf = sdf_module.signed_distance_field( - store_triangles, - np.arange((store_triangles.shape[0])), - points, - include_hit_points=compute_sdf_derivatives, - ) - if compute_sdf_derivatives: - sdf, hit_points = sdf - - sdf = sdf.numpy() # [N] - sdf = np.expand_dims(max_dis * sdf, axis=1) # [N, 1] - - if compute_sdf_derivatives: - hit_points = hit_points.numpy() # [N, 3] - # Gradient of SDF is the unit vector from the query point to the hit point. - sdf_derives = hit_points - points - sdf_derives /= np.linalg.norm(sdf_derives, axis=1, keepdims=True) - return sdf, sdf_derives - - return sdf - - def is_inside(self, x): - # NOTE: point on boundary is included - return np.less(self.sdf_func(x), 0.0).ravel() - - def on_boundary(self, x: np.ndarray, normal: np.ndarray) -> np.ndarray: - x_plus = x + self.eps * normal - x_minus = x - self.eps * normal - - sdf_x_plus = self.sdf_func(x_plus) - sdf_x_minus = self.sdf_func(x_minus) - mask_on_boundary = np.less_equal(sdf_x_plus * sdf_x_minus, 0) - return mask_on_boundary.ravel() - - def translate(self, translation: np.ndarray) -> "SDFMesh": - """Translate by given offsets. - - NOTE: This API generate a completely new Mesh object with translated geometry, - without modifying original Mesh object inplace. - - Args: - translation (np.ndarray): Translation offsets, numpy array of shape (3,): - [offset_x, offset_y, offset_z]. - - Returns: - Mesh: Translated Mesh object. - - Examples: - >>> import ppsci - >>> import pymesh # doctest: +SKIP - >>> mesh = ppsci.geometry.SDFMesh.from_stl('/path/to/mesh.stl') # doctest: +SKIP - >>> mesh = mesh.translate(np.array([1, -1, 2])) # doctest: +SKIP - """ - new_vectors = self.vectors + translation.reshape([1, 1, 3]) - - return SDFMesh( - new_vectors, - self.face_normal, - make_sdf(new_vectors), - ) - - def scale(self, scale: float) -> "SDFMesh": - """Scale by given scale coefficient and center coordinate. - - NOTE: This API generate a completely new Mesh object with scaled geometry, - without modifying original Mesh object inplace. - - Args: - scale (float): Scale coefficient. - - Returns: - Mesh: Scaled Mesh object. - - Examples: - >>> import ppsci - >>> import pymesh # doctest: +SKIP - >>> mesh = ppsci.geometry.SDFMesh.from_stl('/path/to/mesh.stl') # doctest: +SKIP - >>> mesh = mesh.scale(np.array([1.3, 1.5, 2.0])) # doctest: +SKIP - """ - new_vectors = self.vectors * scale - return SDFMesh( - new_vectors, - self.face_normal, - make_sdf(new_vectors), - ) - - def uniform_boundary_points(self, n: int): - """Compute the equi-spaced points on the boundary.""" - raise NotImplementedError( - "'uniform_boundary_points' is not available in SDFMesh." - ) - - def inflated_random_points(self, n, distance, random="pseudo", criteria=None): - raise NotImplementedError( - "'inflated_random_points' is not available in SDFMesh." - ) - - def _approximate_area( - self, - random: Literal["pseudo"] = "pseudo", - criteria: Optional[Callable] = None, - n_appr: int = 10000, - ) -> float: - """Approximate area with given `criteria` and `n_appr` points by Monte Carlo - algorithm. - - Args: - random (str, optional): Random method. Defaults to "pseudo". - criteria (Optional[Callable]): Criteria function. Defaults to None. - n_appr (int): Number of points for approximating area. Defaults to 10000. - - Returns: - float: Approximation area with given criteria. - """ - triangle_areas = area_of_triangles(self.v0, self.v1, self.v2) - triangle_probabilities = triangle_areas / np.linalg.norm(triangle_areas, ord=1) - triangle_index = np.arange(triangle_probabilities.shape[0]) - npoint_per_triangle = np.random.choice( - triangle_index, n_appr, p=triangle_probabilities - ) - npoint_per_triangle, _ = np.histogram( - npoint_per_triangle, - np.arange(triangle_probabilities.shape[0] + 1) - 0.5, - ) - - aux_points = [] - aux_normals = [] - appr_areas = [] - - for i, npoint in enumerate(npoint_per_triangle): - if npoint == 0: - continue - # sample points for computing criteria mask if criteria is given - points_at_triangle_i = sample_in_triangle( - self.v0[i], self.v1[i], self.v2[i], npoint, random - ) - normal_at_triangle_i = np.tile( - self.face_normal[i].reshape(1, 3), (npoint, 1) - ) - aux_points.append(points_at_triangle_i) - aux_normals.append(normal_at_triangle_i) - appr_areas.append( - np.full( - (npoint, 1), triangle_areas[i] / npoint, paddle.get_default_dtype() - ) - ) - - aux_points = np.concatenate(aux_points, axis=0) # [n_appr, 3] - aux_normals = np.concatenate(aux_normals, axis=0) # [n_appr, 3] - appr_areas = np.concatenate(appr_areas, axis=0) # [n_appr, 1] - valid_mask = self.on_boundary(aux_points, aux_normals)[:, None] - # set invalid area to 0 by computing criteria mask with auxiliary points - if criteria is not None: - criteria_mask = criteria(*np.split(aux_points, self.ndim, 1)) - assert valid_mask.shape == criteria_mask.shape - valid_mask = np.logical_and(valid_mask, criteria_mask) - - appr_areas *= valid_mask - - return appr_areas.sum() - - def random_boundary_points( - self, n, random="pseudo" - ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: - triangle_area = area_of_triangles(self.v0, self.v1, self.v2) - triangle_prob = triangle_area / np.linalg.norm(triangle_area, ord=1) - npoint_per_triangle = np.random.choice( - np.arange(len(triangle_prob)), n, p=triangle_prob - ) - npoint_per_triangle, _ = np.histogram( - npoint_per_triangle, np.arange(len(triangle_prob) + 1) - 0.5 - ) - - points = [] - normal = [] - areas = [] - for i, npoint in enumerate(npoint_per_triangle): - if npoint == 0: - continue - points_at_triangle_i = sample_in_triangle( - self.v0[i], self.v1[i], self.v2[i], npoint, random - ) - normal_at_triangle_i = np.tile(self.face_normal[i], (npoint, 1)).astype( - paddle.get_default_dtype() - ) - areas_at_triangle_i = np.full( - (npoint, 1), - triangle_area[i] / npoint, - dtype=paddle.get_default_dtype(), - ) - - points.append(points_at_triangle_i) - normal.append(normal_at_triangle_i) - areas.append(areas_at_triangle_i) - - points = np.concatenate(points, axis=0) - normal = np.concatenate(normal, axis=0) - areas = np.concatenate(areas, axis=0) - - return points, normal, areas - - def sample_boundary( - self, - n: int, - random: Literal["pseudo"] = "pseudo", - criteria: Optional[Callable[..., np.ndarray]] = None, - evenly: bool = False, - inflation_dist: Union[float, Tuple[float, ...]] = None, - ) -> Dict[str, np.ndarray]: - # TODO(sensen): Support for time-dependent points(repeat data in time) - if inflation_dist is not None: - raise NotImplementedError("Not implemented yet") - else: - if evenly: - raise ValueError( - "Can't sample evenly on mesh now, please set evenly=False." - ) - _size, _ntry, _nsuc = 0, 0, 0 - all_points = [] - all_normal = [] - while _size < n: - points, normal, _ = self.random_boundary_points(n, random) - valid_mask = self.on_boundary(points, normal) - - if criteria is not None: - criteria_mask = criteria( - *np.split(points, self.ndim, axis=1) - ).ravel() - assert valid_mask.shape == criteria_mask.shape - valid_mask = np.logical_and(valid_mask, criteria_mask) - - points = points[valid_mask] - normal = normal[valid_mask] - - if len(points) > n - _size: - points = points[: n - _size] - normal = normal[: n - _size] - - all_points.append(points) - all_normal.append(normal) - - _size += len(points) - _ntry += 1 - if len(points) > 0: - _nsuc += 1 - - if _ntry >= 1000 and _nsuc == 0: - raise ValueError( - "Sample boundary points failed, " - "please check correctness of geometry and given criteria." - ) - - all_points = np.concatenate(all_points, axis=0) - all_normal = np.concatenate(all_normal, axis=0) - _appr_area = self._approximate_area(random, criteria) - all_areas = np.full((n, 1), _appr_area / n, paddle.get_default_dtype()) - - x_dict = misc.convert_to_dict(all_points, self.dim_keys) - normal_dict = misc.convert_to_dict( - all_normal, [f"normal_{key}" for key in self.dim_keys if key != "t"] - ) - area_dict = misc.convert_to_dict(all_areas, ["area"]) - return {**x_dict, **normal_dict, **area_dict} - - def random_points(self, n, random="pseudo", criteria=None): - _size = 0 - all_points = [] - cuboid = geometry_3d.Cuboid( - [bound[0] for bound in self.bounds], - [bound[1] for bound in self.bounds], - ) - _nsample, _nvalid = 0, 0 - while _size < n: - random_points = cuboid.random_points(n, random) - valid_mask = self.is_inside(random_points) - - if criteria: - criteria_mask = criteria( - *np.split(random_points, self.ndim, axis=1) - ).ravel() - assert valid_mask.shape == criteria_mask.shape - valid_mask = np.logical_and(valid_mask, criteria_mask) - - valid_points = random_points[valid_mask] - _nvalid += len(valid_points) - - if len(valid_points) > n - _size: - valid_points = valid_points[: n - _size] - - all_points.append(valid_points) - _size += len(valid_points) - _nsample += n - - all_points = np.concatenate(all_points, axis=0) - cuboid_volume = np.prod([b[1] - b[0] for b in self.bounds]) - all_areas = np.full( - (n, 1), cuboid_volume * (_nvalid / _nsample) / n, paddle.get_default_dtype() - ) - return all_points, all_areas - - def sample_interior( - self, - n: int, - random: Literal["pseudo"] = "pseudo", - criteria: Optional[Callable[..., np.ndarray]] = None, - evenly: bool = False, - compute_sdf_derivatives: bool = False, - ): - """Sample random points in the geometry and return those meet criteria.""" - if evenly: - # TODO(sensen): Implement uniform sample for mesh interior. - raise NotImplementedError( - "uniformly sample for interior in mesh is not support yet, " - "you may need to set evenly=False in config dict of constraint" - ) - points, areas = self.random_points(n, random, criteria) - - x_dict = misc.convert_to_dict(points, self.dim_keys) - area_dict = misc.convert_to_dict(areas, ("area",)) - - sdf = self.sdf_func(points, compute_sdf_derivatives) - if compute_sdf_derivatives: - sdf, sdf_derives = sdf - - # NOTE: Negate sdf because weight should be positive. - sdf_dict = misc.convert_to_dict(-sdf, ("sdf",)) - - sdf_derives_dict = {} - if compute_sdf_derivatives: - # NOTE: Negate sdf derivatives - sdf_derives_dict = misc.convert_to_dict( - -sdf_derives, tuple(f"sdf__{key}" for key in self.dim_keys) - ) - - return {**x_dict, **area_dict, **sdf_dict, **sdf_derives_dict} - - def union(self, other: "SDFMesh"): - new_vectors = np.concatenate([self.vectors, other.vectors], axis=0) - new_normals = np.concatenate([self.face_normal, other.face_normal], axis=0) - - def make_union_new_sdf(sdf_func1, sdf_func2): - def new_sdf_func(points: np.ndarray, compute_sdf_derivatives: bool = False): - # Invert definition of sdf to make boolean operation accurate - # see: https://iquilezles.org/articles/interiordistance/ - sdf_self = sdf_func1(points, compute_sdf_derivatives) - sdf_other = sdf_func2(points, compute_sdf_derivatives) - if compute_sdf_derivatives: - sdf_self, sdf_derives_self = sdf_self - sdf_other, sdf_derives_other = sdf_other - - computed_sdf = -np.maximum(-sdf_self, -sdf_other) - - if compute_sdf_derivatives: - computed_sdf_derives = -np.where( - sdf_self < sdf_other, - sdf_derives_self, - sdf_derives_other, - ) - return computed_sdf, computed_sdf_derives - - return computed_sdf - - return new_sdf_func - - return SDFMesh( - new_vectors, - new_normals, - make_union_new_sdf(self.sdf_func, other.sdf_func), - ) - - def __or__(self, other: "SDFMesh"): - return self.union(other) - - def __add__(self, other: "SDFMesh"): - return self.union(other) - - def difference(self, other: "SDFMesh"): - new_vectors = np.concatenate([self.vectors, other.vectors], axis=0) - new_normals = np.concatenate([self.face_normal, -other.face_normal], axis=0) - - def make_difference_new_sdf(sdf_func1, sdf_func2): - def new_sdf_func(points: np.ndarray, compute_sdf_derivatives: bool = False): - # Invert definition of sdf to make boolean operation accurate - # see: https://iquilezles.org/articles/interiordistance/ - sdf_self = sdf_func1(points, compute_sdf_derivatives) - sdf_other = sdf_func2(points, compute_sdf_derivatives) - if compute_sdf_derivatives: - sdf_self, sdf_derives_self = sdf_self - sdf_other, sdf_derives_other = sdf_other - - computed_sdf = -np.minimum(-sdf_self, sdf_other) - - if compute_sdf_derivatives: - computed_sdf_derives = np.where( - -sdf_self < sdf_other, - -sdf_derives_self, - sdf_derives_other, - ) - return computed_sdf, computed_sdf_derives - - return computed_sdf - - return new_sdf_func - - return SDFMesh( - new_vectors, - new_normals, - make_difference_new_sdf(self.sdf_func, other.sdf_func), - ) - - def __sub__(self, other: "SDFMesh"): - return self.difference(other) - - def intersection(self, other: "SDFMesh"): - new_vectors = np.concatenate([self.vectors, other.vectors], axis=0) - new_normals = np.concatenate([self.face_normal, other.face_normal], axis=0) - - def make_intersection_new_sdf(sdf_func1, sdf_func2): - def new_sdf_func(points: np.ndarray, compute_sdf_derivatives: bool = False): - # Invert definition of sdf to make boolean operation accurate - # see: https://iquilezles.org/articles/interiordistance/ - sdf_self = sdf_func1(points, compute_sdf_derivatives) - sdf_other = sdf_func2(points, compute_sdf_derivatives) - if compute_sdf_derivatives: - sdf_self, sdf_derives_self = sdf_self - sdf_other, sdf_derives_other = sdf_other - - computed_sdf = -np.minimum(-sdf_self, -sdf_other) - - if compute_sdf_derivatives: - computed_sdf_derives = np.where( - sdf_self > sdf_other, - -sdf_derives_self, - -sdf_derives_other, - ) - return computed_sdf, computed_sdf_derives - - return computed_sdf - - return new_sdf_func - - return SDFMesh( - new_vectors, - new_normals, - make_intersection_new_sdf(self.sdf_func, other.sdf_func), - ) - - def __and__(self, other: "SDFMesh"): - return self.intersection(other) - - def __str__(self) -> str: - """Return the name of class""" - return ", ".join( - [ - self.__class__.__name__, - f"num_faces = {self.vectors.shape[0]}", - f"bounds = {self.bounds}", - f"dim_keys = {self.dim_keys}", - ] - ) - - -def area_of_triangles(v0, v1, v2): - """Ref https://math.stackexchange.com/questions/128991/how-to-calculate-the-area-of-a-3d-triangle - - Args: - v0 (np.ndarray): Coordinates of the first vertex of the triangle surface with shape of [N, 3]. - v1 (np.ndarray): Coordinates of the second vertex of the triangle surface with shape of [N, 3]. - v2 (np.ndarray): Coordinates of the third vertex of the triangle surface with shape of [N, 3]. - - Returns: - np.ndarray: Area of each triangle with shape of [N, ]. - """ - a = np.sqrt( - (v0[:, 0] - v1[:, 0]) ** 2 - + (v0[:, 1] - v1[:, 1]) ** 2 - + (v0[:, 2] - v1[:, 2]) ** 2 - + 1e-10 - ) - b = np.sqrt( - (v1[:, 0] - v2[:, 0]) ** 2 - + (v1[:, 1] - v2[:, 1]) ** 2 - + (v1[:, 2] - v2[:, 2]) ** 2 - + 1e-10 - ) - c = np.sqrt( - (v0[:, 0] - v2[:, 0]) ** 2 - + (v0[:, 1] - v2[:, 1]) ** 2 - + (v0[:, 2] - v2[:, 2]) ** 2 - + 1e-10 - ) - p = (a + b + c) / 2 - area = np.sqrt(p * (p - a) * (p - b) * (p - c) + 1e-10) - return area - - -def sample_in_triangle(v0, v1, v2, n, random="pseudo", criteria=None): - """ - Uniformly sample n points in an 3D triangle defined by 3 vertices v0, v1, v2 - https://math.stackexchange.com/questions/18686/uniform-random-point-in-triangle - - Args: - v0 (np.ndarray): Coordinates of the first vertex of an triangle with shape of [3, ]. - v1 (np.ndarray): Coordinates of the second vertex of an triangle with shape of [3, ]. - v2 (np.ndarray): Coordinates of the third vertex of an triangle with shape of [3, ]. - n (int): Number of points to be sampled. - - Returns: - np.ndarray: Coordinates of sampled n points with shape of [n, 3]. - """ - xs, ys, zs = [], [], [] - _size = 0 - while _size < n: - r1 = sampler.sample(n, 1, random).ravel() - r2 = sampler.sample(n, 1, random).ravel() - s1 = np.sqrt(r1) - x = v0[0] * (1.0 - s1) + v1[0] * (1.0 - r2) * s1 + v2[0] * r2 * s1 - y = v0[1] * (1.0 - s1) + v1[1] * (1.0 - r2) * s1 + v2[1] * r2 * s1 - z = v0[2] * (1.0 - s1) + v1[2] * (1.0 - r2) * s1 + v2[2] * r2 * s1 - - if criteria is not None: - criteria_mask = criteria(x, y, z).ravel() - x = x[criteria_mask] - y = y[criteria_mask] - z = z[criteria_mask] - - if len(x) > n - _size: - x = x[: n - _size] - y = y[: n - _size] - z = z[: n - _size] - - xs.append(x) - ys.append(y) - zs.append(z) - _size += len(x) - - xs = np.concatenate(xs, axis=0) - ys = np.concatenate(ys, axis=0) - zs = np.concatenate(zs, axis=0) - - return np.stack([xs, ys, zs], axis=1) - - -def make_sdf(vectors: np.ndarray): - def sdf_func(points: np.ndarray, compute_sdf_derivatives=False): - points = points.copy() - x_min, y_min, z_min = np.min(points, axis=0) - x_max, y_max, z_max = np.max(points, axis=0) - max_dis = max(max((x_max - x_min), (y_max - y_min)), (z_max - z_min)) - store_triangles = vectors.copy() - store_triangles[:, :, 0] -= x_min - store_triangles[:, :, 1] -= y_min - store_triangles[:, :, 2] -= z_min - store_triangles *= 1 / max_dis - store_triangles = store_triangles.reshape([-1, 3]) - points[:, 0] -= x_min - points[:, 1] -= y_min - points[:, 2] -= z_min - points *= 1 / max_dis - points = points.astype(np.float64).ravel() - - # compute sdf values - sdf = sdf_module.signed_distance_field( - store_triangles, - np.arange((store_triangles.shape[0])), - points, - include_hit_points=compute_sdf_derivatives, - ) - if compute_sdf_derivatives: - sdf, sdf_derives = sdf - - sdf = sdf.numpy() - sdf = np.expand_dims(max_dis * sdf, axis=1) - - if compute_sdf_derivatives: - sdf_derives = sdf_derives.numpy().reshape(-1) - sdf_derives = -(sdf_derives - points) - sdf_derives = np.reshape(sdf_derives, (sdf_derives.shape[0] // 3, 3)) - sdf_derives = sdf_derives / np.linalg.norm( - sdf_derives, axis=1, keepdims=True - ) - return sdf, sdf_derives - - return sdf - - return sdf_func diff --git a/examples/smc_reac/ppsci/geometry/pointcloud.py b/examples/smc_reac/ppsci/geometry/pointcloud.py deleted file mode 100644 index ae4d4fb4cc..0000000000 --- a/examples/smc_reac/ppsci/geometry/pointcloud.py +++ /dev/null @@ -1,312 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import Dict -from typing import Optional -from typing import Tuple - -import numpy as np - -from ppsci.geometry import geometry -from ppsci.utils import misc - - -class PointCloud(geometry.Geometry): - """Class for point cloud geometry, i.e. a set of points from given file or array. - - Args: - interior (Dict[str, np.ndarray]): Filepath or dict data, which store interior points of a point cloud, such as {"x": np.ndarray, "y": np.ndarray}. - coord_keys (Tuple[str, ...]): Tuple of coordinate keys, such as ("x", "y"). - boundary (Dict[str, np.ndarray]): Boundary points of a point cloud. Defaults to None. - boundary_normal (Dict[str, np.ndarray]): Boundary normal points of a point cloud. Defaults to None. - - Examples: - >>> import ppsci - >>> import numpy as np - >>> interior_points = {"x": np.linspace(-1, 1, dtype="float32").reshape((-1, 1))} - >>> geom = ppsci.geometry.PointCloud(interior_points, ("x",)) - """ - - def __init__( - self, - interior: Dict[str, np.ndarray], - coord_keys: Tuple[str, ...], - boundary: Optional[Dict[str, np.ndarray]] = None, - boundary_normal: Optional[Dict[str, np.ndarray]] = None, - ): - # Interior points - self.interior = misc.convert_to_array(interior, coord_keys) - self.len = self.interior.shape[0] - - # Boundary points - self.boundary = boundary - if self.boundary is not None: - self.boundary = misc.convert_to_array(self.boundary, coord_keys) - - # Boundary normal points - self.normal = boundary_normal - if self.normal is not None: - self.normal = misc.convert_to_array( - self.normal, tuple(f"{key}_normal" for key in coord_keys) - ) - if list(self.normal.shape) != list(self.boundary.shape): - raise ValueError( - f"boundary's shape({self.boundary.shape}) must equal " - f"to normal's shape({self.normal.shape})" - ) - - self.input_keys = coord_keys - super().__init__( - len(coord_keys), - (np.amin(self.interior, axis=0), np.amax(self.interior, axis=0)), - np.inf, - ) - - @property - def dim_keys(self): - return self.input_keys - - def is_inside(self, x): - # NOTE: point on boundary is included - return ( - np.isclose((x[:, None, :] - self.interior[None, :, :]), 0, atol=1e-6) - .all(axis=2) - .any(axis=1) - ) - - def on_boundary(self, x): - if not self.boundary: - raise ValueError( - "self.boundary must be initialized" " when call 'on_boundary' function" - ) - return ( - np.isclose( - (x[:, None, :] - self.boundary[None, :, :]), - 0, - atol=1e-6, - ) - .all(axis=2) - .any(axis=1) - ) - - def translate(self, translation: np.ndarray) -> "PointCloud": - """ - Translate the geometry by the given offset. - - Args: - translation (np.ndarray): Translation offset.The shape of translation must be the same as the shape of the interior points. - - Returns: - PointCloud: Translated point cloud. - - Examples: - >>> import ppsci - >>> import numpy as np - >>> interior_points = {"x": np.linspace(0, 2, 5, dtype="float32").reshape((-1, 1))} - >>> geom = ppsci.geometry.PointCloud(interior_points, ("x",)) - >>> translation = np.array([1.0]) - >>> print(geom.translate(translation).interior) - [[1. ] - [1.5] - [2. ] - [2.5] - [3. ]] - >>> interior_points_2d = {"x": np.linspace(0, 2, 5, dtype="float32").reshape((-1, 1)), - ... "y": np.linspace(0, 2, 5, dtype="float32").reshape((-1, 1))} - >>> geom_2d = ppsci.geometry.PointCloud(interior_points_2d, ("x", "y")) - >>> translation_2d = np.array([1.0, 3.0]) - >>> print(geom_2d.translate(translation_2d).interior) - [[1. 3. ] - [1.5 3.5] - [2. 4. ] - [2.5 4.5] - [3. 5. ]] - """ - for i, offset in enumerate(translation): - self.interior[:, i] += offset - if self.boundary: - self.boundary += offset - return self - - def scale(self, scale: np.ndarray) -> "PointCloud": - """ - Scale the geometry by the given factor. - - Args: - scale (np.ndarray): Scale factor.The shape of scale must be the same as the shape of the interior points. - - Returns: - PointCloud: Scaled point cloud. - - Examples: - >>> import ppsci - >>> import numpy as np - >>> interior_points = {"x": np.linspace(0, 2, 5, dtype="float32").reshape((-1, 1))} - >>> geom = ppsci.geometry.PointCloud(interior_points, ("x",)) - >>> scale = np.array([2.0]) - >>> print(geom.scale(scale).interior) - [[0.] - [1.] - [2.] - [3.] - [4.]] - >>> interior_points_2d = {"x": np.linspace(0, 2, 5, dtype="float32").reshape((-1, 1)), - ... "y": np.linspace(0, 2, 5, dtype="float32").reshape((-1, 1))} - >>> geom_2d = ppsci.geometry.PointCloud(interior_points_2d, ("x", "y")) - >>> scale_2d = np.array([2.0, 0.5]) - >>> print(geom_2d.scale(scale_2d).interior) - [[0. 0. ] - [1. 0.25] - [2. 0.5 ] - [3. 0.75] - [4. 1. ]] - """ - for i, _scale in enumerate(scale): - self.interior[:, i] *= _scale - if self.boundary: - self.boundary[:, i] *= _scale - if self.normal: - self.normal[:, i] *= _scale - return self - - def uniform_boundary_points(self, n: int): - """Compute the equi-spaced points on the boundary.""" - raise NotImplementedError( - "PointCloud do not have 'uniform_boundary_points' method" - ) - - def random_boundary_points(self, n: int, random: str = "pseudo") -> np.ndarray: - """Randomly sample points on the boundary. - - Args: - n (int): Number of sample points. - random (str): Random method. Defaults to "pseudo". - - Returns: - np.ndarray: Randomly sampled points on the boundary.The shape of the returned array is (n, ndim). - - Examples: - >>> import ppsci - >>> import numpy as np - >>> np.random.seed(0) - >>> interior_points = {"x": np.linspace(0, 2, 5, dtype="float32").reshape((-1, 1))} - >>> boundary_points = {"x": np.array([0.0, 2.0], dtype="float32").reshape((-1, 1))} - >>> geom = ppsci.geometry.PointCloud(interior_points, ("x",), boundary_points) - >>> print(geom.random_boundary_points(1)) - [[2.]] - """ - assert self.boundary is not None, ( - "boundary points can't be empty when call " - "'random_boundary_points' method" - ) - assert n <= len(self.boundary), ( - f"number of sample points({n}) " - f"can't be more than that in boundary({len(self.boundary)})" - ) - return self.boundary[ - np.random.choice(len(self.boundary), size=n, replace=False) - ] - - def random_points(self, n: int, random: str = "pseudo") -> np.ndarray: - """Randomly sample points in the geometry. - - Args: - n (int): Number of sample points. - random (str): Random method. Defaults to "pseudo". - - Returns: - np.ndarray: Randomly sampled points in the geometry.The shape of the returned array is (n, ndim). - - Examples: - >>> import ppsci - >>> import numpy as np - >>> np.random.seed(0) - >>> interior_points = {"x": np.linspace(0, 2, 5, dtype="float32").reshape((-1, 1))} - >>> geom = ppsci.geometry.PointCloud(interior_points, ("x",)) - >>> print(geom.random_points(2)) - [[1.] - [0.]] - """ - assert n <= len(self.interior), ( - f"number of sample points({n}) " - f"can't be more than that in points({len(self.interior)})" - ) - return self.interior[ - np.random.choice(len(self.interior), size=n, replace=False) - ] - - def uniform_points(self, n: int, boundary: bool = True) -> np.ndarray: - """Compute the equi-spaced points in the geometry. - - Args: - n (int): Number of sample points. - boundary (bool): Whether to include boundary points. Defaults to True. - - Returns: - np.ndarray: Equi-spaced points in the geometry.The shape of the returned array is (n, ndim). - - Examples: - >>> import ppsci - >>> import numpy as np - >>> interior_points = {"x": np.linspace(0, 2, 5, dtype="float32").reshape((-1, 1))} - >>> geom = ppsci.geometry.PointCloud(interior_points, ("x",)) - >>> print(geom.uniform_points(2)) - [[0. ] - [0.5]] - """ - return self.interior[:n] - - def union(self, other): - raise NotImplementedError( - "Union operation for PointCloud is not supported yet." - ) - - def __or__(self, other): - raise NotImplementedError( - "Union operation for PointCloud is not supported yet." - ) - - def difference(self, other): - raise NotImplementedError( - "Subtraction operation for PointCloud is not supported yet." - ) - - def __sub__(self, other): - raise NotImplementedError( - "Subtraction operation for PointCloud is not supported yet." - ) - - def intersection(self, other): - raise NotImplementedError( - "Intersection operation for PointCloud is not supported yet." - ) - - def __and__(self, other): - raise NotImplementedError( - "Intersection operation for PointCloud is not supported yet." - ) - - def __str__(self) -> str: - """Return the name of class.""" - return ", ".join( - [ - self.__class__.__name__, - f"num_points = {len(self.interior)}", - f"ndim = {self.ndim}", - f"bbox = {self.bbox}", - f"dim_keys = {self.dim_keys}", - ] - ) diff --git a/examples/smc_reac/ppsci/geometry/sampler.py b/examples/smc_reac/ppsci/geometry/sampler.py deleted file mode 100644 index a6de5015ff..0000000000 --- a/examples/smc_reac/ppsci/geometry/sampler.py +++ /dev/null @@ -1,92 +0,0 @@ -# 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. - -""" -Code below is heavily based on [https://github.com/lululxvi/deepxde](https://github.com/lululxvi/deepxde) -""" - -from __future__ import annotations - -import numpy as np -import paddle -import skopt -from typing_extensions import Literal - - -def sample( - n_samples: int, ndim: int, method: Literal["pseudo", "Halton", "LHS"] = "pseudo" -) -> np.ndarray: - """Generate pseudorandom or quasi-random samples in [0, 1]^ndim. - - Args: - n_samples (int): The number of samples. - ndim (int): Number of dimension. - method (str): One of the following: "pseudo" (pseudorandom), "LHS" (Latin - hypercube sampling), "Halton" (Halton sequence), "Hammersley" (Hammersley - sequence), or "Sobol" (Sobol sequence). - - Returns: - np.ndarray: Generated random samples with shape of [n_samples, ndim]. - """ - if method == "pseudo": - return pseudorandom(n_samples, ndim) - if method in ["LHS", "Halton", "Hammersley", "Sobol"]: - return quasirandom(n_samples, ndim, method) - raise ValueError(f"Sampling method({method}) is not available.") - - -def pseudorandom(n_samples: int, ndim: int) -> np.ndarray: - """Pseudo random.""" - # If random seed is set, then the rng based code always returns the same random - # number, which may not be what we expect. - # rng = np.random.default_rng(config.random_seed) - # return rng.random(size=(n_samples, ndim), dtype=dtype=paddle.get_default_dtype()) - return np.random.random(size=(n_samples, ndim)).astype( - dtype=paddle.get_default_dtype() - ) - - -def quasirandom( - n_samples: int, ndim: int, method: Literal["pseudo", "LHS"] -) -> np.ndarray: - """Quasi random""" - # Certain points should be removed: - # - Boundary points such as [..., 0, ...] - # - Special points [0, 0, 0, ...] and [0.5, 0.5, 0.5, ...], which cause error in - # Hypersphere.random_points() and Hypersphere.random_boundary_points() - skip = 0 - if method == "LHS": - sampler = skopt.sampler.Lhs() - elif method == "Halton": - # 1st point: [0, 0, ...] - sampler = skopt.sampler.Halton(min_skip=1, max_skip=1) - elif method == "Hammersley": - # 1st point: [0, 0, ...] - if ndim == 1: - sampler = skopt.sampler.Hammersly(min_skip=1, max_skip=1) - else: - sampler = skopt.sampler.Hammersly() - skip = 1 - elif method == "Sobol": - # 1st point: [0, 0, ...], 2nd point: [0.5, 0.5, ...] - sampler = skopt.sampler.Sobol(randomize=False) - if ndim < 3: - skip = 1 - else: - skip = 2 - space = [(0.0, 1.0)] * ndim - return np.asarray( - sampler.generate(space, n_samples + skip)[skip:], - dtype=paddle.get_default_dtype(), - ) diff --git a/examples/smc_reac/ppsci/geometry/sdf.py b/examples/smc_reac/ppsci/geometry/sdf.py deleted file mode 100644 index bb3e260540..0000000000 --- a/examples/smc_reac/ppsci/geometry/sdf.py +++ /dev/null @@ -1,198 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. -# SPDX-FileCopyrightText: All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# 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. - -# ruff: noqa: F401 - -# modified from: https://github.com/NVIDIA/modulus/blob/main/modulus/utils/sdf.py - -from __future__ import annotations - -import importlib.util -from typing import Tuple -from typing import overload - -from numpy import ndarray - -try: - import warp as wp - - @wp.kernel - def _bvh_query_distance( - mesh: wp.uint64, - points: wp.array(dtype=wp.vec3f), - max_dist: wp.float32, - sdf: wp.array(dtype=wp.float32), - sdf_hit_point: wp.array(dtype=wp.vec3f), - sdf_hit_point_id: wp.array(dtype=wp.int32), - ): - - """ - Computes the signed distance from each point in the given array `points` - to the mesh represented by `mesh`,within the maximum distance `max_dist`, - and stores the result in the array `sdf`. - - Parameters: - mesh (wp.uint64): The identifier of the mesh. - points (wp.array): An array of 3D points for which to compute the - signed distance. - max_dist (wp.float32): The maximum distance within which to search - for the closest point on the mesh. - sdf (wp.array): An array to store the computed signed distances. - sdf_hit_point (wp.array): An array to store the computed hit points. - sdf_hit_point_id (wp.array): An array to store the computed hit point ids. - - Returns: - None - """ - tid = wp.tid() - - res = wp.mesh_query_point_sign_normal(mesh, points[tid], max_dist) - - mesh_ = wp.mesh_get(mesh) - - p0 = mesh_.points[mesh_.indices[3 * res.face + 0]] - p1 = mesh_.points[mesh_.indices[3 * res.face + 1]] - p2 = mesh_.points[mesh_.indices[3 * res.face + 2]] - - p_closest = res.u * p0 + res.v * p1 + (1.0 - res.u - res.v) * p2 - - sdf[tid] = res.sign * wp.abs(wp.length(points[tid] - p_closest)) - sdf_hit_point[tid] = p_closest - sdf_hit_point_id[tid] = res.face - -except ModuleNotFoundError: - pass -except Exception: - raise - - -@overload -def signed_distance_field( - mesh_vertices: list[tuple[float, float, float]], - mesh_indices: ndarray, - input_points: list[tuple[float, float, float]], - max_dist: float = 1e8, - include_hit_points: bool = False, - include_hit_points_id: bool = False, -) -> wp.array: - ... - - -@overload -def signed_distance_field( - mesh_vertices: list[tuple[float, float, float]], - mesh_indices: ndarray, - input_points: list[tuple[float, float, float]], - max_dist: float = 1e8, - include_hit_points: bool = True, - include_hit_points_id: bool = False, -) -> Tuple[wp.array, wp.array]: - ... - - -@overload -def signed_distance_field( - mesh_vertices: list[tuple[float, float, float]], - mesh_indices: ndarray, - input_points: list[tuple[float, float, float]], - max_dist: float = 1e8, - include_hit_points: bool = False, - include_hit_points_id: bool = True, -) -> Tuple[wp.array, wp.array]: - ... - - -@overload -def signed_distance_field( - mesh_vertices: list[tuple[float, float, float]], - mesh_indices: ndarray, - input_points: list[tuple[float, float, float]], - max_dist: float = 1e8, - include_hit_points: bool = True, - include_hit_points_id: bool = True, -) -> Tuple[wp.array, wp.array, wp.array]: - ... - - -def signed_distance_field( - mesh_vertices: list[tuple[float, float, float]], - mesh_indices: ndarray, - input_points: list[tuple[float, float, float]], - max_dist: float = 1e8, - include_hit_points: bool = False, - include_hit_points_id: bool = False, -) -> wp.array: - """ - Computes the signed distance field (SDF) for a given mesh and input points. - - Args: - mesh_vertices (list[tuple[float, float, float]]): List of vertices defining the mesh. - mesh_indices (list[tuple[int, int, int]]): List of indices defining the triangles of the mesh. - input_points (list[tuple[float, float, float]]): List of input points for which to compute the SDF. - max_dist (float, optional): Maximum distance within which to search for - the closest point on the mesh. Default is 1e8. - include_hit_points (bool, optional): Whether to include hit points in - the output. Default is False. - include_hit_points_id (bool, optional): Whether to include hit point - IDs in the output. Default is False. - - Returns: - wp.array: An array containing the computed signed distance field. - - Example: - >>> mesh_vertices = [(0, 0, 0), (1, 0, 0), (0, 1, 0)] - >>> mesh_indices = np.array((0, 1, 2)) - >>> input_points = [(0.5, 0.5, 0.5)] - >>> signed_distance_field(mesh_vertices, mesh_indices, input_points).numpy() - Module modulus.utils.sdf load on device 'cuda:0' took ... - array([0.5], dtype=float32) - """ - if not importlib.util.find_spec("warp"): - raise ModuleNotFoundError("Please install warp with: pip install warp-lang") - - wp.init() - mesh = wp.Mesh( - wp.array(mesh_vertices, dtype=wp.vec3), - wp.array(mesh_indices, dtype=wp.int32), - ) - - sdf_points = wp.array(input_points, dtype=wp.vec3) - - sdf = wp.zeros(shape=sdf_points.shape, dtype=wp.float32) - sdf_hit_point = wp.zeros(shape=sdf_points.shape, dtype=wp.vec3f) - sdf_hit_point_id = wp.zeros(shape=sdf_points.shape, dtype=wp.int32) - - wp.launch( - kernel=_bvh_query_distance, - dim=len(sdf_points), - inputs=[ - mesh.id, - sdf_points, - max_dist, - sdf, - sdf_hit_point, - sdf_hit_point_id, - ], - ) - - if include_hit_points and include_hit_points_id: - return (sdf, sdf_hit_point, sdf_hit_point_id) - elif include_hit_points: - return (sdf, sdf_hit_point) - elif include_hit_points_id: - return (sdf, sdf_hit_point_id) - else: - return sdf diff --git a/examples/smc_reac/ppsci/geometry/timedomain.py b/examples/smc_reac/ppsci/geometry/timedomain.py deleted file mode 100644 index 909944a9b4..0000000000 --- a/examples/smc_reac/ppsci/geometry/timedomain.py +++ /dev/null @@ -1,793 +0,0 @@ -# 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. - -""" -Code below is heavily based on [https://github.com/lululxvi/deepxde](https://github.com/lululxvi/deepxde) -""" - -from __future__ import annotations - -import itertools -from typing import Callable -from typing import Dict -from typing import Optional -from typing import Tuple - -import numpy as np -import paddle - -from ppsci.geometry import geometry -from ppsci.geometry import geometry_1d -from ppsci.geometry import geometry_2d -from ppsci.geometry import geometry_3d -from ppsci.geometry import geometry_nd -from ppsci.geometry import mesh -from ppsci.utils import misc - - -class TimeDomain(geometry_1d.Interval): - """Class for timedomain, an special interval geometry. - - Args: - t0 (float): Start of time. - t1 (float): End of time. - time_step (Optional[float]): Step interval of time. Defaults to None. - timestamps (Optional[Tuple[float, ...]]): List of timestamps. - Defaults to None. - - Examples: - >>> import ppsci - >>> geom = ppsci.geometry.TimeDomain(0, 1) - """ - - def __init__( - self, - t0: float, - t1: float, - time_step: Optional[float] = None, - timestamps: Optional[Tuple[float, ...]] = None, - ): - super().__init__(t0, t1) - self.t0 = t0 - self.t1 = t1 - self.time_step = time_step - if timestamps is None: - self.timestamps = None - else: - self.timestamps = np.array( - timestamps, dtype=paddle.get_default_dtype() - ).reshape([-1]) - if time_step is not None: - if time_step <= 0: - raise ValueError(f"time_step({time_step}) must be larger than 0.") - self.num_timestamps = int(np.ceil((t1 - t0) / time_step)) + 1 - elif timestamps is not None: - self.num_timestamps = len(timestamps) - - def on_initial(self, t: np.ndarray) -> np.ndarray: - """Check if a specific time is on the initial time point. - - Args: - t (np.ndarray): The time to be checked. - - Returns: - np.ndarray: Bool numpy array of whether the specific time is on the initial time point. - - Examples: - >>> import paddle - >>> import ppsci - >>> geom = ppsci.geometry.TimeDomain(0, 1) - >>> T = [0, 0.01, 0.126, 0.2, 0.3] - >>> check = geom.on_initial(T) - >>> print(check) - [ True False False False False] - """ - return np.isclose(t, self.t0).flatten() - - -class TimeXGeometry(geometry.Geometry): - """Class for combination of time and geometry. - - Args: - timedomain (TimeDomain): TimeDomain object. - geometry (geometry.Geometry): Geometry object. - - Examples: - >>> import ppsci - >>> timedomain = ppsci.geometry.TimeDomain(0, 1) - >>> geom = ppsci.geometry.Rectangle((0, 0), (1, 1)) - >>> time_geom = ppsci.geometry.TimeXGeometry(timedomain, geom) - """ - - def __init__(self, timedomain: TimeDomain, geometry: geometry.Geometry): - self.timedomain = timedomain - self.geometry = geometry - self.ndim = geometry.ndim + timedomain.ndim - - @property - def dim_keys(self): - return ("t",) + self.geometry.dim_keys - - def on_boundary(self, x): - # [N, ndim(txyz)] - return self.geometry.on_boundary(x[:, 1:]) - - def on_initial(self, x): - # [N, 1(t)] - return self.timedomain.on_initial(x[:, :1]) - - def boundary_normal(self, x): - # x: [N, ndim(txyz)] - normal = self.geometry.boundary_normal(x[:, 1:]) - return np.hstack((x[:, :1], normal)) - - def uniform_points(self, n: int, boundary: bool = True) -> np.ndarray: - """Uniform points on the spatial-temporal domain. - Geometry volume ~ bbox. - Time volume ~ diam. - - Args: - n (int): The total number of sample points to be generated. - boundary (bool): Indicates whether boundary points are included, default is True. - - Returns: - np.ndarray: a set of spatial-temporal coordinate points 'tx' that represent sample points evenly distributed within the spatial-temporal domain. - - Examples: - >>> import ppsci - >>> timedomain = ppsci.geometry.TimeDomain(0, 1, 0.001) - >>> geom = ppsci.geometry.Rectangle((0, 0), (1, 1)) - >>> time_geom = ppsci.geometry.TimeXGeometry(timedomain, geom) - >>> ts = time_geom.uniform_points(1000) - >>> print(ts.shape) - (1000, 3) - """ - if self.timedomain.time_step is not None: - # exclude start time t0 - nt = int(np.ceil(self.timedomain.diam / self.timedomain.time_step)) - nx = int(np.ceil(n / nt)) - elif self.timedomain.timestamps is not None: - # exclude start time t0 - nt = self.timedomain.num_timestamps - 1 - nx = int(np.ceil(n / nt)) - else: - nx = int( - np.ceil( - ( - n - * np.prod(self.geometry.bbox[1] - self.geometry.bbox[0]) - / self.timedomain.diam - ) - ** 0.5 - ) - ) - nt = int(np.ceil(n / nx)) - x = self.geometry.uniform_points(nx, boundary=boundary) - nx = len(x) - if boundary and ( - self.timedomain.time_step is None and self.timedomain.timestamps is None - ): - t = self.timedomain.uniform_points(nt, boundary=True) - else: - if self.timedomain.time_step is not None: - t = np.linspace( - self.timedomain.t1, - self.timedomain.t0, - num=nt, - endpoint=boundary, - dtype=paddle.get_default_dtype(), - )[:, None][::-1] - else: - t = self.timedomain.timestamps[1:] - tx = [] - for ti in t: - tx.append( - np.hstack((np.full([nx, 1], ti, dtype=paddle.get_default_dtype()), x)) - ) - tx = np.vstack(tx) - if len(tx) > n: - tx = tx[:n] - return tx - - def random_points( - self, n: int, random: str = "pseudo", criteria: Optional[Callable] = None - ) -> np.ndarray: - """Generate random points on the spatial-temporal domain. - - Args: - n (int): The total number of random points to generate. - random (str): Specifies the way to generate random points, default is "pseudo" , which means that a pseudo-random number generator is used. - criteria (Optional[Callable]): A method that filters on the generated random points. Defaults to None. - - Returns: - np.ndarray: A set of random spatial-temporal points. - - Examples: - >>> import ppsci - >>> timedomain = ppsci.geometry.TimeDomain(0, 1, 0.001) - >>> geom = ppsci.geometry.Rectangle((0, 0), (1, 1)) - >>> time_geom = ppsci.geometry.TimeXGeometry(timedomain, geom) - >>> ts = time_geom.random_points(1000) - >>> print(ts.shape) - (1000, 3) - """ - if self.timedomain.time_step is None and self.timedomain.timestamps is None: - raise ValueError("Either time_step or timestamps must be provided.") - # time evenly and geometry random, if time_step if specified - if self.timedomain.time_step is not None: - nt = int(np.ceil(self.timedomain.diam / self.timedomain.time_step)) - t = np.linspace( - self.timedomain.t1, - self.timedomain.t0, - num=nt, - endpoint=False, - dtype=paddle.get_default_dtype(), - )[:, None][ - ::-1 - ] # [nt, 1] - # 1. sample nx points in static geometry with criteria - nx = int(np.ceil(n / nt)) - _size, _ntry, _nsuc = 0, 0, 0 - x = np.empty( - shape=(nx, self.geometry.ndim), dtype=paddle.get_default_dtype() - ) - while _size < nx: - _x = self.geometry.random_points(nx, random) - if criteria is not None: - # fix arg 't' to None in criteria there - criteria_mask = criteria( - None, *np.split(_x, self.geometry.ndim, axis=1) - ).flatten() - _x = _x[criteria_mask] - if len(_x) > nx - _size: - _x = _x[: nx - _size] - x[_size : _size + len(_x)] = _x - - _size += len(_x) - _ntry += 1 - if len(_x) > 0: - _nsuc += 1 - - if _ntry >= 1000 and _nsuc == 0: - raise ValueError( - "Sample points failed, " - "please check correctness of geometry and given criteria." - ) - - # 2. repeat spatial points along time - tx = [] - for ti in t: - tx.append( - np.hstack( - (np.full([nx, 1], ti, dtype=paddle.get_default_dtype()), x) - ) - ) - tx = np.vstack(tx) - if len(tx) > n: - tx = tx[:n] - return tx - elif self.timedomain.timestamps is not None: - nt = self.timedomain.num_timestamps - 1 - t = self.timedomain.timestamps[1:] - nx = int(np.ceil(n / nt)) - - _size, _ntry, _nsuc = 0, 0, 0 - x = np.empty( - shape=(nx, self.geometry.ndim), dtype=paddle.get_default_dtype() - ) - while _size < nx: - _x = self.geometry.random_points(nx, random) - if criteria is not None: - # fix arg 't' to None in criteria there - criteria_mask = criteria( - None, *np.split(_x, self.geometry.ndim, axis=1) - ).flatten() - _x = _x[criteria_mask] - if len(_x) > nx - _size: - _x = _x[: nx - _size] - x[_size : _size + len(_x)] = _x - - _size += len(_x) - _ntry += 1 - if len(_x) > 0: - _nsuc += 1 - - if _ntry >= 1000 and _nsuc == 0: - raise ValueError( - "Sample interior points failed, " - "please check correctness of geometry and given criteria." - ) - - tx = [] - for ti in t: - tx.append( - np.hstack( - (np.full([nx, 1], ti, dtype=paddle.get_default_dtype()), x) - ) - ) - tx = np.vstack(tx) - if len(tx) > n: - tx = tx[:n] - return tx - - if isinstance(self.geometry, geometry_1d.Interval): - geom = geometry_2d.Rectangle( - [self.timedomain.t0, self.geometry.l], - [self.timedomain.t1, self.geometry.r], - ) - return geom.random_points(n, random=random) - - if isinstance(self.geometry, geometry_2d.Rectangle): - geom = geometry_3d.Cuboid( - [self.timedomain.t0, self.geometry.xmin[0], self.geometry.xmin[1]], - [self.timedomain.t1, self.geometry.xmax[0], self.geometry.xmax[1]], - ) - return geom.random_points(n, random=random) - - if isinstance(self.geometry, (geometry_3d.Cuboid, geometry_nd.Hypercube)): - geom = geometry_nd.Hypercube( - np.append(self.timedomain.t0, self.geometry.xmin), - np.append(self.timedomain.t1, self.geometry.xmax), - ) - return geom.random_points(n, random=random) - - x = self.geometry.random_points(n, random=random) - t = self.timedomain.random_points(n, random=random) - t = np.random.permutation(t) - return np.hstack((t, x)) - - def uniform_boundary_points( - self, n: int, criteria: Optional[Callable] = None - ) -> np.ndarray: - """Uniform boundary points on the spatial-temporal domain. - Geometry surface area ~ bbox. - Time surface area ~ diam. - - Args: - n (int): The total number of boundary points on the spatial-temporal domain to be generated that are evenly distributed across geometry boundaries. - criteria (Optional[Callable]): Used to filter the generated boundary points, only points that meet certain conditions are retained. Default is None. - - Returns: - np.ndarray: A set of point coordinates evenly distributed across geometry boundaries on the spatial-temporal domain. - - Examples: - >>> import ppsci - >>> timedomain = ppsci.geometry.TimeDomain(0, 1) - >>> geom = ppsci.geometry.Rectangle((0, 0), (1, 1)) - >>> time_geom = ppsci.geometry.TimeXGeometry(timedomain, geom) - >>> ts = time_geom.uniform_boundary_points(1000) - >>> print(ts.shape) - (1000, 3) - """ - if self.geometry.ndim == 1: - nx = 2 - else: - s = 2 * sum( - map( - lambda l: l[0] * l[1], - itertools.combinations( - self.geometry.bbox[1] - self.geometry.bbox[0], 2 - ), - ) - ) - nx = int((n * s / self.timedomain.diam) ** 0.5) - nt = int(np.ceil(n / nx)) - - _size, _ntry, _nsuc = 0, 0, 0 - x = np.empty(shape=(nx, self.geometry.ndim), dtype=paddle.get_default_dtype()) - while _size < nx: - _x = self.geometry.uniform_boundary_points(nx) - if criteria is not None: - # fix arg 't' to None in criteria there - criteria_mask = criteria( - None, *np.split(_x, self.geometry.ndim, axis=1) - ).flatten() - _x = _x[criteria_mask] - if len(_x) > nx - _size: - _x = _x[: nx - _size] - x[_size : _size + len(_x)] = _x - - _size += len(_x) - _ntry += 1 - if len(_x) > 0: - _nsuc += 1 - - if _ntry >= 1000 and _nsuc == 0: - raise ValueError( - "Sample boundary points failed, " - "please check correctness of geometry and given criteria." - ) - - nx = len(x) - t = np.linspace( - self.timedomain.t1, - self.timedomain.t0, - num=nt, - endpoint=False, - dtype=paddle.get_default_dtype(), - )[:, None][::-1] - tx = [] - for ti in t: - tx.append( - np.hstack((np.full([nx, 1], ti, dtype=paddle.get_default_dtype()), x)) - ) - tx = np.vstack(tx) - if len(tx) > n: - tx = tx[:n] - return tx - - def random_boundary_points( - self, n: int, random: str = "pseudo", criteria: Optional[Callable] = None - ) -> np.ndarray: - """Random boundary points on the spatial-temporal domain. - - Args: - n (int): The total number of spatial-temporal points generated on a given geometry boundary. - random (str): Controls the way to generate random points. Default is "pseudo". - criteria (Optional[Callable]): Used to filter the generated boundary points, only points that meet certain conditions are retained. Default is None. - - Returns: - np.ndarray: A set of point coordinates randomly distributed across geometry boundaries on the spatial-temporal domain. - - Examples: - >>> import ppsci - >>> timedomain = ppsci.geometry.TimeDomain(0, 1, 0.001) - >>> geom = ppsci.geometry.Rectangle((0, 0), (1, 1)) - >>> time_geom = ppsci.geometry.TimeXGeometry(timedomain, geom) - >>> ts = time_geom.random_boundary_points(1000) - >>> print(ts.shape) - (1000, 3) - """ - if self.timedomain.time_step is None and self.timedomain.timestamps is None: - raise ValueError("Either time_step or timestamps must be provided.") - if self.timedomain.time_step is not None: - # exclude start time t0 - nt = int(np.ceil(self.timedomain.diam / self.timedomain.time_step)) - t = np.linspace( - self.timedomain.t1, - self.timedomain.t0, - num=nt, - endpoint=False, - dtype=paddle.get_default_dtype(), - )[:, None][::-1] - nx = int(np.ceil(n / nt)) - - if isinstance(self.geometry, mesh.Mesh): - x, _n, a = self.geometry.random_boundary_points(nx, random=random) - else: - _size, _ntry, _nsuc = 0, 0, 0 - x = np.empty( - shape=(nx, self.geometry.ndim), dtype=paddle.get_default_dtype() - ) - while _size < nx: - _x = self.geometry.random_boundary_points(nx, random) - if criteria is not None: - # fix arg 't' to None in criteria there - criteria_mask = criteria( - None, *np.split(_x, self.geometry.ndim, axis=1) - ).flatten() - _x = _x[criteria_mask] - if len(_x) > nx - _size: - _x = _x[: nx - _size] - x[_size : _size + len(_x)] = _x - - _size += len(_x) - _ntry += 1 - if len(_x) > 0: - _nsuc += 1 - - if _ntry >= 1000 and _nsuc == 0: - raise ValueError( - "Sample boundary points failed, " - "please check correctness of geometry and given criteria." - ) - - t_x = [] - if isinstance(self.geometry, mesh.Mesh): - t_normal = [] - t_area = [] - - for ti in t: - t_x.append( - np.hstack( - (np.full([nx, 1], ti, dtype=paddle.get_default_dtype()), x) - ) - ) - if isinstance(self.geometry, mesh.Mesh): - t_normal.append( - np.hstack( - (np.full([nx, 1], ti, dtype=paddle.get_default_dtype()), _n) - ) - ) - t_area.append( - np.hstack( - (np.full([nx, 1], ti, dtype=paddle.get_default_dtype()), a) - ) - ) - - t_x = np.vstack(t_x) - if isinstance(self.geometry, mesh.Mesh): - t_normal = np.vstack(t_normal) - t_area = np.vstack(t_area) - - if len(t_x) > n: - t_x = t_x[:n] - if isinstance(self.geometry, mesh.Mesh): - t_normal = t_normal[:n] - t_area = t_area[:n] - - if isinstance(self.geometry, mesh.Mesh): - return t_x, t_normal, t_area - else: - return t_x - elif self.timedomain.timestamps is not None: - # exclude start time t0 - nt = self.timedomain.num_timestamps - 1 - t = self.timedomain.timestamps[1:] - nx = int(np.ceil(n / nt)) - - if isinstance(self.geometry, mesh.Mesh): - x, _n, a = self.geometry.random_boundary_points(nx, random=random) - else: - _size, _ntry, _nsuc = 0, 0, 0 - x = np.empty( - shape=(nx, self.geometry.ndim), dtype=paddle.get_default_dtype() - ) - while _size < nx: - _x = self.geometry.random_boundary_points(nx, random) - if criteria is not None: - # fix arg 't' to None in criteria there - criteria_mask = criteria( - None, *np.split(_x, self.geometry.ndim, axis=1) - ).flatten() - _x = _x[criteria_mask] - if len(_x) > nx - _size: - _x = _x[: nx - _size] - x[_size : _size + len(_x)] = _x - - _size += len(_x) - _ntry += 1 - if len(_x) > 0: - _nsuc += 1 - - if _ntry >= 1000 and _nsuc == 0: - raise ValueError( - "Sample boundary points failed, " - "please check correctness of geometry and given criteria." - ) - - t_x = [] - if isinstance(self.geometry, mesh.Mesh): - t_normal = [] - t_area = [] - - for ti in t: - t_x.append( - np.hstack( - (np.full([nx, 1], ti, dtype=paddle.get_default_dtype()), x) - ) - ) - if isinstance(self.geometry, mesh.Mesh): - t_normal.append( - np.hstack( - (np.full([nx, 1], ti, dtype=paddle.get_default_dtype()), _n) - ) - ) - t_area.append( - np.hstack( - (np.full([nx, 1], ti, dtype=paddle.get_default_dtype()), a) - ) - ) - - t_x = np.vstack(t_x) - if isinstance(self.geometry, mesh.Mesh): - t_normal = np.vstack(t_normal) - t_area = np.vstack(t_area) - - if len(t_x) > n: - t_x = t_x[:n] - if isinstance(self.geometry, mesh.Mesh): - t_normal = t_normal[:n] - t_area = t_area[:n] - - if isinstance(self.geometry, mesh.Mesh): - return t_x, t_normal, t_area - else: - return t_x - else: - if isinstance(self.geometry, mesh.Mesh): - x, _n, a = self.geometry.random_boundary_points(n, random=random) - else: - x = self.geometry.random_boundary_points(n, random=random) - - t = self.timedomain.random_points(n, random=random) - t = np.random.permutation(t) - - t_x = np.hstack((t, x)) - - if isinstance(self.geometry, mesh.Mesh): - t_normal = np.hstack((_n, t)) - t_area = np.hstack((_n, t)) - return t_x, t_normal, t_area - else: - return t_x - - def uniform_initial_points(self, n: int) -> np.ndarray: - """Generate evenly distributed point coordinates on the spatial-temporal domain at the initial moment. - - Args: - n (int): The total number of generated points. - - Returns: - np.ndarray: A set of point coordinates evenly distributed on the spatial-temporal domain at the initial moment. - - Examples: - >>> import ppsci - >>> timedomain = ppsci.geometry.TimeDomain(0, 1) - >>> geom = ppsci.geometry.Rectangle((0, 0), (1, 1)) - >>> time_geom = ppsci.geometry.TimeXGeometry(timedomain, geom) - >>> ts = time_geom.uniform_initial_points(1000) - >>> print(ts.shape) - (1000, 3) - """ - x = self.geometry.uniform_points(n, True) - t = self.timedomain.t0 - if len(x) > n: - x = x[:n] - return np.hstack((np.full([n, 1], t, dtype=paddle.get_default_dtype()), x)) - - def random_initial_points(self, n: int, random: str = "pseudo") -> np.ndarray: - """Generate randomly distributed point coordinates on the spatial-temporal domain at the initial moment. - - Args: - n (int): The total number of generated points. - random (str): Controls the way to generate random points. Default is "pseudo". - - Returns: - np.ndarray: A set of point coordinates randomly distributed on the spatial-temporal domain at the initial moment. - - Examples: - >>> import ppsci - >>> timedomain = ppsci.geometry.TimeDomain(0, 1) - >>> geom = ppsci.geometry.Rectangle((0, 0), (1, 1)) - >>> time_geom = ppsci.geometry.TimeXGeometry(timedomain, geom) - >>> ts = time_geom.random_initial_points(1000) - >>> print(ts.shape) - (1000, 3) - """ - x = self.geometry.random_points(n, random=random) - t = self.timedomain.t0 - return np.hstack((np.full([n, 1], t, dtype=paddle.get_default_dtype()), x)) - - def periodic_point( - self, x: Dict[str, np.ndarray], component: int - ) -> Dict[str, np.ndarray]: - """Process given point coordinates to satisfy the periodic boundary conditions of the geometry. - - Args: - x (Dict[str, np.ndarray]): Contains the coordinates and timestamps of the points. It represents the coordinates of the point to be processed. - component (int): Specifies the components or dimensions of specific spatial coordinates that are periodically processed. - - Returns: - Dict[str, np.ndarray] : contains the original timestamps and the coordinates of the spatial point after periodic processing. - - Examples: - >>> import ppsci - >>> timedomain = ppsci.geometry.TimeDomain(0, 1, 0.1) - >>> geom = ppsci.geometry.Rectangle((0, 0), (1, 1)) - >>> time_geom = ppsci.geometry.TimeXGeometry(timedomain, geom) - >>> ts = time_geom.sample_boundary(1000) - >>> result = time_geom.periodic_point(ts, 0) - >>> for k,v in result.items(): - ... print(k, v.shape) - t (1000, 1) - x (1000, 1) - y (1000, 1) - normal_x (1000, 1) - normal_y (1000, 1) - """ - xp = self.geometry.periodic_point(x, component) - txp = {"t": x["t"], **xp} - return txp - - def sample_initial_interior( - self, - n: int, - random: str = "pseudo", - criteria: Optional[Callable] = None, - evenly: bool = False, - compute_sdf_derivatives: bool = False, - ) -> Dict[str, np.ndarray]: - """Sample random points in the time-geometry and return those meet criteria. - - Args: - n (int): The total number of interior points generated. - random (str): The method used to specify the initial point of generation. Default is "pseudo". - criteria (Optional[Callable]): Used to filter the generated interior points, only points that meet certain conditions are retained. Default is None. - evenly (bool): Indicates whether the initial points are generated evenly. Default is False. - compute_sdf_derivatives (bool): Indicates whether to calculate the derivative of signed distance function or not. Default is False. - - Returns: - np.ndarray: Contains the coordinates of the initial internal point generated, as well as the potentially computed signed distance function and its derivative. - - Examples: - >>> import ppsci - >>> timedomain = ppsci.geometry.TimeDomain(0, 1) - >>> geom = ppsci.geometry.Rectangle((0, 0), (1, 1)) - >>> time_geom = ppsci.geometry.TimeXGeometry(timedomain, geom) - >>> ts = time_geom.sample_initial_interior(1000) - >>> for k,v in ts.items(): - ... print(k, v.shape) - t (1000, 1) - x (1000, 1) - y (1000, 1) - sdf (1000, 1) - """ - x = np.empty(shape=(n, self.ndim), dtype=paddle.get_default_dtype()) - _size, _ntry, _nsuc = 0, 0, 0 - while _size < n: - if evenly: - points = self.uniform_initial_points(n) - else: - points = self.random_initial_points(n, random) - - if criteria is not None: - criteria_mask = criteria(*np.split(points, self.ndim, axis=1)).flatten() - points = points[criteria_mask] - - if len(points) > n - _size: - points = points[: n - _size] - x[_size : _size + len(points)] = points - - _size += len(points) - _ntry += 1 - if len(points) > 0: - _nsuc += 1 - - if _ntry >= 1000 and _nsuc == 0: - raise ValueError( - "Sample initial interior points failed, " - "please check correctness of geometry and given criteria." - ) - - # if sdf_func added, return x_dict and sdf_dict, else, only return the x_dict - if hasattr(self.geometry, "sdf_func"): - # compute sdf excluding time t - sdf = -self.geometry.sdf_func(x[..., 1:]) - sdf_dict = misc.convert_to_dict(sdf, ("sdf",)) - sdf_derives_dict = {} - if compute_sdf_derivatives: - # compute sdf derivatives excluding time t - sdf_derives = -self.geometry.sdf_derivatives(x[..., 1:]) - sdf_derives_dict = misc.convert_to_dict( - sdf_derives, tuple(f"sdf__{key}" for key in self.geometry.dim_keys) - ) - else: - sdf_dict = {} - sdf_derives_dict = {} - x_dict = misc.convert_to_dict(x, self.dim_keys) - - return {**x_dict, **sdf_dict, **sdf_derives_dict} - - def __str__(self) -> str: - """Return the name of class""" - return ", ".join( - [ - self.__class__.__name__, - f"ndim = {self.ndim}", - f"bbox = (time){self.timedomain.bbox} x (space){self.geometry.bbox}", - f"diam = (time){self.timedomain.diam} x (space){self.geometry.diam}", - f"dim_keys = {self.dim_keys}", - ] - ) diff --git a/examples/smc_reac/ppsci/loss/__init__.py b/examples/smc_reac/ppsci/loss/__init__.py deleted file mode 100644 index d9502e8248..0000000000 --- a/examples/smc_reac/ppsci/loss/__init__.py +++ /dev/null @@ -1,67 +0,0 @@ -# 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. - -import copy - -from ppsci.loss import mtl -from ppsci.loss.base import Loss -from ppsci.loss.chamfer import ChamferLoss -from ppsci.loss.func import FunctionalLoss -from ppsci.loss.integral import IntegralLoss -from ppsci.loss.kl import KLLoss -from ppsci.loss.l1 import L1Loss -from ppsci.loss.l1 import PeriodicL1Loss -from ppsci.loss.l2 import L2Loss -from ppsci.loss.l2 import L2RelLoss -from ppsci.loss.l2 import PeriodicL2Loss -from ppsci.loss.mae import MAELoss -from ppsci.loss.mse import CausalMSELoss -from ppsci.loss.mse import MSELoss -from ppsci.loss.mse import MSELossWithL2Decay -from ppsci.loss.mse import PeriodicMSELoss - -__all__ = [ - "Loss", - "FunctionalLoss", - "IntegralLoss", - "L1Loss", - "PeriodicL1Loss", - "L2Loss", - "L2RelLoss", - "PeriodicL2Loss", - "MAELoss", - "CausalMSELoss", - "ChamferLoss", - "MSELoss", - "MSELossWithL2Decay", - "PeriodicMSELoss", - "KLLoss", - "mtl", -] - - -def build_loss(cfg): - """Build loss. - - Args: - cfg (DictConfig): Loss config. - - Returns: - Loss: Callable loss object. - """ - cfg = copy.deepcopy(cfg) - - loss_cls = cfg.pop("name") - loss = eval(loss_cls)(**cfg) - return loss diff --git a/examples/smc_reac/ppsci/loss/base.py b/examples/smc_reac/ppsci/loss/base.py deleted file mode 100644 index 378013bb9e..0000000000 --- a/examples/smc_reac/ppsci/loss/base.py +++ /dev/null @@ -1,38 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import Dict -from typing import Optional -from typing import Union - -from paddle import nn -from typing_extensions import Literal - - -class Loss(nn.Layer): - """Base class for loss.""" - - def __init__( - self, - reduction: Literal["mean", "sum"], - weight: Optional[Union[float, Dict[str, float]]] = None, - ): - super().__init__() - self.reduction = reduction - self.weight = weight - - def __str__(self): - return f"{self.__class__.__name__}(reduction={self.reduction}, weight={self.weight})" diff --git a/examples/smc_reac/ppsci/loss/chamfer.py b/examples/smc_reac/ppsci/loss/chamfer.py deleted file mode 100644 index 8740f2a18b..0000000000 --- a/examples/smc_reac/ppsci/loss/chamfer.py +++ /dev/null @@ -1,92 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import Dict -from typing import Optional -from typing import Union - -import paddle - -from ppsci.loss import base - - -class ChamferLoss(base.Loss): - r"""Class for Chamfe distance loss. - - $$ - L = \dfrac{1}{S_1} \sum_{x \in S_1} \min_{y \in S_2} \Vert x - y \Vert_2^2 + \dfrac{1}{S_2} \sum_{y \in S_2} \min_{x \in S_1} \Vert y - x \Vert_2^2 - $$ - - $$ - \text{where } S_1 \text{ and } S_2 \text{ is the coordinate matrix of two point clouds}. - $$ - - Args: - weight (Optional[Union[float, Dict[str, float]]]): Weight for loss. Defaults to None. - - Examples: - >>> import paddle - >>> from ppsci.loss import ChamferLoss - >>> _ = paddle.seed(42) - >>> batch_point_cloud1 = paddle.rand([2, 100, 3]) - >>> batch_point_cloud2 = paddle.rand([2, 50, 3]) - >>> output_dict = {"s1": batch_point_cloud1} - >>> label_dict = {"s1": batch_point_cloud2} - >>> weight = {"s1": 0.8} - >>> loss = ChamferLoss(weight=weight) - >>> result = loss(output_dict, label_dict) - >>> print(result) - {'s1': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, - 0.04415882)} - """ - - def __init__( - self, - weight: Optional[Union[float, Dict[str, float]]] = None, - ): - super().__init__("mean", weight) - - def forward( - self, output_dict, label_dict, weight_dict=None - ) -> Dict[str, "paddle.Tensor"]: - losses = {} - - for key in label_dict: - s1 = output_dict[key] - s2 = label_dict[key] - N1, N2 = s1.shape[1], s2.shape[1] - - # [B, N1, N2, 3] - s1_expand = paddle.expand(s1.reshape([-1, N1, 1, 3]), shape=[-1, N1, N2, 3]) - # [B, N1, N2, 3] - s2_expand = paddle.expand(s2.reshape([-1, 1, N2, 3]), shape=[-1, N1, N2, 3]) - - dis = ((s1_expand - s2_expand) ** 2).sum(axis=3) # [B, N1, N2] - loss_s12 = dis.min(axis=2) # [B, N1] - loss_s21 = dis.min(axis=1) # [B, N2] - loss = loss_s12.mean() + loss_s21.mean() - - if weight_dict and key in weight_dict: - loss *= weight_dict[key] - - if isinstance(self.weight, (float, int)): - loss *= self.weight - elif isinstance(self.weight, dict) and key in self.weight: - loss *= self.weight[key] - - losses[key] = loss - - return losses diff --git a/examples/smc_reac/ppsci/loss/func.py b/examples/smc_reac/ppsci/loss/func.py deleted file mode 100644 index 49992c5f7f..0000000000 --- a/examples/smc_reac/ppsci/loss/func.py +++ /dev/null @@ -1,94 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import Callable -from typing import Dict -from typing import Optional -from typing import Union - -import paddle - -from ppsci.loss import base - - -class FunctionalLoss(base.Loss): - r"""Functional loss class, which allows to use custom loss computing function from given loss_expr for complex computation cases. - - $$ - L = f(\mathbf{x}, \mathbf{y}) - $$ - - $$ - \mathbf{x}, \mathbf{y} \in \mathcal{R}^{N} - $$ - - Args: - loss_expr (Callable[..., paddle.Tensor]): Function for custom loss computation. - weight (Optional[Union[float, Dict[str, float]]]): Weight for loss. Defaults to None. - - Examples: - >>> import paddle - >>> from ppsci.loss import FunctionalLoss - >>> import paddle.nn.functional as F - >>> def mse_sum_loss(output_dict, label_dict, weight_dict=None): - ... losses = 0 - ... for key in output_dict.keys(): - ... loss = F.mse_loss(output_dict[key], label_dict[key], "sum") - ... if weight_dict: - ... loss *= weight_dict[key] - ... losses += loss - ... return {"mse_loss": losses} - >>> loss = FunctionalLoss(mse_sum_loss) - >>> output_dict = {'u': paddle.to_tensor([[0.5, 0.9], [1.1, -1.3]]), - ... 'v': paddle.to_tensor([[0.5, 0.9], [1.1, -1.3]])} - >>> label_dict = {'u': paddle.to_tensor([[-1.8, 1.0], [-0.2, 2.5]]), - ... 'v': paddle.to_tensor([[0.1, 0.1], [0.1, 0.1]])} - >>> weight_dict = {'u': 0.8, 'v': 0.2} - >>> result = loss(output_dict, label_dict, weight_dict) - >>> print(result) - {'mse_loss': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, - 17.89600182)} - """ - - def __init__( - self, - loss_expr: Callable[..., paddle.Tensor], - weight: Optional[Union[float, Dict[str, float]]] = None, - ): - super().__init__(None, weight) - self.loss_expr = loss_expr - - def forward( - self, output_dict, label_dict=None, weight_dict=None - ) -> Dict[str, "paddle.Tensor"]: - losses = self.loss_expr(output_dict, label_dict, weight_dict) - - assert isinstance(losses, dict), ( - "Loss computed by custom function should be type of 'dict', " - f"but got {type(losses)}." - " Please check the return type of custom loss function." - ) - - for key in losses: - assert isinstance( - losses[key], (paddle.Tensor, paddle.static.Variable, paddle.pir.Value) - ), ( - "Loss computed by custom function should be type of 'paddle.Tensor', " - f"'paddle.static.Variable' or 'paddle.pir.Value', but got {type(losses[key])}." - " Please check the return type of custom loss function." - ) - - return losses diff --git a/examples/smc_reac/ppsci/loss/integral.py b/examples/smc_reac/ppsci/loss/integral.py deleted file mode 100644 index 74223c73fc..0000000000 --- a/examples/smc_reac/ppsci/loss/integral.py +++ /dev/null @@ -1,112 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import TYPE_CHECKING -from typing import Dict -from typing import Optional -from typing import Union - -import paddle.nn.functional as F -from typing_extensions import Literal - -from ppsci.loss import base - -if TYPE_CHECKING: - import paddle - - -class IntegralLoss(base.Loss): - r"""Class for integral loss with Monte-Carlo integration algorithm. - - $$ - L = - \begin{cases} - \dfrac{1}{N} \Vert \displaystyle\sum_{i=1}^{M}{\mathbf{s}_i \cdot \mathbf{x}_i} - \mathbf{y} \Vert_2^2, & \text{if reduction='mean'} \\ - \Vert \displaystyle\sum_{i=0}^{M}{\mathbf{s}_i \cdot \mathbf{x}_i} - \mathbf{y} \Vert_2^2, & \text{if reduction='sum'} - \end{cases} - $$ - - $$ - \mathbf{x}, \mathbf{s} \in \mathcal{R}^{M \times N}, \mathbf{y} \in \mathcal{R}^{N} - $$ - - Args: - reduction (Literal["mean", "sum"], optional): Reduction method. Defaults to "mean". - weight (Optional[Union[float, Dict[str, float]]]): Weight for loss. Defaults to None. - - Examples: - >>> import paddle - >>> from ppsci.loss import IntegralLoss - - >>> output_dict = {'u': paddle.to_tensor([[0.5, 2.2, 0.9], [1.1, 0.8, -1.3]]), - ... 'v': paddle.to_tensor([[0.5, 2.2, 0.9], [1.1, 0.8, -1.3]]), - ... 'area': paddle.to_tensor([[0.01, 0.02, 0.03], [0.01, 0.02, 0.03]])} - >>> label_dict = {'u': paddle.to_tensor([-1.8, 0.0]), - ... 'v': paddle.to_tensor([0.1, 0.1])} - >>> weight = {'u': 0.8, 'v': 0.2} - >>> loss = IntegralLoss(weight=weight) - >>> result = loss(output_dict, label_dict) - >>> print(result) - {'u': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, - 1.40780795), 'v': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, - 0.00131200)} - - >>> loss = IntegralLoss(reduction="sum", weight=weight) - >>> result = loss(output_dict, label_dict) - >>> print(result) - {'u': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, - 2.81561589), 'v': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, - 0.00262400)} - """ - - def __init__( - self, - reduction: Literal["mean", "sum"] = "mean", - weight: Optional[Union[float, Dict[str, float]]] = None, - ): - if reduction not in ["mean", "sum"]: - raise ValueError( - f"reduction should be 'mean' or 'sum', but got {reduction}" - ) - super().__init__(reduction, weight) - - def forward( - self, output_dict, label_dict, weight_dict=None - ) -> Dict[str, "paddle.Tensor"]: - losses = {} - - for key in label_dict: - loss = F.mse_loss( - (output_dict[key] * output_dict["area"]).sum(axis=1), - label_dict[key], - "none", - ) - if weight_dict and key in weight_dict: - loss *= weight_dict[key] - - if self.reduction == "sum": - loss = loss.sum() - elif self.reduction == "mean": - loss = loss.mean() - - if isinstance(self.weight, (float, int)): - loss *= self.weight - elif isinstance(self.weight, dict) and key in self.weight: - loss *= self.weight[key] - - losses[key] = loss - - return losses diff --git a/examples/smc_reac/ppsci/loss/kl.py b/examples/smc_reac/ppsci/loss/kl.py deleted file mode 100644 index c07c3ed2c6..0000000000 --- a/examples/smc_reac/ppsci/loss/kl.py +++ /dev/null @@ -1,51 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import Dict -from typing import Optional -from typing import Union - -import paddle -from typing_extensions import Literal - -from ppsci.loss import base - - -class KLLoss(base.Loss): - def __init__( - self, - reduction: Literal["mean", "sum"] = "mean", - weight: Optional[Union[float, Dict[str, float]]] = None, - ): - if reduction not in ["mean", "sum"]: - raise ValueError( - f"reduction should be 'mean' or 'sum', but got {reduction}" - ) - super().__init__(reduction, weight) - - def forward( - self, output_dict, label_dict=None, weight_dict=None - ) -> Dict[str, "paddle.Tensor"]: - losses = {} - - mu, log_sigma = output_dict["mu"], output_dict["log_sigma"] - - base = paddle.exp(2.0 * log_sigma) + paddle.pow(mu, 2) - 1.0 - 2.0 * log_sigma - loss = 0.5 * paddle.sum(base) / mu.shape[0] - - losses["kl_loss"] = loss - - return loss diff --git a/examples/smc_reac/ppsci/loss/l1.py b/examples/smc_reac/ppsci/loss/l1.py deleted file mode 100644 index 3edbc2e102..0000000000 --- a/examples/smc_reac/ppsci/loss/l1.py +++ /dev/null @@ -1,219 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import TYPE_CHECKING -from typing import Dict -from typing import Optional -from typing import Union - -import paddle.nn.functional as F -from typing_extensions import Literal - -from ppsci.loss import base - -if TYPE_CHECKING: - - import paddle - - -class L1Loss(base.Loss): - r"""Class for l1 loss. - - $$ - L = \Vert \mathbf{x} - \mathbf{y} \Vert_1 - $$ - - $$ - \mathbf{x}, \mathbf{y} \in \mathcal{R}^{N} - $$ - - when `reduction` is set to "mean" - - $$ - L = MEAN \left( \Vert \mathbf{x} - \mathbf{y} \Vert_1 \right) - $$ - - when `reduction` is set to "sum" - - $$ - L = SUM \left( \Vert \mathbf{x} - \mathbf{y} \Vert_1 \right) - $$ - - Args: - reduction (Literal["mean", "sum"], optional): Reduction method. Defaults to "mean". - weight (Optional[Union[float, Dict[str, float]]]): Weight for loss. Defaults to None. - - Examples: - >>> import paddle - >>> from ppsci.loss import L1Loss - >>> output_dict = {"u": paddle.to_tensor([[0.5, 0.9], [1.1, -1.3]]), - ... "v": paddle.to_tensor([[0.5, 0.9], [1.1, -1.3]])} - >>> label_dict = {"u": paddle.to_tensor([[-1.8, 1.0], [-0.2, 2.5]]), - ... "v": paddle.to_tensor([[0.1, 0.1], [0.1, 0.1]])} - >>> weight = {"u": 0.8, "v": 0.2} - >>> loss = L1Loss(weight=weight) - >>> result = loss(output_dict, label_dict) - >>> print(result) - {'u': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, - 3.), 'v': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, - 0.35999998)} - - >>> loss = L1Loss(reduction="sum", weight=weight) - >>> result = loss(output_dict, label_dict) - >>> print(result) - {'u': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, - 6.), 'v': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, - 0.71999997)} - """ - - def __init__( - self, - reduction: Literal["mean", "sum"] = "mean", - weight: Optional[Union[float, Dict[str, float]]] = None, - ): - if reduction not in ["mean", "sum"]: - raise ValueError( - f"reduction should be 'mean' or 'sum', but got {reduction}" - ) - super().__init__(reduction, weight) - - def forward( - self, output_dict, label_dict, weight_dict=None - ) -> Dict[str, "paddle.Tensor"]: - losses = {} - - for key in label_dict: - loss = F.l1_loss(output_dict[key], label_dict[key], "none") - if weight_dict and key in weight_dict: - loss *= weight_dict[key] - - if "area" in output_dict: - loss *= output_dict["area"] - - loss = loss.sum(axis=1) - - if self.reduction == "sum": - loss = loss.sum() - elif self.reduction == "mean": - loss = loss.mean() - - if isinstance(self.weight, (float, int)): - loss *= self.weight - elif isinstance(self.weight, dict) and key in self.weight: - loss *= self.weight[key] - - losses[key] = loss - - return losses - - -class PeriodicL1Loss(base.Loss): - r"""Class for periodic l1 loss. - - $$ - L = \Vert \mathbf{x_l}-\mathbf{x_r} \Vert_1 - $$ - - $\mathbf{x_l} \in \mathcal{R}^{N}$ is the first half of batch output, - $\mathbf{x_r} \in \mathcal{R}^{N}$ is the second half of batch output. - - when `reduction` is set to "mean" - - $$ - L = MEAN \left( \Vert \mathbf{x_l}-\mathbf{x_r} \Vert_1 \right) - $$ - - when `reduction` is set to "sum" - - $$ - L = SUM \left( \Vert \mathbf{x_l}-\mathbf{x_r} \Vert_1 \right) - $$ - - Args: - reduction (Literal["mean", "sum"], optional): Reduction method. Defaults to "mean". - weight (Optional[Union[float, Dict[str, float]]]): Weight for loss. Defaults to None. - - Examples: - >>> import paddle - >>> from ppsci.loss import PeriodicL1Loss - - >>> output_dict = {'u': paddle.to_tensor([[0.5, 2.2, 0.9], [1.1, 0.8, -1.3]]), - ... 'v': paddle.to_tensor([[0.5, 2.2, 0.9], [1.1, 0.8, -1.3]])} - >>> label_dict = {'u': paddle.to_tensor([[-1.8, 0.0, 1.0], [-0.2, 0.2, 2.5]]), - ... 'v': paddle.to_tensor([[0.1, 0.1, 0.1], [0.1, 0.1, 0.1]])} - >>> weight = {'u': 0.8, 'v': 0.2} - >>> loss = PeriodicL1Loss(weight=weight) - >>> result = loss(output_dict, label_dict) - >>> print(result) - {'u': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, - 3.35999990), 'v': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, - 0.83999997)} - - >>> loss = PeriodicL1Loss(reduction="sum", weight=weight) - >>> result = loss(output_dict, label_dict) - >>> print(result) - {'u': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, - 3.35999990), 'v': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, - 0.83999997)} - """ - - def __init__( - self, - reduction: Literal["mean", "sum"] = "mean", - weight: Optional[Union[float, Dict[str, float]]] = None, - ): - if reduction not in ["mean", "sum"]: - raise ValueError( - f"reduction should be 'mean' or 'sum', but got {reduction}" - ) - super().__init__(reduction, weight) - - def forward( - self, output_dict, label_dict, weight_dict=None - ) -> Dict[str, "paddle.Tensor"]: - losses = {} - - for key in label_dict: - n_output = len(output_dict[key]) - if n_output % 2 > 0: - raise ValueError( - f"Length of output({n_output}) of key({key}) should be even." - ) - - n_output //= 2 - loss = F.l1_loss( - output_dict[key][:n_output], output_dict[key][n_output:], "none" - ) - if weight_dict and key in weight_dict: - loss *= weight_dict[key] - if "area" in output_dict: - loss *= output_dict["area"] - - loss = loss.sum(axis=1) - - if self.reduction == "sum": - loss = loss.sum() - elif self.reduction == "mean": - loss = loss.mean() - - if isinstance(self.weight, (float, int)): - loss *= self.weight - elif isinstance(self.weight, dict) and key in self.weight: - loss *= self.weight[key] - - losses[key] = loss - - return losses diff --git a/examples/smc_reac/ppsci/loss/l2.py b/examples/smc_reac/ppsci/loss/l2.py deleted file mode 100644 index 7b65a937c6..0000000000 --- a/examples/smc_reac/ppsci/loss/l2.py +++ /dev/null @@ -1,310 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import Dict -from typing import Optional -from typing import Union - -import paddle -import paddle.nn.functional as F -from typing_extensions import Literal - -from ppsci.loss import base - - -class L2Loss(base.Loss): - r"""Class for l2 loss. - - $$ - L =\Vert \mathbf{x} - \mathbf{y} \Vert_2 - $$ - - $$ - \mathbf{x}, \mathbf{y} \in \mathcal{R}^{N} - $$ - - when `reduction` is set to "mean" - - $$ - L = MEAN \left( \Vert \mathbf{x} - \mathbf{y} \Vert_2 \right) - $$ - - when `reduction` is set to "sum" - - $$ - L = SUM \left( \Vert \mathbf{x} - \mathbf{y} \Vert_2 \right) - $$ - - Args: - reduction (Literal["mean", "sum"], optional): Reduction method. Defaults to "mean". - weight (Optional[Union[float, Dict[str, float]]]): Weight for loss. Defaults to None. - - Examples: - >>> import paddle - >>> from ppsci.loss import L2Loss - >>> output_dict = {'u': paddle.to_tensor([[0.5, 0.9], [1.1, -1.3]]), - ... 'v': paddle.to_tensor([[0.5, 0.9], [1.1, -1.3]])} - >>> label_dict = {'u': paddle.to_tensor([[-1.8, 1.0], [-0.2, 2.5]]), - ... 'v': paddle.to_tensor([[0.1, 0.1], [0.1, 0.1]])} - >>> weight = {'u': 0.8, 'v': 0.2} - >>> loss = L2Loss(weight=weight) - >>> result = loss(output_dict, label_dict) - >>> print(result) - {'u': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, - 2.52735591), 'v': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, - 0.26148924)} - >>> loss = L2Loss(reduction="sum", weight=weight) - >>> result = loss(output_dict, label_dict) - >>> print(result) - {'u': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, - 5.05471182), 'v': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, - 0.52297848)} - """ - - def __init__( - self, - reduction: Literal["mean", "sum"] = "mean", - weight: Optional[Union[float, Dict[str, float]]] = None, - ): - if reduction not in ["mean", "sum"]: - raise ValueError( - f"reduction should be 'mean' or 'sum', but got {reduction}" - ) - super().__init__(reduction, weight) - - def forward( - self, output_dict, label_dict, weight_dict=None - ) -> Dict[str, "paddle.Tensor"]: - losses = {} - - for key in label_dict: - loss = F.mse_loss(output_dict[key], label_dict[key], "none") - if weight_dict and key in weight_dict: - loss *= weight_dict[key] - - if "area" in output_dict: - loss *= output_dict["area"] - - loss = loss.sum(axis=1).sqrt() - - if self.reduction == "sum": - loss = loss.sum() - elif self.reduction == "mean": - loss = loss.mean() - - if isinstance(self.weight, (float, int)): - loss *= self.weight - elif isinstance(self.weight, dict) and key in self.weight: - loss *= self.weight[key] - - losses[key] = loss - - return losses - - -class PeriodicL2Loss(base.Loss): - r"""Class for Periodic l2 loss. - - $$ - L = \Vert \mathbf{x_l}-\mathbf{x_r} \Vert_2 - $$ - - $\mathbf{x_l} \in \mathcal{R}^{N}$ is the first half of batch output, - $\mathbf{x_r} \in \mathcal{R}^{N}$ is the second half of batch output. - - when `reduction` is set to "mean" - - $$ - L = MEAN \left( \Vert \mathbf{x_l}-\mathbf{x_r} \Vert_2 \right) - $$ - - when `reduction` is set to "sum" - - $$ - L = SUM \left( \Vert \mathbf{x_l}-\mathbf{x_r} \Vert_2 \right) - $$ - - Args: - reduction (Literal["mean", "sum"], optional): Reduction method. Defaults to "mean". - weight (Optional[Union[float, Dict[str, float]]]): Weight for loss. Defaults to None. - - Examples: - >>> import paddle - >>> from ppsci.loss import PeriodicL2Loss - - >>> output_dict = {'u': paddle.to_tensor([[0.5, 2.2, 0.9], [1.1, 0.8, -1.3]]), - ... 'v': paddle.to_tensor([[0.5, 2.2, 0.9], [1.1, 0.8, -1.3]])} - >>> label_dict = {'u': paddle.to_tensor([[-1.8, 0.0, 1.0], [-0.2, 0.2, 2.5]]), - ... 'v': paddle.to_tensor([[0.1, 0.1, 0.1], [0.1, 0.1, 0.1]])} - >>> weight = {'u': 0.8, 'v': 0.2} - >>> loss = PeriodicL2Loss(weight=weight) - >>> result = loss(output_dict, label_dict) - >>> print(result) - {'u': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, - 2.14065409), 'v': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, - 0.53516352)} - - >>> loss = PeriodicL2Loss(reduction="sum", weight=weight) - >>> result = loss(output_dict, label_dict) - >>> print(result) - {'u': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, - 2.14065409), 'v': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, - 0.53516352)} - """ - - def __init__( - self, - reduction: Literal["mean", "sum"] = "mean", - weight: Optional[Union[float, Dict[str, float]]] = None, - ): - if reduction not in ["mean", "sum"]: - raise ValueError( - f"reduction should be 'mean' or 'sum', but got {reduction}" - ) - super().__init__(reduction, weight) - - def forward( - self, output_dict, label_dict, weight_dict=None - ) -> Dict[str, "paddle.Tensor"]: - losses = {} - - for key in label_dict: - n_output = len(output_dict[key]) - if n_output % 2 > 0: - raise ValueError( - f"Length of output({n_output}) of key({key}) should be even." - ) - n_output //= 2 - - loss = F.mse_loss( - output_dict[key][:n_output], output_dict[key][n_output:], "none" - ) - if weight_dict and key in weight_dict: - loss *= weight_dict[key] - - if "area" in output_dict: - loss *= output_dict["area"] - - loss = loss.sum(axis=1).sqrt() - - if self.reduction == "sum": - loss = loss.sum() - elif self.reduction == "mean": - loss = loss.mean() - - if isinstance(self.weight, (float, int)): - loss *= self.weight - elif isinstance(self.weight, dict) and key in self.weight: - loss *= self.weight[key] - - losses[key] = loss - - return losses - - -class L2RelLoss(base.Loss): - r"""Class for l2 relative loss. - - $$ - L = \dfrac{\Vert \mathbf{x} - \mathbf{y} \Vert_2}{\Vert \mathbf{y} \Vert_2} - $$ - - $$ - \mathbf{x}, \mathbf{y} \in \mathcal{R}^{N} - $$ - - when `reduction` is set to "mean" - - $$ - L = MEAN \left( \dfrac{\Vert \mathbf{x} - \mathbf{y} \Vert_2}{\Vert \mathbf{y} \Vert_2} \right) - $$ - - when `reduction` is set to "sum" - - $$ - L = SUM \left( \dfrac{\Vert \mathbf{x} - \mathbf{y} \Vert_2}{\Vert \mathbf{y} \Vert_2} \right) - $$ - - Args: - reduction (Literal["mean", "sum"], optional): Specifies the reduction to apply to the output: 'mean' | 'sum'. Defaults to "mean". - weight (Optional[Union[float, Dict[str, float]]]): Weight for loss. Defaults to None. - - Examples: - >>> import paddle - >>> from ppsci.loss import L2RelLoss - - >>> output_dict = {'u': paddle.to_tensor([[0.5, 0.9], [1.1, -1.3]]), - ... 'v': paddle.to_tensor([[0.5, 0.9], [1.1, -1.3]])} - >>> label_dict = {'u': paddle.to_tensor([[-1.8, 1.0], [-0.2, 2.5]]), - ... 'v': paddle.to_tensor([[0.1, 0.1], [0.1, 0.1]])} - >>> weight = {'u': 0.8, 'v': 0.2} - >>> loss = L2RelLoss(weight=weight) - >>> result = loss(output_dict, label_dict) - >>> print(result) - {'u': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, - 1.08776188), 'v': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, - 1.84900820)} - - >>> loss = L2RelLoss(reduction="sum", weight=weight) - >>> result = loss(output_dict, label_dict) - >>> print(result) - {'u': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, - 2.17552376), 'v': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, - 3.69801641)} - """ - - def __init__( - self, - reduction: Literal["mean", "sum"] = "mean", - weight: Optional[Union[float, Dict[str, float]]] = None, - ): - if reduction not in ["mean", "sum"]: - raise ValueError( - f"reduction should be 'mean' or 'sum', but got {reduction}" - ) - super().__init__(reduction, weight) - - def rel_loss(self, x, y): - batch_size = x.shape[0] - x_ = x.reshape((batch_size, -1)) - y_ = y.reshape((batch_size, -1)) - diff_norms = paddle.norm(x_ - y_, p=2, axis=1) - y_norms = paddle.norm(y_, p=2, axis=1) - return diff_norms / y_norms - - def forward( - self, output_dict, label_dict, weight_dict=None - ) -> Dict[str, "paddle.Tensor"]: - losses = {} - - for key in label_dict: - loss = self.rel_loss(output_dict[key], label_dict[key]) - if weight_dict: - loss *= weight_dict[key] - - if self.reduction == "sum": - loss = loss.sum() - elif self.reduction == "mean": - loss = loss.mean() - - if isinstance(self.weight, float): - loss *= self.weight - elif isinstance(self.weight, dict) and key in self.weight: - loss *= self.weight[key] - - losses[key] = loss - - return losses diff --git a/examples/smc_reac/ppsci/loss/mae.py b/examples/smc_reac/ppsci/loss/mae.py deleted file mode 100644 index ff7869535c..0000000000 --- a/examples/smc_reac/ppsci/loss/mae.py +++ /dev/null @@ -1,109 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import TYPE_CHECKING -from typing import Dict -from typing import Optional -from typing import Union - -import paddle.nn.functional as F -from typing_extensions import Literal - -from ppsci.loss import base - -if TYPE_CHECKING: - import paddle - - -class MAELoss(base.Loss): - r"""Class for mean absolute error loss. - - $$ - L = - \begin{cases} - \dfrac{1}{N} \Vert {\mathbf{x}-\mathbf{y}} \Vert_1, & \text{if reduction='mean'} \\ - \Vert {\mathbf{x}-\mathbf{y}} \Vert_1, & \text{if reduction='sum'} - \end{cases} - $$ - - $$ - \mathbf{x}, \mathbf{y} \in \mathcal{R}^{N} - $$ - - Args: - reduction (Literal["mean", "sum"], optional): Reduction method. Defaults to "mean". - weight (Optional[Union[float, Dict[str, float]]]): Weight for loss. Defaults to None. - - Examples: - >>> import paddle - >>> from ppsci.loss import MAELoss - - >>> output_dict = {'u': paddle.to_tensor([[0.5, 0.9], [1.1, -1.3]]), - ... 'v': paddle.to_tensor([[0.5, 0.9], [1.1, -1.3]])} - >>> label_dict = {'u': paddle.to_tensor([[-1.8, 1.0], [-0.2, 2.5]]), - ... 'v': paddle.to_tensor([[0.1, 0.1], [0.1, 0.1]])} - >>> weight = {'u': 0.8, 'v': 0.2} - >>> loss = MAELoss(weight=weight) - >>> result = loss(output_dict, label_dict) - >>> print(result) - {'u': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, - 1.50000000), 'v': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, - 0.17999999)} - - >>> loss = MAELoss(reduction="sum", weight=weight) - >>> result = loss(output_dict, label_dict) - >>> print(result) - {'u': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, - 6.), 'v': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, - 0.71999997)} - """ - - def __init__( - self, - reduction: Literal["mean", "sum"] = "mean", - weight: Optional[Union[float, Dict[str, float]]] = None, - ): - if reduction not in ["mean", "sum"]: - raise ValueError( - f"reduction should be 'mean' or 'sum', but got {reduction}" - ) - super().__init__(reduction, weight) - - def forward( - self, output_dict, label_dict, weight_dict=None - ) -> Dict[str, "paddle.Tensor"]: - losses = {} - - for key in label_dict: - loss = F.l1_loss(output_dict[key], label_dict[key], "none") - if weight_dict and key in weight_dict: - loss *= weight_dict[key] - - if "area" in output_dict: - loss *= output_dict["area"] - - if self.reduction == "sum": - loss = loss.sum() - elif self.reduction == "mean": - loss = loss.mean() - if isinstance(self.weight, (float, int)): - loss *= self.weight - elif isinstance(self.weight, dict) and key in self.weight: - loss *= self.weight[key] - - losses[key] = loss - - return losses diff --git a/examples/smc_reac/ppsci/loss/mse.py b/examples/smc_reac/ppsci/loss/mse.py deleted file mode 100644 index 184f11cc40..0000000000 --- a/examples/smc_reac/ppsci/loss/mse.py +++ /dev/null @@ -1,355 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import Dict -from typing import Optional -from typing import Union - -import paddle -import paddle.nn.functional as F -from typing_extensions import Literal - -from ppsci.loss import base - - -class MSELoss(base.Loss): - r"""Class for mean squared error loss. - - $$ - L = - \begin{cases} - \dfrac{1}{N} \Vert {\mathbf{x}-\mathbf{y}} \Vert_2^2, & \text{if reduction='mean'} \\ - \Vert {\mathbf{x}-\mathbf{y}} \Vert_2^2, & \text{if reduction='sum'} - \end{cases} - $$ - - $$ - \mathbf{x}, \mathbf{y} \in \mathcal{R}^{N} - $$ - - Args: - reduction (Literal["mean", "sum"], optional): Reduction method. Defaults to "mean". - weight (Optional[Union[float, Dict[str, float]]]): Weight for loss. Defaults to None. - - Examples: - >>> import paddle - >>> from ppsci.loss import MSELoss - - >>> output_dict = {'u': paddle.to_tensor([[0.5, 0.9], [1.1, -1.3]]), - ... 'v': paddle.to_tensor([[0.5, 0.9], [1.1, -1.3]])} - >>> label_dict = {'u': paddle.to_tensor([[-1.8, 1.0], [-0.2, 2.5]]), - ... 'v': paddle.to_tensor([[0.1, 0.1], [0.1, 0.1]])} - >>> weight = {'u': 0.8, 'v': 0.2} - >>> loss = MSELoss(weight=weight) - >>> result = loss(output_dict, label_dict) - >>> print(result) - {'u': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, - 4.28600025), 'v': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, - 0.18800001)} - - >>> loss = MSELoss(reduction="sum", weight=weight) - >>> result = loss(output_dict, label_dict) - >>> print(result) - {'u': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, - 17.14400101), 'v': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, - 0.75200003)} - """ - - def __init__( - self, - reduction: Literal["mean", "sum"] = "mean", - weight: Optional[Union[float, Dict[str, float]]] = None, - ): - if reduction not in ["mean", "sum"]: - raise ValueError( - f"reduction should be 'mean' or 'sum', but got {reduction}" - ) - super().__init__(reduction, weight) - - def forward( - self, output_dict, label_dict, weight_dict=None - ) -> Dict[str, "paddle.Tensor"]: - losses = {} - - for key in label_dict: - loss = F.mse_loss(output_dict[key], label_dict[key], "none") - if weight_dict and key in weight_dict: - loss *= weight_dict[key] - - if "area" in output_dict: - loss *= output_dict["area"] - - if self.reduction == "sum": - loss = loss.sum() - elif self.reduction == "mean": - loss = loss.mean() - if isinstance(self.weight, (float, int)): - loss *= self.weight - elif isinstance(self.weight, dict) and key in self.weight: - loss *= self.weight[key] - - losses[key] = loss - - return losses - - -class CausalMSELoss(base.Loss): - r"""Class for mean squared error loss. - - $$ - L = \frac{1}{M} \displaystyle\sum_{i=1}^M{w_i} \mathcal{L}_r^i, - $$ - - where $w_i=\exp (-\epsilon \displaystyle\sum_{k=1}^{i-1} \mathcal{L}_r^k), i=2,3, \ldots, M.$ - - Args: - n_chunks (int): $M$, Number of split time windows. - reduction (Literal["mean", "sum"], optional): Reduction method. Defaults to "mean". - weight (Optional[Union[float, Dict[str, float]]]): Weight for loss. Defaults to None. - tol (float, optional): Causal tolerance, i.e. $\epsilon$ in paper. Defaults to 1.0. - - Examples: - >>> import paddle - >>> from ppsci.loss import CausalMSELoss - - >>> output_dict = {'u': paddle.to_tensor([[0.5, 0.9, 1.0], [1.1, -1.3, 0.0]])} - >>> label_dict = {'u': paddle.to_tensor([[-1.8, 1.0, -0.1], [-0.2, 2.5, 2.0]])} - >>> loss = CausalMSELoss(n_chunks=3) - >>> result = loss(output_dict, label_dict) - >>> print(result) - {'u': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, - 0.96841478)} - """ - - def __init__( - self, - n_chunks: int, - reduction: Literal["mean", "sum"] = "mean", - weight: Optional[Union[float, Dict[str, float]]] = None, - tol: float = 1.0, - ): - if n_chunks <= 0: - raise ValueError(f"n_chunks should be positive, but got {n_chunks}") - if reduction not in ["mean", "sum"]: - raise ValueError( - f"reduction should be 'mean' or 'sum', but got {reduction}" - ) - super().__init__(reduction, weight) - self.n_chunks = n_chunks - self.tol = tol - self.register_buffer( - "acc_mat", paddle.tril(paddle.ones([n_chunks, n_chunks]), -1) - ) - - def forward( - self, output_dict, label_dict, weight_dict=None - ) -> Dict[str, "paddle.Tensor"]: - losses = {} - - for key in label_dict: - loss = F.mse_loss(output_dict[key], label_dict[key], "none") - if weight_dict and key in weight_dict: - loss *= weight_dict[key] - - if "area" in output_dict: - loss *= output_dict["area"] - - # causal weighting - loss_t = loss.reshape([self.n_chunks, -1]) # [nt, nx] - weight_t = paddle.exp( - -self.tol * (self.acc_mat @ loss_t.mean(-1, keepdim=True)) - ) # [nt, nt] x [nt, 1] ==> [nt, 1] - assert weight_t.shape[0] == self.n_chunks - loss = loss_t * weight_t.detach() - - if self.reduction == "sum": - loss = loss.sum() - elif self.reduction == "mean": - loss = loss.mean() - if isinstance(self.weight, (float, int)): - loss *= self.weight - elif isinstance(self.weight, dict) and key in self.weight: - loss *= self.weight[key] - - losses[key] = loss - - return losses - - -class MSELossWithL2Decay(MSELoss): - r"""MSELoss with L2 decay. - - $$ - L = - \begin{cases} - \dfrac{1}{N} \Vert {\mathbf{x}-\mathbf{y}} \Vert_2^2 + \displaystyle\sum_{i=1}^{M}{\Vert \mathbf{K_i} \Vert_F^2}, & \text{if reduction='mean'} \\ - \Vert {\mathbf{x}-\mathbf{y}} \Vert_2^2 + \displaystyle\sum_{i=1}^{M}{\Vert \mathbf{K_i} \Vert_F^2}, & \text{if reduction='sum'} - \end{cases} - $$ - - $$ - \mathbf{x}, \mathbf{y} \in \mathcal{R}^{N}, \mathbf{K_i} \in \mathcal{R}^{O_i \times P_i} - $$ - - $M$ is the number of which apply regularization on. - - Args: - reduction (Literal["mean", "sum"], optional): Specifies the reduction to apply to the output: 'mean' | 'sum'. Defaults to "mean". - regularization_dict (Optional[Dict[str, float]]): Regularization dictionary. Defaults to None. - weight (Optional[Union[float, Dict[str, float]]]): Weight for loss. Defaults to None. - - Raises: - ValueError: reduction should be 'mean' or 'sum'. - - Examples: - >>> import paddle - >>> from ppsci.loss import MSELossWithL2Decay - - >>> output_dict = {'u': paddle.to_tensor([[0.5, 0.9], [1.1, -1.3]]), - ... 'v': paddle.to_tensor([[0.5, 0.9], [1.1, -1.3]])} - >>> label_dict = {'u': paddle.to_tensor([[-1.8, 1.0], [-0.2, 2.5]]), - ... 'v': paddle.to_tensor([[0.1, 0.1], [0.1, 0.1]])} - >>> weight = {'u': 0.8, 'v': 0.2} - >>> regularization_dict = {'u': 2.0} - >>> loss = MSELossWithL2Decay(regularization_dict=regularization_dict, weight=weight) - >>> result = loss(output_dict, label_dict) - >>> print(result) - {'u': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, - 7.91999960), 'v': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, - 0.18800001)} - - >>> regularization_dict = {'v': 1.0} - >>> loss = MSELossWithL2Decay(reduction="sum", regularization_dict=regularization_dict, weight=weight) - >>> result = loss(output_dict, label_dict) - >>> print(result) - {'u': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, - 17.14400101), 'v': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, - 3.95999980)} - """ - - def __init__( - self, - reduction: Literal["mean", "sum"] = "mean", - regularization_dict: Optional[Dict[str, float]] = None, - weight: Optional[Union[float, Dict[str, float]]] = None, - ): - if reduction not in ["mean", "sum"]: - raise ValueError( - f"reduction should be 'mean' or 'sum', but got {reduction}" - ) - super().__init__(reduction, weight) - self.regularization_dict = regularization_dict - - def forward( - self, output_dict, label_dict, weight_dict=None - ) -> Dict[str, "paddle.Tensor"]: - losses = super().forward(output_dict, label_dict, weight_dict) - - if self.regularization_dict is not None: - for reg_key, reg_weight in self.regularization_dict.items(): - loss = output_dict[reg_key].pow(2).sum() - losses[reg_key] = loss * reg_weight - - return losses - - -class PeriodicMSELoss(base.Loss): - r"""Class for periodic mean squared error loss. - - $$ - L = - \begin{cases} - \dfrac{1}{N} \Vert \mathbf{x_l}-\mathbf{x_r} \Vert_2^2, & \text{if reduction='mean'} \\ - \Vert \mathbf{x_l}-\mathbf{x_r} \Vert_2^2, & \text{if reduction='sum'} - \end{cases} - $$ - - $\mathbf{x_l} \in \mathcal{R}^{N}$ is the first half of batch output, - $\mathbf{x_r} \in \mathcal{R}^{N}$ is the second half of batch output. - - Args: - reduction (Literal["mean", "sum"], optional): Reduction method. Defaults to "mean". - weight (Optional[Union[float, Dict[str, float]]]): Weight for loss. Defaults to None. - - Examples: - >>> import paddle - >>> from ppsci.loss import PeriodicMSELoss - - >>> output_dict = {'u': paddle.to_tensor([[0.5, 0.9], [1.1, -1.3]]), - ... 'v': paddle.to_tensor([[0.5, 0.9], [1.1, -1.3]])} - >>> label_dict = {'u': paddle.to_tensor([[-1.8, 1.0], [-0.2, 2.5]]), - ... 'v': paddle.to_tensor([[0.1, 0.1], [0.1, 0.1]])} - >>> weight = {'u': 0.8, 'v': 0.2} - >>> loss = PeriodicMSELoss(weight=weight) - >>> result = loss(output_dict, label_dict) - >>> print(result) - {'u': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, - 2.07999969), 'v': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, - 0.51999992)} - - >>> loss = PeriodicMSELoss(reduction="sum", weight=weight) - >>> result = loss(output_dict, label_dict) - >>> print(result) - {'u': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, - 4.15999937), 'v': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, - 1.03999984)} - """ - - def __init__( - self, - reduction: Literal["mean", "sum"] = "mean", - weight: Optional[Union[float, Dict[str, float]]] = None, - ): - if reduction not in ["mean", "sum"]: - raise ValueError( - f"reduction should be 'mean' or 'sum', but got {reduction}" - ) - super().__init__(reduction, weight) - - def forward( - self, output_dict, label_dict, weight_dict=None - ) -> Dict[str, "paddle.Tensor"]: - losses = {} - - for key in label_dict: - n_output = len(output_dict[key]) - if n_output % 2 > 0: - raise ValueError( - f"Length of output({n_output}) of key({key}) should be even." - ) - - n_output //= 2 - loss = F.mse_loss( - output_dict[key][:n_output], output_dict[key][n_output:], "none" - ) - if weight_dict: - loss *= weight_dict[key] - if "area" in output_dict: - loss *= output_dict["area"] - - if self.reduction == "sum": - loss = loss.sum() - elif self.reduction == "mean": - loss = loss.mean() - - if isinstance(self.weight, (float, int)): - loss *= self.weight - elif isinstance(self.weight, dict) and key in self.weight: - loss *= self.weight[key] - - losses[key] = loss - - return losses diff --git a/examples/smc_reac/ppsci/loss/mtl/__init__.py b/examples/smc_reac/ppsci/loss/mtl/__init__.py deleted file mode 100644 index 3bff2aaa7f..0000000000 --- a/examples/smc_reac/ppsci/loss/mtl/__init__.py +++ /dev/null @@ -1,49 +0,0 @@ -# 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. - -import copy - -from ppsci.loss.mtl.agda import AGDA -from ppsci.loss.mtl.base import LossAggregator -from ppsci.loss.mtl.grad_norm import GradNorm -from ppsci.loss.mtl.ntk import NTK -from ppsci.loss.mtl.pcgrad import PCGrad -from ppsci.loss.mtl.relobralo import Relobralo -from ppsci.loss.mtl.sum import Sum - -__all__ = [ - "AGDA", - "GradNorm", - "LossAggregator", - "PCGrad", - "Relobralo", - "Sum", - "NTK", -] - - -def build_mtl_aggregator(cfg): - """Build loss aggregator with multi-task learning method. - - Args: - cfg (DictConfig): Aggregator config. - - Returns: - Loss: Callable loss aggregator object. - """ - cfg = copy.deepcopy(cfg) - - aggregator_cls = cfg.pop("name") - aggregator = eval(aggregator_cls)(**cfg) - return aggregator diff --git a/examples/smc_reac/ppsci/loss/mtl/agda.py b/examples/smc_reac/ppsci/loss/mtl/agda.py deleted file mode 100644 index 31193f5b76..0000000000 --- a/examples/smc_reac/ppsci/loss/mtl/agda.py +++ /dev/null @@ -1,161 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import ClassVar -from typing import List - -import paddle -from paddle import nn - -from ppsci.loss.mtl import base - - -class AGDA(base.LossAggregator): - r""" - **A**daptive **G**radient **D**escent **A**lgorithm - - [Physics-informed neural network based on a new adaptive gradient descent algorithm for solving partial differential equations of flow problems](https://pubs.aip.org/aip/pof/article-abstract/35/6/063608/2899773/Physics-informed-neural-network-based-on-a-new) - - NOTE: This loss aggregator is only suitable for two-task learning and the first task loss must be PDE loss. - - Attributes: - should_persist(bool): Whether to persist the loss aggregator when saving. - Those loss aggregators with parameters and/or buffers should be persisted. - - Args: - model (nn.Layer): Training model. - M (int, optional): Smoothing period. Defaults to 100. - gamma (float, optional): Smooth factor. Defaults to 0.999. - - Examples: - >>> import paddle - >>> from ppsci.loss import mtl - >>> model = paddle.nn.Linear(3, 4) - >>> loss_aggregator = mtl.AGDA(model) - >>> for i in range(5): - ... x1 = paddle.randn([8, 3]) - ... x2 = paddle.randn([8, 3]) - ... y1 = model(x1) - ... y2 = model(x2) - ... pde_loss = paddle.sum(y1) - ... bc_loss = paddle.sum((y2 - 2) ** 2) - ... loss_aggregator({'pde_loss': pde_loss, 'bc_loss': bc_loss}).backward() - """ - should_persist: ClassVar[bool] = False - - def __init__(self, model: nn.Layer, M: int = 100, gamma: float = 0.999) -> None: - super().__init__(model) - self.M = M - self.gamma = gamma - self.Lf_smooth = 0 - self.Lu_smooth = 0 - self.Lf_tilde_acc = 0.0 - self.Lu_tilde_acc = 0.0 - - def __call__(self, losses, step: int = 0) -> "AGDA": - if len(losses) != 2: - raise ValueError( - f"Number of losses(tasks) for AGDA should be 2, but got {len(losses)}" - ) - return super().__call__(losses, step) - - def backward(self) -> None: - grads_list = self._compute_grads() - with paddle.no_grad(): - refined_grads = self._refine_grads(grads_list) - self._set_grads(refined_grads) - - def _compute_grads(self) -> List[paddle.Tensor]: - # compute all gradients derived by each loss - grads_list = [] # num_params x num_losses - for key in self.losses: - # backward with current loss - self.losses[key].backward() - grads_list.append( - paddle.concat( - [ - param.grad.clone().reshape([-1]) - for param in self.model.parameters() - if param.grad is not None - ], - axis=0, - ) - ) - # clear gradients for current loss for not affecting other loss - self.model.clear_gradients() - - return grads_list - - def _refine_grads(self, grads_list: List[paddle.Tensor]) -> List[paddle.Tensor]: - # compute moving average of L^smooth_i(n) - eq.(16) - losses_seq = list(self.losses.values()) - self.Lf_smooth = ( - self.gamma * self.Lf_smooth + (1 - self.gamma) * losses_seq[0].item() - ) - self.Lu_smooth = ( - self.gamma * self.Lu_smooth + (1 - self.gamma) * losses_seq[1].item() - ) - - # compute L^smooth_i(kM) - eq.(17) - if self.step % self.M == 0: - Lf_smooth_kM = self.Lf_smooth - Lu_smooth_kM = self.Lu_smooth - Lf_tilde = self.Lf_smooth / Lf_smooth_kM - Lu_tilde = self.Lu_smooth / Lu_smooth_kM - - # compute r_i(n) - eq.(18) - self.Lf_tilde_acc += Lf_tilde - self.Lu_tilde_acc += Lu_tilde - rf = Lf_tilde / self.Lf_tilde_acc - ru = Lu_tilde / self.Lu_tilde_acc - - # compute E(g(n)) - step1(1) - gf_magn = (grads_list[0] * grads_list[0]).sum().sqrt() - gu_magn = (grads_list[1] * grads_list[1]).sum().sqrt() - Eg = (gf_magn + gu_magn) / 2 - - # compute \omega_f(n) - step1(2) - omega_f = (rf * (Eg - gf_magn) + gf_magn) / gf_magn - omega_u = (ru * (Eg - gu_magn) + gu_magn) / gu_magn - - # compute g_bar(n) - step1(3) - gf_bar = omega_f * grads_list[0] - gu_bar = omega_u * grads_list[1] - - # compute gradient projection - step2(1) - dot_product = (gf_bar * gu_bar).sum() - if dot_product < 0: - gu_bar = gu_bar - (dot_product / (gf_bar * gf_bar).sum()) * gf_bar - grads_list = [gf_bar, gu_bar] - - proj_grads: List[paddle.Tensor] = [] - for j in range(len(self.losses)): - start_idx = 0 - for idx, var in enumerate(self.model.parameters()): - grad_shape = var.shape - flatten_dim = var.numel() - refined_grad = grads_list[j][start_idx : start_idx + flatten_dim] - refined_grad = paddle.reshape(refined_grad, grad_shape) - if len(proj_grads) < self.param_num: - proj_grads.append(refined_grad) - else: - proj_grads[idx] += refined_grad - start_idx += flatten_dim - return proj_grads - - def _set_grads(self, grads_list: List[paddle.Tensor]) -> None: - for i, param in enumerate(self.model.parameters()): - param.grad = grads_list[i] diff --git a/examples/smc_reac/ppsci/loss/mtl/base.py b/examples/smc_reac/ppsci/loss/mtl/base.py deleted file mode 100644 index eec88c9c00..0000000000 --- a/examples/smc_reac/ppsci/loss/mtl/base.py +++ /dev/null @@ -1,68 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import TYPE_CHECKING -from typing import ClassVar -from typing import Dict -from typing import Union - -from paddle import nn - -if TYPE_CHECKING: - import paddle - - -class LossAggregator(nn.Layer): - """Base class of loss aggregator mainly for multitask learning. - - Attributes: - should_persist(bool): Whether to persist the loss aggregator when saving. - Those loss aggregators with parameters and/or buffers should be persisted. - - Args: - model (nn.Layer): Training model. - """ - - should_persist: ClassVar[bool] = False - - def __init__(self, model: nn.Layer) -> None: - super().__init__() - self.model = model - self.step = 0 - self.param_num = 0 - for param in self.model.parameters(): - if not param.stop_gradient: - self.param_num += 1 - - def forward( - self, losses: Dict[str, "paddle.Tensor"], step: int = 0 - ) -> Union["paddle.Tensor", "LossAggregator"]: - self.losses = losses - self.loss_num = len(losses) - self.step = step - return self - - def backward(self) -> None: - raise NotImplementedError( - f"'backward' should be implemented in subclass {self.__class__.__name__}" - ) - - def state_dict(self): - agg_state = super().state_dict() - model_state = self.model.state_dict() - # remove model parameters from state dict for already in pdparams - agg_state = {k: v for k, v in agg_state.items() if k not in model_state} - return agg_state diff --git a/examples/smc_reac/ppsci/loss/mtl/grad_norm.py b/examples/smc_reac/ppsci/loss/mtl/grad_norm.py deleted file mode 100644 index 309d3c257c..0000000000 --- a/examples/smc_reac/ppsci/loss/mtl/grad_norm.py +++ /dev/null @@ -1,145 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import ClassVar -from typing import Dict -from typing import List - -import paddle -from paddle import nn - -from ppsci.loss.mtl import base - -# from ppsci.utils import logger - - -class GradNorm(base.LossAggregator): - r"""GradNorm loss weighting algorithm. - - reference: [https://github.com/PredictiveIntelligenceLab/jaxpi/blob/main/jaxpi/models.py#L132-L146](https://github.com/PredictiveIntelligenceLab/jaxpi/blob/main/jaxpi/models.py#L132-L146) - - $$ - \begin{align*} - L^t &= \sum_{i=1}^{N}{\tilde{w}_i^t\cdot L_i^t}, \\ - \text{where } \\ - \tilde{w}_i^0&=1, \\ - \tilde{w}_i^t&=\tilde{w}_i^{t-1}\cdot m+w_i^t\cdot (1-m), t\ge1\\ - w_i^t&=\dfrac{\overline{\Vert \nabla_{\theta}{L_i^t} \Vert_2}}{\Vert \nabla_{\theta}{L_i^t} \Vert_2}, \\ - \overline{\Vert \nabla_{\theta}{L_i^t} \Vert_2}&=\dfrac{1}{N}\sum_{i=1}^N{\Vert \nabla_{\theta}{L_i^t} \Vert_2}, \\ - &t \text{ is the training step started from 0}. - \end{align*} - $$ - - Attributes: - should_persist(bool): Whether to persist the loss aggregator when saving. - Those loss aggregators with parameters and/or buffers should be persisted. - - Args: - model (nn.Layer): Training model. - num_losses (int, optional): Number of losses. Defaults to 1. - update_freq (int, optional): Weight updating frequency. Defaults to 1000. - momentum (float, optional): Momentum $m$ for moving weight. Defaults to 0.9. - init_weights (List[float]): Initial weights list. Defaults to None. - - Examples: - >>> import paddle - >>> from ppsci.loss import mtl - >>> model = paddle.nn.Linear(3, 4) - >>> loss_aggregator = mtl.GradNorm(model, num_losses=2) - >>> for i in range(5): - ... x1 = paddle.randn([8, 3]) - ... x2 = paddle.randn([8, 3]) - ... y1 = model(x1) - ... y2 = model(x2) - ... loss1 = paddle.sum(y1) - ... loss2 = paddle.sum((y2 - 2) ** 2) - ... loss_aggregator({'loss1': loss1, 'loss2': loss2}).backward() - """ - should_persist: ClassVar[bool] = True - weight: paddle.Tensor - - def __init__( - self, - model: nn.Layer, - num_losses: int = 1, - update_freq: int = 1000, - momentum: float = 0.9, - init_weights: List[float] = None, - ) -> None: - super().__init__(model) - self.step = 0 - self.num_losses = num_losses - self.update_freq = update_freq - self.momentum = momentum - if init_weights is not None and num_losses != len(init_weights): - raise ValueError( - f"Length of init_weights({len(init_weights)}) should be equal to " - f"num_losses({num_losses})." - ) - self.register_buffer( - "weight", - paddle.to_tensor(init_weights, dtype="float32") - if init_weights is not None - else paddle.ones([num_losses]), - ) - - def _compute_weight(self, losses: List["paddle.Tensor"]) -> List["paddle.Tensor"]: - grad_norms = [] - for loss in losses: - loss.backward(retain_graph=True) # NOTE: Keep graph for loss backward - with paddle.no_grad(): - grad_vector = paddle.concat( - [ - p.grad.reshape([-1]) - for p in self.model.parameters() - if p.grad is not None - ] - ) - grad_norms.append(paddle.linalg.norm(grad_vector, p=2)) - self.model.clear_gradients() - - mean_grad_norm = paddle.mean(paddle.stack(grad_norms)) - weight = [(mean_grad_norm / x) for x in grad_norms] - - return weight - - def __call__( - self, losses: Dict[str, "paddle.Tensor"], step: int = 0 - ) -> "paddle.Tensor": - assert len(losses) == self.num_losses, ( - f"Length of given losses({len(losses)}) should be equal to " - f"num_losses({self.num_losses})." - ) - self.step = step - - # compute current loss with moving weights - loss = 0.0 - for i, key in enumerate(losses): - if i == 0: - loss = self.weight[i] * losses[key] - else: - loss += self.weight[i] * losses[key] - - # update moving weights every 'update_freq' steps - if self.step % self.update_freq == 0: - weight = self._compute_weight(list(losses.values())) - for i in range(self.num_losses): - self.weight[i].set_value( - self.momentum * self.weight[i] + (1 - self.momentum) * weight[i] - ) - # logger.message(f"weight at step {self.step}: {self.weight.numpy()}") - - return loss diff --git a/examples/smc_reac/ppsci/loss/mtl/ntk.py b/examples/smc_reac/ppsci/loss/mtl/ntk.py deleted file mode 100644 index b2dab91fc7..0000000000 --- a/examples/smc_reac/ppsci/loss/mtl/ntk.py +++ /dev/null @@ -1,118 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import TYPE_CHECKING -from typing import ClassVar -from typing import Dict - -import paddle - -from ppsci.loss.mtl import base - -if TYPE_CHECKING: - from paddle import nn - - -class NTK(base.LossAggregator): - r"""Weighted Neural Tangent Kernel. - - reference: [https://github.com/PredictiveIntelligenceLab/jaxpi/blob/main/jaxpi/models.py#L148-L158](https://github.com/PredictiveIntelligenceLab/jaxpi/blob/main/jaxpi/models.py#L148-L158) - - Attributes: - should_persist(bool): Whether to persist the loss aggregator when saving. - Those loss aggregators with parameters and/or buffers should be persisted. - - Args: - model (nn.Layer): Training model. - num_losses (int, optional): Number of losses. Defaults to 1. - update_freq (int, optional): Weight updating frequency. Defaults to 1000. - - Examples: - >>> import paddle - >>> from ppsci.loss import mtl - >>> model = paddle.nn.Linear(3, 4) - >>> loss_aggregator = mtl.NTK(model, num_losses=2) - >>> for i in range(5): - ... x1 = paddle.randn([8, 3]) - ... x2 = paddle.randn([8, 3]) - ... y1 = model(x1) - ... y2 = model(x2) - ... loss1 = paddle.sum(y1) - ... loss2 = paddle.sum((y2 - 2) ** 2) - ... loss_aggregator({'loss1': loss1, 'loss2': loss2}).backward() - """ - should_persist: ClassVar[bool] = True - weight: paddle.Tensor - - def __init__( - self, - model: nn.Layer, - num_losses: int = 1, - update_freq: int = 1000, - ) -> None: - super().__init__(model) - self.step = 0 - self.num_losses = num_losses - self.update_freq = update_freq - self.register_buffer("weight", paddle.ones([num_losses])) - - def _compute_weight(self, losses: Dict[str, paddle.Tensor]): - ntk_sum = 0 - ntk_value = [] - for loss in losses.values(): - grads = paddle.grad( - loss, - self.model.parameters(), - create_graph=False, - retain_graph=True, - allow_unused=True, - ) - with paddle.no_grad(): - grad = paddle.concat( - [grad.reshape([-1]) for grad in grads if grad is not None] - ) - ntk_value.append( - paddle.sqrt( - paddle.sum(grad.detach() ** 2), - ) - ) - - ntk_sum += paddle.sum(paddle.stack(ntk_value, axis=0)) - ntk_weight = [(ntk_sum / x) for x in ntk_value] - - return ntk_weight - - def __call__( - self, losses: Dict[str, "paddle.Tensor"], step: int = 0 - ) -> "paddle.Tensor": - assert len(losses) == self.num_losses, ( - f"Length of given losses({len(losses)}) should be equal to " - f"num_losses({self.num_losses})." - ) - self.step = step - - # compute current loss with moving weights - loss = 0 - for i, (k, v) in enumerate(losses.items()): - loss = loss + self.weight[i] * v - - # update moving weights every 'update_freq' steps - if self.step % self.update_freq == 0: - computed_weight = self._compute_weight(losses) - for i in range(self.num_losses): - self.weight[i].set_value(computed_weight[i]) - - return loss diff --git a/examples/smc_reac/ppsci/loss/mtl/pcgrad.py b/examples/smc_reac/ppsci/loss/mtl/pcgrad.py deleted file mode 100644 index 45b5923110..0000000000 --- a/examples/smc_reac/ppsci/loss/mtl/pcgrad.py +++ /dev/null @@ -1,124 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import ClassVar -from typing import List - -import numpy as np -import paddle -from paddle import nn - -from ppsci.loss.mtl import base - - -class PCGrad(base.LossAggregator): - r""" - **P**rojecting **C**onflicting Gradients - - [Gradient Surgery for Multi-Task Learning](https://papers.nips.cc/paper/2020/hash/3fe78a8acf5fda99de95303940a2420c-Abstract.html) - - Code reference: [https://github.com/tianheyu927/PCGrad/blob/master/PCGrad_tf.py](https://github.com/tianheyu927/PCGrad/blob/master/PCGrad_tf.py) - - Attributes: - should_persist(bool): Whether to persist the loss aggregator when saving. - Those loss aggregators with parameters and/or buffers should be persisted. - - Args: - model (nn.Layer): Training model. - - Examples: - >>> import paddle - >>> from ppsci.loss import mtl - >>> model = paddle.nn.Linear(3, 4) - >>> loss_aggregator = mtl.PCGrad(model) - >>> for i in range(5): - ... x1 = paddle.randn([8, 3]) - ... x2 = paddle.randn([8, 3]) - ... y1 = model(x1) - ... y2 = model(x2) - ... loss1 = paddle.sum(y1) - ... loss2 = paddle.sum((y2 - 2) ** 2) - ... loss_aggregator({'loss1': loss1, 'loss2': loss2}).backward() - """ - should_persist: ClassVar[bool] = False - - def __init__(self, model: nn.Layer) -> None: - super().__init__(model) - self._zero = paddle.zeros([]) - - def backward(self) -> None: - # shuffle order of losses - keys = list(self.losses.keys()) - np.random.shuffle(keys) - self.losses = {key: self.losses[key] for key in keys} - - grads_list = self._compute_grads() - with paddle.no_grad(): - refined_grads = self._refine_grads(grads_list) - self._set_grads(refined_grads) - - def _compute_grads(self) -> List[paddle.Tensor]: - # compute all gradients derived by each loss - grads_list = [] # num_params x num_losses - for key in self.losses: - # backward with current loss - self.losses[key].backward() - grads_list.append( - paddle.concat( - [ - param.grad.clone().reshape([-1]) - for param in self.model.parameters() - if param.grad is not None - ], - axis=0, - ) - ) - # clear gradients for current loss for not affecting other loss - self.model.clear_gradients() - - return grads_list - - def _refine_grads(self, grads_list: List[paddle.Tensor]) -> List[paddle.Tensor]: - def proj_grad(grad: paddle.Tensor): - for k in range(self.loss_num): - inner_product = paddle.sum(grad * grads_list[k]) - proj_direction = inner_product / paddle.sum( - grads_list[k] * grads_list[k] - ) - grad = grad - paddle.minimum(proj_direction, self._zero) * grads_list[k] - return grad - - grads_list = [proj_grad(grad) for grad in grads_list] - - # Unpack flattened projected gradients back to their original shapes. - proj_grads: List[paddle.Tensor] = [] - for j in range(self.loss_num): - start_idx = 0 - for idx, var in enumerate(self.model.parameters()): - grad_shape = var.shape - flatten_dim = var.numel() - refined_grad = grads_list[j][start_idx : start_idx + flatten_dim] - refined_grad = paddle.reshape(refined_grad, grad_shape) - if len(proj_grads) < self.param_num: - proj_grads.append(refined_grad) - else: - proj_grads[idx] += refined_grad - start_idx += flatten_dim - return proj_grads - - def _set_grads(self, grads_list: List[paddle.Tensor]) -> None: - for i, param in enumerate(self.model.parameters()): - param.grad = grads_list[i] diff --git a/examples/smc_reac/ppsci/loss/mtl/relobralo.py b/examples/smc_reac/ppsci/loss/mtl/relobralo.py deleted file mode 100644 index 02ec8f1339..0000000000 --- a/examples/smc_reac/ppsci/loss/mtl/relobralo.py +++ /dev/null @@ -1,127 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import ClassVar -from typing import Dict - -import paddle -from paddle import nn - - -class Relobralo(nn.Layer): - r""" - **Re**lative **Lo**ss **B**alancing with **Ra**ndom **Lo**okback - - [Multi-Objective Loss Balancing for Physics-Informed Deep Learning](https://arxiv.org/abs/2110.09813) - - Attributes: - should_persist(bool): Whether to persist the loss aggregator when saving. - Those loss aggregators with parameters and/or buffers should be persisted. - - Args: - num_losses (int): Number of losses. - alpha (float, optional): Ability for remembering past in paper. Defaults to 0.95. - beta (float, optional): Parameter for generating $\rho$ from bernoulli distribution, - and $E[\rho](=\beta)$ should be close to 1. Defaults to 0.99. - tau (float, optional): Temperature factor. Equivalent to softmax when $\tau$=1.0, - equivalent to argmax when $\tau$=0. Defaults to 1.0. - eps (float, optional): $\epsilon$ to avoid divided by 0 in losses. Defaults to 1e-8. - - Examples: - >>> import paddle - >>> from ppsci.loss import mtl - >>> model = paddle.nn.Linear(3, 4) - >>> loss_aggregator = mtl.Relobralo(num_losses=2) - >>> for i in range(5): - ... x1 = paddle.randn([8, 3]) - ... x2 = paddle.randn([8, 3]) - ... y1 = model(x1) - ... y2 = model(x2) - ... loss1 = paddle.sum(y1) - ... loss2 = paddle.sum((y2 - 2) ** 2) - ... loss_aggregator({'loss1': loss1, 'loss2': loss2}).backward() - """ - should_persist: ClassVar[bool] = True - - def __init__( - self, - num_losses: int, - alpha: float = 0.95, - beta: float = 0.99, - tau: float = 1.0, - eps: float = 1e-8, - ) -> None: - super().__init__() - self.step = 0 - self.num_losses: int = num_losses - self.alpha: float = alpha - self.beta: float = beta - self.tau: float = tau - self.eps: float = eps - self.register_buffer("losses_init", paddle.zeros([self.num_losses])) - self.register_buffer("losses_prev", paddle.zeros([self.num_losses])) - self.register_buffer("lmbda", paddle.ones([self.num_losses])) - - def _softmax(self, vec: "paddle.Tensor") -> "paddle.Tensor": - max_item = vec.max() - result = paddle.exp(vec - max_item) / paddle.exp(vec - max_item).sum() - return result - - def _compute_bal( - self, losses_vec1: "paddle.Tensor", losses_vec2: "paddle.Tensor" - ) -> "paddle.Tensor": - return self.num_losses * ( - self._softmax(losses_vec1 / (self.tau * losses_vec2 + self.eps)) - ) - - def __call__( - self, losses: Dict[str, "paddle.Tensor"], step: int = 0 - ) -> "paddle.Tensor": - assert len(losses) == self.num_losses, ( - f"Length of given losses({len(losses)}) should be equal to " - f"num_losses({self.num_losses})." - ) - self.step = step - losses_stacked = paddle.stack(list(losses.values())) # [num_losses, ] - - if self.step == 0: - loss = losses_stacked.sum() - with paddle.no_grad(): - paddle.assign(losses_stacked.detach(), self.losses_init) - else: - with paddle.no_grad(): - # 1. update lambda_hist - rho = paddle.bernoulli(paddle.to_tensor(self.beta)) - lmbda_hist = rho * self.lmbda + (1 - rho) * self._compute_bal( - losses_stacked, self.losses_init - ) - - # 2. update lambda - paddle.assign( - self.alpha * lmbda_hist - + (1 - self.alpha) - * self._compute_bal(losses_stacked, self.losses_prev), - self.lmbda, - ) - - # 3. compute reweighted total loss with lambda - loss = (losses_stacked * self.lmbda).sum() - - # update losses_prev at the end of each step - with paddle.no_grad(): - paddle.assign(losses_stacked.detach(), self.losses_prev) - - return loss diff --git a/examples/smc_reac/ppsci/loss/mtl/sum.py b/examples/smc_reac/ppsci/loss/mtl/sum.py deleted file mode 100644 index d2c9a7bd50..0000000000 --- a/examples/smc_reac/ppsci/loss/mtl/sum.py +++ /dev/null @@ -1,60 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import TYPE_CHECKING -from typing import Dict - -if TYPE_CHECKING: - import paddle - -from typing import ClassVar - -from ppsci.loss.mtl.base import LossAggregator - - -class Sum(LossAggregator): - r""" - **Default loss aggregator** which do simple summation for given losses as below. - - $$ - loss = \sum_i^N losses_i - $$ - - Attributes: - should_persist(bool): Whether to persist the loss aggregator when saving. - Those loss aggregators with parameters and/or buffers should be persisted. - """ - should_persist: ClassVar[bool] = False - - def __init__(self) -> None: - self.step = 0 - - def __call__( - self, losses: Dict[str, "paddle.Tensor"], step: int = 0 - ) -> "paddle.Tensor": - assert ( - len(losses) > 0 - ), f"Number of given losses({len(losses)}) can not be empty." - self.step = step - - total_loss = 0.0 - for i, key in enumerate(losses): - if i == 0: - total_loss = losses[key] - else: - total_loss += losses[key] - - return total_loss diff --git a/examples/smc_reac/ppsci/metric/__init__.py b/examples/smc_reac/ppsci/metric/__init__.py deleted file mode 100644 index 0a1d069aa9..0000000000 --- a/examples/smc_reac/ppsci/metric/__init__.py +++ /dev/null @@ -1,63 +0,0 @@ -# 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. - -import copy - -from ppsci.metric.anomaly_coef import LatitudeWeightedACC -from ppsci.metric.base import Metric -from ppsci.metric.func import FunctionalMetric -from ppsci.metric.l2_rel import L2Rel -from ppsci.metric.l2_rel import MeanL2Rel -from ppsci.metric.mae import MAE -from ppsci.metric.max_ae import MaxAE -from ppsci.metric.mse import MSE -from ppsci.metric.r2_score import R2Score -from ppsci.metric.rmse import RMSE -from ppsci.metric.rmse import LatitudeWeightedRMSE -from ppsci.utils import misc - -__all__ = [ - "LatitudeWeightedACC", - "Metric", - "FunctionalMetric", - "L2Rel", - "MeanL2Rel", - "MAE", - "MaxAE", - "MSE", - "RMSE", - "LatitudeWeightedRMSE", - "R2Score", - "build_metric", -] - - -def build_metric(cfg): - """Build metric. - - Args: - cfg (List[DictConfig]): List of metric config. - - Returns: - Dict[str, Metric]: Dict of callable metric object. - """ - cfg = copy.deepcopy(cfg) - - metric_dict = misc.PrettyOrderedDict() - for _item in cfg: - metric_cls = next(iter(_item.keys())) - metric_cfg = _item.pop(metric_cls) - metric = eval(metric_cls)(**metric_cfg) - metric_dict[metric_cls] = metric - return metric_dict diff --git a/examples/smc_reac/ppsci/metric/anomaly_coef.py b/examples/smc_reac/ppsci/metric/anomaly_coef.py deleted file mode 100644 index 33633228ca..0000000000 --- a/examples/smc_reac/ppsci/metric/anomaly_coef.py +++ /dev/null @@ -1,122 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import Dict -from typing import Optional -from typing import Tuple -from typing import Union - -import numpy as np -import paddle - -from ppsci.metric import base - - -class LatitudeWeightedACC(base.Metric): - r"""Latitude weighted anomaly correlation coefficient. - - $$ - metric = - \dfrac{\sum\limits_{m,n}{L_mX_{mn}Y_{mn}}}{\sqrt{\sum\limits_{m,n}{L_mX_{mn}^{2}}\sum\limits_{m,n}{L_mY_{mn}^{2}}}} - $$ - - $$ - L_m = N_{lat}\dfrac{\cos(lat_m)}{\sum\limits_{j=1}^{N_{lat}}\cos(lat_j)} - $$ - - $lat_m$ is the latitude at m. - $N_{lat}$ is the number of latitude set by `num_lat`. - - Args: - num_lat (int): Number of latitude. - mean (Optional[Union[np.array, Tuple[float, ...]]]): Mean of training data. Defaults to None. - keep_batch (bool, optional): Whether keep batch axis. Defaults to False. - variable_dict (Optional[Dict[str, int]]): Variable dictionary, the key is the name of a variable and - the value is its index. Defaults to None. - unlog (bool, optional): Whether calculate expm1 for all elements in the array. Defaults to False. - scale (float, optional): The scale value used after expm1. Defaults to 1e-5. - - Examples: - >>> import numpy as np - >>> import ppsci - >>> mean = np.random.randn(20, 720, 1440) - >>> metric = ppsci.metric.LatitudeWeightedACC(720, mean=mean) - """ - - def __init__( - self, - num_lat: int, - mean: Optional[Union[np.array, Tuple[float, ...]]], - keep_batch: bool = False, - variable_dict: Optional[Dict[str, int]] = None, - unlog: bool = False, - scale: float = 1e-5, - ): - super().__init__(keep_batch) - self.num_lat = num_lat - self.mean = ( - None if mean is None else paddle.to_tensor(mean, paddle.get_default_dtype()) - ) - self.variable_dict = variable_dict - self.unlog = unlog - self.scale = scale - - self.weight = self.get_latitude_weight(num_lat) - - def get_latitude_weight(self, num_lat: int = 720): - lat_t = paddle.linspace(start=0, stop=1, num=num_lat) - lat_t = paddle.cos(3.1416 * (0.5 - lat_t)) - weight = num_lat * lat_t / paddle.sum(lat_t) - weight = weight.reshape((1, 1, -1, 1)) - return weight - - def scale_expm1(self, x: paddle.Tensor): - return self.scale * paddle.expm1(x) - - @paddle.no_grad() - def forward(self, output_dict, label_dict) -> Dict[str, "paddle.Tensor"]: - metric_dict = {} - - for key in label_dict: - output = ( - self.scale_expm1(output_dict[key]) if self.unlog else output_dict[key] - ) - label = self.scale_expm1(label_dict[key]) if self.unlog else label_dict[key] - - if self.mean is not None: - output = output - self.mean - label = label - self.mean - - rmse = paddle.sum( - self.weight * output * label, axis=(-1, -2) - ) / paddle.sqrt( - paddle.sum(self.weight * output**2, axis=(-1, -2)) - * paddle.sum(self.weight * label**2, axis=(-1, -2)) - ) - - if self.variable_dict is not None: - for variable_name, idx in self.variable_dict.items(): - if self.keep_batch: - metric_dict[f"{key}.{variable_name}"] = rmse[:, idx] - else: - metric_dict[f"{key}.{variable_name}"] = rmse[:, idx].mean() - else: - if self.keep_batch: - metric_dict[key] = rmse.mean(axis=1) - else: - metric_dict[key] = rmse.mean() - - return metric_dict diff --git a/examples/smc_reac/ppsci/metric/base.py b/examples/smc_reac/ppsci/metric/base.py deleted file mode 100644 index 750e629882..0000000000 --- a/examples/smc_reac/ppsci/metric/base.py +++ /dev/null @@ -1,25 +0,0 @@ -# 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. - -from __future__ import annotations - -from paddle import nn - - -class Metric(nn.Layer): - """Base class for metric.""" - - def __init__(self, keep_batch: bool = False): - super().__init__() - self.keep_batch = keep_batch diff --git a/examples/smc_reac/ppsci/metric/func.py b/examples/smc_reac/ppsci/metric/func.py deleted file mode 100644 index bee646b656..0000000000 --- a/examples/smc_reac/ppsci/metric/func.py +++ /dev/null @@ -1,66 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import TYPE_CHECKING -from typing import Callable -from typing import Dict - -from ppsci.metric import base - -if TYPE_CHECKING: - import paddle - - -class FunctionalMetric(base.Metric): - r"""Functional metric class, which allows to use custom metric computing function from given metric_expr for complex computation cases. - - Args: - metric_expr (Callable): Expression of metric calculation. - keep_batch (bool, optional): Whether keep batch axis. Defaults to False. - - Examples: - >>> import paddle - >>> from ppsci.metric import FunctionalMetric - >>> def metric_expr(output_dict, *args): - ... rel_l2 = 0 - ... for key in output_dict: - ... length = int(len(output_dict[key])/2) - ... out_dict = output_dict[key][:length] - ... label_dict = output_dict[key][length:] - ... rel_l2 += paddle.norm(out_dict - label_dict) / paddle.norm(label_dict) - ... return {"rel_l2": rel_l2} - >>> metric_dict = FunctionalMetric(metric_expr) - >>> output_dict = {'u': paddle.to_tensor([[0.5, 0.9], [1.1, -1.3], [-0.2, 1.5], [-0.1, -0.3]]), - ... 'v': paddle.to_tensor([[0.5, 0.9], [1.1, -1.3], [-1.8, 1.0], [-0.2, 2.5]])} - >>> result = metric_dict(output_dict) - >>> print(result) - {'rel_l2': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, - 2.59985542)} - """ - - def __init__( - self, - metric_expr: Callable[ - [Dict["str", "paddle.Tensor"], Dict["str", "paddle.Tensor"]], - Dict["str", "paddle.Tensor"], - ], - keep_batch: bool = False, - ): - super().__init__(keep_batch) - self.metric_expr = metric_expr - - def forward(self, output_dict, label_dict=None) -> Dict[str, "paddle.Tensor"]: - return self.metric_expr(output_dict, label_dict) diff --git a/examples/smc_reac/ppsci/metric/l2_rel.py b/examples/smc_reac/ppsci/metric/l2_rel.py deleted file mode 100644 index 2a64e9befc..0000000000 --- a/examples/smc_reac/ppsci/metric/l2_rel.py +++ /dev/null @@ -1,139 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import Dict - -import numpy as np -import paddle - -from ppsci.metric import base - - -class L2Rel(base.Metric): - r"""Class for l2 relative error. - - NOTE: This metric API is slightly different from `MeanL2Rel`, difference is as below: - - - `L2Rel` regards the input sample as a whole and calculates the l2 relative error of the whole; - - `MeanL2Rel` will calculate L2Rel separately for each input sample and return the average of l2 relative error for all samples. - - $$ - metric = \dfrac{\Vert \mathbf{x} - \mathbf{y} \Vert_2}{\max(\Vert \mathbf{y} \Vert_2, \epsilon)} - $$ - - $$ - \mathbf{x}, \mathbf{y} \in \mathcal{R}^{N} - $$ - - Args: - keep_batch (bool, optional): Whether keep batch axis. Defaults to False. - - Examples: - >>> import paddle - >>> from ppsci.metric import L2Rel - >>> output_dict = {'u': paddle.to_tensor([[0.5, 0.9], [1.1, -1.3]]), - ... 'v': paddle.to_tensor([[0.5, 0.9], [1.1, -1.3]])} - >>> label_dict = {'u': paddle.to_tensor([[-1.8, 1.0], [-0.2, 2.5]]), - ... 'v': paddle.to_tensor([[0.1, 0.1], [0.1, 0.1]])} - >>> loss = L2Rel() - >>> result = loss(output_dict, label_dict) - >>> print(result) - {'u': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, - 1.42658269), 'v': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, - 9.69535923)} - """ - - # NOTE: Avoid divide by zero in result - # see https://github.com/scikit-learn/scikit-learn/pull/15007 - EPS: float = np.finfo(np.float32).eps - - def __init__(self, keep_batch: bool = False): - if keep_batch: - raise ValueError(f"keep_batch should be False, but got {keep_batch}.") - super().__init__(keep_batch) - - @paddle.no_grad() - def forward(self, output_dict, label_dict) -> Dict[str, "paddle.Tensor"]: - metric_dict = {} - for key in label_dict: - rel_l2 = paddle.norm(label_dict[key] - output_dict[key], p=2) / paddle.norm( - label_dict[key], p=2 - ).clip(min=self.EPS) - metric_dict[key] = rel_l2 - - return metric_dict - - -class MeanL2Rel(base.Metric): - r"""Class for mean l2 relative error. - - NOTE: This metric API is slightly different from `L2Rel`, difference is as below: - - - `MeanL2Rel` will calculate L2Rel separately for each input sample and return the average of l2 relative error for all samples. - - `L2Rel` regards the input sample as a whole and calculates the l2 relative error of the whole; - - $$ - metric = \dfrac{1}{M} \sum_{i=1}^{M}\dfrac{\Vert \mathbf{x_i} - \mathbf{y_i} \Vert_2}{\max(\Vert \mathbf{y_i} \Vert_2, \epsilon) } - $$ - - $$ - \mathbf{x_i}, \mathbf{y_i} \in \mathcal{R}^{N} - $$ - - Args: - keep_batch (bool, optional): Whether keep batch axis. Defaults to False. - - Examples: - >>> import paddle - >>> from ppsci.metric import MeanL2Rel - >>> output_dict = {'u': paddle.to_tensor([[0.5, 0.9], [1.1, -1.3]]), - ... 'v': paddle.to_tensor([[0.5, 0.9], [1.1, -1.3]])} - >>> label_dict = {'u': paddle.to_tensor([[-1.8, 1.0], [-0.2, 2.5]]), - ... 'v': paddle.to_tensor([[0.1, 0.1], [0.1, 0.1]])} - >>> loss = MeanL2Rel() - >>> result = loss(output_dict, label_dict) - >>> print(result) - {'u': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, - 1.35970235), 'v': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, - 9.24504089)} - >>> loss = MeanL2Rel(keep_batch=True) - >>> result = loss(output_dict, label_dict) - >>> print(result) - {'u': Tensor(shape=[2], dtype=float32, place=Place(gpu:0), stop_gradient=True, - [1.11803389, 1.60137081]), 'v': Tensor(shape=[2], dtype=float32, place=Place(gpu:0), stop_gradient=True, - [6.32455540 , 12.16552544])} - """ - - # NOTE: Avoid divide by zero in result - # see https://github.com/scikit-learn/scikit-learn/pull/15007 - EPS: float = np.finfo(np.float32).eps - - def __init__(self, keep_batch: bool = False): - super().__init__(keep_batch) - - @paddle.no_grad() - def forward(self, output_dict, label_dict) -> Dict[str, "paddle.Tensor"]: - metric_dict = {} - for key in label_dict: - rel_l2 = paddle.norm( - label_dict[key] - output_dict[key], p=2, axis=1 - ) / paddle.norm(label_dict[key], p=2, axis=1).clip(min=self.EPS) - if self.keep_batch: - metric_dict[key] = rel_l2 - else: - metric_dict[key] = rel_l2.mean() - - return metric_dict diff --git a/examples/smc_reac/ppsci/metric/mae.py b/examples/smc_reac/ppsci/metric/mae.py deleted file mode 100644 index 3b6ebdedbb..0000000000 --- a/examples/smc_reac/ppsci/metric/mae.py +++ /dev/null @@ -1,73 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import Dict - -import paddle -import paddle.nn.functional as F - -from ppsci.metric import base - - -class MAE(base.Metric): - r"""Mean absolute error. - - $$ - metric = \dfrac{1}{N} \Vert \mathbf{x} - \mathbf{y} \Vert_1 - $$ - - $$ - \mathbf{x}, \mathbf{y} \in \mathcal{R}^{N} - $$ - - Args: - keep_batch (bool, optional): Whether keep batch axis. Defaults to False. - - Examples: - >>> import paddle - >>> from ppsci.metric import MAE - >>> output_dict = {'u': paddle.to_tensor([[0.5, 0.9], [1.1, -1.3]]), - ... 'v': paddle.to_tensor([[0.5, 0.9], [1.1, -1.3]])} - >>> label_dict = {'u': paddle.to_tensor([[-1.8, 1.0], [-0.2, 2.5]]), - ... 'v': paddle.to_tensor([[0.1, 0.1], [0.1, 0.1]])} - >>> loss = MAE() - >>> result = loss(output_dict, label_dict) - >>> print(result) - {'u': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, - 1.87500000), 'v': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, - 0.89999998)} - >>> loss = MAE(keep_batch=True) - >>> result = loss(output_dict, label_dict) - >>> print(result) - {'u': Tensor(shape=[2], dtype=float32, place=Place(gpu:0), stop_gradient=True, - [1.20000005, 2.54999995]), 'v': Tensor(shape=[2], dtype=float32, place=Place(gpu:0), stop_gradient=True, - [0.59999996, 1.20000005])} - """ - - def __init__(self, keep_batch: bool = False): - super().__init__(keep_batch) - - @paddle.no_grad() - def forward(self, output_dict, label_dict) -> Dict[str, "paddle.Tensor"]: - metric_dict = {} - for key in label_dict: - mae = F.l1_loss(output_dict[key], label_dict[key], "none") - if self.keep_batch: - metric_dict[key] = mae.mean(axis=tuple(range(1, mae.ndim))) - else: - metric_dict[key] = mae.mean() - - return metric_dict diff --git a/examples/smc_reac/ppsci/metric/max_ae.py b/examples/smc_reac/ppsci/metric/max_ae.py deleted file mode 100644 index 7d94544b97..0000000000 --- a/examples/smc_reac/ppsci/metric/max_ae.py +++ /dev/null @@ -1,77 +0,0 @@ -# Copyright (c) 2024 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. - -from __future__ import annotations - -from typing import Dict - -import paddle - -from ppsci.metric import base - - -class MaxAE(base.Metric): - r"""Maximum Absolute Error (MaxAE). - - $$ - \text{MaxAE} = \max_i \left( |x_i - y_i| \right) - $$ - - $$ - \mathbf{x}, \mathbf{y} \in \mathcal{R}^{N} - $$ - - Args: - keep_batch (bool, optional): Whether keep batch axis. Defaults to False. - - Examples: - >>> import paddle - >>> from ppsci.metric import MaxAE - >>> output_dict = {'u': paddle.to_tensor([[0.5, 0.9], [1.1, -1.3]]), - ... 'v': paddle.to_tensor([[0.5, 0.9], [1.1, -1.3]])} - >>> label_dict = {'u': paddle.to_tensor([[-1.8, 1.0], [-0.2, 2.5]]), - ... 'v': paddle.to_tensor([[0.1, 0.1], [0.1, 0.1]])} - >>> metric = MaxAE() - >>> result = metric(output_dict, label_dict) - >>> print(result) - {'u': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, - 3.80000007), 'v': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, - 0.80000001)} - >>> metric = MaxAE(keep_batch=True) - >>> result = metric(output_dict, label_dict) - >>> print(result) - {'u': Tensor(shape=[2], dtype=float32, place=Place(gpu:0), stop_gradient=True, - [1.30000002, 3.80000007]), 'v': Tensor(shape=[2], dtype=float32, place=Place(gpu:0), stop_gradient=True, - [0.40000001, 0.80000001])} - """ - - def __init__(self, keep_batch: bool = False): - super().__init__(keep_batch) - - @paddle.no_grad() - def forward(self, output_dict, label_dict) -> Dict[str, "paddle.Tensor"]: - maxae_dict = {} - - for key in label_dict: - # Calculate absolute error - ae = paddle.abs(output_dict[key] - label_dict[key]) - - if self.keep_batch: - # Take the maximum AE within each batch - maxae_dict[key] = paddle.amax(ae, axis=tuple(range(1, ae.ndim))) - else: - # Take the global maximum AE across all elements - maxae_dict[key] = paddle.amax(ae) - - return maxae_dict diff --git a/examples/smc_reac/ppsci/metric/mse.py b/examples/smc_reac/ppsci/metric/mse.py deleted file mode 100644 index 9e47a7cf5a..0000000000 --- a/examples/smc_reac/ppsci/metric/mse.py +++ /dev/null @@ -1,73 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import Dict - -import paddle -import paddle.nn.functional as F - -from ppsci.metric import base - - -class MSE(base.Metric): - r"""Mean square error - - $$ - metric = \dfrac{1}{N} \Vert \mathbf{x} - \mathbf{y} \Vert_2^2 - $$ - - $$ - \mathbf{x}, \mathbf{y} \in \mathcal{R}^{N} - $$ - - Args: - keep_batch (bool, optional): Whether keep batch axis. Defaults to False. - - Examples: - >>> import paddle - >>> from ppsci.metric import MSE - >>> output_dict = {'u': paddle.to_tensor([[0.5, 0.9], [1.1, -1.3]]), - ... 'v': paddle.to_tensor([[0.5, 0.9], [1.1, -1.3]])} - >>> label_dict = {'u': paddle.to_tensor([[-1.8, 1.0], [-0.2, 2.5]]), - ... 'v': paddle.to_tensor([[0.1, 0.1], [0.1, 0.1]])} - >>> loss = MSE() - >>> result = loss(output_dict, label_dict) - >>> print(result) - {'u': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, - 5.35750008), 'v': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, - 0.94000000)} - >>> loss = MSE(keep_batch=True) - >>> result = loss(output_dict, label_dict) - >>> print(result) - {'u': Tensor(shape=[2], dtype=float32, place=Place(gpu:0), stop_gradient=True, - [2.65000010, 8.06499958]), 'v': Tensor(shape=[2], dtype=float32, place=Place(gpu:0), stop_gradient=True, - [0.39999998, 1.48000002])} - """ - - def __init__(self, keep_batch: bool = False): - super().__init__(keep_batch) - - @paddle.no_grad() - def forward(self, output_dict, label_dict) -> Dict[str, "paddle.Tensor"]: - metric_dict = {} - for key in label_dict: - mse = F.mse_loss(output_dict[key], label_dict[key], "none") - if self.keep_batch: - metric_dict[key] = mse.mean(axis=tuple(range(1, mse.ndim))) - else: - metric_dict[key] = mse.mean() - - return metric_dict diff --git a/examples/smc_reac/ppsci/metric/r2_score.py b/examples/smc_reac/ppsci/metric/r2_score.py deleted file mode 100644 index 4ad3eb608d..0000000000 --- a/examples/smc_reac/ppsci/metric/r2_score.py +++ /dev/null @@ -1,97 +0,0 @@ -# Copyright (c) 2024 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. - -from __future__ import annotations - -from typing import Dict - -import paddle - -from ppsci.metric import base - - -class R2Score(base.Metric): - r"""Coefficient of Determination (R^2 Score). - - $$ - R^2 = 1 - \frac{\sum_{i=1}^{N} (y_i - \hat{y}_i)^2}{\sum_{i=1}^{N} (y_i - \bar{y})^2} - $$ - - $$ - \mathbf{y}, \mathbf{\hat{y}} \in \mathcal{R}^{N} - $$ - - Args: - keep_batch (bool, optional): Whether keep batch axis. Defaults to False. - - Examples: - >>> import paddle - >>> from ppsci.metric import R2Score - >>> output_dict = {'u': paddle.to_tensor([[0.5, 0.9], [1.1, -1.3]]), - ... 'v': paddle.to_tensor([[0.5, 0.9], [1.1, -1.3]])} - >>> label_dict = {'u': paddle.to_tensor([[-1.8, 1.0], [-0.2, 2.5]]), - ... 'v': paddle.to_tensor([[0.1, 0.1], [0.1, 0.1]])} - >>> metric = R2Score() - >>> result = metric(output_dict, label_dict) - >>> print(result) - {'u': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, - -3.75000000), 'v': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, - 0.00000000)} - >>> metric = R2Score(keep_batch=True) - >>> result = metric(output_dict, label_dict) - >>> print(result) - {'u': Tensor(shape=[2], dtype=float32, place=Place(gpu:0), stop_gradient=True, - [-0.64000008, -6.88000011]), 'v': Tensor(shape=[2], dtype=float32, place=Place(gpu:0), stop_gradient=True, - [0.00000000, 0.00000000])} - """ - - def __init__(self, keep_batch: bool = False): - super().__init__(keep_batch) - - @paddle.no_grad() - def forward(self, output_dict, label_dict) -> Dict[str, "paddle.Tensor"]: - r2score_dict = {} - - for key in label_dict: - - output = output_dict[key] - target = label_dict[key] - - # Ensure the shapes of output and target match - if output.shape != target.shape: - raise ValueError( - f"Output and target shapes do not match for key '{key}'. Output shape: {output.shape}, Target shape: {target.shape}" - ) - - output = output.flatten() - target = target.flatten() - # Calculate mean of target along the last dimension (features) - target_mean = target.mean(axis=-1, keepdim=True) - - # Calculate total sum of squares (TSS) - ss_tot = paddle.sum(x=(target - target_mean) ** 2) - - # Calculate residual sum of squares (RSS) - ss_res = paddle.sum(x=(target - output) ** 2) - - # Calculate R^2 score with numerical stability - r2 = 1 - (ss_res / (ss_tot + 1e-8)) - - # Handle batch-wise or mean R^2 based on self.keep_batch - if self.keep_batch: - r2score_dict[key] = r2.unsqueeze(axis=0) - else: - r2score_dict[key] = r2 - - return r2score_dict diff --git a/examples/smc_reac/ppsci/metric/rmse.py b/examples/smc_reac/ppsci/metric/rmse.py deleted file mode 100644 index 55be8e9102..0000000000 --- a/examples/smc_reac/ppsci/metric/rmse.py +++ /dev/null @@ -1,155 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import Dict -from typing import Optional -from typing import Tuple -from typing import Union - -import numpy as np -import paddle -import paddle.nn.functional as F - -from ppsci.metric import base - - -class RMSE(base.Metric): - r"""Root mean square error - - $$ - metric = \sqrt{\dfrac{1}{N} \Vert \mathbf{x} - \mathbf{y} \Vert_2^2} - $$ - - $$ - \mathbf{x}, \mathbf{y} \in \mathcal{R}^{N} - $$ - - Args: - keep_batch (bool, optional): Whether keep batch axis. Defaults to False. - - Examples: - >>> import paddle - >>> from ppsci.metric import RMSE - >>> output_dict = {'u': paddle.to_tensor([[0.5, 0.9], [1.1, -1.3]]), - ... 'v': paddle.to_tensor([[0.5, 0.9], [1.1, -1.3]])} - >>> label_dict = {'u': paddle.to_tensor([[-1.8, 1.0], [-0.2, 2.5]]), - ... 'v': paddle.to_tensor([[0.1, 0.1], [0.1, 0.1]])} - >>> loss = RMSE() - >>> result = loss(output_dict, label_dict) - >>> print(result) - {'u': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, - 2.31462741), 'v': Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True, - 0.96953595)} - """ - - def __init__(self, keep_batch: bool = False): - if keep_batch: - raise ValueError(f"keep_batch should be False, but got {keep_batch}.") - super().__init__(keep_batch) - - @paddle.no_grad() - def forward(self, output_dict, label_dict) -> Dict[str, "paddle.Tensor"]: - metric_dict = {} - for key in label_dict: - rmse = F.mse_loss(output_dict[key], label_dict[key], "mean") ** 0.5 - metric_dict[key] = rmse - - return metric_dict - - -class LatitudeWeightedRMSE(base.Metric): - r"""Latitude weighted root mean square error. - - $$ - metric =\sqrt{\dfrac{1}{MN}\sum\limits_{m=1}^{M}\sum\limits_{n=1}^{N}L_m(X_{mn}-Y_{mn})^{2}} - $$ - - $$ - L_m = N_{lat}\dfrac{\cos(lat_m)}{\sum\limits_{j=1}^{N_{lat}}\cos(lat_j)} - $$ - - $lat_m$ is the latitude at m. - $N_{lat}$ is the number of latitude set by `num_lat`. - - Args: - num_lat (int): Number of latitude. - std (Optional[Union[np.array, Tuple[float, ...]]]): Standard Deviation of training dataset. Defaults to None. - keep_batch (bool, optional): Whether keep batch axis. Defaults to False. - variable_dict (Optional[Dict[str, int]]): Variable dictionary, the key is the name of a variable and - the value is its index. Defaults to None. - unlog (bool, optional): Whether calculate expm1 for all elements in the array. Defaults to False. - scale (float, optional): The scale value used after expm1. Defaults to 1e-5. - - Examples: - >>> import numpy as np - >>> import ppsci - >>> std = np.random.randn(20, 1, 1) - >>> metric = ppsci.metric.LatitudeWeightedRMSE(720, std=std) - """ - - def __init__( - self, - num_lat: int, - std: Optional[Union[np.array, Tuple[float, ...]]] = None, - keep_batch: bool = False, - variable_dict: Dict[str, int] = None, - unlog: bool = False, - scale: float = 1e-5, - ): - super().__init__(keep_batch) - self.num_lat = num_lat - self.std = ( - None - if std is None - else paddle.to_tensor(std, paddle.get_default_dtype()).reshape((1, -1)) - ) - self.variable_dict = variable_dict - self.unlog = unlog - self.scale = scale - self.weight = self.get_latitude_weight(num_lat) - - def get_latitude_weight(self, num_lat: int = 720): - lat_t = paddle.linspace(start=0, stop=1, num=num_lat) - lat_t = paddle.cos(3.1416 * (0.5 - lat_t)) - weight = num_lat * lat_t / paddle.sum(lat_t) - weight = weight.reshape((1, 1, -1, 1)) - return weight - - def scale_expm1(self, x: paddle.Tensor): - return self.scale * paddle.expm1(x) - - @paddle.no_grad() - def forward(self, output_dict, label_dict) -> Dict[str, "paddle.Tensor"]: - metric_dict = {} - for key in label_dict: - output = ( - self.scale_expm1(output_dict[key]) if self.unlog else output_dict[key] - ) - label = self.scale_expm1(label_dict[key]) if self.unlog else label_dict[key] - - mse = F.mse_loss(output, label, "none") - rmse = (mse * self.weight).mean(axis=(-1, -2)) ** 0.5 - if self.std is not None: - rmse = rmse * self.std - if self.variable_dict is not None: - for variable_name, idx in self.variable_dict.items(): - metric_dict[f"{key}.{variable_name}"] = ( - rmse[:, idx] if self.keep_batch else rmse[:, idx].mean() - ) - else: - metric_dict[key] = rmse.mean(axis=1) if self.keep_batch else rmse.mean() - - return metric_dict diff --git a/examples/smc_reac/ppsci/optimizer/__init__.py b/examples/smc_reac/ppsci/optimizer/__init__.py deleted file mode 100644 index c03b0717ee..0000000000 --- a/examples/smc_reac/ppsci/optimizer/__init__.py +++ /dev/null @@ -1,84 +0,0 @@ -# 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. - -import copy - -from ppsci.optimizer import lr_scheduler -from ppsci.optimizer.optimizer import LBFGS -from ppsci.optimizer.optimizer import SGD -from ppsci.optimizer.optimizer import SOAP -from ppsci.optimizer.optimizer import Adam -from ppsci.optimizer.optimizer import AdamW -from ppsci.optimizer.optimizer import Momentum -from ppsci.optimizer.optimizer import OptimizerList -from ppsci.optimizer.optimizer import RMSProp - -__all__ = [ - "LBFGS", - "SGD", - "Adam", - "AdamW", - "Momentum", - "RMSProp", - "OptimizerList", - "lr_scheduler", - "SOAP", -] - - -def build_lr_scheduler(cfg, epochs, iters_per_epoch): - """Build learning rate scheduler. - - Args: - cfg (DictConfig): Learning rate scheduler config. - epochs (int): Total epochs. - iters_per_epoch (int): Number of iterations of one epoch. - - Returns: - LRScheduler: Learning rate scheduler. - """ - cfg = copy.deepcopy(cfg) - cfg.update({"epochs": epochs, "iters_per_epoch": iters_per_epoch}) - lr_scheduler_cls = cfg.pop("name") - lr_scheduler_ = eval(lr_scheduler_cls)(**cfg) - return lr_scheduler_() - - -def build_optimizer(cfg, model_list, epochs, iters_per_epoch): - """Build optimizer and learning rate scheduler - - Args: - cfg (DictConfig): Learning rate scheduler config. - model_list (Tuple[nn.Layer, ...]): Tuple of model(s). - epochs (int): Total epochs. - iters_per_epoch (int): Number of iterations of one epoch. - - Returns: - Optimizer, LRScheduler: Optimizer and learning rate scheduler. - """ - # build lr_scheduler - cfg = copy.deepcopy(cfg) - lr_cfg = cfg.pop("lr") - if isinstance(lr_cfg, float): - lr_scheduler = lr_cfg - else: - lr_scheduler = build_lr_scheduler(lr_cfg, epochs, iters_per_epoch) - - # build optimizer - opt_cls = cfg.pop("name") - optimizer = eval(opt_cls)(learning_rate=lr_scheduler, **cfg)(model_list) - - if isinstance(lr_scheduler, float): - return optimizer, None - return optimizer, lr_scheduler diff --git a/examples/smc_reac/ppsci/optimizer/lr_scheduler.py b/examples/smc_reac/ppsci/optimizer/lr_scheduler.py deleted file mode 100644 index cad7da3503..0000000000 --- a/examples/smc_reac/ppsci/optimizer/lr_scheduler.py +++ /dev/null @@ -1,911 +0,0 @@ -# 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. - -from __future__ import annotations - -import abc -import math -from typing import Callable -from typing import List -from typing import Tuple -from typing import Union - -from paddle.optimizer import lr - -from ppsci.utils import logger - -__all__ = [ - "Linear", - "Cosine", - "Step", - "Piecewise", - "MultiStepDecay", - "ExponentialDecay", - "CosineWarmRestarts", - "OneCycleLR", - "LambdaDecay", - "ReduceOnPlateau", -] - - -class LRBase: - """Base class for custom learning rates. - - Args: - epochs (int): Total epoch(s). - iters_per_epoch (int): Number of iterations within an epoch. - learning_rate (float): Learning rate. - warmup_epoch (int): Number of warmup epochs. - warmup_start_lr (float): Start learning rate within warmup. - last_epoch (int): Last epoch. - by_epoch (bool): Learning rate decays by epoch when by_epoch is True, else by iter. - verbose (bool): If True, prints a message to stdout for each update. Defaults to False. - """ - - def __init__( - self, - epochs: int, - iters_per_epoch: int, - learning_rate: float, - warmup_epoch: int, - warmup_start_lr: float, - last_epoch: int, - by_epoch: bool, - verbose: bool = False, - ) -> None: - """Initialize and record the necessary parameters.""" - super().__init__() - if warmup_epoch >= epochs: - msg = ( - "When using warm up, the value of 'Global.epochs' should be greater " - "than value of 'Optimizer.lr.warmup_epoch'. The value of " - f"'Optimizer.lr.warmup_epoch' has been set to {epochs}." - ) - logger.warning(msg) - warmup_epoch = epochs - self.epochs = epochs - self.iters_per_epoch = iters_per_epoch - self.learning_rate = learning_rate - self.warmup_epoch = warmup_epoch - self.warmup_steps = ( - self.warmup_epoch - if by_epoch - else round(self.warmup_epoch * self.iters_per_epoch) - ) - self.warmup_start_lr = warmup_start_lr - self.last_epoch = last_epoch - self.by_epoch = by_epoch - self.verbose = verbose - - @abc.abstractmethod - def __call__(self, *args, **kwargs) -> lr.LRScheduler: - """Generate an learning rate scheduler. - - Returns: - lr.LinearWarmup: learning rate scheduler. - """ - pass - - def linear_warmup( - self, learning_rate: Union[float, lr.LRScheduler] - ) -> lr.LinearWarmup: - """Add an Linear Warmup before learning_rate. - - Args: - learning_rate (Union[float, lr.LRScheduler]): Original learning rate without - warmup. - - Returns: - lr.LinearWarmup: learning rate scheduler with warmup. - """ - warmup_lr = lr.LinearWarmup( - learning_rate=learning_rate, - warmup_steps=self.warmup_steps, - start_lr=self.warmup_start_lr, - end_lr=self.learning_rate, - last_epoch=self.last_epoch, - verbose=self.verbose, - ) - return warmup_lr - - -class Constant(lr.LRScheduler): - """Constant learning rate Class implementation. - - Args: - learning_rate (float): The initial learning rate. - last_epoch (int, optional): The index of last epoch. Default: -1. - """ - - def __init__(self, learning_rate: float, last_epoch: int = -1): - self.learning_rate = learning_rate - self.last_epoch = last_epoch - super().__init__() - - def get_lr(self) -> float: - """Always return the same learning rate""" - return self.learning_rate - - -class Linear(LRBase): - """Linear learning rate decay. - - Args: - epochs (int): Total epoch(s). - iters_per_epoch (int): Number of iterations within an epoch. - learning_rate (float): Learning rate. - end_lr (float, optional): The minimum final learning rate. Defaults to 0.0. - power (float, optional): Power of polynomial. Defaults to 1.0. - cycle (bool, optional): Whether the learning rate rises again. If True, then the learning rate will rise when it decrease - to ``end_lr`` . If False, the learning rate is monotone decreasing. Defaults to False. - warmup_epoch (int): Number of warmup epochs. - warmup_start_lr (float): Start learning rate within warmup. - last_epoch (int): Last epoch. - by_epoch (bool): Learning rate decays by epoch when by_epoch is True, else by iter. - - Examples: - >>> import ppsci - >>> lr = ppsci.optimizer.lr_scheduler.Linear(10, 2, 0.001)() - """ - - def __init__( - self, - epochs: int, - iters_per_epoch: int, - learning_rate: float, - end_lr: float = 0.0, - power: float = 1.0, - cycle: bool = False, - warmup_epoch: int = 0, - warmup_start_lr: float = 0.0, - last_epoch: int = -1, - by_epoch: bool = False, - ): - super().__init__( - epochs, - iters_per_epoch, - learning_rate, - warmup_epoch, - warmup_start_lr, - last_epoch, - by_epoch, - ) - self.decay_steps = (epochs - self.warmup_epoch) * iters_per_epoch - self.end_lr = end_lr - self.power = power - self.cycle = cycle - self.warmup_steps = round(self.warmup_epoch * iters_per_epoch) - if self.by_epoch: - self.decay_steps = self.epochs - self.warmup_epoch - - def __call__(self): - learning_rate = ( - lr.PolynomialDecay( - learning_rate=self.learning_rate, - decay_steps=self.decay_steps, - end_lr=self.end_lr, - power=self.power, - cycle=self.cycle, - last_epoch=self.last_epoch, - ) - if self.decay_steps > 0 - else Constant(self.learning_rate) - ) - - if self.warmup_steps > 0: - learning_rate = self.linear_warmup(learning_rate) - - setattr(learning_rate, "by_epoch", self.by_epoch) - return learning_rate - - -class ExponentialDecay(LRBase): - """ExponentialDecay learning rate decay. - - Args: - epochs (int): Total epoch(s). - iters_per_epoch (int): Number of iterations within an epoch. - learning_rate (float): Learning rate. - gamma (float): The decay rate. - decay_steps (int): The number of steps to decay. - warmup_epoch (int): Number of warmup epochs. - warmup_start_lr (float): Start learning rate within warmup. - last_epoch (int): Last epoch. - by_epoch (bool): Learning rate decays by epoch when by_epoch is True, else by iter. - - Examples: - >>> import ppsci - >>> lr = ppsci.optimizer.lr_scheduler.ExponentialDecay(10, 2, 1e-3, 0.95, 3)() - """ - - def __init__( - self, - epochs: int, - iters_per_epoch: int, - learning_rate: float, - gamma: float, - decay_steps: int, - warmup_epoch: int = 0, - warmup_start_lr: float = 0.0, - last_epoch: int = -1, - by_epoch: bool = False, - ): - super().__init__( - epochs, - iters_per_epoch, - learning_rate, - warmup_epoch, - warmup_start_lr, - last_epoch, - by_epoch, - ) - self.decay_steps = decay_steps - self.gamma = gamma - self.warmup_steps = round(self.warmup_epoch * iters_per_epoch) - if self.by_epoch: - self.decay_steps /= iters_per_epoch - - def __call__(self): - learning_rate = lr.ExponentialDecay( - learning_rate=self.learning_rate, - gamma=self.gamma ** (1 / self.decay_steps), - last_epoch=self.last_epoch, - ) - - if self.warmup_steps > 0: - learning_rate = self.linear_warmup(learning_rate) - - setattr(learning_rate, "by_epoch", self.by_epoch) - return learning_rate - - -class Cosine(LRBase): - """Cosine learning rate decay. - - lr = 0.05 * (math.cos(epoch * (math.pi / epochs)) + 1) - - Args: - epochs (int): Total epoch(s). - iters_per_epoch (int): Number of iterations within an epoch. - learning_rate (float): Learning rate. - eta_min (float, optional): Minimum learning rate. Defaults to 0.0. - warmup_epoch (int, optional): The epoch numbers for LinearWarmup. Defaults to 0. - warmup_start_lr (float, optional): Start learning rate within warmup. Defaults to 0.0. - last_epoch (int, optional): Last epoch. Defaults to -1. - by_epoch (bool, optional): Learning rate decays by epoch when by_epoch is True, - else by iter. Defaults to False. - - Examples: - >>> import ppsci - >>> lr = ppsci.optimizer.lr_scheduler.Cosine(10, 2, 1e-3)() - """ - - def __init__( - self, - epochs: int, - iters_per_epoch: int, - learning_rate: float, - eta_min: float = 0.0, - warmup_epoch: int = 0, - warmup_start_lr: float = 0.0, - last_epoch: int = -1, - by_epoch: bool = False, - ): - super().__init__( - epochs, - iters_per_epoch, - learning_rate, - warmup_epoch, - warmup_start_lr, - last_epoch, - by_epoch, - ) - self.T_max = (self.epochs - self.warmup_epoch) * self.iters_per_epoch - self.eta_min = eta_min - if self.by_epoch: - self.T_max = self.epochs - self.warmup_epoch - - def __call__(self): - learning_rate = ( - lr.CosineAnnealingDecay( - learning_rate=self.learning_rate, - T_max=self.T_max, - eta_min=self.eta_min, - last_epoch=self.last_epoch, - ) - if self.T_max > 0 - else Constant(self.learning_rate) - ) - - if self.warmup_steps > 0: - learning_rate = self.linear_warmup(learning_rate) - - setattr(learning_rate, "by_epoch", self.by_epoch) - return learning_rate - - -class Step(LRBase): - """Step learning rate decay. - - Args: - epochs (int): Total epoch(s). - iters_per_epoch (int): Number of iterations within an epoch. - learning_rate (float): Learning rate. - step_size (int): The interval to update. - gamma (float, optional): The Ratio that the learning rate will be reduced. - ``new_lr = origin_lr * gamma``. It should be less than 1.0. Default: 0.1. - warmup_epoch (int, optional): The epoch numbers for LinearWarmup. Defaults to 0. - warmup_start_lr (float, optional): Start learning rate within warmup. Defaults to 0.0. - last_epoch (int, optional): Last epoch. Defaults to -1. - by_epoch (bool, optional): Learning rate decays by epoch when by_epoch is True, - else by iter. Defaults to False. - - Examples: - >>> import ppsci - >>> lr = ppsci.optimizer.lr_scheduler.Step(10, 1, 1e-3, 2, 0.95)() - """ - - def __init__( - self, - epochs: int, - iters_per_epoch: int, - learning_rate: float, - step_size: int, - gamma: float, - warmup_epoch: int = 0, - warmup_start_lr: float = 0.0, - last_epoch: int = -1, - by_epoch: bool = False, - ): - super().__init__( - epochs, - iters_per_epoch, - learning_rate, - warmup_epoch, - warmup_start_lr, - last_epoch, - by_epoch, - ) - self.step_size = step_size * iters_per_epoch - self.gamma = gamma - if self.by_epoch: - self.step_size = step_size - - def __call__(self): - learning_rate = lr.StepDecay( - learning_rate=self.learning_rate, - step_size=self.step_size, - gamma=self.gamma, - last_epoch=self.last_epoch, - ) - - if self.warmup_steps > 0: - learning_rate = self.linear_warmup(learning_rate) - - setattr(learning_rate, "by_epoch", self.by_epoch) - return learning_rate - - -class Piecewise(LRBase): - """Piecewise learning rate decay - - Args: - epochs (int): Total epoch(s) - iters_per_epoch (int): Number of iterations within an epoch - decay_epochs (Tuple[int, ...]): A list of steps numbers. The type of element in the - list is python int. - values (Tuple[float, ...]): Tuple of learning rate values that will be picked during - different epoch boundaries. - warmup_epoch (int, optional): The epoch numbers for LinearWarmup. Defaults to 0. - warmup_start_lr (float, optional): Start learning rate within warmup. Defaults to 0.0. - last_epoch (int, optional): Last epoch. Defaults to -1. - by_epoch (bool, optional): Learning rate decays by epoch when by_epoch is True, - else by iter. Defaults to False. - - Examples: - >>> import ppsci - >>> lr = ppsci.optimizer.lr_scheduler.Piecewise( - ... 10, 1, [2, 4], (1e-3, 1e-4, 1e-5) - ... )() - """ - - def __init__( - self, - epochs: int, - iters_per_epoch: int, - decay_epochs: Tuple[int, ...], - values: Tuple[float, ...], - warmup_epoch: int = 0, - warmup_start_lr: float = 0.0, - last_epoch: int = -1, - by_epoch: bool = False, - ): - super().__init__( - epochs, - iters_per_epoch, - values[0], - warmup_epoch, - warmup_start_lr, - last_epoch, - by_epoch, - ) - self.values = values - self.boundaries_steps = [e * iters_per_epoch for e in decay_epochs] - if self.by_epoch is True: - self.boundaries_steps = decay_epochs - - def __call__(self): - learning_rate = lr.PiecewiseDecay( - boundaries=self.boundaries_steps, - values=self.values, - last_epoch=self.last_epoch, - ) - - if self.warmup_steps > 0: - learning_rate = self.linear_warmup(learning_rate) - - setattr(learning_rate, "by_epoch", self.by_epoch) - return learning_rate - - -class MultiStepDecay(LRBase): - """MultiStepDecay learning rate decay - - Args: - epochs (int): Total epoch(s) - iters_per_epoch (int): Number of iterations within an epoch - learning_rate (float): Learning rate - milestones (Tuple[int, ...]): Tuple of each boundaries. should be increasing. - gamma (float, optional): The Ratio that the learning rate will be reduced. - `new_lr = origin_lr * gamma`. It should be less than 1.0. Defaults to 0.1. - warmup_epoch (int, optional): The epoch numbers for LinearWarmup. Defaults to 0. - warmup_start_lr (float, optional): Start learning rate within warmup. Defaults to 0.0. - last_epoch (int, optional): Last epoch. Defaults to -1. - by_epoch (bool, optional): Learning rate decays by epoch when by_epoch is True, - else by iter. Defaults to False. - - Examples: - >>> import ppsci - >>> lr = ppsci.optimizer.lr_scheduler.MultiStepDecay(10, 1, 1e-3, (4, 5))() - """ - - def __init__( - self, - epochs: int, - iters_per_epoch: int, - learning_rate: float, - milestones: Tuple[int, ...], - gamma: float = 0.1, - warmup_epoch: int = 0, - warmup_start_lr: float = 0.0, - last_epoch: int = -1, - by_epoch: bool = False, - ): - super().__init__( - epochs, - iters_per_epoch, - learning_rate, - warmup_epoch, - warmup_start_lr, - last_epoch, - by_epoch, - ) - self.milestones = [x * iters_per_epoch for x in milestones] - self.gamma = gamma - if self.by_epoch: - self.milestones = milestones - - def __call__(self): - learning_rate = lr.MultiStepDecay( - learning_rate=self.learning_rate, - milestones=self.milestones, - gamma=self.gamma, - last_epoch=self.last_epoch, - ) - - if self.warmup_steps > 0: - learning_rate = self.linear_warmup(learning_rate) - - setattr(learning_rate, "by_epoch", self.by_epoch) - return learning_rate - - -class CosineAnnealingWarmRestarts(lr.LRScheduler): - """The implementation of cosine annealing schedule with warm restarts. - - Args: - learning_rate (float): Learning rate - T_0 (int): Number of iterations for the first restart. - T_mult (int, optional): A factor increases T_i after a restart. Defaults to 1. - eta_min (float, optional): Minimum learning rate. Defaults to 0. - last_epoch (int, optional): The index of last epoch. Defaults to -1. - verbose (bool, optional): If `True`, prints a message to stdout for each update. Defaults to False. - """ - - def __init__( - self, - learning_rate: float, - T_0: int, - T_mult: int = 1, - eta_min: float = 0.0, - last_epoch: int = -1, - verbose: bool = False, - ): - if T_0 <= 0 or not isinstance(T_0, int): - raise ValueError(f"Expected positive integer T_0, but got {T_0}") - if T_mult < 1 or not isinstance(T_mult, int): - raise ValueError(f"Expected integer T_mult >= 1, but got {T_mult}") - self.T_0 = T_0 - self.T_i = T_0 - self.T_mult = T_mult - self.eta_min = eta_min - self.T_cur = last_epoch - super().__init__(learning_rate, last_epoch, verbose) - - def get_lr(self): - return ( - self.eta_min - + (self.base_lr - self.eta_min) - * (1 + math.cos(math.pi * self.T_cur / self.T_i)) - / 2 - ) - - def step(self, epoch=None): - if epoch is None and self.last_epoch < 0: - epoch = 0 - - if epoch is None: - epoch = self.last_epoch + 1 - self.T_cur = self.T_cur + 1 - if self.T_cur >= self.T_i: - self.T_cur = self.T_cur - self.T_i - self.T_i = self.T_i * self.T_mult - else: - if epoch < 0: - raise ValueError(f"Expected non-negative epoch, but got {epoch}") - if epoch >= self.T_0: - if self.T_mult == 1: - self.T_cur = epoch % self.T_0 - else: - n = int( - math.log( - (epoch / self.T_0 * (self.T_mult - 1) + 1), self.T_mult - ) - ) - self.T_cur = epoch - self.T_0 * (self.T_mult**n - 1) / ( - self.T_mult - 1 - ) - self.T_i = self.T_0 * self.T_mult ** (n) - else: - self.T_i = self.T_0 - self.T_cur = epoch - self.last_epoch = math.floor(epoch) - self.last_lr = self.get_lr() - - -class CosineWarmRestarts(LRBase): - """Set the learning rate using a cosine annealing schedule with warm restarts. - - Args: - epochs (int): Total epoch(s) - iters_per_epoch (int): Number of iterations within an epoch - learning_rate (float): Learning rate - T_0 (int): Number of iterations for the first restart. - T_mult (int): A factor increases T_i after a restart - eta_min (float, optional): Minimum learning rate. Defaults to 0.0. - warmup_epoch (int, optional): The epoch numbers for LinearWarmup. Defaults to 0. - warmup_start_lr (float, optional): Start learning rate within warmup. Defaults to 0.0. - last_epoch (int, optional): Last epoch. Defaults to -1. - by_epoch (bool, optional): Learning rate decays by epoch when by_epoch is True, else by iter. Defaults to False. - - Examples: - >>> import ppsci - >>> lr = ppsci.optimizer.lr_scheduler.CosineWarmRestarts(20, 1, 1e-3, 14, 2)() - """ - - def __init__( - self, - epochs: int, - iters_per_epoch: int, - learning_rate: float, - T_0: int, - T_mult: int, - eta_min: float = 0.0, - warmup_epoch: int = 0, - warmup_start_lr: float = 0.0, - last_epoch: int = -1, - by_epoch: bool = False, - ): - super().__init__( - epochs, - iters_per_epoch, - learning_rate, - warmup_epoch, - warmup_start_lr, - last_epoch, - by_epoch, - ) - self.T_0 = T_0 - self.T_mult = T_mult - self.eta_min = eta_min - if self.by_epoch is False: - self.T_0 = T_0 * iters_per_epoch - - def __call__(self): - learning_rate = CosineAnnealingWarmRestarts( - learning_rate=self.learning_rate, - T_0=self.T_0, - T_mult=self.T_mult, - eta_min=self.eta_min, - last_epoch=self.last_epoch, - verbose=self.verbose, - ) - - if self.warmup_steps > 0: - learning_rate = self.linear_warmup(learning_rate) - - setattr(learning_rate, "by_epoch", self.by_epoch) - return learning_rate - - -class OneCycleLR(LRBase): - """Sets the learning rate according to the one cycle learning rate scheduler. - The scheduler adjusts the learning rate from an initial learning rate to the maximum learning rate and then - from that maximum learning rate to the minimum learning rate, which is much less than the initial learning rate. - - It has been proposed in [Super-Convergence: Very Fast Training of Neural Networks Using Large Learning Rates](https://arxiv.org/abs/1708.07120). - - Please note that the default behavior of this scheduler follows the fastai implementation of one cycle, - which claims that **"unpublished work has shown even better results by using only two phases"**. - If you want the behavior of this scheduler to be consistent with the paper, please set `three_phase=True`. - - Args: - epochs (int): Total epoch(s). - iters_per_epoch (int): Number of iterations within an epoch. - max_learning_rate (float): The maximum learning rate. It is a python float number. Functionally, it defines the initial learning rate by `divide_factor` . - divide_factor (float, optional): Initial learning rate will be determined by initial_learning_rate = max_learning_rate / divide_factor. Defaults to 25.0. - end_learning_rate (float, optional): The minimum learning rate during training, it should be much less than initial learning rate. Defaults to 0.0001. - phase_pct (float): The percentage of total steps which used to increasing learning rate. Defaults to 0.3. - anneal_strategy (str, optional): Strategy of adjusting learning rate. "cos" for cosine annealing, "linear" for linear annealing. Defaults to "cos". - three_phase (bool, optional): Whether to use three phase. Defaults to False. - warmup_epoch (int, optional): The epoch numbers for LinearWarmup. Defaults to 0. - warmup_start_lr (float, optional): Start learning rate within warmup. Defaults to 0.0. - last_epoch (int, optional): Last epoch. Defaults to -1. - by_epoch (bool, optional): Learning rate decays by epoch when by_epoch is True, else by iter. Defaults to False. - - Examples: - >>> import ppsci - >>> lr = ppsci.optimizer.lr_scheduler.OneCycleLR(100, 1, 1e-3)() - """ - - def __init__( - self, - epochs: int, - iters_per_epoch: int, - max_learning_rate: float, - divide_factor: float = 25.0, - end_learning_rate: float = 0.0001, - phase_pct: float = 0.3, - anneal_strategy: str = "cos", - three_phase: bool = False, - warmup_epoch: int = 0, - warmup_start_lr: float = 0.0, - last_epoch: int = -1, - by_epoch: bool = False, - ): - super().__init__( - epochs, - iters_per_epoch, - max_learning_rate, - warmup_epoch, - warmup_start_lr, - last_epoch, - by_epoch, - ) - self.total_steps = epochs - if not by_epoch: - self.total_steps *= iters_per_epoch - self.divide_factor = divide_factor - self.end_learning_rate = end_learning_rate - self.phase_pct = phase_pct - self.anneal_strategy = anneal_strategy - self.three_phase = three_phase - - def __call__(self): - learning_rate = lr.OneCycleLR( - max_learning_rate=self.learning_rate, - total_steps=self.total_steps, - divide_factor=self.divide_factor, - end_learning_rate=self.end_learning_rate, - phase_pct=self.phase_pct, - anneal_strategy=self.anneal_strategy, - three_phase=self.three_phase, - last_epoch=self.last_epoch, - verbose=self.verbose, - ) - - if self.warmup_steps > 0: - learning_rate = self.linear_warmup(learning_rate) - - setattr(learning_rate, "by_epoch", self.by_epoch) - return learning_rate - - -class LambdaDecay(LRBase): - """This interface provides a lambda function to set the learning rate strategy. - - Args: - epochs (int): Total epoch(s). - iters_per_epoch (int): Number of iterations within an epoch. - learning_rate (float): Learning rate. - lr_lambda (Callable): A lambda function that calculates a factor through epoch, which is multiplied by the initial learning rate. - warmup_epoch (int, optional): The epoch numbers for LinearWarmup. Defaults to 0. - warmup_start_lr (float, optional): Start learning rate within warmup. Defaults to 0.0. - last_epoch (int, optional): Last epoch. Defaults to -1. - by_epoch (bool, optional): Learning rate decays by epoch when by_epoch is True, - else by iter. Defaults to False. - verbose (bool, optional): If True, prints a message to stdout for each update. Defaults to False. - - Examples: - >>> import ppsci - >>> lr = ppsci.optimizer.lr_scheduler.LambdaDecay(0.5, lr_lambda=lambda x:0.95**x, verbose=True)() - """ - - def __init__( - self, - epochs: int, - iters_per_epoch: int, - learning_rate: float, - lr_lambda: Callable, - warmup_epoch: int = 0, - warmup_start_lr: float = 0.0, - last_epoch: int = -1, - by_epoch: bool = False, - verbose: bool = False, - ): - super().__init__( - epochs, - iters_per_epoch, - learning_rate, - warmup_epoch, - warmup_start_lr, - last_epoch, - by_epoch, - verbose, - ) - self.learning_rate = learning_rate - self.lr_lambda = lr_lambda - self.last_epoch = last_epoch - self.verbose = verbose - self.by_epoch = by_epoch - - def __call__(self): - learning_rate = lr.LambdaDecay( - learning_rate=self.learning_rate, - lr_lambda=self.lr_lambda, - last_epoch=self.last_epoch, - verbose=self.verbose, - ) - - if self.warmup_steps > 0: - learning_rate = self.linear_warmup(learning_rate) - - setattr(learning_rate, "by_epoch", self.by_epoch) - return learning_rate - - -class ReduceOnPlateau(LRBase): - """This interface provides a learning rate scheduler that reduces the learning rate when a metric has stopped improving. - - Args: - epochs (int): Total epoch(s). - iters_per_epoch (int): Number of iterations within an epoch. - learning_rate (float): Initial learning rate. - last_epoch (int, optional): The index of the last epoch. Defaults to -1. - warmup_epoch (int, optional): The epoch numbers for LinearWarmup. Defaults to 0. - warmup_start_lr (float, optional): Start learning rate within warmup. Defaults to 0.0. - mode (str, optional): One of `min` or `max`. In `min` mode, lr will be reduced when the quantity monitored has stopped decreasing; in `max` mode it will be reduced when the quantity monitored has stopped increasing. Defaults to "min". - patience (int, optional): Number of epochs with no improvement after which learning rate will be reduced. Defaults to 20. - factor (float, optional): Factor by which the learning rate will be reduced. new_lr = lr * factor. Defaults to 1e-4. - verbose (bool, optional): If True, prints a message to stdout for each update. Defaults to True. - by_epoch (bool, optional): Learning rate decays by epoch when by_epoch is True, else by iter. Defaults to True. - - Examples: - >>> import ppsci - >>> lr = ppsci.optimizer.lr_scheduler.ReduceOnPlateau(epochs=50, iters_per_epoch=100, learning_rate=0.1, mode='min', patience=10, factor=0.5, verbose=True)() - """ - - def __init__( - self, - epochs: int, - iters_per_epoch: int, - learning_rate: float, - last_epoch: int = -1, - warmup_epoch: int = 0, - warmup_start_lr: float = 0.0, - mode: str = "min", - patience: int = 20, - factor: float = 1e-4, - verbose: bool = True, - by_epoch: bool = True, - ): - super().__init__( - epochs, - iters_per_epoch, - learning_rate, - warmup_epoch, - warmup_start_lr, - last_epoch, - by_epoch, - ) - self.mode = mode - self.patience = patience - self.factor = factor - self.verbose = verbose - self.learning_rate = learning_rate - self.by_epoch = by_epoch - - def __call__(self): - learning_rate = lr.ReduceOnPlateau( - mode=self.mode, - patience=self.patience, - factor=self.factor, - verbose=self.verbose, - learning_rate=self.learning_rate, - ) - - if self.warmup_steps > 0: - learning_rate = self.linear_warmup(learning_rate) - - setattr(learning_rate, "by_epoch", self.by_epoch) - return learning_rate - - -class SchedulerList: - """SchedulerList which wrap more than one scheduler. - - Args: - scheduler_list (Tuple[lr.LRScheduler, ...]): Schedulers listed in a tuple. - - Examples: - >>> import ppsci - >>> sch1 = ppsci.optimizer.lr_scheduler.Linear(10, 2, 0.001)() - >>> sch2 = ppsci.optimizer.lr_scheduler.ExponentialDecay(10, 2, 1e-3, 0.95, 3)() - >>> sch = ppsci.optimizer.lr_scheduler.SchedulerList((sch1, sch2)) - """ - - def __init__(self, scheduler_list: Tuple[lr.LRScheduler, ...]): - super().__init__() - self._sch_list = scheduler_list - self.by_epoch = False - - def step(self): - for sch in self._sch_list: - sch.step() - - def get_lr(self) -> float: - """Return learning rate of first scheduler""" - return self._sch_list[0].get_lr() - - def _state_keys(self) -> List[str]: - return ["last_epoch", "last_lr"] - - def __len__(self) -> int: - return len(self._sch_list) - - def __getitem__(self, idx): - return self._sch_list[idx] - - def __setitem__(self, idx, sch): - raise NotImplementedError("Can not modify any item in SchedulerList.") diff --git a/examples/smc_reac/ppsci/optimizer/optimizer.py b/examples/smc_reac/ppsci/optimizer/optimizer.py deleted file mode 100644 index fd4b3d447a..0000000000 --- a/examples/smc_reac/ppsci/optimizer/optimizer.py +++ /dev/null @@ -1,649 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import TYPE_CHECKING -from typing import Dict -from typing import List -from typing import Optional -from typing import Tuple -from typing import Union - -from paddle import nn -from paddle import optimizer as optim -from paddle import regularizer -from paddle.incubate import optimizer as incubate_optim -from typing_extensions import Literal - -from ppsci.optimizer.soap import SOAP as SOAP_impl -from ppsci.utils import logger -from ppsci.utils import misc - -if TYPE_CHECKING: - import paddle - -__all__ = ["SGD", "Momentum", "Adam", "RMSProp", "AdamW", "LBFGS", "OptimizerList"] - - -class SGD: - """Stochastic Gradient Descent. - - Args: - learning_rate (Union[float, optim.lr.LRScheduler], optional): The learning rate - used to update parameter(s). Defaults to 0.001. - weight_decay (Optional[Union[float, regularizer.L1Decay, regularizer.L2Decay]]): - Regularization strategy. Defaults to None. - grad_clip (Optional[Union[nn.ClipGradByNorm, nn.ClipGradByValue, nn.ClipGradByGlobalNorm]]): - Gradient clipping strategy. Defaults to None. - - Examples: - >>> import ppsci - >>> model = ppsci.arch.MLP(("x",), ("u",), 5, 20) - >>> opt = ppsci.optimizer.SGD(1e-3)(model) - """ - - def __init__( - self, - learning_rate: Union[float, optim.lr.LRScheduler] = 0.001, - weight_decay: Optional[ - Union[float, regularizer.L1Decay, regularizer.L2Decay] - ] = None, - grad_clip: Optional[ - Union[nn.ClipGradByNorm, nn.ClipGradByValue, nn.ClipGradByGlobalNorm] - ] = None, - ): - self.learning_rate = learning_rate - self.weight_decay = weight_decay - self.grad_clip = grad_clip - - def __call__(self, model_list: Union[nn.Layer, Tuple[nn.Layer, ...]]): - # model_list is None in static graph - if not isinstance(model_list, (tuple, list)): - model_list = (model_list,) - parameters = ( - sum([m.parameters() for m in model_list], []) if model_list else None - ) - opt = optim.SGD( - learning_rate=self.learning_rate, - parameters=parameters, - weight_decay=self.weight_decay, - grad_clip=self.grad_clip, - ) - return opt - - -class Momentum: - """Simple Momentum optimizer with velocity state. - - Args: - learning_rate (Union[float, optim.lr.LRScheduler]): The learning rate - used to update parameter(s). - momentum (float): Momentum factor. - weight_decay (Optional[Union[float, regularizer.L1Decay, regularizer.L2Decay]]): - Regularization strategy. Defaults to None. - grad_clip (Optional[Union[nn.ClipGradByNorm, nn.ClipGradByValue, nn.ClipGradByGlobalNorm]]): - Gradient clipping strategy. Defaults to None. - use_nesterov (bool, optional): Whether to use nesterov momentum. Defaults to False. - no_weight_decay_name (Optional[str]): List of names of no weight decay parameters split by white space. Defaults to None. - - Examples: - >>> import ppsci - >>> model = ppsci.arch.MLP(("x",), ("u",), 5, 20) - >>> opt = ppsci.optimizer.Momentum(1e-3, 0.9)(model) - """ - - def __init__( - self, - learning_rate: Union[float, optim.lr.LRScheduler], - momentum: float, - weight_decay: Optional[ - Union[float, regularizer.L1Decay, regularizer.L2Decay] - ] = None, - grad_clip: Optional[ - Union[nn.ClipGradByNorm, nn.ClipGradByValue, nn.ClipGradByGlobalNorm] - ] = None, - use_nesterov: bool = False, - no_weight_decay_name: Optional[str] = None, - ): - super().__init__() - self.learning_rate = learning_rate - self.momentum = momentum - self.weight_decay = weight_decay - self.grad_clip = grad_clip - self.use_nesterov = use_nesterov - self.no_weight_decay_name_list = ( - no_weight_decay_name.split() if no_weight_decay_name else [] - ) - - def __call__(self, model_list: Union[nn.Layer, Tuple[nn.Layer, ...]]): - # model_list is None in static graph - if not isinstance(model_list, (tuple, list)): - model_list = (model_list,) - parameters = None - if len(self.no_weight_decay_name_list) > 0: - params_with_decay = [] - params_without_decay = [] - for m in model_list: - params = [ - p - for n, p in m.named_parameters() - if not any(nd in n for nd in self.no_weight_decay_name_list) - ] - params_with_decay.extend(params) - params = [ - p - for n, p in m.named_parameters() - if any(nd in n for nd in self.no_weight_decay_name_list) - ] - params_without_decay.extend(params) - parameters = [ - {"params": params_with_decay, "weight_decay": self.weight_decay}, - {"params": params_without_decay, "weight_decay": 0.0}, - ] - else: - parameters = ( - sum([m.parameters() for m in model_list], []) if model_list else None - ) - opt = optim.Momentum( - learning_rate=self.learning_rate, - momentum=self.momentum, - weight_decay=self.weight_decay, - grad_clip=self.grad_clip, - use_nesterov=self.use_nesterov, - parameters=parameters, - ) - if hasattr(opt, "_use_multi_tensor"): - opt = optim.Momentum( - learning_rate=self.learning_rate, - momentum=self.momentum, - weight_decay=self.weight_decay, - grad_clip=self.grad_clip, - parameters=parameters, - use_nesterov=self.use_nesterov, - use_multi_tensor=True, - ) - return opt - - -class Adam: - """Adam: A Method for Stochastic Optimization. - - Args: - learning_rate (Union[float, optim.lr.LRScheduler], optional): The learning rate - used to update parameter(s). Defaults to 0.001. - beta1 (float, optional): The exponential decay rate for the 1st moment estimates. Defaults to 0.9. - beta2 (float, optional): The exponential decay rate for the 2nd moment estimates. Defaults to 0.999. - epsilon (float, optional): A small float value for numerical stability. Defaults to 1e-08. - weight_decay (Optional[Union[float, regularizer.L1Decay, regularizer.L2Decay]]): Regularization strategy. Defaults to None. - grad_clip (Optional[Union[nn.ClipGradByNorm, nn.ClipGradByValue, nn.ClipGradByGlobalNorm]]): Gradient clipping strategy. Defaults to None. - lazy_mode (bool, optional): Whether to enable lazy mode for moving-average. Defaults to False. - amsgrad (bool, optional): Whether to use the AMSGrad variant of this algorithm from the paper - `On the Convergence of Adam and Beyond `_. Defaults to False. - - Examples: - >>> import ppsci - >>> model = ppsci.arch.MLP(("x",), ("u",), 5, 20) - >>> opt = ppsci.optimizer.Adam(1e-3)(model) - """ - - def __init__( - self, - learning_rate: Union[float, optim.lr.LRScheduler] = 0.001, - beta1: float = 0.9, - beta2: float = 0.999, - epsilon: float = 1e-08, - weight_decay: Optional[ - Union[float, regularizer.L1Decay, regularizer.L2Decay] - ] = None, - grad_clip: Optional[ - Union[nn.ClipGradByNorm, nn.ClipGradByValue, nn.ClipGradByGlobalNorm] - ] = None, - lazy_mode: bool = False, - amsgrad: bool = False, - ): - self.learning_rate = learning_rate - self.beta1 = beta1 - self.beta2 = beta2 - self.epsilon = epsilon - self.learning_rate = learning_rate - self.weight_decay = weight_decay - self.grad_clip = grad_clip - self.lazy_mode = lazy_mode - self.amsgrad = amsgrad - - def __call__(self, model_list: Union[nn.Layer, Tuple[nn.Layer, ...]]): - # model_list is None in static graph - if not isinstance(model_list, (tuple, list)): - model_list = (model_list,) - parameters = ( - sum([m.parameters() for m in model_list], []) if model_list else None - ) - import inspect - - extra_kwargs = {} - if "amsgrad" in inspect.signature(optim.Adam.__init__).parameters: - extra_kwargs["amsgrad"] = self.amsgrad - opt = optim.Adam( - learning_rate=self.learning_rate, - beta1=self.beta1, - beta2=self.beta2, - epsilon=self.epsilon, - weight_decay=self.weight_decay, - grad_clip=self.grad_clip, - lazy_mode=self.lazy_mode, - parameters=parameters, - **extra_kwargs, - ) - return opt - - -class LBFGS: - """The L-BFGS is a quasi-Newton method for solving an unconstrained optimization - problem over a differentiable function. Closely related is the Newton method for minimization. - - Args: - learning_rate (float, optional): The learning rate - used to update parameter(s). Defaults to 1.0. - max_iter (int, optional): Maximal number of iterations per optimization step. - Defaults to 1. - max_eval (Optional[int]): Maximal number of function evaluations per - optimization step. Defaults to None. - tolerance_grad (float, optional): Termination tolerance on first order optimality. - Defaults to 1e-07. - tolerance_change (float, optional): Termination tolerance on function - value/parameter changes. Defaults to 1e-09. - history_size (int, optional): Update history size. Defaults to 100. - line_search_fn (Optional[Literal["strong_wolfe"]]): Either 'strong_wolfe' or None. - Defaults to "strong_wolfe". - - Examples: - >>> import ppsci - >>> model = ppsci.arch.MLP(("x",), ("u",), 5, 20) - >>> opt = ppsci.optimizer.LBFGS(1e-3)(model) - """ - - def __init__( - self, - learning_rate: float = 1.0, - max_iter: int = 1, - max_eval: Optional[int] = None, - tolerance_grad: float = 1e-07, - tolerance_change: float = 1e-09, - history_size: int = 100, - line_search_fn: Optional[Literal["strong_wolfe"]] = "strong_wolfe", - ): - self.lr = learning_rate - self.max_iter = max_iter - self.max_eval = max_eval - self.tolerance_grad = tolerance_grad - self.tolerance_change = tolerance_change - self.history_size = history_size - self.line_search_fn = line_search_fn - - def __call__(self, model_list: Union[nn.Layer, Tuple[nn.Layer, ...]]): - # model_list is None in static graph - if not isinstance(model_list, (tuple, list)): - model_list = (model_list,) - parameters = ( - sum([m.parameters() for m in model_list], []) if model_list else None - ) - try: - opt = getattr(optim, "LBFGS")( - learning_rate=self.lr, - max_iter=self.max_iter, - max_eval=self.max_eval, - tolerance_grad=self.tolerance_grad, - tolerance_change=self.tolerance_change, - history_size=self.history_size, - line_search_fn=self.line_search_fn, - parameters=parameters, - ) - except AttributeError: - opt = getattr(incubate_optim, "LBFGS")( - learning_rate=self.lr, - max_iter=self.max_iter, - max_eval=self.max_eval, - tolerance_grad=self.tolerance_grad, - tolerance_change=self.tolerance_change, - history_size=self.history_size, - line_search_fn=self.line_search_fn, - parameters=parameters, - ) - return opt - - -class RMSProp: - """Root Mean Squared Propagation (RMSProp) is an unpublished, adaptive learning rate method. - - Args: - learning_rate (Union[float, optim.lr.LRScheduler]): The learning rate - used to update parameter(s) - rho (float, optional): Factor ρ in equation. Defaults to 0.95. - epsilon (float, optional): Factor ϵ in equation as a smoothing term. Defaults to 1e-6. - momentum (float, optional):β in equation is the momentum term. Defaults to 0.0. - weight_decay (Optional[Union[float, regularizer.L1Decay, regularizer.L2Decay]]): - Regularization strategy. Defaults to None. - grad_clip (Optional[Union[nn.ClipGradByNorm, nn.ClipGradByValue, nn.ClipGradByGlobalNorm]]): - Gradient clipping strategy. Defaults to None. - - Examples: - >>> import ppsci - >>> model = ppsci.arch.MLP(("x",), ("u",), 5, 20) - >>> opt = ppsci.optimizer.RMSProp(1e-3)(model) - """ - - def __init__( - self, - learning_rate: Union[float, optim.lr.LRScheduler], - rho: float = 0.95, - epsilon: float = 1e-6, - momentum: float = 0.0, - weight_decay: Optional[ - Union[float, regularizer.L1Decay, regularizer.L2Decay] - ] = None, - grad_clip: Optional[ - Union[nn.ClipGradByNorm, nn.ClipGradByValue, nn.ClipGradByGlobalNorm] - ] = None, - ): - super().__init__() - self.learning_rate = learning_rate - self.momentum = momentum - self.rho = rho - self.epsilon = epsilon - self.weight_decay = weight_decay - self.grad_clip = grad_clip - - def __call__(self, model_list: Union[nn.Layer, Tuple[nn.Layer, ...]]): - # model_list is None in static graph - if not isinstance(model_list, (tuple, list)): - model_list = (model_list,) - parameters = ( - sum([m.parameters() for m in model_list], []) if model_list else None - ) - opt = optim.RMSProp( - learning_rate=self.learning_rate, - momentum=self.momentum, - rho=self.rho, - epsilon=self.epsilon, - weight_decay=self.weight_decay, - grad_clip=self.grad_clip, - parameters=parameters, - ) - return opt - - -class AdamW: - """AdamW is implemented based on DECOUPLED WEIGHT DECAY REGULARIZATION. - - Args: - learning_rate (Union[float, optim.lr.LRScheduler], optional): The learning rate - used to update parameter(s). Defaults to 0.001. - beta1 (float, optional): The exponential decay rate for the 1st moment estimates. Defaults to 0.9. - beta2 (float, optional): The exponential decay rate for the 2nd moment estimates. Defaults to 0.999. - epsilon (float, optional): A small float value for numerical stability. Defaults to 1e-8. - weight_decay (float, optional): Regularization coefficient. Defaults to 0.01. - grad_clip (Optional[Union[nn.ClipGradByNorm, nn.ClipGradByValue, nn.ClipGradByGlobalNorm]]): Gradient clipping strategy. Defaults to None. - no_weight_decay_name (Optional[str]): List of names of no weight decay parameters split by white space. Defaults to None. - one_dim_param_no_weight_decay (bool, optional): Apply no weight decay on 1-D parameter(s). Defaults to False. - amsgrad (bool, optional): Whether to use the AMSGrad variant of this algorithm from the paper - `On the Convergence of Adam and Beyond `_. Defaults to False. - - Examples: - >>> import ppsci - >>> model = ppsci.arch.MLP(("x",), ("u",), 5, 20) - >>> opt = ppsci.optimizer.AdamW(1e-3)(model) - """ - - def __init__( - self, - learning_rate: Union[float, optim.lr.LRScheduler] = 0.001, - beta1: float = 0.9, - beta2: float = 0.999, - epsilon: float = 1e-8, - weight_decay: float = 0.001, - grad_clip: Optional[ - Union[nn.ClipGradByNorm, nn.ClipGradByValue, nn.ClipGradByGlobalNorm] - ] = None, - no_weight_decay_name: Optional[str] = None, - one_dim_param_no_weight_decay: bool = False, - amsgrad: bool = False, - ): - super().__init__() - self.learning_rate = learning_rate - self.beta1 = beta1 - self.beta2 = beta2 - self.epsilon = epsilon - self.grad_clip = grad_clip - self.weight_decay = weight_decay - self.no_weight_decay_name_list = ( - no_weight_decay_name.split() if no_weight_decay_name else [] - ) - self.one_dim_param_no_weight_decay = one_dim_param_no_weight_decay - self.amsgrad = amsgrad - - def __call__(self, model_list: Union[nn.Layer, Tuple[nn.Layer, ...]]): - # model_list is None in static graph - if not isinstance(model_list, (tuple, list)): - model_list = (model_list,) - parameters = ( - sum([m.parameters() for m in model_list], []) if model_list else None - ) - - # TODO(gaotingquan): Model_list is None when in static graph, "no_weight_decay" not work. - if model_list is None: - if ( - self.one_dim_param_no_weight_decay - or len(self.no_weight_decay_name_list) != 0 - ): - msg = '"AdamW" does not support setting "no_weight_decay" in static graph. Please use dynamic graph.' - logger.error(Exception(msg)) - raise Exception(msg) - - self.no_weight_decay_param_name_list = ( - [ - p.name - for model in model_list - for n, p in model.named_parameters() - if any(nd in n for nd in self.no_weight_decay_name_list) - ] - if model_list - else [] - ) - - if self.one_dim_param_no_weight_decay: - self.no_weight_decay_param_name_list += ( - [ - p.name - for model in model_list - for n, p in model.named_parameters() - if len(p.shape) == 1 - ] - if model_list - else [] - ) - import inspect - - extra_kwargs = {} - if "amsgrad" in inspect.signature(optim.AdamW.__init__).parameters: - extra_kwargs["amsgrad"] = self.amsgrad - - opt = optim.AdamW( - learning_rate=self.learning_rate, - beta1=self.beta1, - beta2=self.beta2, - epsilon=self.epsilon, - parameters=parameters, - weight_decay=self.weight_decay, - grad_clip=self.grad_clip, - apply_decay_param_fun=self._apply_decay_param_fun, - **extra_kwargs, - ) - return opt - - def _apply_decay_param_fun(self, name): - return name not in self.no_weight_decay_param_name_list - - -class SOAP: - """ - Improving and Stabilizing Shampoo using Adam. Implements SOAP algorithm (https://arxiv.org/abs/2409.11321). - - Args: - learning_rate (float, optional): - The learning rate to use. defaults to 0.003. - beta1 (float, optional): - Adam's betas parameters beta1. defaults to 0.95. - beta2 (float, optional): - Adam's betas parameters beta2. defaults to 0.95. - shampoo_beta (float, optional): - If >= 0, use this beta for the preconditioner (L and R in paper, state['GG'] below) moving average instead of betas[1]. - defaults to -1. - epsilon (float, optional): - Adam's epsilon for numerical stability. defaults to 1e-08. - weight_decay (float, optional): weight decay coefficient. defaults to 0.01. - precondition_frequency (int, optional): - How often to update the preconditioner. defaults to 10. - max_precond_dim (int, optional): - Maximum dimension of the preconditioner. - Set to 10000, so that we exclude most common vocab sizes while including layers. defaults to 10000. - merge_dims (bool, optional): - Whether or not to merge dimensions of the preconditioner. defaults to `False`. - precondition_1d (bool, optional): - Whether or not to precondition 1D gradients. defaults to `False`. - normalize_grads (bool, optional): - Whether or not to normalize gradients per layer. - Helps at large precondition_frequency (~100 in our experiments), - but hurts performance at small precondition_frequency (~10 in our experiments). defaults to `False`. - data_format (str, optional): - Data format of the input for convolutional layers. - Should be "channels_last" for data_format of NHWC and "channels_first" for NCHW. defaults to `channels_first`. - correct_bias (bool, optional): - Whether or not to use bias correction in Adam. defaults to `True`. - - Examples: - >>> import ppsci - >>> model = ppsci.arch.MLP(("x",), ("u",), 5, 20) - >>> opt = ppsci.optimizer.SOAP(1e-3)(model) - """ - - def __init__( - self, - learning_rate: float = 3e-3, - beta1: float = 0.95, - beta2: float = 0.95, - shampoo_beta: float = -1, - epsilon: float = 1e-8, - weight_decay: float = 0.01, - precondition_frequency: int = 10, - max_precond_dim: int = 10000, # - merge_dims: bool = False, # Merge dimensions till the product of the dimensions is less than or equal to max_precond_dim. - precondition_1d: bool = False, - normalize_grads: bool = False, - data_format: str = "channels_first", - correct_bias: bool = True, - ): - self.learning_rate = learning_rate - self.beta1 = beta1 - self.beta2 = beta2 - self.shampoo_beta = shampoo_beta - self.epsilon = epsilon - self.weight_decay = weight_decay - self.precondition_frequency = precondition_frequency - self.max_precond_dim = max_precond_dim - self.merge_dims = merge_dims - self.precondition_1d = precondition_1d - self.normalize_grads = normalize_grads - self.data_format = data_format - self.correct_bias = correct_bias - - def __call__(self, model_list: Union[nn.Layer, Tuple[nn.Layer, ...]]): - # model_list is None in static graph - if not isinstance(model_list, (tuple, list)): - model_list = (model_list,) - parameters = ( - sum([m.parameters() for m in model_list], []) if model_list else None - ) - opt = SOAP_impl( - parameters=parameters, - learning_rate=self.learning_rate, - beta1=self.beta1, - beta2=self.beta2, - shampoo_beta=self.shampoo_beta, - epsilon=self.epsilon, - weight_decay=self.weight_decay, - precondition_frequency=self.precondition_frequency, - max_precond_dim=self.max_precond_dim, - merge_dims=self.merge_dims, - precondition_1d=self.precondition_1d, - normalize_grads=self.normalize_grads, - data_format=self.data_format, - correct_bias=self.correct_bias, - ) - return opt - - -class OptimizerList: - """OptimizerList which wrap more than one optimizer. - NOTE: LBFGS is not supported yet. - - Args: - optimizer_list (Tuple[optim.Optimizer, ...]): Optimizers listed in a tuple. - - Examples: - >>> import ppsci - >>> model1 = ppsci.arch.MLP(("x",), ("u",), 5, 20) - >>> opt1 = ppsci.optimizer.Adam(1e-3)(model1) - >>> model2 = ppsci.arch.MLP(("y",), ("v",), 5, 20) - >>> opt2 = ppsci.optimizer.Adam(1e-3)(model2) - >>> opt = ppsci.optimizer.OptimizerList((opt1, opt2)) - """ - - def __init__(self, optimizer_list: Tuple[optim.Optimizer, ...]): - super().__init__() - self._opt_list = optimizer_list - if "LBFGS" in set(misc.typename(opt) for opt in optimizer_list): - raise ValueError("LBFGS is not supported in OptimizerList yet.") - - def step(self): - for opt in self._opt_list: - opt.step() - - def clear_grad(self): - for opt in self._opt_list: - opt.clear_grad() - - def get_lr(self) -> float: - """Return learning rate of first optimizer""" - return self._opt_list[0].get_lr() - - def set_state_dict(self, state_dicts: List[Dict[str, "paddle.Tensor"]]): - for i, opt in enumerate(self._opt_list): - opt.set_state_dict(state_dicts[i]) - - def state_dict(self) -> List[Dict[str, "paddle.Tensor"]]: - state_dicts = [opt.state_dict() for opt in self._opt_list] - return state_dicts - - def __len__(self) -> int: - return len(self._opt_list) - - def __getitem__(self, idx): - return self._opt_list[idx] - - def __setitem__(self, idx, opt): - raise NotImplementedError("Can not modify any item in OptimizerList.") - - def __iter__(self): - yield from iter(self._opt_list) diff --git a/examples/smc_reac/ppsci/optimizer/soap.py b/examples/smc_reac/ppsci/optimizer/soap.py deleted file mode 100644 index de239a978b..0000000000 --- a/examples/smc_reac/ppsci/optimizer/soap.py +++ /dev/null @@ -1,558 +0,0 @@ -# Copyright (c) 2025 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. -# -# refs: https://github.com/nikhilvyas/SOAP - -from collections import defaultdict -from itertools import chain - -import paddle -import paddle.optimizer as optim - - -class SOAP(optim.Optimizer): - """ - Improving and Stabilizing Shampoo using Adam. Implements SOAP algorithm (https://arxiv.org/abs/2409.11321). - - Parameters: - parameters (list|tuple): - List/Tuple of ``Tensor`` names to update to minimize ``loss``. - learning_rate (float, optional): - The learning rate to use. defaults to 0.003. - beta1 (float optional): - Adam's betas parameters b1. defaults to 0.95. - beta2 (float optional): - Adam's betas parameters b1. defaults to 0.95. - shampoo_beta (float, optional): - If >= 0, use this beta for the preconditioner (L and R in paper, state['GG'] below) moving average instead of betas[1]. - defaults to -1. - epsilon (float, optional): - Adam's epsilonilon for numerical stability. defaults to 1e-08. - weight_decay (float, optional): weight decay coefficient. defaults to 0.01. - precondition_frequency (int, optional): - How often to update the preconditioner. defaults to 10. - max_precond_dim (int, optional): - Maximum dimension of the preconditioner. - Set to 10000, so that we exclude most common vocab sizes while including layers. defaults to 10000. - merge_dims (bool, optional): - Whether or not to merge dimensions of the preconditioner. defaults to `False`. - precondition_1d (bool, optional): - Whether or not to precondition 1D gradients. defaults to `False`. - normalize_grads (bool, optional): - Whether or not to normalize gradients per layer. - Helps at large precondition_frequency (~100 in our experiments), - but hurts performance at small precondition_frequency (~10 in our experiments). defaults to `False`. - data_format (str, optional): - Data format of the input for convolutional layers. - Should be "channels_last" for data_format of NHWC and "channels_first" for NCHW. defaults to `channels_first`. - correct_bias (bool, optional): - Whether or not to use bias correction in Adam. defaults to `True`. - name (str, optional): Normally there is no need for user to set this property. - For more information, please refer to :ref:`api_guide_Name`. - The default value is None. - - Return: - loss (Tensor): the final loss of closure. - - Examples: - .. code-block:: python - - >>> import paddle - >>> import ppsci - >>> import numpy as np - - >>> np.random.seed(0) - >>> np_w = np.random.rand(1).astype(np.float32) - >>> np_x = np.random.rand(1).astype(np.float32) - - >>> inputs = [np.random.rand(1).astype(np.float32) for i in range(10)] - >>> # y = 2x - >>> targets = [2 * x for x in inputs] - - >>> class Net(paddle.nn.Layer): - ... def __init__(self): - ... super().__init__() - ... w = paddle.to_tensor(np_w) - ... self.w = paddle.create_parameter(shape=w.shape, dtype=w.dtype, default_initializer=paddle.nn.initializer.Assign(w)) - ... - ... def forward(self, x): - ... return self.w * x - ... - >>> net = Net() - >>> opt = ppsci.optimizer.soap.SOAP(parameters=net.parameters()) - >>> def train_step(inputs, targets): - ... def closure(): - ... outputs = net(inputs) - ... loss = paddle.nn.functional.mse_loss(outputs, targets) - ... print('loss: ', loss.item()) - ... opt.clear_grad() - ... loss.backward() - ... return loss - ... opt.step(closure) - ... - >>> for input, target in zip(inputs, targets): - ... input = paddle.to_tensor(input) - ... target = paddle.to_tensor(target) - ... train_step(input, target) - """ - - def __init__( - self, - parameters, - learning_rate: float = 3e-3, - beta1: float = 0.95, - beta2: float = 0.95, - shampoo_beta: float = -1, - epsilon: float = 1e-8, - weight_decay: float = 0.01, - precondition_frequency: int = 10, - max_precond_dim: int = 10000, # - merge_dims: bool = False, # Merge dimensions till the product of the dimensions is less than or equal to max_precond_dim. - precondition_1d: bool = False, - normalize_grads: bool = False, - data_format: str = "channels_first", - correct_bias: bool = True, - name: str = None, - ): - self._betas = paddle.to_tensor((beta1, beta2)) - self._shampoo_beta = shampoo_beta - self._epsilon = epsilon - self._weight_decay = weight_decay - self._precondition_frequency = precondition_frequency - self._max_precond_dim = max_precond_dim - self._merge_dims = merge_dims - self._precondition_1d = precondition_1d - self._normalize_grads = normalize_grads - self._correct_bias = correct_bias - - self.state = defaultdict(dict) - - super().__init__( - learning_rate=learning_rate, - parameters=parameters, - weight_decay=weight_decay, - name=name, - ) - - if isinstance(self._parameter_list[0], dict): - raise TypeError("The parameter groups is not supported on SOAP optimizer.") - - self._data_format = data_format - - def merge_dims(self, grad, max_precond_dim): - """ - Merges dimensions of the gradient tensor till the product of the dimensions is less than or equal to max_precond_dim. - """ - assert self._data_format in ["channels_first", "channels_last"] - if self._data_format == "channels_last" and grad.ndim == 4: - grad = grad.transpose(0, 3, 1, 2) - shape = grad.shape - new_shape = [] - - curr_shape = 1 - for dim_size in shape: - temp_shape = curr_shape * dim_size - if temp_shape > max_precond_dim: - if curr_shape > 1: - new_shape.append(curr_shape) - curr_shape = dim_size - else: - new_shape.append(dim_size) - curr_shape = 1 - else: - curr_shape = temp_shape - - if curr_shape > 1 or len(new_shape) == 0: - new_shape.append(curr_shape) - - new_grad = grad.reshape(new_shape) - return new_grad - - @paddle.base.framework.non_static_only - def step(self, closure=None): - """ - Performs a single optimization step. - - Arguments: - closure (Optional[Callable]): A closure that reevaluates the model and returns the loss. - """ - with paddle.no_grad(): - if closure is None: - loss = None - else: - closure = paddle.enable_grad()(closure) - loss = closure() - - for p in self._parameter_list: - if p.grad is None: - continue - grad = p.grad - - state = self.state[p] - - if "step" not in state: - state["step"] = 0 - - # State initialization - if "exp_avg" not in state: - # Exponential moving average of gradient values - state["exp_avg"] = paddle.zeros_like(grad) - # Exponential moving average of squared gradient values - state["exp_avg_sq"] = paddle.zeros_like(grad) - - if "Q" not in state: - self.init_preconditioner( - grad, - state, - precondition_frequency=self._precondition_frequency, - precondition_1d=self._precondition_1d, - shampoo_beta=( - self._shampoo_beta - if self._shampoo_beta >= 0 - else self._betas[1] - ), - max_precond_dim=self._max_precond_dim, - merge_dims=self._merge_dims, - ) - self.update_preconditioner( - grad, - state, - max_precond_dim=self._max_precond_dim, - merge_dims=self._merge_dims, - precondition_1d=self._precondition_1d, - ) - continue # first step is skipped so that we never use the current gradients in the projection. - - # Projecting gradients to the eigenbases of Shampoo's preconditioner - # i.e. projecting to the eigenbases of matrices in state['GG'] - grad_projected = self.project( - grad, - state, - merge_dims=self._merge_dims, - max_precond_dim=self._max_precond_dim, - ) - - exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"] - beta1, beta2 = self._betas - - state["step"] += 1 - - # Decay the first and second moment running average coefficient - # In-place operations to update the averages at the same time - exp_avg.multiply_(beta1).add_((1.0 - beta1) * grad_projected) - exp_avg_sq.multiply_(beta2).add_( - (1.0 - beta2) * grad_projected.square() - ) - - denom = exp_avg_sq.sqrt().add_( - paddle.full([], self._epsilon, dtype=exp_avg_sq.dtype) - ) - - # Projecting the exponential moving average of gradients to the eigenbases of Shampoo's preconditioner - # i.e. projecting to the eigenbases of matrices in state['GG'] - # exp_avg_projected = self.project(exp_avg, state, merge_dims=self._merge_dims"], - # max_precond_dim=self._max_precond_dim']) - exp_avg_projected = exp_avg - - step_size = self.get_lr() - if self._correct_bias: - bias_correction1 = 1.0 - beta1 ** (state["step"]) - bias_correction2 = 1.0 - beta2 ** (state["step"]) - step_size = step_size * (bias_correction2**0.5) / bias_correction1 - - # Projecting back the preconditioned (by Adam) exponential moving average of gradients - # to the original space - norm_grad = self.project_back( - exp_avg_projected / denom, - state, - merge_dims=self._merge_dims, - max_precond_dim=self._max_precond_dim, - ) - - if self._normalize_grads: - norm_grad = norm_grad / (1e-30 + paddle.mean(norm_grad**2) ** 0.5) - - p.add_(-step_size * norm_grad) - - # From AdamW code: Just adding the square of the weights to the loss function is *not* - # the correct way of using L2 regularization/weight decay with Adam, - # since that will interact with the m and v parameters in strange ways. - # - # Instead we want to decay the weights in a manner that doesn't interact - # with the m/v parameters. This is equivalent to adding the square - # of the weights to the loss with plain (non-momentum) SGD. - # Add weight decay at the end (fixed version) - if self._weight_decay > 0.0: - p.add_((-self.get_lr() * self._weight_decay) * p) - - # Update is done after the gradient step to avoid using current gradients in the projection. - self.update_preconditioner( - grad, - state, - max_precond_dim=self._max_precond_dim, - merge_dims=self._merge_dims, - precondition_1d=self._precondition_1d, - ) - - return loss - - def init_preconditioner( - self, - grad, - state, - precondition_frequency=10, - shampoo_beta=0.95, - max_precond_dim=10000, - precondition_1d=False, - merge_dims=False, - ): - """ - Initializes the preconditioner matrices (L and R in the paper). - """ - state[ - "GG" - ] = [] # Will hold all the preconditioner matrices (L and R in the paper). - if grad.ndim == 1: - if not precondition_1d or grad.shape[0] > max_precond_dim: - state["GG"].append([]) - else: - state["GG"].append(paddle.zeros([grad.shape[0], grad.shape[0]])) - else: - if merge_dims: - grad = self.merge_dims(grad, max_precond_dim) - - for dim_size in grad.shape: - if dim_size > max_precond_dim: - state["GG"].append([]) - else: - state["GG"].append(paddle.zeros([dim_size, dim_size])) - - state["Q"] = None # Will hold all the eigenbases of the preconditioner. - state["precondition_frequency"] = precondition_frequency - state["shampoo_beta"] = shampoo_beta - - def project(self, grad, state, merge_dims=False, max_precond_dim=10000): - """ - Projects the gradient to the eigenbases of the preconditioner. - """ - original_shape = grad.shape - if merge_dims: - if grad.ndim == 4 and self._data_format == "channels_last": - transposed_shape = grad.transpose(0, 3, 1, 2).shape - grad = self.merge_dims(grad, max_precond_dim) - - for mat in state["Q"]: - if len(mat) > 0: - grad = paddle.tensordot( - grad, - mat, - axes=[[0], [0]], - ) - else: - transpose_order = list(range(1, len(grad.shape))) + [0] - grad = grad.transpose(transpose_order) - - if merge_dims: - if self._data_format == "channels_last" and len(original_shape) == 4: - grad = grad.reshape(transposed_shape).transpose(0, 2, 3, 1) - else: - grad = grad.reshape(original_shape) - return grad - - def update_preconditioner( - self, - grad, - state, - max_precond_dim=10000, - merge_dims=False, - precondition_1d=False, - ): - """ - Updates the preconditioner matrices and the eigenbases (L, R, Q_L, Q_R in the paper). - """ - if state["Q"] is not None: - state["exp_avg"] = self.project_back( - state["exp_avg"], - state, - merge_dims=merge_dims, - max_precond_dim=max_precond_dim, - ) - if grad.ndim == 1: - if precondition_1d and grad.shape[0] <= max_precond_dim: - state["GG"][0].lerp_( - grad.unsqueeze(1) @ grad.unsqueeze(0), 1 - state["shampoo_beta"] - ) - else: - if merge_dims: - new_grad = self.merge_dims(grad, max_precond_dim) - for idx, dim_size in enumerate(new_grad.shape): - if dim_size <= max_precond_dim: - outer_product = paddle.tensordot( - new_grad, - new_grad, - axes=[ - [ - *chain( - range(idx), range(idx + 1, len(new_grad.shape)) - ) - ] - ] - * 2, - ) - state["GG"][idx].lerp_(outer_product, 1 - state["shampoo_beta"]) - else: - for idx, dim_size in enumerate(grad.shape): - if dim_size <= max_precond_dim: - outer_product = paddle.tensordot( - grad, - grad, - # Contracts across all dimensions except for k. - axes=[[*chain(range(idx), range(idx + 1, len(grad.shape)))]] - * 2, - ) - state["GG"][idx].lerp_(outer_product, 1 - state["shampoo_beta"]) - - if state["Q"] is None: - state["Q"] = self.get_orthogonal_matrix(state["GG"]) - if state["step"] > 0 and state["step"] % state["precondition_frequency"] == 0: - state["Q"] = self.get_orthogonal_matrix_QR( - state, max_precond_dim, merge_dims - ) - # state['Q'] = self.get_fast_QR(state, max_precond_dim, merge_dims) - - if state["step"] > 0: - state["exp_avg"] = self.project( - state["exp_avg"], - state, - merge_dims=merge_dims, - max_precond_dim=max_precond_dim, - ) - - def project_back(self, grad, state, merge_dims=False, max_precond_dim=10000): - """ - Projects the gradient back to the original space. - """ - original_shape = grad.shape - if merge_dims: - if self._data_format == "channels_last" and grad.ndim == 4: - transposed_shape = grad.transpose(0, 3, 1, 2).shape - grad = self.merge_dims(grad, max_precond_dim) - for mat in state["Q"]: - if len(mat) > 0: - grad = paddle.tensordot( - grad, - mat, - axes=[[0], [1]], - ) - else: - transpose_order = list(range(1, len(grad.shape))) + [0] - grad = grad.transpose(transpose_order) - - if merge_dims: - if self._data_format == "channels_last" and len(original_shape) == 4: - grad = grad.reshape(transposed_shape).transpose(0, 2, 3, 1) - else: - grad = grad.reshape(original_shape) - return grad - - def get_orthogonal_matrix(self, mat): - """ - Computes the eigenbases of the preconditioner using paddle.linalg.eigh decomposition. - """ - matrix = [] - for m in mat: - if len(m) == 0: - matrix.append([]) - continue - if m.dtype != paddle.float32: - float_data = False - original_type = m.dtype - original_device = m.place - matrix.append(m.to(paddle.float32)) - else: - float_data = True - matrix.append(m) - - final = [] - for m in matrix: - if len(m) == 0: - final.append([]) - continue - _, Q = paddle.linalg.eigh(m + 1e-30 * paddle.eye(m.shape[0])) - Q = paddle.flip(Q, [1]) - - if not float_data: - Q = Q.to(original_device, dtype=original_type) - final.append(Q) - return final - - def get_orthogonal_matrix_QR(self, state, max_precond_dim=10000, merge_dims=False): - """ - Computes the eigenbases of the preconditioner using one round of power iteration - followed by paddle.linalg.qr decomposition. - """ - precond_list = state["GG"] - orth_list = state["Q"] - - matrix = [] - orth_matrix = [] - for m, o in zip(precond_list, orth_list): - if len(m) == 0: - matrix.append([]) - orth_matrix.append([]) - continue - if m.dtype != paddle.float32: - float_data = False - original_type = m.dtype - original_device = m.place - matrix.append(m.to(paddle.float32)) - orth_matrix.append(o.to(paddle.float32)) - else: - float_data = True - matrix.append(m.to(paddle.float32)) - orth_matrix.append(o.to(paddle.float32)) - - orig_shape = state["exp_avg_sq"].shape - if self._data_format == "channels_last" and len(orig_shape) == 4: - transposed_shape = state["exp_avg_sq"].transpose(0, 3, 1, 2).shape - if merge_dims: - exp_avg_sq = self.merge_dims(state["exp_avg_sq"], max_precond_dim) - else: - exp_avg_sq = state["exp_avg_sq"] - - final = [] - for ind, (m, o) in enumerate(zip(matrix, orth_matrix)): - if len(m) == 0: - final.append([]) - continue - est_eig = paddle.diag(o.T @ m @ o) - sort_idx = paddle.argsort(est_eig, descending=True) - exp_avg_sq = exp_avg_sq.index_select(sort_idx, ind) - o = o[:, sort_idx] - power_iter = m @ o - Q, _ = paddle.linalg.qr(power_iter) - - if not float_data: - Q = Q.to(original_device, dtype=original_type) - final.append(Q) - - if merge_dims: - if self._data_format == "channels_last" and len(orig_shape) == 4: - exp_avg_sq = exp_avg_sq.reshape(transposed_shape).transpose(0, 2, 3, 1) - else: - exp_avg_sq = exp_avg_sq.reshape(orig_shape) - - state["exp_avg_sq"] = exp_avg_sq - - return final diff --git a/examples/smc_reac/ppsci/probability/__init__.py b/examples/smc_reac/ppsci/probability/__init__.py deleted file mode 100644 index 1068ada02f..0000000000 --- a/examples/smc_reac/ppsci/probability/__init__.py +++ /dev/null @@ -1,19 +0,0 @@ -# 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. - -from ppsci.probability.hmc import HamiltonianMonteCarlo - -__all__ = [ - "HamiltonianMonteCarlo", -] diff --git a/examples/smc_reac/ppsci/probability/hmc.py b/examples/smc_reac/ppsci/probability/hmc.py deleted file mode 100644 index 76a4847f77..0000000000 --- a/examples/smc_reac/ppsci/probability/hmc.py +++ /dev/null @@ -1,175 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import Callable -from typing import Dict - -import paddle -from paddle import distribution - -from ppsci import utils - - -class EnableGradient: - """ - This class is for enabling a dict of tensor for autodiff - """ - - def __init__(self, tensor_dict: Dict[str, paddle.Tensor]): - self.tensor_dict = tensor_dict - - def __enter__(self): - for t in self.tensor_dict.values(): - t.stop_gradient = False - t.clear_grad() - - def __exit__(self, exec_type, exec_val, exec_tb): - for t in self.tensor_dict.values(): - t.stop_gradient = True - - -class HamiltonianMonteCarlo: - """ - Using the HamiltonianMonteCarlo(HMC) to sample from the desired probability distribution. The HMC combine the Hamiltonian Dynamics and Markov Chain Monte Carlo sampling algorithm which is a more efficient way compared to the Metropolis Hasting (MH) method. - - Args: - distribution_fn (paddle.distribution.Distribution): The Log (Posterior) Distribution function that of the parameters needed to be sampled. - path_len (float): The total path length. - step_size (float): Every step size. - num_warmup_steps (int): The number of warm-up steps for the MCMC run. - random_seed (int): Random seed number. - - Examples: - >>> import paddle - >>> from ppsci.probability.hmc import HamiltonianMonteCarlo - >>> def log_posterior(**kwargs): - ... dist = paddle.distribution.Normal(loc=0, scale=1) - ... return dist.log_prob(kwargs['x']) - >>> HMC = HamiltonianMonteCarlo(log_posterior, path_len=1.5, step_size=0.25) - >>> trial = HMC.run_chain(1000, {'x': paddle.to_tensor(0.0)}) - """ - - def __init__( - self, - distribution_fn: Callable, - path_len: float = 1.0, - step_size: float = 0.25, - num_warmup_steps: int = 0, - random_seed: int = 1024, - ): - self.distribution_fn = distribution_fn - self.steps = int(path_len / step_size) - self.step_size = step_size - self.path_len = path_len - self.num_warmup_steps = num_warmup_steps - utils.set_random_seed(random_seed) - self._rv_unif = distribution.Uniform(0, 1) - - def sample( - self, last_position: Dict[str, paddle.Tensor] - ) -> Dict[str, paddle.Tensor]: - """ - Single step for sample - """ - q0 = q1 = last_position - p0 = p1 = self._sample_r(q0) - - for s in range(self.steps): - grad = self._potential_energy_gradient(q1) - for site_name in p1.keys(): - p1[site_name] -= self.step_size * grad[site_name] / 2 - for site_name in q1.keys(): - q1[site_name] += self.step_size * p1[site_name] - - grad = self._potential_energy_gradient(q1) - for site_name in p1.keys(): - p1[site_name] -= self.step_size * grad[site_name] / 2 - - # set the next state in the Markov chain - return q1 if self._check_acceptance(q0, q1, p0, p1) else q0 - - def run_chain( - self, epochs: int, initial_position: Dict[str, paddle.Tensor] - ) -> Dict[str, paddle.Tensor]: - sampling_result: Dict[str, paddle.Tensor] = {} - for k in initial_position.keys(): - sampling_result[k] = [] - pos = initial_position - - # warmup - for _ in range(self.num_warmup_steps): - pos = self.sample(pos) - - # begin collecting sampling result - for e in range(epochs): - pos = self.sample(pos) - for k in pos.keys(): - sampling_result[k].append(pos[k].numpy()) - - for k in initial_position.keys(): - sampling_result[k] = paddle.to_tensor(sampling_result[k]) - - return sampling_result - - def _potential_energy_gradient( - self, pos: Dict[str, paddle.Tensor] - ) -> Dict[str, paddle.Tensor]: - """ - Calculate the gradient of potential energy - """ - grads = {} - with EnableGradient(pos): - (-self.distribution_fn(**pos)).backward() - for k, v in pos.items(): - grads[k] = v.grad.detach() - return grads - - def _k_energy_fn(self, r: Dict[str, paddle.Tensor]) -> paddle.Tensor: - energy = 0.0 - for v in r.values(): - energy = energy + v.dot(v) - return 0.5 * energy - - def _sample_r( - self, params_dict: Dict[str, paddle.Tensor] - ) -> Dict[str, paddle.Tensor]: - # sample r for params - r = {} - for k, v in params_dict.items(): - rv_r = distribution.Normal(paddle.zeros_like(v), paddle.ones_like(v)) - r[k] = rv_r.sample([1]) - if not (v.shape == [] or v.shape == 1): - r[k] = r[k].squeeze() - return r - - def _check_acceptance( - self, - q0: Dict[str, paddle.Tensor], - q1: Dict[str, paddle.Tensor], - p0: Dict[str, paddle.Tensor], - p1: Dict[str, paddle.Tensor], - ) -> bool: - # calculate the Metropolis acceptance probability - energy_current = -self.distribution_fn(**q0) + self._k_energy_fn(p0) - energy_proposed = -self.distribution_fn(**q1) + self._k_energy_fn(p1) - - acceptance = paddle.minimum( - paddle.to_tensor(1.0), paddle.exp(energy_current - energy_proposed) - ) - - # whether accept the proposed state position - event = self._rv_unif.sample([]) - return event <= acceptance diff --git a/examples/smc_reac/ppsci/solver/__init__.py b/examples/smc_reac/ppsci/solver/__init__.py deleted file mode 100644 index 03f97bc2d9..0000000000 --- a/examples/smc_reac/ppsci/solver/__init__.py +++ /dev/null @@ -1,25 +0,0 @@ -# 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. - -from ppsci.solver import eval -from ppsci.solver import train -from ppsci.solver import visu -from ppsci.solver.solver import Solver - -__all__ = [ - "eval", - "train", - "visu", - "Solver", -] diff --git a/examples/smc_reac/ppsci/solver/eval.py b/examples/smc_reac/ppsci/solver/eval.py deleted file mode 100644 index 1af7655901..0000000000 --- a/examples/smc_reac/ppsci/solver/eval.py +++ /dev/null @@ -1,316 +0,0 @@ -# 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. - -from __future__ import annotations - -import time -from typing import TYPE_CHECKING -from typing import Dict -from typing import Optional -from typing import Tuple -from typing import Union - -import paddle -from paddle import io - -from ppsci.solver import printer -from ppsci.solver.train import _compute_batch_size -from ppsci.utils import misc - -if TYPE_CHECKING: - from pgl.utils import data as pgl_data - - from ppsci import solver - - -def _get_dataset_length( - data_loader: Union["io.DataLoader", "pgl_data.Dataloader", "io.IterableDataset"] -) -> int: - """Get full dataset length of given dataloader. - - Args: - data_loader (Union[io.DataLoader, pgl_data.Dataloader, io.IterableDataset]): - Given dataloader. - - Returns: - int: Length of full dataset. - """ - if isinstance(data_loader, io.DataLoader): - num_samples = len(data_loader.dataset) - elif isinstance(data_loader, io.IterableDataset): - num_samples = data_loader.num_samples - elif str(type(data_loader)) == "": - num_samples = len(data_loader.dataset) - else: - raise NotImplementedError( - f"Can not fetch the length of given dataset({type(data_loader)})." - ) - - return num_samples - - -def _eval_by_dataset( - solver: "solver.Solver", epoch_id: Optional[int], log_freq: int -) -> Tuple[float, Dict[str, Dict[str, float]]]: - """Evaluate with computing metric on total samples(default process). - - NOTE: This is the default evaluation method as general for most cases, but may not - memory-efficiency for large dataset or large output. - - Args: - solver (solver.Solver): Main Solver. - epoch_id (Optional[int]): Epoch id. - log_freq (int): Log evaluation information every `log_freq` steps. - - Returns: - Tuple[float, Dict[str, Dict[str, float]]]: Target metric and all metric dicts - computed during evaluation. - """ - target_metric: float = float("inf") - metric_dict_group: Dict[str, Dict[str, float]] = misc.PrettyOrderedDict() - for _, _validator in solver.validator.items(): - all_output = misc.Prettydefaultdict(list) - all_label = misc.Prettydefaultdict(list) - num_samples = _get_dataset_length(_validator.data_loader) - - loss_dict = misc.Prettydefaultdict(float) - reader_tic = time.perf_counter() - batch_tic = time.perf_counter() - for iter_id, batch in enumerate(_validator.data_loader, start=1): - input_dict, label_dict, weight_dict = batch - reader_cost = time.perf_counter() - reader_tic - - for v in input_dict.values(): - if hasattr(v, "stop_gradient"): - v.stop_gradient = False - - # forward - with solver.autocast_context_manager( - solver.use_amp, solver.amp_level - ), solver.no_grad_context_manager(solver.eval_with_no_grad): - output_dict, validator_loss = solver.forward_helper.eval_forward( - _validator.output_expr, - input_dict, - solver.model, - _validator, - label_dict, - weight_dict, - ) - - loss_dict[f"{_validator.name}/loss"] = float( - sum(list(validator_loss.values())) - ) - - for key, output in output_dict.items(): - all_output[key].append( - (output.detach() if hasattr(output, "detach") else output) - if solver.world_size == 1 - else misc.all_gather(output.detach()) - ) - - for key, label in label_dict.items(): - all_label[key].append( - (label.detach() if hasattr(label, "detach") else label) - if solver.world_size == 1 - else misc.all_gather(label.detach()) - ) - - batch_cost = time.perf_counter() - batch_tic - solver.eval_time_info["reader_cost"].update(reader_cost) - solver.eval_time_info["batch_cost"].update(batch_cost) - batch_size = _compute_batch_size(input_dict) - printer.update_eval_loss(solver, loss_dict, batch_size) - if ( - iter_id == 1 - or iter_id % log_freq == 0 - or iter_id == len(_validator.data_loader) - ): - printer.log_eval_info( - solver, - batch_size, - epoch_id, - len(_validator.data_loader), - iter_id, - ) - - reader_tic = time.perf_counter() - batch_tic = time.perf_counter() - - # concatenate all data and discard padded sample(s) - for key in all_output: - if paddle.is_tensor(all_output[key][0]): - all_output[key] = paddle.concat(all_output[key]) - if len(all_output[key]) > num_samples: - all_output[key] = all_output[key][:num_samples] - - for key in all_label: - if paddle.is_tensor(all_label[key][0]): - all_label[key] = paddle.concat(all_label[key]) - if len(all_label[key]) > num_samples: - all_label[key] = all_label[key][:num_samples] - - for metric_name, metric_func in _validator.metric.items(): - # NOTE: compute metric with entire output and label - metric_dict = metric_func(all_output, all_label) - metric_dict_group[metric_name] = { - k: float(v) for k, v in metric_dict.items() - } - for var_name, metric_value in metric_dict.items(): - metric_str = f"{_validator.name}/{metric_name}.{var_name}" - if metric_str not in solver.eval_output_info: - solver.eval_output_info[metric_str] = misc.AverageMeter( - metric_str, ".5f" - ) - solver.eval_output_info[metric_str].update( - float(metric_value), num_samples - ) - - # use the first metric for return value - tmp = metric_dict_group - while isinstance(tmp, dict): - tmp = next(iter(tmp.values())) - # avoid that none of metric is set - if isinstance(tmp, float): - target_metric = float(tmp) - - return target_metric, metric_dict_group - - -def _eval_by_batch( - solver: "solver.Solver", epoch_id: Optional[int], log_freq: int -) -> Tuple[float, Dict[str, Dict[str, float]]]: - """Evaluate with computing metric by batch, which is memory-efficient. - - NOTE: This is a evaluation function for large dataset or large output, as is more - memory-efficiency than evaluating by dataset, but less general because some metric - is not independent among samples, e.g. L2 relative error. - - Args: - solver (solver.Solver): Main Solver. - epoch_id (Optional[int]): Epoch id. - log_freq (int): Log evaluation information every `log_freq` steps. - - Returns: - Tuple[float, Dict[str, Dict[str, float]]]: Target metric and all metric dicts - computed during evaluation. - """ - target_metric: float = float("inf") - metric_dict_group: Dict[str, Dict[str, float]] = misc.PrettyOrderedDict() - for _, _validator in solver.validator.items(): - num_samples = _get_dataset_length(_validator.data_loader) - - loss_dict = misc.Prettydefaultdict(float) - reader_tic = time.perf_counter() - batch_tic = time.perf_counter() - for iter_id, batch in enumerate(_validator.data_loader, start=1): - input_dict, label_dict, weight_dict = batch - reader_cost = time.perf_counter() - reader_tic - - batch_size = _compute_batch_size(input_dict) - for v in input_dict.values(): - if hasattr(v, "stop_gradient"): - v.stop_gradient = False - - # forward - with solver.autocast_context_manager( - solver.use_amp, solver.amp_level - ), solver.no_grad_context_manager(solver.eval_with_no_grad): - output_dict, validator_loss = solver.forward_helper.eval_forward( - _validator.output_expr, - input_dict, - solver.model, - _validator, - label_dict, - weight_dict, - ) - - loss_dict[f"{_validator.name}/loss"] = float( - sum(list(validator_loss.values())) - ) - - # collect batch metric - for metric_name, metric_func in _validator.metric.items(): - metric_dict_group[metric_name] = misc.Prettydefaultdict(list) - metric_dict = metric_func(output_dict, label_dict) - for var_name, metric_value in metric_dict.items(): - metric_dict_group[metric_name][var_name].append( - metric_value - if solver.world_size == 1 - else misc.all_gather(metric_value) - ) - - batch_cost = time.perf_counter() - batch_tic - solver.eval_time_info["reader_cost"].update(reader_cost) - solver.eval_time_info["batch_cost"].update(batch_cost) - printer.update_eval_loss(solver, loss_dict, batch_size) - if ( - iter_id == 1 - or iter_id % log_freq == 0 - or iter_id == len(_validator.data_loader) - ): - printer.log_eval_info( - solver, - batch_size, - epoch_id, - len(_validator.data_loader), - iter_id, - ) - - reader_tic = time.perf_counter() - batch_tic = time.perf_counter() - - # concatenate all metric and discard metric of padded sample(s) - for metric_name, metric_dict in metric_dict_group.items(): - for var_name, metric_value in metric_dict.items(): - # NOTE: concat single metric(scalar) list into metric vector - metric_value = paddle.concat(metric_value)[:num_samples] - # NOTE: compute metric via averaging metric over all samples, - # this might be not general for certain evaluation case - metric_value = float(metric_value.mean()) - metric_dict_group[metric_name][var_name] = metric_value - metric_str = f"{_validator.name}/{metric_name}.{var_name}" - if metric_str not in solver.eval_output_info: - solver.eval_output_info[metric_str] = misc.AverageMeter( - metric_str, ".5f" - ) - solver.eval_output_info[metric_str].update(metric_value, num_samples) - - # use the first metric for return value - tmp = metric_dict_group - while isinstance(tmp, dict): - tmp = next(iter(tmp.values())) - # avoid that none of metric is set - if isinstance(tmp, float): - target_metric = tmp - - return target_metric, metric_dict_group - - -def eval_func( - solver: "solver.Solver", epoch_id: Optional[int], log_freq: int -) -> Tuple[float, Dict[str, Dict[str, float]]]: - """Evaluation function. - - Args: - solver (solver.Solver): Main Solver. - epoch_id (Optional[int]): Epoch id. - log_freq (int): Log evaluation information every `log_freq` steps. - - Returns: - Tuple[float, Dict[str, Dict[str, float]]]: Target metric and all metric dicts - computed during evaluation. - """ - if solver.compute_metric_by_batch: - return _eval_by_batch(solver, epoch_id, log_freq) - return _eval_by_dataset(solver, epoch_id, log_freq) diff --git a/examples/smc_reac/ppsci/solver/printer.py b/examples/smc_reac/ppsci/solver/printer.py deleted file mode 100644 index cedaeab7cc..0000000000 --- a/examples/smc_reac/ppsci/solver/printer.py +++ /dev/null @@ -1,161 +0,0 @@ -# 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. - -import datetime -from typing import TYPE_CHECKING -from typing import Dict -from typing import Optional - -from paddle import device - -from ppsci.utils import logger -from ppsci.utils import misc - -if TYPE_CHECKING: - from ppsci import solver - - -def update_train_loss( - solver: "solver.Solver", loss_dict: Dict[str, float], batch_size: int -): - for key in loss_dict: - if key not in solver.train_output_info: - solver.train_output_info[key] = misc.AverageMeter(key, "7.5f") - solver.train_output_info[key].update(float(loss_dict[key]), batch_size) - if key not in solver.train_loss_info: - solver.train_loss_info[key] = misc.AverageMeter(key, ".5f") - solver.train_loss_info[key].update(float(loss_dict[key])) - - -def update_eval_loss( - solver: "solver.Solver", loss_dict: Dict[str, float], batch_size: int -): - for key in loss_dict: - if key not in solver.eval_output_info: - solver.eval_output_info[key] = misc.AverageMeter(key, "7.5f") - solver.eval_output_info[key].update(float(loss_dict[key]), batch_size) - - -def log_train_info( - solver: "solver.Solver", batch_size: int, epoch_id: int, iter_id: int -): - lr_msg = f"lr: {solver.optimizer.get_lr():.5f}" - - metric_msg = ", ".join( - [ - f"{key}: {solver.train_output_info[key].avg:.5f}" - for key in solver.train_output_info - ] - ) - - time_msg = ", ".join( - [solver.train_time_info[key].mean for key in solver.train_time_info] - ) - - ips_msg = f"ips: {batch_size / solver.train_time_info['batch_cost'].avg:.2f}" - if solver.benchmark_flag: - ips_msg += " samples/s" - - eta_sec = ( - (solver.epochs - epoch_id + 1) * solver.iters_per_epoch - iter_id - ) * solver.train_time_info["batch_cost"].avg - eta_msg = f"eta: {str(datetime.timedelta(seconds=int(eta_sec)))}" - - epoch_width = len(str(solver.epochs)) - iters_width = len(str(solver.iters_per_epoch)) - log_str = ( - f"[Train][Epoch {epoch_id:>{epoch_width}}/{solver.epochs}]" - f"[Iter {iter_id:>{iters_width}}/{solver.iters_per_epoch}] {lr_msg}, " - f"{metric_msg}, {time_msg}, {ips_msg}, {eta_msg}" - ) - if solver.benchmark_flag: - max_mem_reserved_msg = ( - f"max_mem_reserved: {device.cuda.max_memory_reserved() // (1 << 20)} MB" - ) - max_mem_allocated_msg = ( - f"max_mem_allocated: {device.cuda.max_memory_allocated() // (1 << 20)} MB" - ) - log_str += f", {max_mem_reserved_msg}, {max_mem_allocated_msg}" - logger.info(log_str) - - # reset time information after printing - for key in solver.train_time_info: - solver.train_time_info[key].reset() - - logger.scalar( - { - "train/lr": solver.optimizer.get_lr(), - **{ - f"train/{key}": solver.train_output_info[key].avg - for key in solver.train_output_info - }, - }, - step=solver.global_step, - vdl_writer=solver.vdl_writer, - wandb_writer=solver.wandb_writer, - tbd_writer=solver.tbd_writer, - ) - - -def log_eval_info( - solver: "solver.Solver", - batch_size: int, - epoch_id: Optional[int], - iters_per_epoch: int, - iter_id: int, -): - metric_msg = ", ".join( - [ - f"{key}: {solver.eval_output_info[key].avg:.5f}" - for key in solver.eval_output_info - ] - ) - - time_msg = ", ".join( - [solver.eval_time_info[key].mean for key in solver.eval_time_info] - ) - - ips_msg = f"ips: {batch_size / solver.eval_time_info['batch_cost'].avg:.2f}" - - eta_sec = (iters_per_epoch - iter_id) * solver.eval_time_info["batch_cost"].avg - eta_msg = f"eta: {str(datetime.timedelta(seconds=int(eta_sec)))}" - - epoch_width = len(str(solver.epochs)) - iters_width = len(str(iters_per_epoch)) - if isinstance(epoch_id, int): - logger.info( - f"[Eval][Epoch {epoch_id:>{epoch_width}}/{solver.epochs}]" - f"[Iter {iter_id:>{iters_width}}/{iters_per_epoch}] " - f"{metric_msg}, {time_msg}, {ips_msg}, {eta_msg}" - ) - else: - logger.info( - f"[Eval][Iter {iter_id:>{iters_width}}/{iters_per_epoch}] " - f"{metric_msg}, {time_msg}, {ips_msg}, {eta_msg}" - ) - - # reset time information after printing - for key in solver.eval_time_info: - solver.eval_time_info[key].reset() - - # logger.scalar( - # { - # f"eval/{key}": solver.eval_output_info[key].avg - # for key in solver.eval_output_info - # }, - # step=solver.global_step, - # vdl_writer=solver.vdl_writer, - # wandb_writer=solver.wandb_writer, - # tbd_writer=solver.tbd_writer, - # ) diff --git a/examples/smc_reac/ppsci/solver/solver.py b/examples/smc_reac/ppsci/solver/solver.py deleted file mode 100644 index 390211c023..0000000000 --- a/examples/smc_reac/ppsci/solver/solver.py +++ /dev/null @@ -1,1219 +0,0 @@ -# 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. - -from __future__ import annotations - -import contextlib -import functools -import importlib -import itertools -import os -import sys -from os import path as osp -from typing import TYPE_CHECKING -from typing import Callable -from typing import Dict -from typing import List -from typing import Mapping -from typing import Optional -from typing import Tuple -from typing import Union - -import numpy as np -import paddle -import paddle.distributed as dist -import sympy as sp -from omegaconf import DictConfig -from omegaconf import OmegaConf -from packaging import version -from paddle import amp -from paddle import jit -from paddle import nn -from paddle import optimizer as optim -from paddle.distributed import fleet -from paddle.framework import core -from typing_extensions import Literal - -import ppsci -from ppsci.loss import mtl -from ppsci.utils import ema -from ppsci.utils import expression -from ppsci.utils import logger -from ppsci.utils import misc -from ppsci.utils import save_load - -if TYPE_CHECKING: - from types import ModuleType - - from paddle.static import InputSpec - - -class Solver: - """Class for solver. - - Args: - model (nn.Layer): Model. - constraint (Optional[Dict[str, ppsci.constraint.Constraint]]): Constraint(s) applied on model. Defaults to None. - output_dir (Optional[str]): Output directory. Defaults to "./output/". - optimizer (Optional[optimizer.Optimizer]): Optimizer object. Defaults to None. - lr_scheduler (Optional[optimizer.lr.LRScheduler]): Learning rate scheduler. Defaults to None. - epochs (int, optional): Training epoch(s). Defaults to 5. - iters_per_epoch (int, optional): Number of iterations within an epoch. If set to -1, - than will be automatically set to the length of dataloader of given constraint. - Defaults to 20. - update_freq (int, optional): Update frequency of parameters. Defaults to 1. - save_freq (int, optional): Saving frequency for checkpoint. Defaults to 0. - log_freq (int, optional): Logging frequency. Defaults to 10. - eval_during_train (bool, optional): Whether evaluate model during training. Defaults to False. - start_eval_epoch (int, optional): Epoch number evaluation applied begin after. Defaults to 1. - eval_freq (int, optional): Evaluation frequency. Defaults to 1. - seed (int, optional): Random seed. Defaults to 42. - use_vdl (Optional[bool]): Whether use VisualDL to log scalars. Defaults to False. - use_wandb (Optional[bool]): Whether use wandb to log data. Defaults to False. - use_tbd (Optional[bool]): Whether use tensorboardX to log data. Defaults to False. - wandb_config (Optional[Dict[str, str]]): Config dict of WandB. Defaults to None. - device (Literal["cpu", "gpu", "xpu", "sdaa", None], optional): Runtime device. Defaults to None, which means use default device on current platform. - equation (Optional[Dict[str, ppsci.equation.PDE]]): Equation dict. Defaults to None. - geom (Optional[Dict[str, ppsci.geometry.Geometry]]): Geometry dict. Defaults to None. - validator (Optional[Dict[str, ppsci.validate.Validator]]): Validator dict. Defaults to None. - visualizer (Optional[Dict[str, ppsci.visualize.Visualizer]]): Visualizer dict. Defaults to None. - use_amp (bool, optional): Whether use AMP. Defaults to False. - amp_level (Literal["O0", "O1", "O2", "OD"], optional): AMP level. Defaults to "O1". - pretrained_model_path (Optional[str]): Pretrained model path. Defaults to None. - checkpoint_path (Optional[str]): Checkpoint path. Defaults to None. - compute_metric_by_batch (bool, optional): Whether calculate metrics after each batch during evaluation. Defaults to False. - eval_with_no_grad (bool, optional): Whether set `stop_gradient=True` for every Tensor if no differentiation - involved during computation, generally for save GPU memory and accelerate computing. Defaults to False. - to_static (bool, optional): Whether enable to_static for forward pass. Defaults to False. - loss_aggregator (Optional[mtl.LossAggregator]): Loss aggregator, such as a multi-task learning loss aggregator. Defaults to None. - cfg: (Optional[DictConfig]): Running config dict. Defaults to None. NOTE: This will be required in the future. - - Examples: - >>> import ppsci - >>> model = ppsci.arch.MLP(("x",), ("u",), 5, 20) - >>> opt = ppsci.optimizer.AdamW(1e-3)(model) - >>> geom = ppsci.geometry.Rectangle((0, 0), (1, 1)) - >>> pde_constraint = ppsci.constraint.InteriorConstraint( - ... {"u": lambda out: out["u"]}, - ... {"u": 0}, - ... geom, - ... { - ... "dataset": "IterableNamedArrayDataset", - ... "iters_per_epoch": 1, - ... "batch_size": 16, - ... }, - ... ppsci.loss.MSELoss("mean"), - ... name="EQ", - ... ) # doctest: +SKIP - >>> solver = ppsci.solver.Solver( - ... model, - ... {"EQ": pde_constraint}, - ... "./output", - ... opt, - ... None, - ... ) # doctest: +SKIP - """ - - def __init__( - self, - model: nn.Layer, - constraint: Optional[Dict[str, ppsci.constraint.Constraint]] = None, - output_dir: Optional[str] = "./output/", - optimizer: Optional[optim.Optimizer] = None, - lr_scheduler: Optional[optim.lr.LRScheduler] = None, - epochs: int = 5, - iters_per_epoch: int = 20, - update_freq: int = 1, - save_freq: int = 0, - log_freq: int = 10, - eval_during_train: bool = False, - start_eval_epoch: int = 1, - eval_freq: int = 1, - seed: int = 42, - use_vdl: bool = False, - use_wandb: bool = False, - use_tbd: bool = False, - wandb_config: Optional[Mapping] = None, - device: Literal["cpu", "gpu", "xpu", "sdaa", None] = None, - equation: Optional[Dict[str, ppsci.equation.PDE]] = None, - geom: Optional[Dict[str, ppsci.geometry.Geometry]] = None, - validator: Optional[Dict[str, ppsci.validate.Validator]] = None, - visualizer: Optional[Dict[str, ppsci.visualize.Visualizer]] = None, - use_amp: bool = False, - amp_level: Literal["O0", "O1", "O2", "OD"] = "O1", - pretrained_model_path: Optional[str] = None, - checkpoint_path: Optional[str] = None, - compute_metric_by_batch: bool = False, - eval_with_no_grad: bool = False, - to_static: bool = False, - loss_aggregator: Optional[mtl.LossAggregator] = None, - *, - cfg: Optional[DictConfig] = None, - ): - self.cfg = cfg - if isinstance(cfg, DictConfig): - # (Recommended)Params can be passed within cfg - # rather than passed to 'Solver.__init__' one-by-one. - self._parse_params_from_cfg(cfg) - - # set model - self.model = model - # set constraint - self.constraint = constraint - # set output directory - if not cfg: - self.output_dir = output_dir - - # set optimizer - self.optimizer = optimizer - # set learning rate scheduler - if lr_scheduler is not None: - logger.warning( - "The argument: 'lr_scheduler' now automatically retrieves from " - "'optimizer._learning_rate' when 'optimizer' is given, so it is " - "recommended to remove it from the Solver's initialization arguments." - ) - self.lr_scheduler = ( - optimizer._learning_rate - if ( - isinstance(optimizer, optim.Optimizer) - and isinstance(optimizer._learning_rate, optim.lr.LRScheduler) - ) - else None - ) - if isinstance(self.optimizer, ppsci.optimizer.OptimizerList): - self.lr_scheduler = ppsci.optimizer.lr_scheduler.SchedulerList( - tuple( - opt._learning_rate - for opt in self.optimizer - if isinstance(opt._learning_rate, optim.lr.LRScheduler) - ) - ) - - # set training hyper-parameter - if not cfg: - self.epochs = epochs - self.iters_per_epoch = iters_per_epoch - # set update_freq for gradient accumulation - self.update_freq = update_freq - # set checkpoint saving frequency - self.save_freq = save_freq - # set logging frequency - self.log_freq = log_freq - - # set evaluation hyper-parameter - self.eval_during_train = eval_during_train - self.start_eval_epoch = start_eval_epoch - self.eval_freq = eval_freq - - if self.iters_per_epoch == -1 and self.constraint is not None: - if len(self.constraint) != 1: - raise NotImplementedError( - f"Multiple({len(self.constraint)}) constraints are detected, " - "which is not supported yet, when 'iters_per_epoch' is set to -1." - ) - self.iters_per_epoch = len(next(iter(self.constraint.values())).data_loader) - logger.message( - "Detected 'iters_per_epoch' is set to -1, 'iters_per_epoch' is now " - f"reset to the length of dataloader({self.iters_per_epoch}) of given constraint." - ) - - # initialize training log(training loss, time cost, etc.) recorder during one epoch - self.train_output_info: Dict[str, misc.AverageMeter] = {} - self.train_time_info = { - "batch_cost": misc.AverageMeter("batch_cost", ".5f", postfix="s"), - "reader_cost": misc.AverageMeter("reader_cost", ".5f", postfix="s"), - } - self.train_loss_info: Dict[str, misc.AverageMeter] = {} - - # initialize evaluation log(evaluation loss, metric, etc.) recorder. - self.eval_output_info: Dict[str, misc.AverageMeter] = {} - self.eval_time_info = { - "batch_cost": misc.AverageMeter("batch_cost", ".5f", postfix="s"), - "reader_cost": misc.AverageMeter("reader_cost", ".5f", postfix="s"), - } - - # set running device - if not cfg: - self.device = device - if self.device is None: - # set to default device if not specified - self.device: str = paddle.device.get_device() - - if self.device != "cpu" and paddle.device.get_device() == "cpu": - # fall back to cpu if no other device available - logger.warning(f"Set device({device}) to 'cpu' for only cpu available.") - self.device = "cpu" - self.device = paddle.device.set_device(self.device) - - # set equations for physics-driven or data-physics hybrid driven task, such as PINN - self.equation = equation - - # set validator - self.validator = validator - - # set visualizer - self.visualizer = visualizer - - # set automatic mixed precision(AMP) configuration - if not cfg: - self.use_amp = use_amp - self.amp_level = amp_level - self.scaler = amp.GradScaler(True) if self.use_amp else None - - # whether calculate metrics by each batch during evaluation, mainly for memory efficiency - if not cfg: - self.compute_metric_by_batch = compute_metric_by_batch - if validator is not None: - for metric in itertools.chain( - *[_v.metric.values() for _v in self.validator.values()] - ): - if metric.keep_batch ^ self.compute_metric_by_batch: - raise ValueError( - f"{misc.typename(metric)}.keep_batch should be " - f"{self.compute_metric_by_batch} when compute_metric_by_batch=" - f"{self.compute_metric_by_batch}." - ) - # check metric name uniqueness over all validators - _count = {} - for _validator in validator.values(): - for metric_name in _validator.metric: - if metric_name in _count: - logger.warning( - f"Metric name({metric_name}) is duplicated, please ensure " - "all metric names are unique over all given validators." - ) - _count[metric_name] = 1 - del _count - - # whether set `stop_gradient=True` for every Tensor if no differentiation involved during evaluation - if not cfg: - self.eval_with_no_grad = eval_with_no_grad - - self.rank = dist.get_rank() - self.world_size = dist.get_world_size() - # initialize distributed environment - if self.world_size > 1: - # TODO(sensen): Support different kind of DistributedStrategy - fleet.init(is_collective=True) - logger.warning( - f"Detected 'world_size'({self.world_size}) > 1, it is recommended to " - "scale up the learning rate and reduce the 'epochs' or " - "'iters_per_epoch' according to the 'world_size' both linearly if you " - "are training model." - ) - - # set moving average model(optional) - self.ema_model = None - if self.cfg and any(key in self.cfg.TRAIN for key in ["ema", "swa"]): - if "ema" in self.cfg.TRAIN and cfg.TRAIN.ema.get("use_ema", False): - self.ema_model = ema.ExponentialMovingAverage( - self.model, self.cfg.TRAIN.ema.decay - ) - elif "swa" in self.cfg.TRAIN and cfg.TRAIN.swa.get("use_swa", False): - self.ema_model = ema.StochasticWeightAverage(self.model) - - # load pretrained model, usually used for transfer learning - if not cfg: - self.pretrained_model_path = pretrained_model_path - if self.pretrained_model_path is not None: - save_load.load_pretrain( - self.model, self.pretrained_model_path, self.equation - ) - - self.cur_metric = float("inf") - # initialize an dict for tracking best metric during training - self.best_metric = { - "metric": float("inf"), - "epoch": 0, - } - - # use loss aggregator, use Sum if None - if isinstance(loss_aggregator, (mtl.AGDA, mtl.PCGrad)) and self.use_amp: - raise ValueError( - "Auto Mix Precision do not support AGDA, PCGrad loss aggregator yet, " - "please set use_amp=False." - ) - self.loss_aggregator = loss_aggregator or mtl.Sum() - - # load model checkpoint, usually used for resume training - if not cfg: - self.checkpoint_path = checkpoint_path - if self.checkpoint_path is not None: - if self.pretrained_model_path is not None: - logger.warning( - "Detected 'pretrained_model_path' is given, weights in which might be" - "overridden by weights loaded from given 'checkpoint_path'." - ) - loaded_metric = save_load.load_checkpoint( - self.checkpoint_path, - self.model, - self.optimizer, - self.scaler, - self.equation, - self.ema_model, - self.loss_aggregator, - ) - if isinstance(loaded_metric, dict): - self.best_metric.update(loaded_metric) - - # decorate model(s) and optimizer(s) for AMP - if self.use_amp: - self.model, self.optimizer = amp.decorate( - self.model, - self.optimizer, - self.amp_level, - save_dtype="float32", - ) - - # choosing an appropriate training function for different optimizers - if misc.typename(self.optimizer) == "LBFGS": - if self.use_amp: - raise ValueError( - "Auto Mix Precision is not supported for L-BFGS optimizer." - ) - self.train_epoch_func = ppsci.solver.train.train_LBFGS_epoch_func - if self.update_freq != 1: - self.update_freq = 1 - logger.warning("Set 'update_freq' to to 1 when using L-BFGS optimizer.") - else: - self.train_epoch_func = ppsci.solver.train.train_epoch_func - - # wrap model and optimizer to parallel object - if self.world_size > 1: - if isinstance(self.model, paddle.DataParallel): - raise ValueError( - "Given model is already wrapped by paddle.DataParallel." - "Please do not wrap your model with DataParallel " - "before 'Solver.__init__' and keep it's type as 'nn.Layer'." - ) - - def dist_wrapper(model: nn.Layer) -> paddle.DataParallel: - dist_model = fleet.distributed_model(model) - if hasattr(model, "input_keys"): - dist_model.input_keys = dist_model._layers.input_keys - if hasattr(model, "output_keys"): - dist_model.output_keys = dist_model._layers.output_keys - return dist_model - - if isinstance(self.model, ppsci.arch.ModelList): - for i in range(len(self.model.model_list)): - # NOTE: Convert each model in model_list to DataParallel - self.model.model_list[i] = dist_wrapper(self.model.model_list[i]) - else: - self.model = dist_wrapper(self.model) - - if self.optimizer is not None: - self.optimizer = fleet.distributed_optimizer(self.optimizer) - - # set VisualDL tool - self.vdl_writer = None - if not cfg: - self.use_vdl = use_vdl - if self.use_vdl: - try: - import visualdl as vdl - except ModuleNotFoundError: - raise ModuleNotFoundError( - "Please install 'visualdl' with `pip install visualdl` first." - ) - with misc.RankZeroOnly(self.rank) as is_master: - if is_master: - self.vdl_writer = vdl.LogWriter(osp.join(self.output_dir, "vdl")) - logger.info( - "VisualDL is enabled for logging, you can view it by " - f"running:\nvisualdl --logdir {self.vdl_writer._logdir} --port 8080" - ) - - # set WandB tool - self.wandb_writer = None - if not cfg: - self.use_wandb = use_wandb - self.wandb_config = {} - if self.use_wandb: - try: - import wandb - except ModuleNotFoundError: - raise ModuleNotFoundError( - "Please install 'wandb' with `pip install wandb` first." - ) - with misc.RankZeroOnly(self.rank) as is_master: - if is_master: - self.wandb_writer = wandb.init(**self.wandb_config) - - # set TensorBoardX tool - self.tbd_writer = None - if not cfg: - self.use_tbd = use_tbd - if self.use_tbd: - try: - import tensorboardX - except ModuleNotFoundError: - raise ModuleNotFoundError( - "Please install 'tensorboardX' with `pip install tensorboardX` first." - ) - with misc.RankZeroOnly(self.rank) as is_master: - if is_master: - self.tbd_writer = tensorboardX.SummaryWriter( - osp.join(self.output_dir, "tensorboard") - ) - logger.message( - "TensorboardX is enabled for logging, you can view it by " - f"running:\ntensorboard --logdir {self.tbd_writer.logdir}" - ) - - self.global_step = 0 - - # log paddlepaddle's version - if version.Version(paddle.__version__) != version.Version("0.0.0"): - paddle_version = paddle.__version__ - if version.Version(paddle.__version__) < version.Version("2.6.0"): - logger.warning( - f"Detected paddlepaddle version is '{paddle_version}', " - "currently it is recommended to use paddlepaddle >= 2.6 or develop version." - ) - else: - paddle_version = f"develop({paddle.version.commit[:7]})" - - logger.info(f"Using paddlepaddle {paddle_version} on device {self.device}") - - self.forward_helper = expression.ExpressionSolver() - - # whether enable static for forward pass. Defaults to False - if not cfg: - self.to_static = to_static - jit.enable_to_static(self.to_static) - logger.message( - f"Set to_static={self.to_static} for computational optimization." - ) - - # convert sympy to callable object if exist - extra_parameters = [] - if self.equation: - for equation in self.equation.values(): - extra_parameters += list(equation.learnable_parameters) - - def convert_expr( - container_dict: Union[ - Dict[str, ppsci.constraint.Constraint], - Dict[str, ppsci.validate.Validator], - Dict[str, ppsci.visualize.Visualizer], - ] - ) -> None: - for container in container_dict.values(): - exprs = [ - expr - for expr in container.output_expr.values() - if isinstance(expr, sp.Basic) - ] - if len(exprs) > 0: - funcs = ppsci.lambdify( - exprs, - self.model, - extra_parameters=extra_parameters, - # graph_filename=osp.join(self.output_dir, "symbolic_graph_visual"), # HACK: Activate it for DEBUG. - fuse_derivative=True, - ) - ind = 0 - for name in container.output_expr: - if isinstance(container.output_expr[name], sp.Basic): - container.output_expr[name] = funcs[ind] - # FIXME: Equation with parameter not support yet. - # if self.world_size > 1: - # container.output_expr[name] = dist_wrapper( - # container.output_expr[name] - # ) - ind += 1 - - if self.constraint: - convert_expr(self.constraint) - - if self.validator: - convert_expr(self.validator) - - if self.visualizer: - convert_expr(self.visualizer) - - # set up benchmark flag, will print memory stat if enabled - self.benchmark_flag: bool = os.getenv("BENCHMARK_ROOT", None) is not None - - # set up nvtx flag for nsight analysis - self.nvtx_flag: bool = os.getenv("NVTX", None) is not None - self.forward_helper.nvtx_flag = self.nvtx_flag - - # for callbacks - self.callbacks_on_epoch_begin: List[Callable[[Solver]]] = [] - self.callbacks_on_epoch_end: List[Callable[[Solver]]] = [] - self.callbacks_on_iter_begin: List[Callable[[Solver]]] = [] - self.callbacks_on_iter_end: List[Callable[[Solver]]] = [] - - def train(self) -> None: - """Training.""" - self.global_step = self.best_metric["epoch"] * self.iters_per_epoch - self.max_steps = self.epochs * self.iters_per_epoch - - start_epoch = self.best_metric["epoch"] + 1 - - if self.use_tbd and isinstance(self.cfg, DictConfig): - self.tbd_writer.add_text( - "config", f"
{str(OmegaConf.to_yaml(self.cfg))}
" - ) - - if self.nvtx_flag: - core.nvprof_start() - core.nvprof_enable_record_event() - - for epoch_id in range(start_epoch, self.epochs + 1): - self._invoke_callbacks_on_epoch_begin() # [optional] - self.train_epoch_func(self, epoch_id, self.log_freq) - self._invoke_callbacks_on_epoch_end() # [optional] - - self.train_output_info.clear() - - # update average model if exist - if self.ema_model and epoch_id % self.avg_freq == 0: - self.ema_model.update() - - # evaluate during training - if ( - self.eval_during_train - and epoch_id % self.eval_freq == 0 - and epoch_id >= self.start_eval_epoch - ): - self.cur_metric, metric_dict_group = self.eval(epoch_id) - if self.cur_metric < self.best_metric["metric"]: - self.best_metric["metric"] = self.cur_metric - self.best_metric["epoch"] = epoch_id - save_load.save_checkpoint( - self.model, - self.optimizer, - self.best_metric, - self.scaler, - self.output_dir, - "best_model", - self.equation, - aggregator=self.loss_aggregator, - ) - logger.info( - f"[Eval][Epoch {epoch_id}]" - f"[best metric: {self.best_metric['metric']}]" - ) - for metric_name, metric_dict in metric_dict_group.items(): - logger.scalar( - {f"eval/{metric_name}/{k}": v for k, v in metric_dict.items()}, - epoch_id, - self.vdl_writer, - self.wandb_writer, - self.tbd_writer, - ) - - # visualize after evaluation - if self.visualizer is not None: - self.visualize(epoch_id) - - # evaluation for moving average evaluation(almost same procedure) - if self.ema_model and epoch_id % self.avg_freq == 0: - self.ema_model.apply_shadow() - logger.info("Evaluating metric of averaging model...") - cur_metric_ema, metric_dict_group_ema = self.eval(epoch_id) - self.ema_model.restore() - - if cur_metric_ema < self.best_metric["metric"]: - self.best_metric["metric"] = cur_metric_ema - self.best_metric["epoch"] = epoch_id - save_load.save_checkpoint( - self.ema_model, - None, - metric=self.best_metric, - output_dir=self.output_dir, - prefix="best_model_ema", - ) - logger.info( - f"[Eval][Epoch {epoch_id}]" - f"[best metric: {self.best_metric['metric']}]" - ) - for metric_name, metric_dict in metric_dict_group_ema.items(): - logger.scalar( - { - f"eval_ema/{metric_name}/{k}": v - for k, v in metric_dict.items() - }, - epoch_id, - self.vdl_writer, - self.wandb_writer, - self.tbd_writer, - ) - - # update learning rate by epoch - if self.lr_scheduler is not None and self.lr_scheduler.by_epoch: - self.lr_scheduler.step() - - # save epoch model every save_freq epochs - if self.save_freq > 0 and epoch_id % self.save_freq == 0: - save_load.save_checkpoint( - self.model, - self.optimizer, - {"metric": self.cur_metric, "epoch": epoch_id}, - self.scaler, - self.output_dir, - f"epoch_{epoch_id}", - self.equation, - ema_model=self.ema_model, - aggregator=self.loss_aggregator, - ) - - # save the latest model for convenient resume training - save_load.save_checkpoint( - self.model, - self.optimizer, - {"metric": self.cur_metric, "epoch": epoch_id}, - self.scaler, - self.output_dir, - "latest", - self.equation, - print_log=(epoch_id == start_epoch), - ema_model=self.ema_model, - aggregator=self.loss_aggregator, - ) - - def finetune(self, pretrained_model_path: str) -> None: - """Finetune model based on given pretrained model path. - - Args: - pretrained_model_path (str): Pretrained model path or url. - """ - # load pretrained model - save_load.load_pretrain(self.model, pretrained_model_path, self.equation) - - # call train program - self.train() - - @misc.run_on_eval_mode - def eval( - self, epoch_id: Optional[int] = None - ) -> Tuple[float, Dict[str, Dict[str, float]]]: - """Evaluation. - - Args: - epoch_id (Optional[int]): Epoch id. Defaults to None. - - Returns: - Tuple[float, Dict[str, Dict[str, float]]]: A targe metric value(float) and - all metric(s)(dict) of evaluation, used to judge the quality of the model. - """ - # set eval func - self.eval_func = ppsci.solver.eval.eval_func - - result = self.eval_func(self, epoch_id, self.log_freq) - metric_msg = ", ".join( - [self.eval_output_info[key].avg_info for key in self.eval_output_info] - ) - - if isinstance(epoch_id, int): - logger.info(f"[Eval][Epoch {epoch_id}][Avg] {metric_msg}") - else: - logger.info(f"[Eval][Avg] {metric_msg}") - self.eval_output_info.clear() - - return result - - @misc.run_on_eval_mode - def visualize(self, epoch_id: Optional[int] = None): - """Visualization. - - Args: - epoch_id (Optional[int]): Epoch id. Defaults to None. - """ - # set visualize func - self.visu_func = ppsci.solver.visu.visualize_func - - self.visu_func(self, epoch_id) - if isinstance(epoch_id, int): - logger.info(f"[Visualize][Epoch {epoch_id}] Finish visualization") - else: - logger.info("[Visualize] Finish visualization") - - @misc.run_on_eval_mode - def predict( - self, - input_dict: Dict[str, Union[np.ndarray, paddle.Tensor]], - expr_dict: Optional[Dict[str, Callable]] = None, - batch_size: Optional[int] = 64, - no_grad: bool = True, - return_numpy: bool = False, - ) -> Dict[str, Union[paddle.Tensor, np.ndarray]]: - """Pure prediction using model.forward(...) and expression(optional, if given). - - Args: - input_dict (Dict[str, Union[np.ndarray, paddle.Tensor]]): Input data in dict. - expr_dict (Optional[Dict[str, Callable]]): Expression dict, which guide to - compute equation variable with callable function. Defaults to None. - batch_size (Optional[int]): Predicting by batch size. If None, data in - `input_dict` will be used directly for inference without any batch slicing. - Defaults to 64. - no_grad (bool): Whether set stop_gradient=True for entire prediction, mainly - for memory-efficiency. Defaults to True. - return_numpy (bool): Whether convert result from Tensor to numpy ndarray. - Defaults to False. - - Returns: - Dict[str, Union[paddle.Tensor, np.ndarray]]: Prediction in dict. - - Examples: - >>> import paddle - >>> import ppsci - >>> model = ppsci.arch.MLP(('x', 'y'), ('u', 'v'), num_layers=None, hidden_size=[32, 8]) - >>> solver = ppsci.solver.Solver(model) # doctest: +SKIP - >>> input_dict = {'x': paddle.rand([32, 1]), - ... 'y': paddle.rand([32, 1])} - >>> pred = solver.predict(input_dict) # doctest: +SKIP - >>> for k, v in pred.items(): # doctest: +SKIP - ... print(k, v.shape) # doctest: +SKIP - u [32, 1] - v [32, 1] - """ - num_samples = len(next(iter(input_dict.values()))) - num_pad = (self.world_size - num_samples % self.world_size) % self.world_size - # pad with last element if `num_samples` is not divisible by `world_size` - # ensuring every device get same number of data. - if num_pad > 0: - for k, v in input_dict.items(): - repeat_times = (num_pad, *(1 for _ in range(v.ndim - 1))) - if isinstance(v, np.ndarray): - input_dict[k] = np.concatenate( - ( - v, - np.tile(v[num_samples - 1 : num_samples], repeat_times), - ), - ) - elif isinstance(v, paddle.Tensor): - input_dict[k] = paddle.concat( - ( - v, - paddle.tile(v[num_samples - 1 : num_samples], repeat_times), - ), - ) - else: - raise ValueError(f"Unsupported data type {type(v)}.") - - num_samples_pad = num_samples + num_pad - local_num_samples_pad = num_samples_pad // self.world_size - local_input_dict = ( - {k: v[self.rank :: self.world_size] for k, v in input_dict.items()} - if self.world_size > 1 - else input_dict - ) - local_batch_num = ( - (local_num_samples_pad + (batch_size - 1)) // batch_size - if batch_size is not None - else 1 - ) - - pred_dict = misc.Prettydefaultdict(list) - with self.no_grad_context_manager(no_grad), self.no_sync_context_manager( - self.world_size > 1, self.model - ): - for batch_id in range(local_batch_num): - # prepare local batch input - if batch_size is not None: - st = batch_id * batch_size - ed = min(local_num_samples_pad, (batch_id + 1) * batch_size) - batch_input_dict = { - k: v[st:ed] for k, v in local_input_dict.items() - } - else: - batch_input_dict = {**local_input_dict} - # Keep dtype unchanged as all dtype be correct when given into predict function - for key in batch_input_dict: - if not paddle.is_tensor(batch_input_dict[key]): - batch_input_dict[key] = paddle.to_tensor( - batch_input_dict[key], stop_gradient=no_grad - ) - - # forward - with self.autocast_context_manager(self.use_amp, self.amp_level): - batch_output_dict = self.forward_helper.visu_forward( - expr_dict, batch_input_dict, self.model - ) - - # collect local batch output - for key, batch_output in batch_output_dict.items(): - pred_dict[key].append( - batch_output.detach() if no_grad else batch_output - ) - - # concatenate local output - pred_dict = {key: paddle.concat(value) for key, value in pred_dict.items()} - - if self.world_size > 1: - # gather global output from all devices if world_size > 1 - pred_dict = { - key: misc.all_gather(value) for key, value in pred_dict.items() - } - # rearrange output as the same order of input_dict according - # to inverse permutation - perm = np.arange(num_samples_pad, dtype="int64") - perm = np.concatenate( - [perm[rank :: self.world_size] for rank in range(self.world_size)], - axis=0, - ) - perm_inv = np.empty_like(perm) - perm_inv[perm] = np.arange(num_samples_pad, dtype="int64") - perm_inv = paddle.to_tensor(perm_inv) - pred_dict = {key: value[perm_inv] for key, value in pred_dict.items()} - # then discard output of padding data at the end if num_pad > 0 - if num_pad > 0: - pred_dict = { - key: value[:num_samples] for key, value in pred_dict.items() - } - # NOTE: Discard padding data in input_dict for consistency - for k in input_dict: - input_dict[k] = input_dict[k][:num_samples] - - # convert to numpy ndarray if specified - if return_numpy: - pred_dict = { - k: (v.numpy() if paddle.is_tensor(v) else v) - for k, v in pred_dict.items() - } - - return pred_dict - - @misc.run_on_eval_mode - def export( - self, - input_spec: List[Dict[str, InputSpec]], - export_path: str, - with_onnx: bool = False, - skip_prune_program: bool = False, - *, - full_graph: bool = True, - ignore_modules: Optional[List[ModuleType]] = None, - ): - """ - Convert model to static graph model and export to files. - - Args: - input_spec (List[Dict[str, InputSpec]]): InputSpec describes the signature - information of the model input. - export_path (str): The path prefix to save model. - with_onnx (bool, optional): Whether to export model into onnx after - paddle inference models are exported. Defaults to False. - skip_prune_program (bool, optional): Whether prune program, pruning program - may cause unexpectable result, e.g. llm-inference. Defaults to False. - full_graph (bool, optional): Symbolic OpCode Translator(SOT) will be used - when set to True, where otherwise use Abstract Syntax Tree(AST) if False. - Defaults to True. - ignore_modules (List[ModuleType]): Adds modules that should be ignored during - conversion. Builtin modules that have been ignored are collections, pdb, - copy, inspect, re, numpy, logging, six. For example, einops can be added - here. Defaults to None. - """ - if ignore_modules is not None: - jit.ignore_module(ignore_modules) - - jit.enable_to_static(True) - - if self.pretrained_model_path is None: - logger.warning( - "'INFER.pretrained_model_path' is not given, so the weights of exported " - "model will be random initialized." - ) - - # convert model to static graph model - static_model = jit.to_static( - self.model, - input_spec=input_spec, - full_graph=full_graph, - ) - - # save static graph model to disk - if len(osp.dirname(export_path)): - os.makedirs(osp.dirname(export_path), exist_ok=True) - try: - jit.save(static_model, export_path, skip_prune_program=skip_prune_program) - except Exception as e: - raise e - logger.message( - f"Inference model has been exported to: {export_path}, including " - + ( - "*.json, *.pdiparams files." - if paddle.framework.use_pir_api() - else "*.pdmodel, *.pdiparams and *.pdiparams.info files." - ) - ) - jit.enable_to_static(False) - - if with_onnx: - # TODO: support pir + onnx - if not importlib.util.find_spec("paddle2onnx"): - raise ModuleNotFoundError( - "Please install paddle2onnx with `pip install paddle2onnx`" - " before exporting onnx model." - ) - import paddle2onnx - - DEFAULT_OPSET_VERSION = 19 - - paddle2onnx.export( - model_filename=export_path + ".json" - if paddle.framework.use_pir_api() - else ".pdmodel", - params_filename=export_path + ".pdiparams", - save_file=export_path + ".onnx", - opset_version=DEFAULT_OPSET_VERSION, - enable_onnx_checker=True, - ) - logger.message(f"ONNX model has been exported to: {export_path}.onnx") - - def autocast_context_manager( - self, enable: bool, level: Literal["O0", "O1", "O2", "OD"] = "O1" - ) -> contextlib.AbstractContextManager: - """Smart autocast context manager for Auto Mix Precision. - - Args: - enable (bool): Enable autocast. - level (Literal["O0", "O1", "O2", "OD"]): Autocast level. - - Returns: - contextlib.AbstractContextManager: Smart autocast context manager. - """ - if enable: - ctx_manager = amp.auto_cast(level=level) - else: - ctx_manager = ( - contextlib.nullcontext() - if sys.version_info >= (3, 7) - else contextlib.suppress() - ) - return ctx_manager - - @functools.lru_cache() - def no_grad_context_manager( - self, enable: bool - ) -> contextlib.AbstractContextManager: - """Smart no_grad context manager. - - Args: - enable (bool): Enable no_grad. - - Returns: - contextlib.AbstractContextManager: Smart no_grad context manager. - """ - if enable: - ctx_manager = paddle.no_grad() - else: - ctx_manager = ( - contextlib.nullcontext() - if sys.version_info >= (3, 7) - else contextlib.suppress() - ) - return ctx_manager - - def no_sync_context_manager( - self, - enable: bool, - ddp_model: paddle.DataParallel, - ) -> contextlib.AbstractContextManager: - """Smart no_sync context manager for given model. - NOTE: Only `paddle.DataParallel` object has `no_sync` interface. - - Args: - enable (bool): Enable no_sync. - - Returns: - contextlib.AbstractContextManager: Smart no_sync context manager. - """ - if enable: - if isinstance(self.model, ppsci.arch.ModelList): - for model in self.model.model_list: - if not isinstance(model, paddle.DataParallel): - raise TypeError( - "no_sync interface is only for model with type " - "paddle.DataParallel, but got type " - f"{misc.typename(model)}" - ) - ctx_manager = contextlib.ExitStack() - for model in self.model.model_list: - ctx_manager.enter_context(model.no_sync()) - else: - if not isinstance(self.model, paddle.DataParallel): - raise TypeError( - "no_sync interface is only for model with type " - f"paddle.DataParallel, but got type {misc.typename(ddp_model)}" - ) - ctx_manager = ddp_model.no_sync() - else: - ctx_manager = ( - contextlib.nullcontext() - if sys.version_info >= (3, 7) - else contextlib.suppress() - ) - return ctx_manager - - def plot_loss_history( - self, - by_epoch: bool = False, - smooth_step: int = 1, - use_semilogy: bool = True, - ) -> None: - """Plotting iteration/epoch-loss curve. - - Args: - by_epoch (bool, optional): Whether the abscissa axis of the curve is epoch or iteration. Defaults to False. - smooth_step (int, optional): How many steps of loss are squeezed to one point to smooth the curve. Defaults to 1. - use_semilogy (bool, optional): Whether to set non-uniform coordinates for the y-axis. Defaults to True. - """ - loss_dict = {} - for key in self.train_loss_info: - loss_arr = np.asarray(self.train_loss_info[key].history) - if by_epoch: - loss_arr = np.mean( - np.reshape(loss_arr, (-1, self.iters_per_epoch)), - axis=1, - ) - loss_dict[key] = list(loss_arr) - - misc.plot_curve( - data=loss_dict, - xlabel="Epoch" if by_epoch else "Iteration", - ylabel="Loss", - output_dir=self.output_dir, - smooth_step=smooth_step, - use_semilogy=use_semilogy, - ) - - def _parse_params_from_cfg(self, cfg: DictConfig): - """ - Parse hyper-parameters from DictConfig. - """ - self.output_dir = cfg.output_dir - self.log_freq = cfg.log_freq - self.use_tbd = cfg.use_tbd - self.use_vdl = cfg.use_vdl - self.wandb_config = cfg.wandb_config - self.use_wandb = cfg.use_wandb - self.device = cfg.device - self.to_static = cfg.to_static - - self.use_amp = cfg.use_amp - self.amp_level = cfg.amp_level - - self.epochs = cfg.TRAIN.epochs - self.iters_per_epoch = cfg.TRAIN.iters_per_epoch - self.update_freq = cfg.TRAIN.update_freq - self.save_freq = cfg.TRAIN.save_freq - self.eval_during_train = cfg.TRAIN.eval_during_train - self.start_eval_epoch = cfg.TRAIN.start_eval_epoch - self.eval_freq = cfg.TRAIN.eval_freq - self.checkpoint_path = cfg.TRAIN.checkpoint_path - - if "ema" in cfg.TRAIN and cfg.TRAIN.ema.get("use_ema", False): - self.avg_freq = cfg.TRAIN.ema.avg_freq - elif "swa" in cfg.TRAIN and cfg.TRAIN.swa.get("use_swa", False): - self.avg_freq = cfg.TRAIN.swa.avg_freq - - self.compute_metric_by_batch = cfg.EVAL.compute_metric_by_batch - self.eval_with_no_grad = cfg.EVAL.eval_with_no_grad - - if cfg.mode == "train": - self.pretrained_model_path = cfg.TRAIN.pretrained_model_path - elif cfg.mode == "eval": - self.pretrained_model_path = cfg.EVAL.pretrained_model_path - elif cfg.mode in ["export", "infer"]: - self.pretrained_model_path = cfg.INFER.pretrained_model_path - - def register_callback_on_epoch_begin( - self: Solver, callback_fn: Callable[[Solver]] - ) -> None: - """ - Registers a callback function to be executed at the beginning of each training epoch. - - Args: - callback_fn : Callable[[Solver]] - A function that takes a Solver instance as an argument. This function - will be called at the start of every epoch. - """ - self.callbacks_on_epoch_begin.append(callback_fn) - - def register_callback_on_epoch_end( - self: Solver, callback_fn: Callable[[Solver]] - ) -> None: - """ - Registers a callback function to be executed at the end of each training epoch. - - Args: - callback_fn : Callable[[Solver]] - A function that takes a Solver instance as an argument. This function - will be called at the end of every epoch. - """ - self.callbacks_on_epoch_end.append(callback_fn) - - def register_callback_on_iter_begin( - self: Solver, callback_fn: Callable[[Solver]] - ) -> None: - """ - Registers a callback function to be executed at the beginning of each training iteration. - - Args: - callback_fn : Callable[[Solver]] - A function that takes a Solver instance as an argument. This function - will be called at the start of every iteration. - """ - self.callbacks_on_iter_begin.append(callback_fn) - - def register_callback_on_iter_end( - self: Solver, callback_fn: Callable[[Solver]] - ) -> None: - """ - Registers a callback function to be executed at the end of each training iteration. - - Args: - callback_fn : Callable[[Solver]] - A function that takes a Solver instance as an argument. This function - will be called at the end of every iteration. - - Returns: - ------- - None - """ - self.callbacks_on_iter_end.append(callback_fn) - - def _invoke_callbacks_on_epoch_begin(self: Solver) -> None: - """ - Invokes all registered callbacks at the beginning of an epoch. - """ - for callback in self.callbacks_on_epoch_begin: - callback(self) - - def _invoke_callbacks_on_epoch_end(self: Solver) -> None: - """ - Invokes all registered callbacks at the end of an epoch. - """ - for callback in self.callbacks_on_epoch_end: - callback(self) - - def _invoke_callbacks_on_iter_begin(self: Solver) -> None: - """ - Invokes all registered callbacks at the beginning of an iteration. - """ - for callback in self.callbacks_on_iter_begin: - callback(self) - - def _invoke_callbacks_on_iter_end(self: Solver) -> None: - """ - Invokes all registered callbacks at the end of an iteration. - """ - for callback in self.callbacks_on_iter_end: - callback(self) diff --git a/examples/smc_reac/ppsci/solver/train.py b/examples/smc_reac/ppsci/solver/train.py deleted file mode 100644 index aec95a0dc4..0000000000 --- a/examples/smc_reac/ppsci/solver/train.py +++ /dev/null @@ -1,324 +0,0 @@ -# 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. - -from __future__ import annotations - -import sys -import time -from typing import TYPE_CHECKING -from typing import Dict -from typing import Sequence -from typing import Union - -import paddle -from paddle.distributed.fleet.utils import hybrid_parallel_util as hpu -from paddle.framework import core - -from ppsci.solver import printer -from ppsci.utils import misc - -if TYPE_CHECKING: - from ppsci import solver - - -def _compute_batch_size( - input_dict: Dict[str, Union[paddle.Tensor, Sequence[paddle.Tensor]]] -) -> int: - """Compute batch size from given input dict. - - NOTE: Returned `batch_size` might be inaccurate, but it won't affect the correctness - of the training results because `batch_size` is now only used for timing. - - Args: - input_dict (Dict[str, Union[paddle.Tensor, Sequence[paddle.Tensor]]]): Given input dict. - - Returns: - int: Batch size of input dict. - """ - sample = next(iter(input_dict.values())) - if hasattr(sample, "shape"): - return sample.shape[0] - elif hasattr(sample, "__len__"): # Might be inaccurate here. - return len(sample) - else: - raise ValueError("Unsupported type of input dict value.") - - -def train_epoch_func(solver: "solver.Solver", epoch_id: int, log_freq: int): - """Train program for one epoch. - - Args: - solver (solver.Solver): Main solver. - epoch_id (int): Epoch id. - log_freq (int): Log training information every `log_freq` steps. - """ - batch_tic = time.perf_counter() - - for iter_id in range(1, solver.iters_per_epoch + 1): - solver._invoke_callbacks_on_iter_begin() - if solver.nvtx_flag: # only for nsight analysis - core.nvprof_nvtx_push( - f"Training iteration {solver.global_step + 1}" - ) # Training iteration - - total_batch_size = 0 - reader_cost = 0.0 - batch_cost = 0.0 - reader_tic = time.perf_counter() - - input_dicts = [] - label_dicts = [] - weight_dicts = [] - for _, _constraint in solver.constraint.items(): - # fetch data from data loader - if solver.nvtx_flag: # only for nsight analysis - core.nvprof_nvtx_push("Data load") - - try: - input_dict, label_dict, weight_dict = next(_constraint.data_iter) - except StopIteration: - _constraint.data_iter = iter(_constraint.data_loader) - input_dict, label_dict, weight_dict = next(_constraint.data_iter) - - if solver.nvtx_flag: # only for nsight analysis - core.nvprof_nvtx_pop() - - reader_cost += time.perf_counter() - reader_tic - - for v in input_dict.values(): - if hasattr(v, "stop_gradient"): - v.stop_gradient = False - - # gather each constraint's input, label, weight to a list - input_dicts.append(input_dict) - label_dicts.append(label_dict) - weight_dicts.append(weight_dict) - total_batch_size += _compute_batch_size(input_dict) - reader_tic = time.perf_counter() - - loss_dict = misc.Prettydefaultdict(float) - loss_dict["loss"] = 0.0 - # forward for every constraint, including model and equation expression - with solver.no_sync_context_manager(solver.world_size > 1, solver.model): - with solver.autocast_context_manager(solver.use_amp, solver.amp_level): - if solver.nvtx_flag: # only for nsight analysis - core.nvprof_nvtx_push("Loss computation") - - losses_all, losses_constraint = solver.forward_helper.train_forward( - tuple( - _constraint.output_expr - for _constraint in solver.constraint.values() - ), - input_dicts, - solver.model, - solver.constraint, - label_dicts, - weight_dicts, - ) - assert "loss" not in losses_all, ( - "Key 'loss' is not allowed in loss_dict for it is an preserved key" - " representing total loss, please use other name instead." - ) - - if solver.nvtx_flag: # only for nsight analysis - core.nvprof_nvtx_pop() # Loss computation - - # accumulate all losses - if solver.nvtx_flag: # only for nsight analysis - core.nvprof_nvtx_push("Loss aggregator") - - total_loss = solver.loss_aggregator(losses_all, solver.global_step) - if solver.update_freq > 1: - total_loss = total_loss / solver.update_freq - - loss_dict.update(losses_constraint) - loss_dict["loss"] = float(total_loss) - - if solver.nvtx_flag: # only for nsight analysis - core.nvprof_nvtx_pop() # Loss aggregator - - # backward - if solver.nvtx_flag: # only for nsight analysis - core.nvprof_nvtx_push("Loss backward") - - if solver.use_amp: - total_loss_scaled = solver.scaler.scale(total_loss) - total_loss_scaled.backward() - else: - total_loss.backward() - - if solver.nvtx_flag: # only for nsight analysis - core.nvprof_nvtx_pop() # Loss backward - - # update parameters - if iter_id % solver.update_freq == 0 or iter_id == solver.iters_per_epoch: - if solver.nvtx_flag: # only for nsight analysis - core.nvprof_nvtx_push("Optimizer update") - - if solver.world_size > 1: - # fuse + allreduce manually before optimization if use DDP + no_sync - # details in https://github.com/PaddlePaddle/Paddle/issues/48898#issuecomment-1343838622 - hpu.fused_allreduce_gradients(list(solver.model.parameters()), None) - if solver.use_amp: - solver.scaler.minimize(solver.optimizer, total_loss_scaled) - else: - solver.optimizer.step() - - if solver.nvtx_flag: # only for nsight analysis - core.nvprof_nvtx_pop() # Optimizer update - - solver.optimizer.clear_grad() - - # update learning rate by step - if solver.lr_scheduler is not None and not solver.lr_scheduler.by_epoch: - solver.lr_scheduler.step() - - if solver.benchmark_flag: - paddle.device.synchronize() - batch_cost += time.perf_counter() - batch_tic - - # update and log training information - solver.global_step += 1 - solver.train_time_info["reader_cost"].update(reader_cost) - solver.train_time_info["batch_cost"].update(batch_cost) - printer.update_train_loss(solver, loss_dict, total_batch_size) - if ( - solver.global_step % log_freq == 0 - or solver.global_step == 1 - or solver.global_step == solver.max_steps - ): - printer.log_train_info(solver, total_batch_size, epoch_id, iter_id) - - batch_tic = time.perf_counter() - - if solver.nvtx_flag: # only for nsight analysis - core.nvprof_nvtx_pop() # Training iteration - NVTX_STOP_ITER = 25 - if solver.global_step >= NVTX_STOP_ITER: - print( - f"Only run {NVTX_STOP_ITER} steps when 'NVTX' is set in environment" - " for nsight analysis. Exit now ......\n" - ) - core.nvprof_stop() - sys.exit(0) - - solver._invoke_callbacks_on_iter_end() - - -def train_LBFGS_epoch_func(solver: "solver.Solver", epoch_id: int, log_freq: int): - """Train function for one epoch with L-BFGS optimizer. - - NOTE: L-BFGS training program do not support AMP now. - - Args: - solver (solver.Solver): Main solver. - epoch_id (int): Epoch id. - log_freq (int): Log training information every `log_freq` steps. - """ - batch_tic = time.perf_counter() - - for iter_id in range(1, solver.iters_per_epoch + 1): - solver._invoke_callbacks_on_iter_begin() - loss_dict = misc.Prettydefaultdict(float) - loss_dict["loss"] = 0.0 - total_batch_size = 0 - reader_cost = 0.0 - batch_cost = 0.0 - reader_tic = time.perf_counter() - - input_dicts = [] - label_dicts = [] - weight_dicts = [] - for _, _constraint in solver.constraint.items(): - # fetch data from data loader - try: - input_dict, label_dict, weight_dict = next(_constraint.data_iter) - except StopIteration: - _constraint.data_iter = iter(_constraint.data_loader) - input_dict, label_dict, weight_dict = next(_constraint.data_iter) - reader_cost += time.perf_counter() - reader_tic - - for v in input_dict.values(): - if hasattr(v, "stop_gradient"): - v.stop_gradient = False - - # gather each constraint's input, label, weight to a list - input_dicts.append(input_dict) - label_dicts.append(label_dict) - weight_dicts.append(weight_dict) - total_batch_size += _compute_batch_size(input_dict) - reader_tic = time.perf_counter() - - def closure() -> paddle.Tensor: - """Forward-backward closure function for LBFGS optimizer. - - Returns: - paddle.Tensor: Computed loss scalar. - """ - with solver.no_sync_context_manager(solver.world_size > 1, solver.model): - with solver.autocast_context_manager(solver.use_amp, solver.amp_level): - # forward for every constraint, including model and equation expression - losses_all, losses_constraint = solver.forward_helper.train_forward( - tuple( - _constraint.output_expr - for _constraint in solver.constraint.values() - ), - input_dicts, - solver.model, - solver.constraint, - label_dicts, - weight_dicts, - ) - - # accumulate all losses - total_loss = solver.loss_aggregator(losses_all, solver.global_step) - loss_dict.update(losses_constraint) - loss_dict["loss"] = float(total_loss) - - # backward - solver.optimizer.clear_grad() - total_loss.backward() - - if solver.world_size > 1: - # fuse + allreduce manually before optimization if use DDP model - # details in https://github.com/PaddlePaddle/Paddle/issues/48898#issuecomment-1343838622 - hpu.fused_allreduce_gradients(list(solver.model.parameters()), None) - - return total_loss - - # update parameters - solver.optimizer.step(closure) - - # update learning rate by step - if solver.lr_scheduler is not None and not solver.lr_scheduler.by_epoch: - solver.lr_scheduler.step() - - if solver.benchmark_flag: - paddle.device.synchronize() - batch_cost += time.perf_counter() - batch_tic - - # update and log training information - solver.global_step += 1 - solver.train_time_info["reader_cost"].update(reader_cost) - solver.train_time_info["batch_cost"].update(batch_cost) - printer.update_train_loss(solver, loss_dict, total_batch_size) - if ( - solver.global_step % log_freq == 0 - or solver.global_step == 1 - or solver.global_step == solver.max_steps - ): - printer.log_train_info(solver, total_batch_size, epoch_id, iter_id) - - batch_tic = time.perf_counter() - solver._invoke_callbacks_on_iter_end() diff --git a/examples/smc_reac/ppsci/solver/visu.py b/examples/smc_reac/ppsci/solver/visu.py deleted file mode 100644 index 80e11abe75..0000000000 --- a/examples/smc_reac/ppsci/solver/visu.py +++ /dev/null @@ -1,98 +0,0 @@ -# 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. - -from __future__ import annotations - -import os -import os.path as osp -from typing import TYPE_CHECKING -from typing import Optional - -import paddle - -from ppsci.utils import misc - -if TYPE_CHECKING: - from ppsci import solver - - -def visualize_func(solver: "solver.Solver", epoch_id: Optional[int]): - """Visualization program. - - Args: - solver (solver.Solver): Main Solver. - epoch_id (Optional[int]): Epoch id. - """ - for _, _visualizer in solver.visualizer.items(): - all_input = misc.Prettydefaultdict(list) - all_output = misc.Prettydefaultdict(list) - - # NOTE: 'visualize_func' now do not apply data sharding(different from 'Solver.predict'), - # where every rank receive same input data and compute same output data - # (which will cause computational redundancy), - # but only the 0-rank(master) device save the visualization result into disk. - # TODO(HydrogenSulfate): This will be optimized in the future. - - input_dict = _visualizer.input_dict - batch_size = _visualizer.batch_size - num_samples = len(next(iter(input_dict.values()))) - batch_num = (num_samples + (batch_size - 1)) // batch_size - - for batch_id in range(batch_num): - batch_input_dict = {} - st = batch_id * batch_size - ed = min(num_samples, (batch_id + 1) * batch_size) - - # prepare batch input dict - for key in input_dict: - if not paddle.is_tensor(input_dict[key]): - batch_input_dict[key] = paddle.to_tensor( - input_dict[key][st:ed], paddle.get_default_dtype() - ) - else: - batch_input_dict[key] = input_dict[key][st:ed] - batch_input_dict[key].stop_gradient = False - - # forward - with solver.autocast_context_manager( - solver.use_amp, solver.amp_level - ), solver.no_grad_context_manager(solver.eval_with_no_grad): - batch_output_dict = solver.forward_helper.visu_forward( - _visualizer.output_expr, batch_input_dict, solver.model - ) - - # collect batch data with dtype fixed to float32 regardless of the dtypes of - # paddle runtime, which is most compatible with almost visualization tools. - for key, batch_input in batch_input_dict.items(): - all_input[key].append(batch_input.detach().astype("float32")) - for key, batch_output in batch_output_dict.items(): - all_output[key].append(batch_output.detach().astype("float32")) - - # concatenate all data - for key in all_input: - all_input[key] = paddle.concat(all_input[key]) - for key in all_output: - all_output[key] = paddle.concat(all_output[key]) - - # save visualization - with misc.RankZeroOnly(solver.rank) as is_master: - if is_master: - visual_dir = osp.join(solver.output_dir, "visual") - if epoch_id: - visual_dir = osp.join(visual_dir, f"epoch_{epoch_id}") - os.makedirs(visual_dir, exist_ok=True) - _visualizer.save( - osp.join(visual_dir, _visualizer.prefix), - {**all_input, **all_output}, - ) diff --git a/examples/smc_reac/ppsci/utils/__init__.py b/examples/smc_reac/ppsci/utils/__init__.py deleted file mode 100644 index 3382eee856..0000000000 --- a/examples/smc_reac/ppsci/utils/__init__.py +++ /dev/null @@ -1,66 +0,0 @@ -# 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. - -# NOTE: Put config module import at the top level for register default config(s) in -# ConfigStore at the beginning of ppsci -from ppsci.utils import config # isort:skip # noqa: F401 -from ppsci.utils import ema -from ppsci.utils import initializer -from ppsci.utils import logger -from ppsci.utils import misc -from ppsci.utils import reader -from ppsci.utils import writer -from ppsci.utils.checker import dynamic_import_to_globals -from ppsci.utils.checker import run_check -from ppsci.utils.checker import run_check_mesh -from ppsci.utils.expression import ExpressionSolver -from ppsci.utils.misc import AverageMeter -from ppsci.utils.misc import set_random_seed -from ppsci.utils.reader import load_csv_file -from ppsci.utils.reader import load_mat_file -from ppsci.utils.reader import load_npz_file -from ppsci.utils.reader import load_vtk_file -from ppsci.utils.reader import load_vtk_with_time_file -from ppsci.utils.save_load import load_checkpoint -from ppsci.utils.save_load import load_pretrain -from ppsci.utils.save_load import save_checkpoint -from ppsci.utils.symbolic import lambdify -from ppsci.utils.writer import save_csv_file -from ppsci.utils.writer import save_tecplot_file - -__all__ = [ - "AverageMeter", - "ExpressionSolver", - "initializer", - "logger", - "misc", - "ema", - "reader", - "writer", - "load_csv_file", - "load_mat_file", - "load_npz_file", - "load_vtk_file", - "load_vtk_with_time_file", - "save_csv_file", - "save_tecplot_file", - "dynamic_import_to_globals", - "run_check", - "run_check_mesh", - "set_random_seed", - "load_checkpoint", - "load_pretrain", - "save_checkpoint", - "lambdify", -] diff --git a/examples/smc_reac/ppsci/utils/callbacks.py b/examples/smc_reac/ppsci/utils/callbacks.py deleted file mode 100644 index a753432b55..0000000000 --- a/examples/smc_reac/ppsci/utils/callbacks.py +++ /dev/null @@ -1,136 +0,0 @@ -# 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. - -import importlib.util -import inspect -import sys -import traceback -from os import path as osp -from typing import Any - -from hydra.core.hydra_config import HydraConfig -from hydra.experimental.callback import Callback -from omegaconf import DictConfig - -from ppsci.utils import config as config_module -from ppsci.utils import logger -from ppsci.utils import misc - -RUNTIME_EXIT_CODE = 1 # for other errors -VALIDATION_ERROR_EXIT_CODE = 2 # for invalid argument detected in config file - - -class InitCallback(Callback): - """Callback class for: - 1. Parse config dict from given yaml file and check its validity. - 2. Fixing random seed to 'config.seed'. - 3. Initialize logger while creating output directory(if not exist). - 4. Enable prim mode if specified. - - NOTE: This callback is mainly for reducing unnecessary duplicate code in each - examples code when runing with hydra. - - This callback should be added to hydra config file as follows: - - ``` yaml hl_lines="7-11" - # content of example.yaml below - hydra: - run: - ... - job: - ... - callbacks: - init_callback: - _target_: ppsci.utils.callbacks.InitCallback # <-- add callback at here - xxx_callback: - _target_: ppsci.utils.callbacks.XxxCallback # <-- add more callback here - sweep: - ... - ... - ... - ``` - """ - - def on_job_start(self, config: DictConfig, **kwargs: Any) -> None: - if importlib.util.find_spec("pydantic") is not None: - from pydantic import ValidationError - else: - logger.error( - f"ModuleNotFoundError at {__file__}:{inspect.currentframe().f_lineno}\n" - "Please install pydantic with `pip install pydantic` when set callbacks" - " in your config yaml." - ) - sys.exit(RUNTIME_EXIT_CODE) - - # check given cfg using pre-defined pydantic schema in 'SolverConfig', - # error(s) will be printed and exit program if any checking failed at this step - try: - _model_pydantic = config_module.SolverConfig(**dict(config)) - full_cfg = DictConfig(_model_pydantic.model_dump()) - except ValidationError as e: - print(e) - sys.exit(VALIDATION_ERROR_EXIT_CODE) - except Exception as e: - print(e) - sys.exit(RUNTIME_EXIT_CODE) - - # fix random seed for reproducibility - misc.set_random_seed(full_cfg.seed) - - # initialize logger while creating output directory - logger.init_logger( - "ppsci", - osp.join(full_cfg.output_dir, f"{full_cfg.mode}.log") - if full_cfg.output_dir and full_cfg.mode not in ["export", "infer"] - else None, - full_cfg.log_level, - ) - - # set device before running into example function - if "device" in full_cfg: - import paddle - - if isinstance(full_cfg.device, str): - paddle.device.set_device(full_cfg.device) - - try: - if "num" in HydraConfig.get().job: - jobs_id = HydraConfig.get().job.num - else: - jobs_id = None - if "n_jobs" in HydraConfig.get().launcher: - parallel_jobs_num = HydraConfig.get().launcher.n_jobs - else: - parallel_jobs_num = None - - if jobs_id and parallel_jobs_num: - job_device_id = jobs_id % parallel_jobs_num - device_type = paddle.get_device().split(":")[0] - logger.message( - f"Running job {jobs_id} on device {device_type}:{job_device_id}(logical device id)" - ) - paddle.set_device(f"{device_type}:{job_device_id}") - except Exception as e: - print(e) - traceback.print_exc() - sys.exit(RUNTIME_EXIT_CODE) - - # enable prim if specified - if "prim" in full_cfg and bool(full_cfg.prim): - # Mostly for compiler running with dy2st. - from paddle.framework import core - - core.set_prim_eager_enabled(True) - core._set_prim_all_enabled(True) - logger.message("Prim mode is enabled.") diff --git a/examples/smc_reac/ppsci/utils/checker.py b/examples/smc_reac/ppsci/utils/checker.py deleted file mode 100644 index 8991a89e2c..0000000000 --- a/examples/smc_reac/ppsci/utils/checker.py +++ /dev/null @@ -1,287 +0,0 @@ -# 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. - -from __future__ import annotations - -import importlib.util -import traceback -from typing import Dict -from typing import Sequence -from typing import Union - -import paddle - -from ppsci.utils import logger - -__all__ = [ - "run_check", - "run_check_mesh", - "dynamic_import_to_globals", -] - - -def run_check() -> None: - """Check whether PaddleScience is installed correctly and running successfully on - your system. - - Examples: - >>> import ppsci - >>> ppsci.utils.run_check() # doctest: +SKIP - """ - # test demo code below. - import ppsci - - try: - ppsci.utils.set_random_seed(42) - ppsci.utils.logger.init_logger() - model = ppsci.arch.MLP(("x", "y"), ("u", "v", "p"), 3, 16, "tanh") - - equation = {"NavierStokes": ppsci.equation.NavierStokes(0.01, 1.0, 2, False)} - - geom = {"rect": ppsci.geometry.Rectangle((-0.05, -0.05), (0.05, 0.05))} - - ITERS_PER_EPOCH = 5 - train_dataloader_cfg = { - "dataset": "IterableNamedArrayDataset", - "iters_per_epoch": ITERS_PER_EPOCH, - } - - NPOINT_PDE = 8**2 - pde_constraint = 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={ - "continuity": 0.0001, - "momentum_x": 0.0001, - "momentum_y": 0.0001, - }, - name="EQ", - ) - constraint = {pde_constraint.name: pde_constraint} - - residual_validator = ppsci.validate.GeometryValidator( - equation["NavierStokes"].equations, - {"continuity": 0, "momentum_x": 0, "momentum_y": 0}, - geom["rect"], - { - "dataset": "NamedArrayDataset", - "total_size": 8**2, - "batch_size": 32, - "sampler": {"name": "BatchSampler"}, - }, - ppsci.loss.MSELoss("sum"), - evenly=True, - metric={"MSE": ppsci.metric.MSE(False)}, - name="Residual", - ) - validator = {residual_validator.name: residual_validator} - - EPOCHS = 2 - optimizer = ppsci.optimizer.Adam(0.001)(model) - solver = ppsci.solver.Solver( - model, - constraint, - None, - optimizer, - None, - EPOCHS, - ITERS_PER_EPOCH, - device=paddle.device.get_device(), - equation=equation, - validator=validator, - ) - solver.train() - solver.eval(EPOCHS) - except Exception as e: - traceback.print_exc() - logger.error( - f"PaddleScience meets some problem with \n {repr(e)} \nplease check whether " - "Paddle's version and PaddleScience's version are both correct." - ) - else: - logger.message("PaddleScience is installed successfully.✨ 🍰 ✨") - - -def run_check_mesh() -> None: - """Check whether geometry packages is installed correctly and `ppsci.geometry.Mesh` - can running successfully on your system. - - Examples: - >>> import ppsci - >>> ppsci.utils.run_check_mesh() # doctest: +SKIP - """ - # test demo code below. - if importlib.util.find_spec("open3d") is None: - raise ModuleNotFoundError( - "Please install open3d first with: " "`pip install open3d`" - ) - if importlib.util.find_spec("pysdf") is None: - raise ModuleNotFoundError( - "Please install pysdf first with: `pip install pysdf`" - ) - if importlib.util.find_spec("pymesh") is None: - raise ModuleNotFoundError( - "Please install pymesh first as " - "https://paddlescience-docs.readthedocs.io/zh/latest/zh/install_setup/#__tabbed_4_4" - ) - - import numpy as np - import pymesh - - import ppsci - - try: - ppsci.utils.set_random_seed(42) - ppsci.utils.logger.init_logger() - model = ppsci.arch.MLP(("x", "y"), ("u", "v", "p"), 3, 16, "tanh") - - equation = {"NavierStokes": ppsci.equation.NavierStokes(0.01, 1.0, 2, False)} - - # create a 1x1x1 simple cube geometry - vertices = np.array( - [ - [0.0, 0.0, 0.0], - [1.0, 0.0, 0.0], - [0.0, 0.0, 1.0], - [1.0, 0.0, 1.0], - [0.0, 1.0, 0.0], - [1.0, 1.0, 0.0], - [0.0, 1.0, 1.0], - [1.0, 1.0, 1.0], - ] - ) # 8 vertices for mesh - faces = np.array( - [ - [4, 7, 5], - [4, 6, 7], - [0, 2, 4], - [2, 6, 4], - [0, 1, 2], - [1, 3, 2], - [1, 5, 7], - [1, 7, 3], - [2, 3, 7], - [2, 7, 6], - [0, 4, 1], - [1, 4, 5], - ] - ) # 12 triangle faces for mesh - box_mesh = pymesh.form_mesh(vertices, faces) - geom = {"rect": ppsci.geometry.Mesh(box_mesh)} - - ITERS_PER_EPOCH = 5 - train_dataloader_cfg = { - "dataset": "IterableNamedArrayDataset", - "iters_per_epoch": ITERS_PER_EPOCH, - } - - NPOINT_PDE = 8**2 - pde_constraint = 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"), - weight_dict={ - "continuity": "sdf", - "momentum_x": "sdf", - "momentum_y": "sdf", - }, - name="EQ", - ) - constraint = {pde_constraint.name: pde_constraint} - - residual_validator = ppsci.validate.GeometryValidator( - equation["NavierStokes"].equations, - {"continuity": 0, "momentum_x": 0, "momentum_y": 0}, - geom["rect"], - { - "dataset": "NamedArrayDataset", - "total_size": 8**2, - "batch_size": 32, - "sampler": {"name": "BatchSampler"}, - }, - ppsci.loss.MSELoss("sum"), - metric={"MSE": ppsci.metric.MSE(False)}, - name="Residual", - ) - validator = {residual_validator.name: residual_validator} - - EPOCHS = 2 - optimizer = ppsci.optimizer.Adam(0.001)(model) - solver = ppsci.solver.Solver( - model, - constraint, - None, - optimizer, - None, - EPOCHS, - ITERS_PER_EPOCH, - device=paddle.device.get_device(), - equation=equation, - validator=validator, - ) - solver.train() - solver.eval(EPOCHS) - except Exception as e: - traceback.print_exc() - logger.error( - f"PaddleScience meets some problem with \n {repr(e)} \nplease check whether " - "open3d, pysdf, pybind11, PyMesh are all installed correctly." - ) - else: - logger.message("ppsci.geometry.Mesh module running successfully.✨ 🍰 ✨") - - -def dynamic_import_to_globals( - names: Union[str, Sequence[str]], alias: Dict[str, str] = None -) -> bool: - """Import module and add it to globals() by given names dynamically. - - Args: - names (Union[str, Sequence[str]]): Module name or sequence of module names. - alias (Dict[str, str]): Alias name of module when imported into globals(). - - Returns: - bool: Whether given names all exist. - """ - if isinstance(names, str): - names = (names,) - - if alias is None: - alias = {} - - for name in names: - # find module in environment by it's name and alias(if given) - module_spec = importlib.util.find_spec(name) - if module_spec is None and name in alias: - module_spec = importlib.util.find_spec(alias[name]) - - # log error and return False if module do not exist - if not module_spec: - logger.error(f"Module {name} should be installed first.") - return False - - # module exist, add to globals() if not in globals() - add_name = name - if add_name in alias: - add_name = alias[add_name] - if add_name not in globals(): - globals()[add_name] = importlib.import_module(name) - - return True diff --git a/examples/smc_reac/ppsci/utils/config.py b/examples/smc_reac/ppsci/utils/config.py deleted file mode 100644 index d887a75d3d..0000000000 --- a/examples/smc_reac/ppsci/utils/config.py +++ /dev/null @@ -1,457 +0,0 @@ -# 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. - -from __future__ import annotations - -import importlib.util -from typing import Mapping -from typing import Optional -from typing import Tuple - -from typing_extensions import Literal - -__all__ = [] - -if importlib.util.find_spec("pydantic") is not None: - try: - from hydra.core.config_store import ConfigStore - from omegaconf import OmegaConf - from pydantic import BaseModel - from pydantic import field_validator - from pydantic import model_validator - from pydantic_core.core_schema import ValidationInfo - - __all__.append("SolverConfig") - - class EMAConfig(BaseModel): - use_ema: bool = False - decay: float = 0.9 - avg_freq: int = 1 - - @field_validator("decay") - def decay_check(cls, v): - if v <= 0 or v >= 1: - raise ValueError( - f"'ema.decay' should be in (0, 1) when is type of float, but got {v}" - ) - return v - - @field_validator("avg_freq") - def avg_freq_check(cls, v): - if v <= 0: - raise ValueError( - "'ema.avg_freq' should be a positive integer when is type of int, " - f"but got {v}" - ) - return v - - class SWAConfig(BaseModel): - use_swa: bool = False - avg_freq: int = 1 - avg_range: Optional[Tuple[int, int]] = None - - @field_validator("avg_range") - def avg_range_check(cls, v, info: ValidationInfo): - if isinstance(v, tuple) and v[0] > v[1]: - raise ValueError( - f"'swa.avg_range' should be a valid range, but got {v}." - ) - if isinstance(v, tuple) and v[0] < 0: - raise ValueError( - "The start epoch of 'swa.avg_range' should be a non-negtive integer" - f" , but got {v[0]}." - ) - return v - - @field_validator("avg_freq") - def avg_freq_check(cls, v): - if v <= 0: - raise ValueError( - "'swa.avg_freq' should be a positive integer when is type of int, " - f"but got {v}" - ) - return v - - class TrainConfig(BaseModel): - """ - Schema of training config for pydantic validation. - """ - - epochs: int = 1 - iters_per_epoch: int = 20 - update_freq: int = 1 - save_freq: int = 0 - eval_during_train: bool = False - start_eval_epoch: int = 1 - eval_freq: int = 1 - checkpoint_path: Optional[str] = None - pretrained_model_path: Optional[str] = None - ema: Optional[EMAConfig] = None - swa: Optional[SWAConfig] = None - - # Fine-grained validator(s) below - @field_validator("epochs") - def epochs_check(cls, v): - if v <= 0: - raise ValueError( - "'TRAIN.epochs' should be a positive integer when is type of int, " - f"but got {v}" - ) - return v - - @field_validator("iters_per_epoch") - def iters_per_epoch_check(cls, v): - if v <= 0 and v != -1: - raise ValueError( - f"'TRAIN.iters_per_epoch' received an invalid value({v}), " - "but is expected one of: \n" - "* A positive integer, to manually specify the number of iterations per epoch, " - "which is commonly used in PINN training.\n" - "* -1, to automatically set the number of iterations per epoch to " - "the length of dataloader of given constraint, which is commonly " - f"used in data-driven training.\n" - ) - return v - - @field_validator("update_freq") - def update_freq_check(cls, v): - if v <= 0: - raise ValueError( - "'TRAIN.update_freq' should be a positive integer when is type of int" - f", but got {v}" - ) - return v - - @field_validator("save_freq") - def save_freq_check(cls, v): - if v < 0: - raise ValueError( - "'TRAIN.save_freq' should be a non-negtive integer when is type of int" - f", but got {v}" - ) - return v - - @field_validator("start_eval_epoch") - def start_eval_epoch_check(cls, v, info: ValidationInfo): - if info.data["eval_during_train"]: - if v <= 0: - raise ValueError( - f"'TRAIN.start_eval_epoch' should be a positive integer when " - f"'TRAIN.eval_during_train' is True, but got {v}" - ) - return v - - @field_validator("eval_freq") - def eval_freq_check(cls, v, info: ValidationInfo): - if info.data["eval_during_train"]: - if v <= 0: - raise ValueError( - f"'TRAIN.eval_freq' should be a positive integer when " - f"'TRAIN.eval_during_train' is True, but got {v}" - ) - return v - - @model_validator(mode="after") - def ema_swa_checker(self): - if (self.ema and self.swa) and (self.ema.use_ema and self.swa.use_swa): - raise ValueError( - "Cannot enable both EMA and SWA at the same time, " - "please disable at least one of them." - ) - return self - - @model_validator(mode="after") - def swa_avg_range_checker(self): - if ( - self.swa - and self.swa.use_swa - and self.swa.avg_range[1] > self.epochs - ): - raise ValueError( - "The end epoch of 'swa.avg_range' should not be lager than " - f"'epochs'({self.epochs}), but got {self.swa.avg_range[1]}." - ) - return self - - class EvalConfig(BaseModel): - """ - Schema of evaluation config for pydantic validation. - """ - - pretrained_model_path: Optional[str] = None - eval_with_no_grad: bool = False - compute_metric_by_batch: bool = False - batch_size: Optional[int] = 256 - - @field_validator("batch_size") - def batch_size_check(cls, v): - if isinstance(v, int) and v <= 0: - raise ValueError( - f"'EVAL.batch_size' should be greater than 0 or None, but got {v}" - ) - return v - - class InferConfig(BaseModel): - """ - Schema of inference config for pydantic validation. - """ - - pretrained_model_path: Optional[str] = None - export_path: str = "./inference" - pdmodel_path: Optional[str] = None - pdiparams_path: Optional[str] = None - onnx_path: Optional[str] = None - device: Literal["cpu", "gpu", "npu", "xpu", "sdaa"] = "cpu" - engine: Literal["native", "tensorrt", "onnx", "mkldnn"] = "native" - precision: Literal["fp32", "fp16", "int8"] = "fp32" - ir_optim: bool = True - min_subgraph_size: int = 30 - gpu_mem: int = 2000 - gpu_id: int = 0 - max_batch_size: int = 1024 - num_cpu_threads: int = 10 - batch_size: Optional[int] = 256 - - # Fine-grained validator(s) below - @field_validator("engine") - def engine_check(cls, v, info: ValidationInfo): - if v == "tensorrt" and info.data["device"] != "gpu": - raise ValueError( - "'INFER.device' should be 'gpu' when 'INFER.engine' is 'tensorrt', " - f"but got '{info.data['device']}'" - ) - if v == "mkldnn" and info.data["device"] != "cpu": - raise ValueError( - "'INFER.device' should be 'cpu' when 'INFER.engine' is 'mkldnn', " - f"but got '{info.data['device']}'" - ) - - return v - - @field_validator("min_subgraph_size") - def min_subgraph_size_check(cls, v): - if v <= 0: - raise ValueError( - "'INFER.min_subgraph_size' should be greater than 0, " - f"but got {v}" - ) - return v - - @field_validator("gpu_mem") - def gpu_mem_check(cls, v): - if v <= 0: - raise ValueError( - "'INFER.gpu_mem' should be greater than 0, " f"but got {v}" - ) - return v - - @field_validator("gpu_id") - def gpu_id_check(cls, v): - if v < 0: - raise ValueError( - "'INFER.gpu_id' should be greater than or equal to 0, " - f"but got {v}" - ) - return v - - @field_validator("max_batch_size") - def max_batch_size_check(cls, v): - if v <= 0: - raise ValueError( - "'INFER.max_batch_size' should be greater than 0, " - f"but got {v}" - ) - return v - - @field_validator("num_cpu_threads") - def num_cpu_threads_check(cls, v): - if v < 0: - raise ValueError( - "'INFER.num_cpu_threads' should be greater than or equal to 0, " - f"but got {v}" - ) - return v - - @field_validator("batch_size") - def batch_size_check(cls, v): - if isinstance(v, int) and v <= 0: - raise ValueError( - f"'INFER.batch_size' should be greater than 0 or None, but got {v}" - ) - return v - - class SolverConfig(BaseModel): - """ - Schema of global config for pydantic validation. - """ - - # Global settings config - mode: Literal["train", "eval", "export", "infer"] = "train" - output_dir: Optional[str] = None - log_freq: int = 20 - seed: int = 42 - use_vdl: bool = False - use_tbd: bool = False - wandb_config: Mapping = {} - use_wandb: bool = False - device: Literal["cpu", "gpu", "xpu", "sdaa", None] = None - use_amp: bool = False - amp_level: Literal["O0", "O1", "O2", "OD"] = "O1" - to_static: bool = False - prim: bool = False - log_level: Literal["debug", "info", "warning", "error"] = "info" - - # Training related config - TRAIN: Optional[TrainConfig] = None - - # Evaluation related config - EVAL: Optional[EvalConfig] = None - - # Inference related config - INFER: Optional[InferConfig] = None - - # Fine-grained validator(s) below - @field_validator("log_freq") - def log_freq_check(cls, v): - if v <= 0: - raise ValueError( - "'log_freq' should be a non-negtive integer when is type of int" - f", but got {v}" - ) - return v - - @field_validator("seed") - def seed_check(cls, v): - if v < 0: - raise ValueError( - f"'seed' should be a non-negtive integer, but got {v}" - ) - return v - - @field_validator("use_wandb") - def use_wandb_check(cls, v, info: ValidationInfo): - if v and not isinstance(info.data["wandb_config"], dict): - raise ValueError( - "'wandb_config' should be a dict when 'use_wandb' is True, " - f"but got {info.data['wandb_config'].__class__.__name__}" - ) - return v - - # Register 'XXXConfig' as default node, so as to be used as default config in *.yaml - """ - #### xxx.yaml #### - defaults: - - ppsci_default <-- 'ppsci_default' used here - - TRAIN: train_default <-- 'train_default' used here - - TRAIN/ema: ema_default <-- 'ema_default' used here - - TRAIN/swa: swa_default <-- 'swa_default' used here - - EVAL: eval_default <-- 'eval_default' used here - - INFER: infer_default <-- 'infer_default' used here - - _self_ <-- config defined in current yaml - - mode: train - seed: 42 - ... - ... - ################## - """ - - cs = ConfigStore.instance() - - global_default_cfg = SolverConfig().model_dump() - omegaconf_dict_config = OmegaConf.create(global_default_cfg) - cs.store(name="ppsci_default", node=omegaconf_dict_config) - - train_default_cfg = TrainConfig().model_dump() - train_omegaconf_dict_config = OmegaConf.create(train_default_cfg) - cs.store(group="TRAIN", name="train_default", node=train_omegaconf_dict_config) - - ema_default_cfg = EMAConfig().model_dump() - ema_omegaconf_dict_config = OmegaConf.create(ema_default_cfg) - cs.store(group="TRAIN/ema", name="ema_default", node=ema_omegaconf_dict_config) - - swa_default_cfg = SWAConfig().model_dump() - swa_omegaconf_dict_config = OmegaConf.create(swa_default_cfg) - cs.store(group="TRAIN/swa", name="swa_default", node=swa_omegaconf_dict_config) - - eval_default_cfg = EvalConfig().model_dump() - eval_omegaconf_dict_config = OmegaConf.create(eval_default_cfg) - cs.store(group="EVAL", name="eval_default", node=eval_omegaconf_dict_config) - - infer_default_cfg = InferConfig().model_dump() - infer_omegaconf_dict_config = OmegaConf.create(infer_default_cfg) - cs.store(group="INFER", name="infer_default", node=infer_omegaconf_dict_config) - - exclude_keys_default = [ - "mode", - "output_dir", - "log_freq", - "seed", - "use_vdl", - "use_tbd", - "wandb_config", - "use_wandb", - "device", - "use_amp", - "amp_level", - "to_static", - "prim", - "log_level", - "TRAIN.save_freq", - "TRAIN.eval_during_train", - "TRAIN.start_eval_epoch", - "TRAIN.eval_freq", - "TRAIN.checkpoint_path", - "TRAIN.pretrained_model_path", - "EVAL.pretrained_model_path", - "EVAL.eval_with_no_grad", - "EVAL.compute_metric_by_batch", - "EVAL.batch_size", - "INFER.pretrained_model_path", - "INFER.export_path", - "INFER.pdmodel_path", - "INFER.pdiparams_path", - "INFER.onnx_path", - "INFER.device", - "INFER.engine", - "INFER.precision", - "INFER.ir_optim", - "INFER.min_subgraph_size", - "INFER.gpu_mem", - "INFER.gpu_id", - "INFER.max_batch_size", - "INFER.num_cpu_threads", - "INFER.batch_size", - ] - cs.store( - group="hydra/job/config/override_dirname/exclude_keys", - name="exclude_keys_default", - node=exclude_keys_default, - ) - except ImportError as e: - from ppsci.utils import logger - - logger.error(e) - logger.error( - "paddlesci requires pydantic>=2.5.0; otherwise, built-in examples may not run properly." - ) - except Exception as e: - raise e - -else: - from ppsci.utils import logger - - logger.error( - "paddlesci requires pydantic>=2.5.0; otherwise, built-in examples may not run properly." - ) diff --git a/examples/smc_reac/ppsci/utils/download.py b/examples/smc_reac/ppsci/utils/download.py deleted file mode 100644 index 291703e2d2..0000000000 --- a/examples/smc_reac/ppsci/utils/download.py +++ /dev/null @@ -1,285 +0,0 @@ -# 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. - -from __future__ import annotations - -import hashlib -import os -import os.path as osp -import shutil -import tarfile -import time -import zipfile - -import requests -import tqdm - -from ppsci.utils import logger -from ppsci.utils import misc - -__all__ = ["get_weights_path_from_url"] - -WEIGHTS_HOME = osp.expanduser("~/.paddlesci/weights") - -DOWNLOAD_RETRY_LIMIT = 3 - - -def is_url(path): - """ - Whether path is URL. - - Args: - path (str): URL string or not. - """ - return path.startswith("http://") or path.startswith("https://") - - -def get_weights_path_from_url(url, md5sum=None): - """Get weights path from WEIGHT_HOME, if not exists, - download it from url. - - Args: - url (str): Download url - md5sum (str): md5 sum of download package - - Returns: - str: a local path to save downloaded weights. - """ - path = get_path_from_url(url, WEIGHTS_HOME, md5sum) - return path - - -def _map_path(url, root_dir): - # parse path after download under root_dir - fname = osp.split(url)[-1] - fpath = fname - return osp.join(root_dir, fpath) - - -def get_path_from_url(url, root_dir, md5sum=None, check_exist=True, decompress=True): - """Download from given url to root_dir. - if file or directory specified by url is exists under - root_dir, return the path directly, otherwise download - from url and decompress it, return the path. - - Args: - url (str): Download url - root_dir (str): Root dir for downloading, it should be - WEIGHTS_HOME or DATASET_HOME - md5sum (str): md5 sum of download package - - Returns: - str: a local path to save downloaded models & weights & datasets. - """ - if not is_url(url): - raise ValueError(f"Given url({url}) is not valid") - # parse path after download to decompress under root_dir - fullpath = _map_path(url, root_dir) - # Mainly used to solve the problem of downloading data from different - # machines in the case of multiple machines. Different nodes will download - # data, and the same node will only download data once. - rank_id_curr_node = int(os.environ.get("PADDLE_RANK_IN_NODE", 0)) - - if osp.exists(fullpath) and check_exist and _md5check(fullpath, md5sum): - logger.message(f"Found {fullpath} already in {WEIGHTS_HOME}, skip downloading.") - else: - with misc.RankZeroOnly(rank_id_curr_node) as is_master: - if is_master: - fullpath = _download(url, root_dir, md5sum) - - if decompress and (tarfile.is_tarfile(fullpath) or zipfile.is_zipfile(fullpath)): - with misc.RankZeroOnly(rank_id_curr_node) as is_master: - if is_master: - fullpath = _decompress(fullpath) - - return fullpath - - -def _download(url, path, md5sum=None): - """ - Download from url, save to path. - - url (str): Download url - path (str): Download to given path - """ - if not osp.exists(path): - os.makedirs(path) - - fname = osp.split(url)[-1] - fullname = osp.join(path, fname) - retry_cnt = 0 - - while not (osp.exists(fullname) and _md5check(fullname, md5sum)): - if retry_cnt < DOWNLOAD_RETRY_LIMIT: - retry_cnt += 1 - else: - raise RuntimeError(f"Download from {url} failed. " "Retry limit reached") - - logger.message(f"Downloading {fname} from {url}") - - try: - req = requests.get(url, stream=True) - except Exception as e: # requests.exceptions.ConnectionError - logger.warning( - f"Downloading {fname} from {url} failed {retry_cnt + 1} times with exception {str(e)}" - ) - time.sleep(1) - continue - - if req.status_code != 200: - raise RuntimeError( - f"Downloading from {url} failed with code " f"{req.status_code}!" - ) - - # For protecting download interrupted, download to - # tmp_fullname firstly, move tmp_fullname to fullname - # after download finished - tmp_fullname = fullname + "_tmp" - total_size = req.headers.get("content-length") - with open(tmp_fullname, "wb") as f: - if total_size: - with tqdm.tqdm(total=(int(total_size) + 1023) // 1024) as pbar: - for chunk in req.iter_content(chunk_size=1024): - f.write(chunk) - pbar.update(1) - else: - for chunk in req.iter_content(chunk_size=1024): - if chunk: - f.write(chunk) - shutil.move(tmp_fullname, fullname) - logger.message(f"Finish downloading pretrained model and saved to {fullname}") - - return fullname - - -def _md5check(fullname, md5sum=None): - if md5sum is None: - return True - - logger.message(f"File {fullname} md5 checking...") - md5 = hashlib.md5() - with open(fullname, "rb") as f: - for chunk in iter(lambda: f.read(4096), b""): - md5.update(chunk) - calc_md5sum = md5.hexdigest() - - if calc_md5sum != md5sum: - logger.error( - f"File {fullname} md5 check failed, {calc_md5sum}(calc) != " - f"{md5sum}(base)" - ) - return False - return True - - -def _decompress(fname): - """ - Decompress for zip and tar file - """ - logger.message(f"Decompressing {fname}...") - - # For protecting decompressing interrupted, - # decompress to fpath_tmp directory firstly, if decompress - # succeed, move decompress files to fpath and delete - # fpath_tmp and remove download compress file. - - if tarfile.is_tarfile(fname): - uncompressed_path = _uncompress_file_tar(fname) - elif zipfile.is_zipfile(fname): - uncompressed_path = _uncompress_file_zip(fname) - else: - raise TypeError(f"Unsupported compress file type {fname}") - - return uncompressed_path - - -def _uncompress_file_zip(filepath): - with zipfile.ZipFile(filepath, "r") as files: - file_list = files.namelist() - - file_dir = os.path.dirname(filepath) - - if _is_a_single_file(file_list): - rootpath = file_list[0] - uncompressed_path = os.path.join(file_dir, rootpath) - - for item in file_list: - files.extract(item, file_dir) - - elif _is_a_single_dir(file_list): - rootpath = os.path.splitext(file_list[0])[0].split(os.sep)[-1] - uncompressed_path = os.path.join(file_dir, rootpath) - - for item in file_list: - files.extract(item, file_dir) - - else: - rootpath = os.path.splitext(filepath)[0].split(os.sep)[-1] - uncompressed_path = os.path.join(file_dir, rootpath) - if not os.path.exists(uncompressed_path): - os.makedirs(uncompressed_path) - for item in file_list: - files.extract(item, os.path.join(file_dir, rootpath)) - - return uncompressed_path - - -def _uncompress_file_tar(filepath, mode="r:*"): - with tarfile.open(filepath, mode) as files: - file_list = files.getnames() - - file_dir = os.path.dirname(filepath) - - if _is_a_single_file(file_list): - rootpath = file_list[0] - uncompressed_path = os.path.join(file_dir, rootpath) - for item in file_list: - files.extract(item, file_dir) - elif _is_a_single_dir(file_list): - rootpath = os.path.splitext(file_list[0])[0].split(os.sep)[-1] - uncompressed_path = os.path.join(file_dir, rootpath) - for item in file_list: - files.extract(item, file_dir) - else: - rootpath = os.path.splitext(filepath)[0].split(os.sep)[-1] - uncompressed_path = os.path.join(file_dir, rootpath) - if not os.path.exists(uncompressed_path): - os.makedirs(uncompressed_path) - - for item in file_list: - files.extract(item, os.path.join(file_dir, rootpath)) - - return uncompressed_path - - -def _is_a_single_file(file_list): - if len(file_list) == 1 and file_list[0].find(os.sep) < -1: - return True - return False - - -def _is_a_single_dir(file_list): - new_file_list = [] - for file_path in file_list: - if "/" in file_path: - file_path = file_path.replace("/", os.sep) - elif "\\" in file_path: - file_path = file_path.replace("\\", os.sep) - new_file_list.append(file_path) - - file_name = new_file_list[0].split(os.sep)[0] - for i in range(1, len(new_file_list)): - if file_name != new_file_list[i].split(os.sep)[0]: - return False - return True diff --git a/examples/smc_reac/ppsci/utils/ema.py b/examples/smc_reac/ppsci/utils/ema.py deleted file mode 100644 index 690ee6fda8..0000000000 --- a/examples/smc_reac/ppsci/utils/ema.py +++ /dev/null @@ -1,172 +0,0 @@ -# Copyright (c) 2024 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. - -from __future__ import annotations - -import itertools -from typing import Dict -from typing import Optional - -import paddle -from paddle import nn - -__all__ = [ - "AveragedModel", - "ExponentialMovingAverage", - "StochasticWeightAverage", -] - - -class AveragedModel(nn.Layer): - """Base class for Averaged Model. - - Args: - model (nn.Layer): The model to be averaged. - decay (float): The decay rate for averaging. - """ - - def __init__(self, model: nn.Layer, decay: Optional[float] = None): - super().__init__() - self.model = model # As a quick reference to online model - self.decay = decay - - self.params_shadow: Dict[str, paddle.Tensor] = {} # ema param or buffer - self.params_backup: Dict[str, paddle.Tensor] = {} # used for apply and restore - for name, param_or_buffer in itertools.chain( - self.model.named_parameters(), self.model.named_buffers() - ): - self.params_shadow[name] = param_or_buffer.clone().detach() - - self.register_buffer("n_avg", paddle.to_tensor(0, "int64"), True) - - def _update_fn_( - self, - shadow_param: paddle.Tensor, - model_param: paddle.Tensor, - step: paddle.Tensor, - ): - raise NotImplementedError("AveragedModel._update_fn_ should be implemented.") - - def update(self): - for name, param_or_buffer in itertools.chain( - self.model.named_parameters(), self.model.named_buffers() - ): - if not param_or_buffer.stop_gradient: - assert ( - name in self.params_shadow - ), f"Parameter: {name} should be in params_shadow dict, but not found." - - # only update floating and complex data - if paddle.is_floating_point(param_or_buffer) or paddle.is_complex( - param_or_buffer - ): - with paddle.no_grad(): - self._update_fn_( - self.params_shadow[name], - param_or_buffer, - self.n_avg, - ) - self.n_avg += 1 - - def apply_shadow(self): - """Set averaged model parameters to online model.""" - for name, param_or_buffer in itertools.chain( - self.model.named_parameters(), self.model.named_buffers() - ): - if name in self.params_shadow: - stop_gradient = param_or_buffer.stop_gradient - with paddle.no_grad(): - self.params_backup[name] = paddle.assign(param_or_buffer) - paddle.assign(self.params_shadow[name], param_or_buffer) - param_or_buffer.stop_gradient = stop_gradient - - def restore(self): - """Restore online model parameters from backup parameter dict.""" - assert self.params_backup, ( - "params_backup should not be empty, may be caused by calling 'restore' " - "before 'apply_shadow'." - ) - for name, param_or_buffer in itertools.chain( - self.model.named_parameters(), self.model.named_buffers() - ): - if name in self.params_backup: - assert name in self.params_shadow - stop_gradient = param_or_buffer.stop_gradient - with paddle.no_grad(): - paddle.assign(self.params_backup[name], param_or_buffer) - param_or_buffer.stop_gradient = stop_gradient - - self.params_backup = {} - - def set_state_dict(self, state_dict: Dict[str, paddle.Tensor]): - assert ( - "n_avg" in state_dict - ), "state_dict should contain 'n_avg' key, but not found." - self.n_avg.set_value(state_dict.pop("n_avg")) - self.params_shadow.update(state_dict) - - def state_dict(self) -> Dict[str, paddle.Tensor]: - return { - **self.params_shadow, - "n_avg": self.n_avg, - } - - -class ExponentialMovingAverage(AveragedModel): - r"""Implements the exponential moving average (EMA) of the model. - - All parameters are updated by the formula as below: - - $$ - \mathbf{\theta}_{EMA}^{t+1} = \alpha \mathbf{\theta}_{EMA}^{t} + (1 - \alpha) \mathbf{\theta}^{t} - $$ - - Where $\alpha$ is the decay rate, $\theta_{EMA}^{t}$ is the moving average parameters and $\theta^{t}$ is the online parameters at step $t$. - - Args: - model (nn.Layer): The model to be averaged. - decay (float): The decay rate for averaging. - """ - - def __init__(self, model: nn.Layer, decay: float = 0.9): - super().__init__(model, decay) - - def _update_fn_(self, shadow_param, model_param, step): - shadow_param.lerp_(model_param, 1.0 - self.decay) - - -class StochasticWeightAverage(AveragedModel): - r"""Implements the stochastic weight averaging (SWA) of the model. - - Stochastic Weight Averaging was proposed in [Averaging Weights Leads to Wider Optima and Better Generalization](https://arxiv.org/abs/1803.05407), - - All parameters are updated by the formula as below: - - $$ - \mathbf{\theta}_{SWA}^{t} = \frac{1}{t-t_0+1}\sum_{i=t_0}^t{\mathbf{\theta}^{i}} - $$ - - Where $\theta_{SWA}^{t}$ is the average parameters between step $t_0$ and $t$, $\theta^{i}$ is the online parameters at step $i$. - - Args: - model (nn.Layer): The model to be averaged. - """ - - def __init__(self, model: nn.Layer): - super().__init__(model, None) - self.n_avg += 1 # Set to 1 for model already initialized - - def _update_fn_(self, shadow_param, model_param, step): - dynamic_decay = step / (step + 1) - shadow_param.lerp_(model_param, 1.0 - dynamic_decay) diff --git a/examples/smc_reac/ppsci/utils/expression.py b/examples/smc_reac/ppsci/utils/expression.py deleted file mode 100644 index 6bfcddd214..0000000000 --- a/examples/smc_reac/ppsci/utils/expression.py +++ /dev/null @@ -1,212 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import TYPE_CHECKING -from typing import Callable -from typing import Dict -from typing import Optional -from typing import Tuple - -from paddle import jit -from paddle import nn -from paddle.framework import core - -if TYPE_CHECKING: - import paddle - from ppsci import constraint - from ppsci import validate - from ppsci import arch - -from ppsci.autodiff import clear - -__all__ = [ - "ExpressionSolver", -] - - -class ExpressionSolver(nn.Layer): - """Expression computing helper, which compute named result according to corresponding - function and related inputs. - - Examples: - >>> import ppsci - >>> model = ppsci.arch.MLP(("x", "y"), ("u", "v"), 5, 128) - >>> expr_solver = ExpressionSolver() - """ - - nvtx_flag: bool # only for nsight analysis - - def __init__(self): - super().__init__() - - def forward(self, *args, **kwargs): - raise NotImplementedError( - "Use train_forward/eval_forward/visu_forward instead of forward." - ) - - @jit.to_static - def train_forward( - self, - expr_dicts: Tuple[Dict[str, Callable], ...], - input_dicts: Tuple[Dict[str, "paddle.Tensor"], ...], - model: arch.Arch, - constraint: Dict[str, "constraint.Constraint"], - label_dicts: Tuple[Dict[str, "paddle.Tensor"], ...], - weight_dicts: Tuple[Dict[str, "paddle.Tensor"], ...], - ) -> Tuple[Dict[str, "paddle.Tensor"], Dict[str, float]]: - """Forward computation for training, including model forward and equation - forward. - - Args: - expr_dicts (Tuple[Dict[str, Callable], ...]): Tuple of expression dicts. - input_dicts (Tuple[Dict[str, paddle.Tensor], ...]): Tuple of input dicts. - model (arch.Arch): NN model. - constraint (Dict[str, "constraint.Constraint"]): Constraint dict. - label_dicts (Tuple[Dict[str, paddle.Tensor], ...]): Tuple of label dicts. - weight_dicts (Tuple[Dict[str, paddle.Tensor], ...]): Tuple of weight dicts. - - Returns: - Tuple[Dict[str, "paddle.Tensor"], Dict[str, float]]: - all_losses: A loss dictionary containing the output terms of all constraints, - constraint_losses: The loss values of all constraints. - """ - losses_all: Dict[str, "paddle.Tensor"] = {} - losses_constraint: Dict[str, float] = {} - - for i, cst_name in enumerate(constraint): - cst_obj = constraint[cst_name] - - # model forward - if self.nvtx_flag: # only for nsight analysis - core.nvprof_nvtx_push(f"Constraint {cst_name}") - - output_dict = model(input_dicts[i]) - - # equation forward - data_dict = {k: v for k, v in input_dicts[i].items()} - data_dict.update(output_dict) - for name, expr in expr_dicts[i].items(): - output_dict[name] = expr(data_dict) - - # put field 'area' into output_dict - if "area" in input_dicts[i]: - output_dict["area"] = input_dicts[i]["area"] - - # clear differentiation cache - clear() - - # compute loss for each constraint according to its' own output, label and weight - losses: Dict[str, "paddle.Tensor"] = cst_obj.loss( - output_dict, - label_dicts[i], - weight_dicts[i], - ) - # update losses into 'losses_all' and 'losses_constraint' - # 'losses_all': Will be send to loss aggregator for further computing final loss(scalar) - # 'losses_constraint': Will be used in logging - losses_constraint[cst_name] = 0.0 - for key in losses: - losses_constraint[cst_name] += losses[key].item() - if key in losses_all: - losses_all[key] += losses[key] - else: - losses_all[key] = losses[key] - - if self.nvtx_flag: # only for nsight analysis - core.nvprof_nvtx_pop() - - return losses_all, losses_constraint - - @jit.to_static - def eval_forward( - self, - expr_dict: Dict[str, Callable], - input_dict: Dict[str, "paddle.Tensor"], - model: arch.Arch, - validator: "validate.Validator", - label_dict: Dict[str, "paddle.Tensor"], - weight_dict: Dict[str, "paddle.Tensor"], - ) -> Tuple[Dict[str, "paddle.Tensor"], Dict[str, "paddle.Tensor"]]: - """Forward computation for evaluation, including model forward and equation - forward. - - Args: - expr_dict (Dict[str, Callable]): Expression dict. - input_dict (Dict[str, paddle.Tensor]): Input dict. - model (arch.Arch): NN model. - validator (validate.Validator): Validator. - label_dict (Dict[str, paddle.Tensor]): Label dict. - weight_dict (Dict[str, paddle.Tensor]): Weight dict. - - Returns: - Tuple[Dict[str, paddle.Tensor], Dict[str, paddle.Tensor]]: Result dict and loss for - given validator. - """ - # model forward - output_dict = model(input_dict) - - # equation forward - data_dict = {k: v for k, v in input_dict.items()} - data_dict.update(output_dict) - for name, expr in expr_dict.items(): - output_dict[name] = expr(data_dict) - - # put field 'area' into output_dict - if "area" in input_dict: - output_dict["area"] = input_dict["area"] - - # clear differentiation cache - clear() - - # compute loss for each validator according to its' own output, label and weight - validator_losses = validator.loss( - output_dict, - label_dict, - weight_dict, - ) - return output_dict, validator_losses - - def visu_forward( - self, - expr_dict: Optional[Dict[str, Callable]], - input_dict: Dict[str, "paddle.Tensor"], - model: arch.Arch, - ) -> Dict[str, "paddle.Tensor"]: - """Forward computation for visualization, including model forward and equation - forward. - - Args: - expr_dict (Optional[Dict[str, Callable]]): Expression dict. - input_dict (Dict[str, paddle.Tensor]): Input dict. - model (arch.Arch): NN model. - - Returns: - Dict[str, paddle.Tensor]: Result dict for given expression dict. - """ - # model forward - output_dict = model(input_dict) - - if isinstance(expr_dict, dict): - # equation forward - data_dict = {k: v for k, v in input_dict.items()} - data_dict.update(output_dict) - for name, expr in expr_dict.items(): - output_dict[name] = expr(data_dict) - - # clear differentiation cache - clear() - - return output_dict diff --git a/examples/smc_reac/ppsci/utils/initializer.py b/examples/smc_reac/ppsci/utils/initializer.py deleted file mode 100644 index 0a5ececf84..0000000000 --- a/examples/smc_reac/ppsci/utils/initializer.py +++ /dev/null @@ -1,498 +0,0 @@ -# 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. - -""" -The initialization method under this module is aligned with pytorch initialization. -If you need to use the initialization method of PaddlePaddle, please refer to -[paddle.nn.initializer](https://github.com/PaddlePaddle/Paddle/tree/develop/python/paddle/nn/initializer) - -This code is based on [torch.nn.init](https://github.com/pytorch/pytorch/blob/main/torch/nn/init.py) -The copyright of pytorch/pytorch is a BSD-style license, as found in the LICENSE file. -""" - -from __future__ import annotations - -import math - -import numpy as np -import paddle -from paddle import nn -from typing_extensions import Literal - -from ppsci.utils import logger - -__all__ = [ - "uniform_", - "normal_", - "trunc_normal_", - "glorot_normal_", - "constant_", - "ones_", - "zeros_", - "xavier_uniform_", - "xavier_normal_", - "kaiming_uniform_", - "kaiming_normal_", - "linear_init_", - "conv_init_", -] - - -def _no_grad_uniform_(tensor, a, b): - with paddle.no_grad(): - tensor.set_value( - paddle.uniform(shape=tensor.shape, dtype=tensor.dtype, min=a, max=b) - ) - return tensor - - -def _no_grad_normal_(tensor, mean=0.0, std=1.0): - with paddle.no_grad(): - tensor.set_value(paddle.normal(mean=mean, std=std, shape=tensor.shape)) - return tensor - - -def _no_grad_trunc_normal_(tensor, mean=0.0, std=1.0, a=-2.0, b=2.0): - # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf - def norm_cdf(x): - # Computes standard normal cumulative distribution function - return (1.0 + math.erf(x / math.sqrt(2.0))) / 2.0 - - if (mean < a - 2 * std) or (mean > b + 2 * std): - logger.warning( - f"mean({mean}) is more than 2 std({std}) from [a, b]([{a}, {b}]) in _no_grad_trunc_normal_. " - "The distribution of values may be incorrect." - ) - with paddle.no_grad(): - # Values are generated by using a truncated uniform distribution and - # then using the inverse CDF for the normal distribution. - # Get upper and lower cdf values - l = norm_cdf((a - mean) / std) - u = norm_cdf((b - mean) / std) - - # Uniformly fill tensor with values from [l, u], then translate to - # [2l-1, 2u-1]. - _tensor = paddle.uniform( - shape=tensor.shape, dtype=tensor.dtype, min=2 * l - 1, max=2 * u - 1 - ) - - # Use inverse cdf transform for normal distribution to get truncated - # standard normal - _tensor.erfinv_() - - # Transform to proper mean, std - _tensor = paddle.multiply( - _tensor, paddle.to_tensor(std * math.sqrt(2.0), tensor.dtype) - ) - _tensor = paddle.add(_tensor, paddle.to_tensor(mean, tensor.dtype)) - - # Clamp to ensure it"s in the proper range - _tensor = paddle.clip(_tensor, min=a, max=b) - tensor.set_value(_tensor) - return tensor - - -def _no_grad_fill_(tensor, value=0.0): - with paddle.no_grad(): - tensor.set_value(paddle.full_like(tensor, value, dtype=tensor.dtype)) - return tensor - - -def uniform_(tensor: paddle.Tensor, a: float, b: float) -> paddle.Tensor: - """Modify tensor inplace using uniform_. - - Args: - tensor (paddle.Tensor): Paddle Tensor. - a (float): Min value. - b (float): Max value. - - Returns: - paddle.Tensor: Initialized tensor. - - Examples: - >>> import paddle - >>> import ppsci - >>> param = paddle.empty((128, 256), "float32") - >>> param = ppsci.utils.initializer.uniform_(param, -1, 1) - """ - return _no_grad_uniform_(tensor, a, b) - - -def normal_( - tensor: paddle.Tensor, mean: float = 0.0, std: float = 1.0 -) -> paddle.Tensor: - """Modify tensor inplace using normal_. - - Args: - tensor (paddle.Tensor): Paddle Tensor. - mean (float, optional): Mean value. Defaults to 0.0. - std (float, optional): Std value. Defaults to 1.0. - - Returns: - paddle.Tensor: Initialized tensor. - - Examples: - >>> import paddle - >>> import ppsci - >>> param = paddle.empty((128, 256), "float32") - >>> param = ppsci.utils.initializer.normal_(param, 0, 1) - """ - return _no_grad_normal_(tensor, mean, std) - - -def trunc_normal_( - tensor: paddle.Tensor, - mean: float = 0.0, - std: float = 1.0, - a: float = -2.0, - b: float = 2.0, -) -> paddle.Tensor: - """Modify tensor inplace using trunc_normal_. - - Args: - tensor (paddle.Tensor): Paddle Tensor. - mean (float, optional): The mean of the normal distribution. Defaults to 0.0. - std (float, optional): The standard deviation of the normal distribution. Defaults to 1.0. - a (float, optional): The minimum cutoff value. Defaults to -2.0. - b (float, optional): The maximum cutoff value. Defaults to 2.0. - - Returns: - paddle.Tensor: Initialized tensor. - - Examples: - >>> import paddle - >>> import ppsci - >>> param = paddle.empty((128, 256), "float32") - >>> param = ppsci.utils.initializer.trunc_normal_(param, 0.0, 1.0) - """ - return _no_grad_trunc_normal_(tensor, mean, std, a, b) - - -def constant_(tensor: paddle.Tensor, value: float = 0.0) -> paddle.Tensor: - """Modify tensor inplace using constant_. - - Args: - tensor (paddle.Tensor): Paddle Tensor. - value (float, optional): Value to fill tensor. Defaults to 0.0. - - Returns: - paddle.Tensor: Initialized tensor. - - Examples: - >>> import paddle - >>> import ppsci - >>> param = paddle.empty((128, 256), "float32") - >>> param = ppsci.utils.initializer.constant_(param, 2) - """ - return _no_grad_fill_(tensor, value) - - -def ones_(tensor: paddle.Tensor) -> paddle.Tensor: - """Modify tensor inplace using ones_. - - Args: - tensor (paddle.Tensor): Paddle Tensor. - - Returns: - paddle.Tensor: Initialized tensor. - - Examples: - >>> import paddle - >>> import ppsci - >>> param = paddle.empty((128, 256), "float32") - >>> param = ppsci.utils.initializer.ones_(param) - """ - return _no_grad_fill_(tensor, 1) - - -def zeros_(tensor: paddle.Tensor) -> paddle.Tensor: - """Modify tensor inplace using zeros_. - - Args: - tensor (paddle.Tensor): Paddle Tensor. - - Returns: - paddle.Tensor: Initialized tensor. - - Examples: - >>> import paddle - >>> import ppsci - >>> param = paddle.empty((128, 256), "float32") - >>> param = ppsci.utils.initializer.zeros_(param) - """ - return _no_grad_fill_(tensor, 0) - - -def _calculate_fan_in_and_fan_out(tensor, reverse=False): - """ - Calculate (fan_in, _fan_out) for tensor. - - Args: - tensor (paddle.Tensor): paddle.Tensor. - reverse (bool): Tensor data format order, False by default as [fout, fin, ...]. - e.g. : conv.weight [cout, cin, kh, kw] is False; linear.weight [cin, cout] - is True. - - Return: - Tuple[float, float]: (fan_in, fan_out). - """ - if tensor.ndim < 2: - raise ValueError( - f"tensor.ndim should be no less than 2, but got {tensor.ndim}." - ) - - if reverse: - num_input_fmaps, num_output_fmaps = tensor.shape[0], tensor.shape[1] - else: - num_input_fmaps, num_output_fmaps = tensor.shape[1], tensor.shape[0] - - receptive_field_size = 1 - if tensor.ndim > 2: - receptive_field_size = np.prod(tensor.shape[2:]) - - fan_in = num_input_fmaps * receptive_field_size - fan_out = num_output_fmaps * receptive_field_size - - return fan_in, fan_out - - -def xavier_uniform_( - tensor: paddle.Tensor, gain: float = 1.0, reverse: bool = False -) -> paddle.Tensor: - """Modify tensor inplace using xavier_uniform_. - - Args: - tensor (paddle.Tensor): Paddle Tensor. - gain (float, optional): Hyperparameter. Defaults to 1.0. - reverse (bool, optional): Tensor data format order, False by default as - [fout, fin, ...].. Defaults to False. - - Returns: - paddle.Tensor: Initialized tensor. - - Examples: - >>> import paddle - >>> import ppsci - >>> param = paddle.empty((128, 256), "float32") - >>> param = ppsci.utils.initializer.xavier_uniform_(param) - """ - fan_in, fan_out = _calculate_fan_in_and_fan_out(tensor, reverse=reverse) - std = gain * math.sqrt(2.0 / float(fan_in + fan_out)) - k = math.sqrt(3.0) * std - return _no_grad_uniform_(tensor, -k, k) - - -def xavier_normal_( - tensor: paddle.Tensor, gain: float = 1.0, reverse: bool = False -) -> paddle.Tensor: - """Modify tensor inplace using xavier_normal_. - - Args: - tensor (paddle.Tensor): Paddle Tensor. - gain (float, optional): Hyperparameter. Defaults to 1.0. - reverse (bool, optional): Tensor data format order, False by - default as [fout, fin, ...]. Defaults to False. - - Returns: - paddle.Tensor: Initialized tensor. - - Examples: - >>> import paddle - >>> import ppsci - >>> param = paddle.empty((128, 256), "float32") - >>> param = ppsci.utils.initializer.xavier_normal_(param) - """ - fan_in, fan_out = _calculate_fan_in_and_fan_out(tensor, reverse=reverse) - std = gain * math.sqrt(2.0 / float(fan_in + fan_out)) - return _no_grad_normal_(tensor, 0, std) - - -# reference: https://pytorch.org/docs/stable/_modules/torch/nn/init.html -def _calculate_correct_fan(tensor, mode, reverse=False): - mode = mode.lower() - valid_modes = ["fan_in", "fan_out"] - if mode not in valid_modes: - raise ValueError(f"Mode {mode} not supported, please use one of {valid_modes}") - - fan_in, fan_out = _calculate_fan_in_and_fan_out(tensor, reverse) - - return fan_in if mode == "fan_in" else fan_out - - -def _calculate_gain(nonlinearity, param=None): - linear_fns = [ - "linear", - "conv1d", - "conv2d", - "conv3d", - "conv_transpose1d", - "conv_transpose2d", - "conv_transpose3d", - ] - if nonlinearity in linear_fns or nonlinearity == "sigmoid": - return 1 - elif nonlinearity == "tanh": - return 5.0 / 3 - elif nonlinearity == "relu": - return math.sqrt(2.0) - elif nonlinearity == "leaky_relu": - if param is None: - negative_slope = 0.01 - elif ( - not isinstance(param, bool) - and isinstance(param, int) - or isinstance(param, float) - ): - # True/False are instances of int, hence check above - negative_slope = param - else: - raise ValueError(f"negative_slope {param} not a valid number") - return math.sqrt(2.0 / (1 + negative_slope**2)) - elif nonlinearity == "selu": - return 3.0 / 4 - else: - raise ValueError(f"Unsupported nonlinearity {nonlinearity}") - - -def kaiming_uniform_( - tensor: paddle.Tensor, - a: float = 0, - mode: Literal["fan_in", "fan_out"] = "fan_in", - nonlinearity: str = "leaky_relu", - reverse: bool = False, -) -> paddle.Tensor: - """Modify tensor inplace using kaiming_uniform method. - - Args: - tensor (paddle.Tensor): Paddle Tensor. - a (float, optional): The negative slope of the rectifier used after this layer. - Defaults to 0. - mode (Literal["fan_in", "fan_out"], optional): - ["fan_in", "fan_out"]. Defaults to "fan_in". - nonlinearity (str, optional): Nonlinearity method name. Defaults to "leaky_relu". - reverse (bool, optional): Tensor data format order, False by default as - [fout, fin, ...].. Defaults to False. - - Returns: - paddle.Tensor: Initialized tensor. - - Examples: - >>> import paddle - >>> import ppsci - >>> param = paddle.empty((128, 256), "float32") - >>> param = ppsci.utils.initializer.kaiming_uniform_(param) - """ - fan = _calculate_correct_fan(tensor, mode, reverse) - gain = _calculate_gain(nonlinearity, a) - std = gain / math.sqrt(fan) - k = math.sqrt(3.0) * std - return _no_grad_uniform_(tensor, -k, k) - - -def kaiming_normal_( - tensor: paddle.Tensor, - a: float = 0, - mode: Literal["fan_in", "fan_out"] = "fan_in", - nonlinearity: str = "leaky_relu", - reverse: bool = False, -) -> paddle.Tensor: - """Modify tensor inplace using kaiming_normal_. - - Args: - tensor (paddle.Tensor): Paddle Tensor. - a (float, optional): The negative slope of the rectifier used after this layer. - Defaults to 0. - mode (Literal["fan_in", "fan_out"], optional): Either - 'fan_in' (default) or 'fan_out'. Defaults to "fan_in". - nonlinearity (str, optional): Nonlinearity method name. Defaults to "leaky_relu". - reverse (bool, optional): Tensor data format order. Defaults to False. - - Returns: - paddle.Tensor: Initialized tensor. - - Examples: - >>> import paddle - >>> import ppsci - >>> param = paddle.empty((128, 256), "float32") - >>> param = ppsci.utils.initializer.kaiming_normal_(param) - """ - fan = _calculate_correct_fan(tensor, mode, reverse) - gain = _calculate_gain(nonlinearity, a) - std = gain / math.sqrt(fan) - return _no_grad_normal_(tensor, 0, std) - - -def linear_init_(module: nn.Layer) -> None: - """Initialize module's weight and bias as it is a linear layer. - - Args: - module (nn.Layer): Linear Layer to be initialized. - - Examples: - >>> import paddle - >>> import ppsci - >>> layer = paddle.nn.Linear(128, 256) - >>> ppsci.utils.initializer.linear_init_(layer) - """ - kaiming_uniform_(module.weight, a=math.sqrt(5), reverse=True) - if module.bias is not None: - fan_in, _ = _calculate_fan_in_and_fan_out(module.weight, reverse=True) - bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0 - uniform_(module.bias, -bound, bound) - - -def conv_init_(module: nn.Layer) -> None: - """Initialize module's weight and bias as it is a conv layer. - - Args: - module (nn.Layer): Convolution Layer to be initialized. - - Examples: - >>> import paddle - >>> import ppsci - >>> layer = paddle.nn.Conv2D(4, 16, 2) - >>> ppsci.utils.initializer.conv_init_(layer) - """ - kaiming_uniform_(module.weight, a=math.sqrt(5)) - if module.bias is not None: - fan_in, _ = _calculate_fan_in_and_fan_out(module.weight, reverse=False) - if fan_in != 0: - bound = 1 / math.sqrt(fan_in) - uniform_(module.bias, -bound, bound) - - -def glorot_normal_(tensor: paddle.Tensor) -> paddle.Tensor: - """Modify tensor inplace using jax-style glorot_normal. - - Args: - tensor (paddle.Tensor): Paddle Tensor/Parameter. - - Returns: - paddle.Tensor: Initialized tensor. - - Examples: - >>> import paddle - >>> import ppsci - >>> param = paddle.empty((128, 256), "float32") - >>> param = ppsci.utils.initializer.glorot_normal_(param) - """ - assert ( - tensor.ndim == 2 - ), f"glorot_normal_ only support 2D tensor now, but got ndim={tensor.ndim}" - fin, fout = tensor.shape - var = 2.0 / (fin + fout) - stddev = math.sqrt(var) * 0.87962566103423978 - trunc_normal_(tensor) - tensor.set_value(tensor * stddev) - return tensor diff --git a/examples/smc_reac/ppsci/utils/logger.py b/examples/smc_reac/ppsci/utils/logger.py deleted file mode 100644 index 46ca57bd46..0000000000 --- a/examples/smc_reac/ppsci/utils/logger.py +++ /dev/null @@ -1,264 +0,0 @@ -# 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. - -from __future__ import annotations - -import functools -import logging -import os -import sys -from typing import TYPE_CHECKING -from typing import Callable -from typing import Dict -from typing import Optional - -import colorlog -import paddle.distributed as dist - -from ppsci.utils import misc - -if TYPE_CHECKING: - import visualdl # isort:skip - import wandb # isort:skip - import tensorboardX as tbd - -_logger: logging.Logger = None - -# INFO(20) is white(no color) -# use custom log level `MESSAGE` for printing message in color -_MESSAGE_LEVEL = 25 - -_COLORLOG_CONFIG = { - "DEBUG": "green", - "WARNING": "yellow", - "ERROR": "red", - "MESSAGE": "bold_cyan", -} - -__all__ = [ - "init_logger", - "set_log_level", - "info", - "message", - "debug", - "warning", - "error", - "scalar", -] - - -def init_logger( - name: str = "ppsci", - log_file: Optional[str] = None, - log_level: int = logging.INFO, -) -> None: - """Initialize and get a logger by name. - - If the logger has not been initialized, this method will initialize the logger by - adding one or two handlers, otherwise the initialized logger will be directly - returned. During initialization, a StreamHandler will always be added. If `log_file` - is specified a FileHandler will also be added. - - Args: - name (str, optional): Logger name. Defaults to "ppsci". - log_file (Optional[str]): The log filename. If specified, a FileHandler - will be added to the logger. Defaults to None. - log_level (int, optional): The logger level. Note that only the process of - rank 0 is affected, and other processes will set the level to - "Error" thus be silent most of the time. Defaults to logging.INFO. - """ - # Add custom log level MESSAGE(25), between WARNING(30) and INFO(20) - logging.addLevelName(_MESSAGE_LEVEL, "MESSAGE") - - if isinstance(log_level, str): - log_level = getattr(logging, log_level.upper()) - - global _logger - - # get a clean logger - _logger = logging.getLogger(name) - _logger.handlers.clear() - - # add stream_handler, output to stdout such as terminal - stream_formatter = colorlog.ColoredFormatter( - "%(log_color)s[%(asctime)s] %(name)s %(levelname)s: %(message)s", - datefmt="%Y/%m/%d %H:%M:%S", - log_colors=_COLORLOG_CONFIG, - ) - stream_handler = logging.StreamHandler(stream=sys.stdout) - stream_handler.setFormatter(stream_formatter) - stream_handler._name = "stream_handler" - _logger.addHandler(stream_handler) - - # add file_handler, output to log_file(if specified), only for rank 0 device - if log_file is not None and dist.get_rank() == 0: - log_file_folder = os.path.dirname(log_file) - if len(log_file_folder): - os.makedirs(log_file_folder, exist_ok=True) - file_formatter = logging.Formatter( - "[%(asctime)s] %(name)s %(levelname)s: %(message)s", - datefmt="%Y/%m/%d %H:%M:%S", - ) - file_handler = logging.FileHandler(log_file, "a") # append mode - file_handler.setFormatter(file_formatter) - file_handler._name = "file_handler" - _logger.addHandler(file_handler) - - if dist.get_rank() == 0: - _logger.setLevel(log_level) - else: - _logger.setLevel(logging.ERROR) - - _logger.propagate = False - - -def set_log_level(log_level: int): - """Set logger level, only message of level >= `log_level` will be printed. - - Built-in log level are below: - - CRITICAL = 50, - FATAL = 50, - ERROR = 40, - WARNING = 30, - WARN = 30, - INFO = 20, - DEBUG = 10, - NOTSET = 0. - - Args: - log_level (int): Log level. - """ - if dist.get_rank() == 0: - _logger.setLevel(log_level) - else: - _logger.setLevel(logging.ERROR) - - -def ensure_logger(log_func: Callable) -> Callable: - """ - A decorator which automatically initialize `logger` by default arguments - when init_logger() is not called manually. - """ - - @functools.wraps(log_func) - def wrapped_log_func(msg, *args): - if _logger is None: - init_logger() - _logger.warning( - "Logger has already been automatically initialized as `log_file` is " - "set to None by default, information will only be printed to terminal " - "without writing to any file." - ) - - log_func(msg, *args) - - return wrapped_log_func - - -@ensure_logger -@misc.run_at_rank0 -def info(msg, *args): - _logger.info(msg, *args) - - -@ensure_logger -@misc.run_at_rank0 -def message(msg, *args): - _logger.log(_MESSAGE_LEVEL, msg, *args) - - -@ensure_logger -@misc.run_at_rank0 -def debug(msg, *args): - _logger.debug(msg, *args) - - -@ensure_logger -@misc.run_at_rank0 -def warning(msg, *args): - _logger.warning(msg, *args) - - -@ensure_logger -@misc.run_at_rank0 -def error(msg, *args): - _logger.error(msg, *args) - - -def scalar( - metric_dict: Dict[str, float], - step: int, - vdl_writer: Optional["visualdl.LogWriter"] = None, - wandb_writer: Optional["wandb.run"] = None, - tbd_writer: Optional["tbd.SummaryWriter"] = None, -): - """This function will add scalar data to VisualDL or WandB for plotting curve(s). - - Args: - metric_dict (Dict[str, float]): Metrics dict with metric name and value. - step (int): The step of the metric. - vdl_writer (Optional[visualdl.LogWriter]): VisualDL writer to record metrics. Defaults to None. - wandb_writer (Optional[wandb.run]): Run object of WandB to record metrics. Defaults to None. - tbd_writer (Optional[tbd.SummaryWriter]): Run object of WandB to record metrics. Defaults to None. - """ - if vdl_writer is not None: - with misc.RankZeroOnly() as is_master: - if is_master: - for name, value in metric_dict.items(): - vdl_writer.add_scalar(name, value, step) - - if wandb_writer is not None: - with misc.RankZeroOnly() as is_master: - if is_master: - wandb_writer.log({"step": step, **metric_dict}) - - if tbd_writer is not None: - with misc.RankZeroOnly() as is_master: - if is_master: - for name, value in metric_dict.items(): - tbd_writer.add_scalar(name, value, global_step=step) - - -def advertise(): - """ - Show the advertising message like the following: - - =========================================================== - == PaddleScience is powered by PaddlePaddle ! == - =========================================================== - == == - == For more info please go to the following website. == - == == - == https://github.com/PaddlePaddle/PaddleScience == - =========================================================== - """ - - _copyright = "PaddleScience is powered by PaddlePaddle !" - ad = "Please refer to the following website for more info." - website = "https://github.com/PaddlePaddle/PaddleScience" - AD_LEN = 6 + len(max([_copyright, ad, website], key=len)) - - info( - "\n{0}\n{1}\n{2}\n{3}\n{4}\n{5}\n{6}\n{7}\n".format( - "=" * (AD_LEN + 4), - "=={}==".format(_copyright.center(AD_LEN)), - "=" * (AD_LEN + 4), - "=={}==".format(" " * AD_LEN), - "=={}==".format(ad.center(AD_LEN)), - "=={}==".format(" " * AD_LEN), - "=={}==".format(website.center(AD_LEN)), - "=" * (AD_LEN + 4), - ) - ) diff --git a/examples/smc_reac/ppsci/utils/misc.py b/examples/smc_reac/ppsci/utils/misc.py deleted file mode 100644 index 7874290d4e..0000000000 --- a/examples/smc_reac/ppsci/utils/misc.py +++ /dev/null @@ -1,684 +0,0 @@ -# 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. - -from __future__ import annotations - -import collections -import functools -import os -import random -import time -from contextlib import ContextDecorator -from typing import Callable -from typing import Dict -from typing import List -from typing import Optional -from typing import Sequence -from typing import Tuple -from typing import Union - -import numpy as np -import paddle -from matplotlib import pyplot as plt -from paddle import distributed as dist - -from ppsci.utils import logger - -__all__ = [ - "AverageMeter", - "PrettyOrderedDict", - "Prettydefaultdict", - "RankZeroOnly", - "RankZeroFirst", - "Timer", - "all_gather", - "concat_dict_list", - "convert_to_array", - "convert_to_dict", - "stack_dict_list", - "cartesian_product", - "combine_array_with_time", - "set_random_seed", - "run_on_eval_mode", - "run_at_rank0", - "plot_curve", - "check_flag_enabled", -] - - -class AverageMeter: - """ - Computes and stores the average and current value - Code was based on https://github.com/pytorch/examples/blob/master/imagenet/main.py - """ - - def __init__(self, name="", fmt="f", postfix="", need_avg=True): - self.name = name - self.fmt = fmt - self.postfix = postfix - self.need_avg = need_avg - self.reset() - - def reset(self): - """Reset.""" - self.val = 0 - self.avg = 0 - self.sum = 0 - self.count = 0 - self.history = [] - - def update(self, val, n=1): - """Update.""" - self.val = val - self.sum += val * n - self.count += n - self.avg = self.sum / self.count - self.history.append(val) - - @property - def avg_info(self): - if isinstance(self.avg, paddle.Tensor): - self.avg = float(self.avg) - return f"{self.name}: {self.avg:.5f}" - - @property - def total(self): - return f"{self.name}_sum: {self.sum:{self.fmt}}{self.postfix}" - - @property - def total_minute(self): - return f"{self.name} {self.sum / 60:{self.fmt}}{self.postfix} min" - - @property - def mean(self): - return ( - f"{self.name}: {self.avg:{self.fmt}}{self.postfix}" if self.need_avg else "" - ) - - @property - def value(self): - return f"{self.name}: {self.val:{self.fmt}}{self.postfix}" - - -class PrettyOrderedDict(collections.OrderedDict): - """ - The ordered dict which can be prettily printed. - - Examples: - >>> import ppsci - >>> dic = ppsci.utils.misc.PrettyOrderedDict() - >>> dic.update({'a':1, 'b':2, 'c':3}) - >>> print(dic) - ('a', 1)('b', 2)('c', 3) - """ - - def __str__(self): - return "".join([str((k, v)) for k, v in self.items()]) - - -class Prettydefaultdict(collections.defaultdict): - """ - The default dict which can be prettily printed. - - Examples: - >>> import ppsci - >>> dic = ppsci.utils.misc.Prettydefaultdict() - >>> dic.update({'a':1, 'b':2, 'c':3}) - >>> print(dic) - ('a', 1)('b', 2)('c', 3) - """ - - def __str__(self): - return "".join([str((k, v)) for k, v in self.items()]) - - -class RankZeroOnly: - """ - A context manager that ensures the code inside it is only executed by the process - with rank zero. All rank will be synchronized by `dist.barrier` in - distributed environment. - - NOTE: Always used for time consuming code blocks, such as initialization of log - writer, saving result to disk, etc. - - Args: - rank (Optional[int]): The rank of the current process. If not provided, - it will be obtained from `dist.get_rank()`. - - Examples: - >>> import paddle.distributed as dist - >>> with RankZeroOnly(dist.get_rank()) as is_master: - ... if is_master: - ... # code here which should only be executed in the master process - ... pass - """ - - def __init__(self, rank: Optional[int] = None): - """ - Enter the context and check if the current process is the master. - - Args: - rank (Optional[int]): The rank of the current process. If not provided, - it will be obtained from `dist.get_rank()`. - """ - super().__init__() - self.rank = rank if (rank is not None) else dist.get_rank() - self.is_master = self.rank == 0 - - def __enter__(self) -> bool: - """ - Enter the context and check if the current process is the master. - - Returns: - bool: True if the current process is the master (rank zero), False otherwise. - """ - return self.is_master - - def __exit__(self, exc_type, exc_value, traceback): - if dist.get_world_size() > 1: - dist.barrier() - - -class RankZeroFirst(ContextDecorator): - """ - A context manager that ensures the code inside it is only executed by the process - with rank zero first. All ranks will be synchronized by `dist.barrier()`. - - Args: - rank (Optional[int]): The rank of the current process. If not provided, - it will be obtained from `dist.get_rank()`. - - Examples: - >>> import paddle.distributed as dist - >>> with RankZeroFirst(dist.get_rank()): - ... # code here which should be executed first in the master(rank-0) process - ... pass - """ - - def __init__(self, rank: Optional[int] = None): - if dist.is_initialized(): - self.rank = rank if rank is not None else dist.get_rank() - self.world_size = dist.get_world_size() - else: - self.rank = 0 - self.world_size = 1 - self.is_master = self.rank == 0 - - def __enter__(self): - if self.world_size > 1 and not self.is_master: - dist.barrier() # Non-master processs wait for master to finish - - def __exit__(self, type, value, traceback): - if self.world_size > 1 and self.is_master: - dist.barrier() # Allow others to proceed - - -class Timer(ContextDecorator): - """Count time cost for code block within context. - - Args: - name (str, optional): Name of timer discriminate different code block. - Defaults to "Timer". - auto_print (bool, optional): Whether print time cost when exit context. - Defaults to True. - - Examples: - >>> import paddle - >>> from ppsci.utils import misc - >>> with misc.Timer("test1", auto_print=False) as timer: - ... w = sum(range(0, 10)) - >>> print(f"time cost of 'sum(range(0, 10))' is {timer.interval:.2f}") # doctest: +SKIP - time cost of 'sum(range(0, 10))' is 0.00 - - >>> @misc.Timer("test2", auto_print=True) - ... def func(): - ... w = sum(range(0, 10)) - >>> func() # doctest: +SKIP - - >>> timer = misc.Timer("cost_of_func", auto_print=False) - >>> timer.start() - >>> def func(): - ... w = sum(range(0, 10)) - >>> func() - >>> timer.end() - >>> print(f"time cost of 'cost_of_func' is {timer.interval:.2f}") # doctest: +SKIP - time cost of 'cost_of_func' is 0.00 - """ - - interval: float # Time cost for code within Timer context - - def __init__(self, name: str = "Timer", auto_print: bool = True): - super().__init__() - self.name = name - self.auto_print = auto_print - - def __enter__(self): - paddle.device.synchronize() - self.start_time = time.perf_counter() - return self - - def __exit__(self, type, value, traceback): - paddle.device.synchronize() - self.end_time = time.perf_counter() - self.interval = self.end_time - self.start_time - if self.auto_print: - logger.message(f"{self.name}.time_cost = {self.interval * 1000:.2f} ms") - - def start(self, name: str = "Timer"): - """Push a new timer context. - - Args: - name (str, optional): Name of code block to be clocked. Defaults to "Timer". - """ - paddle.device.synchronize() - self.start_time = time.perf_counter() - - def end(self): - """End current timer context and print time cost.""" - paddle.device.synchronize() - self.end_time = time.perf_counter() - self.interval = self.end_time - self.start_time - if self.auto_print: - logger.message(f"{self.name}.time_cost = {self.interval:.2f} s") - - -def convert_to_dict(array: np.ndarray, keys: Tuple[str, ...]) -> Dict[str, np.ndarray]: - """Split given array into single channel array at axis -1 in order of given keys. - - Args: - array (np.ndarray): Array to be split. - keys (Tuple[str, ...]): Keys used in split. - - Returns: - Dict[str, np.ndarray]: Split dict. - - Examples: - >>> import numpy as np - >>> import ppsci - >>> arr = np.array([[1., 2., 3.], [4., 5., 6.]]) - >>> result = ppsci.utils.misc.convert_to_dict(arr, ("x", "y", "z")) - >>> print(arr.shape) - (2, 3) - >>> for k, v in result.items(): - ... print(k, v.shape) - x (2, 1) - y (2, 1) - z (2, 1) - """ - if array.shape[-1] != len(keys): - raise ValueError( - f"dim of array({array.shape[-1]}) must equal to " f"len(keys)({len(keys)})" - ) - - split_array = np.split(array, len(keys), axis=-1) - return {key: split_array[i] for i, key in enumerate(keys)} - - -def all_gather( - tensor: paddle.Tensor, concat: bool = True, axis: int = 0 -) -> Union[paddle.Tensor, List[paddle.Tensor]]: - """Gather tensor from all devices, concatenate them along given axis if specified. - - Args: - tensor (paddle.Tensor): Tensor to be gathered from all GPUs. - concat (bool, optional): Whether to concatenate gathered Tensors. Defaults to True. - axis (int, optional): Axis which concatenated along. Defaults to 0. - - Returns: - Union[paddle.Tensor, List[paddle.Tensor]]: Gathered Tensors. - - Examples: - >>> import paddle - >>> import ppsci - >>> import paddle.distributed as dist - >>> dist.init_parallel_env() # doctest: +SKIP - >>> if dist.get_rank() == 0: # doctest: +SKIP - ... data = paddle.to_tensor([[1, 2, 3], [4, 5, 6]]) - ... else: - ... data = paddle.to_tensor([[7, 8, 9], [10, 11, 12]]) - >>> result = ppsci.utils.misc.all_gather(data) # doctest: +SKIP - >>> print(result.numpy()) # doctest: +SKIP - [[ 1 2 3] - [ 4 5 6] - [ 7 8 9] - [10 11 12]] - """ - result: List[paddle.Tensor] = [] - - # NOTE: Put tensor to CUDAPlace from CUDAPinnedPlace to use communication. - if tensor.place.is_cuda_pinned_place(): - tensor = tensor.cuda() - - # TODO(HydrogenSulfate): As non-contiguous(strided) tensor is not supported in - # dist.all_gather, manually convert given Tensor to contiguous below. Strided tensor - # will be supported in future. - dist.all_gather(result, tensor.contiguous()) - - if concat: - return paddle.concat(result, axis) - return result - - -def convert_to_array(dict_: Dict[str, np.ndarray], keys: Tuple[str, ...]) -> np.ndarray: - """Concatenate arrays in axis -1 in order of given keys. - - Args: - dict_ (Dict[str, np.ndarray]): Dict contains arrays. - keys (Tuple[str, ...]): Concatenate keys used in concatenation. - - Returns: - np.ndarray: Concatenated array. - - Examples: - >>> import numpy as np - >>> import ppsci - >>> dic = {"x": np.array([[1., 2.], [3., 4.]]), - ... "y": np.array([[5., 6.], [7., 8.]]), - ... "z": np.array([[9., 10.], [11., 12.]])} - >>> result = ppsci.utils.misc.convert_to_array(dic, ("x", "z")) - >>> print(result) - [[ 1. 2. 9. 10.] - [ 3. 4. 11. 12.]] - """ - return np.concatenate([dict_[key] for key in keys], axis=-1) - - -def concat_dict_list( - dict_list: Sequence[Dict[str, np.ndarray]] -) -> Dict[str, np.ndarray]: - """Concatenate arrays in tuple of dicts at axis 0. - - Args: - dict_list (Sequence[Dict[str, np.ndarray]]): Sequence of dicts. - - Returns: - Dict[str, np.ndarray]: A dict with concatenated arrays for each key. - - Examples: - >>> import numpy as np - >>> import ppsci - >>> dic1 = {"x": np.array([[1., 2.], [3., 4.]]), "y": np.array([[5., 6.], [7., 8.]])} - >>> dic2 = {"x": np.array([[1., 2.], [3., 4.]]), "y": np.array([[5., 6.], [7., 8.]])} - >>> result = ppsci.utils.misc.concat_dict_list((dic1, dic2)) - >>> print(result) - {'x': array([[1., 2.], - [3., 4.], - [1., 2.], - [3., 4.]]), 'y': array([[5., 6.], - [7., 8.], - [5., 6.], - [7., 8.]])} - """ - ret = {} - for key in dict_list[0].keys(): - ret[key] = np.concatenate([_dict[key] for _dict in dict_list], axis=0) - return ret - - -def stack_dict_list( - dict_list: Sequence[Dict[str, np.ndarray]] -) -> Dict[str, np.ndarray]: - """Stack arrays in tuple of dicts at axis 0. - - Args: - dict_list (Sequence[Dict[str, np.ndarray]]): Sequence of dicts. - - Returns: - Dict[str, np.ndarray]: A dict with stacked arrays for each key. - - Examples: - >>> import numpy as np - >>> import ppsci - >>> dic1 = {"x": np.array([[1., 2.], [3., 4.]]), "y": np.array([[5., 6.], [7., 8.]])} - >>> dic2 = {"x": np.array([[1., 2.], [3., 4.]]), "y": np.array([[5., 6.], [7., 8.]])} - >>> result = ppsci.utils.misc.stack_dict_list((dic1, dic2)) - >>> for k, v in result.items(): - ... print(k, v.shape) - x (2, 2, 2) - y (2, 2, 2) - """ - ret = {} - for key in dict_list[0].keys(): - ret[key] = np.stack([_dict[key] for _dict in dict_list], axis=0) - return ret - - -def typename(obj: object) -> str: - """Return type name of given object. - - Args: - obj (object): Python object which is instantiated from a class. - - Returns: - str: Class name of given object. - """ - return obj.__class__.__name__ - - -def combine_array_with_time(x: np.ndarray, t: Tuple[int, ...]) -> np.ndarray: - """Combine given data x with time sequence t. - Given x with shape (N, D) and t with shape (T, ), - this function will repeat t_i for N times and will concat it with data x for each t_i in t, - finally return the stacked result, which is of shape (N×T, D+1). - - Args: - x (np.ndarray): Points data with shape (N, D). - t (Tuple[int, ...]): Time sequence with shape (T, ). - - Returns: - np.ndarray: Combined data with shape of (N×T, D+1). - - Examples: - >>> import numpy as np - >>> import ppsci - >>> data_point = np.arange(10).reshape((2, 5)) - >>> time = (1, 2, 3) - >>> result = ppsci.utils.misc.combine_array_with_time(data_point, time) - >>> print(result) - [[1. 0. 1. 2. 3. 4.] - [1. 5. 6. 7. 8. 9.] - [2. 0. 1. 2. 3. 4.] - [2. 5. 6. 7. 8. 9.] - [3. 0. 1. 2. 3. 4.] - [3. 5. 6. 7. 8. 9.]] - """ - nx = len(x) - tx = [] - for ti in t: - tx.append( - np.hstack( - (np.full([nx, 1], float(ti), dtype=paddle.get_default_dtype()), x) - ) - ) - tx = np.vstack(tx) - return tx - - -def cartesian_product(*arrays: np.ndarray) -> np.ndarray: - """Cartesian product for input sequence of array(s). - - Reference: https://stackoverflow.com/questions/11144513/cartesian-product-of-x-and-y-array-points-into-single-array-of-2d-points - - Assume shapes of input arrays are: $(N_1,), (N_2,), (N_3,), ..., (N_M,)$, - then the cartesian product result will be shape of $(N_1xN_2xN_3x...xN_M, M)$. - - Args: - arrays (np.ndarray): Input arrays. - - Returns: - np.ndarray: Cartesian product result of shape $(N_1xN_2xN_3x...xN_M, M)$. - - Examples: - >>> t = np.array([1, 2]) - >>> x = np.array([10, 20]) - >>> y = np.array([100, 200]) - >>> txy = cartesian_product(t, x, y) - >>> print(txy) - [[ 1 10 100] - [ 1 10 200] - [ 1 20 100] - [ 1 20 200] - [ 2 10 100] - [ 2 10 200] - [ 2 20 100] - [ 2 20 200]] - """ - la = len(arrays) - dtype = np.result_type(*arrays) - arr = np.empty([len(a) for a in arrays] + [la], dtype=dtype) - for i, a in enumerate(np.ix_(*arrays)): - arr[..., i] = a - return arr.reshape(-1, la) - - -def set_random_seed(seed: int): - """Set numpy, random, paddle random_seed to given seed. - - Args: - seed (int): Random seed. - """ - paddle.seed(seed) - np.random.seed(seed) - random.seed(seed) - - -def run_on_eval_mode(func: Callable) -> Callable: - """A decorator automatically running given class method in eval mode and keep - training state unchanged after function finished. - - Args: - func (Callable): Class method which is expected running in eval mode. - - Returns: - Callable: Decorated class method. - """ - - @functools.wraps(func) - def function_with_eval_state(self, *args, **kwargs): - # log original state - train_state = self.model.training - - # switch to eval mode - if train_state: - self.model.eval() - - # run func in eval mode - result = func(self, *args, **kwargs) - - # restore state - if train_state: - self.model.train() - - return result - - return function_with_eval_state - - -def run_at_rank0(func: Callable) -> Callable: - """A decorator that allow given function run only at rank 0 to avoid - multiple logs or other events. Usually effected in distributed environment. - - Args: - func (Callable): Given function. - - Returns: - Callable: Wrapped function which will only run at at rank 0, - skipped at other rank. - - Examples: - >>> import paddle - >>> from ppsci.utils import misc - >>> @misc.run_at_rank0 - ... def func(): - ... print(f"now_rank is {paddle.distributed.get_rank()}") - >>> func() - now_rank is 0 - """ - - @functools.wraps(func) - def wrapped_func(*args, **kwargs): - if dist.get_rank() == 0: - return func(*args, **kwargs) - - return wrapped_func - - -def plot_curve( - data: Dict[str, List], - xlabel: str = "X", - ylabel: str = "Y", - output_dir: str = "./output/", - smooth_step: int = 1, - use_semilogy: bool = False, -) -> None: - """Plotting curve. - - Args: - data (Dict[str, List]): Dict of all data, keys are curves' name. - xlabel (str, optional): Label of x-axis. Defaults to "X". - ylabel (str, optional): Label of y-axis. Defaults to "Y". - output_dir (str, optional): Output directory of figure. Defaults to "./output/". - smooth_step (int, optional): How many points are squeezed to one point to smooth the curve. Defaults to 1. - use_semilogy (bool, optional): Whether to set non-uniform coordinates for the y-axis. Defaults to False. - """ - data_arr = np.concatenate( - [np.asarray(arr).reshape(-1, 1) for arr in data.values()], axis=1 - ) - - # smooth - if data_arr.shape[0] % smooth_step != 0: - data_arr = np.reshape( - data_arr[: -(data_arr.shape[0] % smooth_step), :], - (-1, smooth_step, data_arr.shape[1]), - ) - else: - data_arr = np.reshape(data_arr, (-1, smooth_step, data_arr.shape[1])) - data_arr = np.mean(data_arr, axis=1) - - # plot - plt.figure() - if use_semilogy: - plt.yscale("log") - plt.xscale("log") - plt.plot(np.arange(data_arr.shape[0]) * smooth_step, data_arr) - plt.legend( - list(data.keys()), - loc="upper left", - bbox_to_anchor=(1, 1), - ) - plt.xlabel(xlabel) - plt.ylabel(ylabel) - plt.grid() - plt.yticks(size=10) - plt.xticks(size=10) - plt.tight_layout() - - plt.savefig(os.path.join(output_dir, f"{xlabel}-{ylabel}_curve.jpg"), dpi=200) - plt.clf() - plt.close() - - -def check_flag_enabled(flag_name: str) -> bool: - """Check whether the flag is enabled. - - Args: - flag_name (str): Flag name to be checked whether enabled or disabled. - - Returns: - bool: Whether given flag name is enabled in environment. - """ - value = os.getenv(flag_name, False) - if isinstance(value, str): - return value.lower() in ["true", "1"] - return False diff --git a/examples/smc_reac/ppsci/utils/reader.py b/examples/smc_reac/ppsci/utils/reader.py deleted file mode 100644 index ef0fb8f191..0000000000 --- a/examples/smc_reac/ppsci/utils/reader.py +++ /dev/null @@ -1,266 +0,0 @@ -# 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. - -from __future__ import annotations - -import collections -import csv -import pickle -from typing import Dict -from typing import Optional -from typing import Tuple - -import meshio -import numpy as np -import paddle -import scipy.io as sio - -__all__ = [ - "load_csv_file", - "load_mat_file", - "load_npz_file", - "load_vtk_file", - "load_vtk_with_time_file", - "load_dat_file", -] - - -def load_csv_file( - file_path: str, - keys: Tuple[str, ...], - alias_dict: Optional[Dict[str, str]] = None, - delimiter: str = ",", - encoding: str = "utf-8", -) -> Dict[str, np.ndarray]: - """Load *.csv file and fetch data as given keys. - - Args: - file_path (str): CSV file path. - keys (Tuple[str, ...]): Required fetching keys. - alias_dict (Optional[Dict[str, str]]): Alias for keys, - i.e. {inner_key: outer_key}. Defaults to None. - encoding (str, optional): Encoding code when open file. Defaults to "utf-8". - - Returns: - Dict[str, np.ndarray]: Loaded data in dict. - """ - if alias_dict is None: - alias_dict = {} - - try: - # read all data from csv file - with open(file_path, "r", encoding=encoding) as csv_file: - reader = csv.DictReader(csv_file, delimiter=delimiter) - raw_data = collections.defaultdict(list) - for _, line_dict in enumerate(reader): - for key, value in line_dict.items(): - raw_data[key].append(value) - except FileNotFoundError as e: - raise e - - # convert to numpy array - data_dict = {} - for key in keys: - fetch_key = alias_dict[key] if key in alias_dict else key - if fetch_key not in raw_data: - raise KeyError(f"fetch_key({fetch_key}) do not exist in raw_data.") - data_dict[key] = np.asarray(raw_data[fetch_key]) - if not np.issubdtype(data_dict[key].dtype, np.integer): - data_dict[key] = data_dict[key].astype(paddle.get_default_dtype()) - data_dict[key] = data_dict[key].reshape([-1, 1]) - - return data_dict - - -def load_mat_file( - file_path: str, keys: Tuple[str, ...], alias_dict: Optional[Dict[str, str]] = None -) -> Dict[str, np.ndarray]: - """Load *.mat file and fetch data as given keys. - - Args: - file_path (str): Mat file path. - keys (Tuple[str, ...]): Required fetching keys. - alias_dict (Optional[Dict[str, str]]): Alias for keys, - i.e. {original_key: original_key}. Defaults to None. - - Returns: - Dict[str, np.ndarray]: Loaded data in dict. - """ - - if alias_dict is None: - alias_dict = {} - - try: - # read all data from mat file - raw_data = sio.loadmat(file_path) - except FileNotFoundError as e: - raise e - - # convert to numpy array - data_dict = {} - for key in keys: - fetch_key = alias_dict[key] if key in alias_dict else key - if fetch_key not in raw_data: - raise KeyError(f"fetch_key({fetch_key}) do not exist in raw_data.") - data_dict[key] = np.asarray(raw_data[fetch_key]) - if not np.issubdtype(data_dict[key].dtype, np.integer): - data_dict[key] = data_dict[key].astype(paddle.get_default_dtype()) - data_dict[key] = data_dict[key].reshape([-1, 1]) - - return data_dict - - -def load_npz_file( - file_path: str, keys: Tuple[str, ...], alias_dict: Optional[Dict[str, str]] = None -) -> Dict[str, np.ndarray]: - """Load *.npz file and fetch data as given keys. - - Args: - file_path (str): Npz file path. - keys (Tuple[str, ...]): Required fetching keys. - alias_dict (Optional[Dict[str, str]]): Alias for keys, - i.e. {original_key: original_key}. Defaults to None. - - Returns: - Dict[str, np.ndarray]: Loaded data in dict. - """ - - if alias_dict is None: - alias_dict = {} - - try: - # read all data from npz file - raw_data = np.load(file_path, allow_pickle=True) - except FileNotFoundError as e: - raise e - - # convert to numpy array - data_dict = {} - for key in keys: - fetch_key = alias_dict[key] if key in alias_dict else key - if fetch_key not in raw_data: - raise KeyError(f"fetch_key({fetch_key}) do not exist in raw_data.") - data_dict[key] = np.asarray(raw_data[fetch_key]) - if data_dict[key].dtype in (np.float16, np.float32, np.float64): - data_dict[key] = data_dict[key].astype(paddle.get_default_dtype()) - - return data_dict - - -def load_vtk_file( - filename_without_timeid: str, - time_step: float, - time_index: Tuple[int, ...], - input_keys: Tuple[str, ...], - label_keys: Optional[Tuple[str, ...]], -) -> Dict[str, np.ndarray]: - """Load coordinates and attached label from the *.vtu file. - - Args: - filename_without_timeid (str): File name without time id. - time_step (float): Physical time step. - time_index (Tuple[int, ...]): Physical time indexes. - input_keys (Tuple[str, ...]): Input coordinates name keys. - label_keys (Optional[Tuple[str, ...]]): Input label name keys. - - Returns: - Dict[str, np.ndarray]: Input coordinates dict, label coordinates dict - """ - input_dict = {var: [] for var in input_keys} - label_dict = {var: [] for var in label_keys} - for index in time_index: - file = filename_without_timeid + f"{index}.vtu" - mesh = meshio.read(file) - n = mesh.points.shape[0] - i = 0 - for key in input_dict: - if key == "t": - input_dict[key].append( - np.full((n, 1), index * time_step, paddle.get_default_dtype()) - ) - else: - input_dict[key].append( - mesh.points[:, i].reshape(n, 1).astype(paddle.get_default_dtype()) - ) - i += 1 - for i, key in enumerate(label_dict): - label_dict[key].append( - np.array(mesh.point_data[key], paddle.get_default_dtype()) - ) - for key in input_dict: - input_dict[key] = np.concatenate(input_dict[key]) - for key in label_dict: - label_dict[key] = np.concatenate(label_dict[key]) - - return input_dict, label_dict - - -def load_vtk_with_time_file(file: str) -> Dict[str, np.ndarray]: - """Temporary interface for points cloud, will be banished sooner. - - Args: - file (str): Input file name. - - Returns: - Dict[str, np.ndarray]: Input coordinates dict. - """ - mesh = meshio.read(file) - n = mesh.points.shape[0] - t = np.array(mesh.point_data["time"]) - x = mesh.points[:, 0].reshape(n, 1) - y = mesh.points[:, 1].reshape(n, 1) - z = mesh.points[:, 2].reshape(n, 1) - input_dict = {"t": t, "x": x, "y": y, "z": z} - return input_dict - - -def load_dat_file( - file_path: str, - keys: Tuple[str, ...] = None, - alias_dict: Optional[Dict[str, str]] = None, -) -> Dict[str, np.ndarray]: - """Load *.dat file and fetch data as given keys. - - Args: - file_path (str): Dat file path. - keys (Tuple[str, ...]): Required fetching keys. - alias_dict (Optional[Dict[str, str]]): Alias for keys, - i.e. {original_key: original_key}. Defaults to None. - - Returns: - Dict[str, np.ndarray]: Loaded data in dict. - """ - - if alias_dict is None: - alias_dict = {} - - try: - # read all data from .dat file - raw_data = pickle.load(open(file_path, "rb")) - except FileNotFoundError as e: - raise e - - # convert to numpy array - data_dict = {} - if keys is None: - keys = raw_data.keys() - for key in keys: - fetch_key = alias_dict[key] if key in alias_dict else key - if fetch_key not in raw_data: - raise KeyError(f"fetch_key({fetch_key}) do not exist in raw_data.") - data_dict[key] = np.asarray(raw_data[fetch_key]) - if data_dict[key].dtype in (np.float16, np.float32, np.float64): - data_dict[key] = data_dict[key].astype(paddle.get_default_dtype()) - - return data_dict diff --git a/examples/smc_reac/ppsci/utils/save_load.py b/examples/smc_reac/ppsci/utils/save_load.py deleted file mode 100644 index cdf14cce87..0000000000 --- a/examples/smc_reac/ppsci/utils/save_load.py +++ /dev/null @@ -1,300 +0,0 @@ -# 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. - -from __future__ import annotations - -import os -from typing import TYPE_CHECKING -from typing import Any -from typing import Dict -from typing import Optional - -import paddle - -from ppsci.utils import download -from ppsci.utils import logger - -if TYPE_CHECKING: - from paddle import amp - from paddle import nn - from paddle import optimizer - - from ppsci import equation - from ppsci.loss import mtl - from ppsci.utils import ema - - -__all__ = [ - "load_checkpoint", - "save_checkpoint", - "load_pretrain", -] - - -def _load_pretrain_from_path( - path: str, - model: nn.Layer, - equation: Optional[Dict[str, equation.PDE]] = None, -): - """Load pretrained model from given path. - - Args: - path (str): File path of pretrained model, i.e. `/path/to/model.pdparams`. - model (nn.Layer): Model with parameters. - equation (Optional[Dict[str, equation.PDE]]): Equations. Defaults to None. - """ - if not (os.path.isdir(path) or os.path.exists(f"{path}.pdparams")): - raise FileNotFoundError( - f"Pretrained model path {path}.pdparams does not exists." - ) - - param_state_dict = paddle.load(f"{path}.pdparams") - model.set_state_dict(param_state_dict) - logger.message(f"Finish loading pretrained model from: {path}.pdparams") - if equation is not None: - if not os.path.exists(f"{path}.pdeqn"): - num_learnable_params = sum( - [len(eq.learnable_parameters) for eq in equation.values()] - ) - if num_learnable_params > 0: - logger.warning( - f"There are a total of {num_learnable_params} learnable parameters" - f" in the equation, but {path}.pdeqn not found." - ) - else: - equation_dict = paddle.load(f"{path}.pdeqn") - for name, _equation in equation.items(): - _equation.set_state_dict(equation_dict[name]) - logger.message( - f"Finish loading pretrained equation parameters from: {path}.pdeqn" - ) - - -def load_pretrain( - model: nn.Layer, - path: str, - equation: Optional[Dict[str, equation.PDE]] = None, -): - """ - Load pretrained model from given path or url. - - Args: - model (nn.Layer): Model with parameters. - path (str): File path or url of pretrained model, i.e. `/path/to/model.pdparams` - or `http://xxx.com/model.pdparams`. - equation (Optional[Dict[str, equation.PDE]]): Equations. Defaults to None. - - Examples: - >>> import ppsci - >>> from ppsci.utils import save_load - >>> model = ppsci.arch.MLP(("x", "y"), ("u", "v", "p"), 9, 50, "tanh") - >>> save_load.load_pretrain( - ... model=model, - ... path="path/to/pretrain_model") # doctest: +SKIP - """ - if path.startswith("http"): - # download from path(url) and get its' physical path - eqn_path = path.replace(".pdparams", ".pdeqn", 1) - path = download.get_weights_path_from_url(path) - - # automatically download additional equation weights if available - def is_url_accessible(url: str): - try: - import requests - - response = requests.head(url, timeout=5) - return response.status_code == requests.codes.ok - except requests.RequestException: - return False - except Exception: - return False - - if is_url_accessible(eqn_path): - download.get_weights_path_from_url(eqn_path) - - # remove ".pdparams" in suffix of path for convenient - if path.endswith(".pdparams"): - path = path[:-9] - _load_pretrain_from_path(path, model, equation) - - -def load_checkpoint( - path: str, - model: nn.Layer, - optimizer: optimizer.Optimizer, - grad_scaler: Optional[amp.GradScaler] = None, - equation: Optional[Dict[str, equation.PDE]] = None, - ema_model: Optional[ema.AveragedModel] = None, - aggregator: Optional[mtl.LossAggregator] = None, -) -> Dict[str, Any]: - """Load from checkpoint. - - Args: - path (str): Path for checkpoint. - model (nn.Layer): Model with parameters. - optimizer (optimizer.Optimizer): Optimizer for model. - grad_scaler (Optional[amp.GradScaler]): GradScaler for AMP. Defaults to None. - equation (Optional[Dict[str, equation.PDE]]): Equations. Defaults to None. - ema_model: Optional[ema.AveragedModel]: Average model. Defaults to None. - aggregator: Optional[mtl.LossAggregator]: Loss aggregator. Defaults to None. - - Returns: - Dict[str, Any]: Loaded metric information. - """ - if not os.path.exists(f"{path}.pdparams"): - raise FileNotFoundError(f"{path}.pdparams not exist.") - if not os.path.exists(f"{path}.pdopt"): - raise FileNotFoundError(f"{path}.pdopt not exist.") - if grad_scaler is not None and not os.path.exists(f"{path}.pdscaler"): - raise FileNotFoundError(f"{path}.scaler not exist.") - - # load state dict - model_dict = paddle.load(f"{path}.pdparams") - optim_dict = paddle.load(f"{path}.pdopt") - metric_dict = {} - if os.path.exists(f"{path}.pdstates"): - metric_dict = paddle.load(f"{path}.pdstates") - if grad_scaler is not None: - scaler_dict = paddle.load(f"{path}.pdscaler") - if equation is not None: - if not os.path.exists(f"{path}.pdeqn"): - logger.warning(f"{path}.pdeqn not found.") - equation_dict = None - else: - equation_dict = paddle.load(f"{path}.pdeqn") - - # set model state dict - logger.message(f"* Loading model checkpoint from {path}.pdparams") - missing_keys, unexpected_keys = model.set_state_dict(model_dict) - if missing_keys: - logger.warning( - f"There are missing keys when loading checkpoint: {missing_keys}, " - "and corresponding parameters will be initialized by default." - ) - if unexpected_keys: - logger.warning( - f"There are redundant keys: {unexpected_keys}, " - "and corresponding weights will be ignored." - ) - - # set optimizer state dict - logger.message(f"* Loading optimizer checkpoint from {path}.pdopt") - optimizer.set_state_dict(optim_dict) - - if grad_scaler is not None: - logger.message(f"* Loading grad scaler checkpoint from {path}.pdscaler") - grad_scaler.load_state_dict(scaler_dict) - - if equation is not None and equation_dict is not None: - logger.message(f"* Loading equation checkpoint from {path}.pdeqn") - for name, _equation in equation.items(): - _equation.set_state_dict(equation_dict[name]) - - if ema_model is not None: - logger.message(f"* Loading EMA checkpoint from {path}_ema.pdparams") - avg_model_dict = paddle.load(f"{path}_ema.pdparams") - ema_model.set_state_dict(avg_model_dict) - - if aggregator is not None and aggregator.should_persist: - logger.message(f"* Loading loss aggregator checkpoint from {path}.pdagg") - aggregator_dict = paddle.load(f"{path}.pdagg") - aggregator.set_state_dict(aggregator_dict) - - logger.message(f"Finish loading checkpoint from {path}") - return metric_dict - - -def save_checkpoint( - model: nn.Layer, - optimizer: Optional[optimizer.Optimizer], - metric: Optional[Dict[str, float]] = None, - grad_scaler: Optional[amp.GradScaler] = None, - output_dir: Optional[str] = None, - prefix: str = "model", - equation: Optional[Dict[str, equation.PDE]] = None, - print_log: bool = True, - ema_model: Optional[ema.AveragedModel] = None, - aggregator: Optional[mtl.LossAggregator] = None, -): - """ - Save checkpoint, including model params, optimizer params, metric information. - - Args: - model (nn.Layer): Model with parameters. - optimizer (Optional[optimizer.Optimizer]): Optimizer for model. - metric (Optional[Dict[str, float]]): Metric information, such as {"RMSE": 0.1, "MAE": 0.2}. Defaults to None. - grad_scaler (Optional[amp.GradScaler]): GradScaler for AMP. Defaults to None. - output_dir (Optional[str]): Directory for checkpoint storage. - prefix (str, optional): Prefix for storage. Defaults to "model". - equation (Optional[Dict[str, equation.PDE]]): Equations. Defaults to None. - print_log (bool, optional): Whether print saving log information, mainly for - keeping log tidy without duplicate 'Finish saving checkpoint ...' log strings. - Defaults to True. - ema_model: Optional[ema.AveragedModel]: Average model. Defaults to None. - aggregator: Optional[mtl.LossAggregator]: Loss aggregator. Defaults to None. - - Examples: - >>> import ppsci - >>> import paddle - >>> from ppsci.utils import save_load - >>> model = ppsci.arch.MLP(("x", "y", "z"), ("u", "v", "w"), 5, 64, "tanh") - >>> optimizer = ppsci.optimizer.Adam(0.001)(model) - >>> save_load.save_checkpoint(model, optimizer, {"RMSE": 0.1}, output_dir="path/to/output/dir") # doctest: +SKIP - """ - if paddle.distributed.get_rank() != 0: - return - - if output_dir is None: - logger.warning("output_dir is None, skip save_checkpoint") - return - - ckpt_dir = os.path.join(output_dir, "checkpoints") - ckpt_path = os.path.join(ckpt_dir, prefix) - os.makedirs(ckpt_dir, exist_ok=True) - - paddle.save(model.state_dict(), f"{ckpt_path}.pdparams") - - if optimizer is not None: - paddle.save(optimizer.state_dict(), f"{ckpt_path}.pdopt") - - if metric is not None and len(metric) > 0: - paddle.save(metric, f"{ckpt_path}.pdstates") - - if grad_scaler is not None: - paddle.save(grad_scaler.state_dict(), f"{ckpt_path}.pdscaler") - - if equation is not None: - num_learnable_params = sum( - [len(eq.learnable_parameters) for eq in equation.values()] - ) - if num_learnable_params > 0: - paddle.save( - {key: eq.state_dict() for key, eq in equation.items()}, - f"{ckpt_path}.pdeqn", - ) - - if ema_model is not None: - paddle.save(ema_model.state_dict(), f"{ckpt_path}_ema.pdparams") - - if aggregator is not None and aggregator.should_persist: - paddle.save(aggregator.state_dict(), f"{ckpt_path}.pdagg") - - if print_log: - log_str = f"Finish saving checkpoint to: {ckpt_path}" - if prefix == "latest": - log_str += ( - "(latest checkpoint will be saved every epoch as expected, " - "but this log will be printed only once for tidy logging)" - ) - logger.message(log_str) diff --git a/examples/smc_reac/ppsci/utils/symbolic.py b/examples/smc_reac/ppsci/utils/symbolic.py deleted file mode 100644 index dcd089d017..0000000000 --- a/examples/smc_reac/ppsci/utils/symbolic.py +++ /dev/null @@ -1,981 +0,0 @@ -# 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. - -""" -Sympy to python function conversion module -""" - -from __future__ import annotations - -import functools -import os -from typing import Dict -from typing import List -from typing import Optional -from typing import Sequence -from typing import Tuple -from typing import Union - -import paddle -import sympy as sp -from paddle import nn -from typing_extensions import TypeAlias - -from ppsci import arch -from ppsci import equation -from ppsci.autodiff import hessian -from ppsci.autodiff import jacobian -from ppsci.utils import logger - -__all__ = [ - "lambdify", - "_cvt_to_key", -] - - -DATA_DICT: TypeAlias = Dict[str, paddle.Tensor] - -SYMPY_BUILTIN_FUNC: TypeAlias = Union[ - sp.sin, - sp.sinh, - sp.asin, - sp.cos, - sp.acos, - sp.cosh, - sp.tan, - sp.atan, - sp.atan2, - sp.acosh, - sp.asinh, - sp.tanh, - sp.atanh, - sp.erf, - sp.loggamma, - sp.exp, - sp.Pow, - sp.log, - sp.Max, - sp.Min, - sp.Abs, - sp.Heaviside, - sp.sign, - sp.ceiling, - sp.floor, - sp.Add, - sp.Mul, -] - -SYMPY_TO_PADDLE = { - sp.sin: paddle.sin, - sp.sinh: paddle.sinh, - sp.asin: paddle.asin, - sp.cos: paddle.cos, - sp.acos: paddle.acos, - sp.cosh: paddle.cosh, - sp.tan: paddle.tan, - sp.atan: paddle.atan, - sp.atan2: paddle.atan2, - sp.acosh: paddle.acosh, - sp.asinh: paddle.asinh, - sp.tanh: paddle.tanh, - sp.atanh: paddle.atanh, - sp.erf: paddle.erf, - sp.loggamma: paddle.lgamma, - sp.exp: paddle.exp, - sp.Pow: paddle.pow, - sp.log: paddle.log, - sp.Max: paddle.maximum, - sp.Min: paddle.minimum, - sp.Abs: paddle.abs, - sp.Heaviside: paddle.heaviside, - sp.sign: paddle.sign, - sp.ceiling: paddle.ceil, - sp.floor: paddle.floor, - # NOTE: sp.Add and sp.Mul is not included here for un-alignment with paddle - # and are implemented manually in 'OperatorNode._add_operator_func' and - # 'OperatorNode._mul_operator_func' -} - - -def _cvt_to_key(expr: sp.Basic) -> str: - """Convert sympy expression to a string key, mainly as retrieval key in dict. - - Args: - expr (sp.Basic): Sympy expression. - - Returns: - str: Converted string key. - """ - if isinstance(expr, sp.Function) and str(expr.func) == equation.DETACH_FUNC_NAME: - return f"{_cvt_to_key(expr.args[0])}_{equation.DETACH_FUNC_NAME}" - - if isinstance(expr, (sp.Symbol, sp.core.function.UndefinedFunction, sp.Function)): - # use name of custom function(e.g. "f") instead of itself(e.g. "f(x, y)") - # for simplicity. - if hasattr(expr, "name"): - return expr.name - else: - return str(expr) - elif isinstance(expr, sp.Derivative): - # convert "Derivative(u(x,y),(x,2),(y,2))" to "u__x__x__y__y" - expr_str = expr.args[0].name - for symbol, order in expr.args[1:]: - expr_str += f"__{symbol}" * order - return expr_str - else: - return str(expr) - - -class Node(nn.Layer): - """The base class of the node in expression tree. - - Args: - expr (sp.Basic): Sympy expression. - """ - - def __init__(self, expr: sp.Basic): - super().__init__() - self.expr = expr - self.key = _cvt_to_key(self.expr) - - def forward(self, **kwargs): - raise NotImplementedError("Node.forward is not implemented") - - def __str__(self): - return ( - f"{self.__class__.__name__}(expr: {self.expr}, " - f"expr_type: {type(self.expr)})" - ) - - def __repr__(self): - return f"{self.__class__.__name__}(expr: {self.expr})" - - -class DetachNode(Node): - """Class for detach operation in converted expression tree. - - Args: - expr (sp.Basic): Sympy expression. - """ - - def __init__(self, expr: sp.Basic): - super().__init__(expr) - self.child = _cvt_to_key(self.expr.args[0]) - - def forward(self, data_dict: DATA_DICT): - if self.key in data_dict: - return data_dict - - data_dict[self.key] = data_dict[self.child].detach() - return data_dict - - -class OperatorNode(Node): - """Class for operator node in converted expression tree. - - Args: - expr (SYMPY_BUILTIN_FUNC): Sympy expression. - """ - - def __init__( - self, - expr: SYMPY_BUILTIN_FUNC, - ): - super().__init__(expr) - # preprocess children's key instead of processing at run-time in forward - # which can reduce considerable overhead of time for calling "_cvt_to_key" - self.childs = [_cvt_to_key(arg) for arg in self.expr.args] - - if self.expr.func == sp.Add: - self._apply_func = self._add_operator_func - elif self.expr.func == sp.Mul: - self._apply_func = self._mul_operator_func - elif self.expr.func == sp.Heaviside: - self._apply_func = self._heaviside_operator_func - self._auxiliary_func = SYMPY_TO_PADDLE[sp.Heaviside] - self._auxiliary_func = functools.partial( - self._auxiliary_func, y=paddle.zeros([]) - ) - elif self.expr.func == sp.Min: - self._apply_func = self._minimum_operator_func - elif self.expr.func == sp.Max: - self._apply_func = self._maximum_operator_func - else: - self._apply_func = self._vanilla_operator_func - self._auxiliary_func = SYMPY_TO_PADDLE[self.expr.func] - - def forward(self, data_dict: DATA_DICT): - # use cache - if self.key in data_dict: - return data_dict - - return self._apply_func(data_dict) - - def _add_operator_func(self, data_dict: DATA_DICT) -> DATA_DICT: - data_dict[self.key] = data_dict[self.childs[0]] - for p in self.childs[1:]: - data_dict[self.key] += data_dict[p] - return data_dict - - def _mul_operator_func(self, data_dict: DATA_DICT) -> DATA_DICT: - data_dict[self.key] = data_dict[self.childs[0]] - for child in self.childs[1:]: - data_dict[self.key] *= data_dict[child] - return data_dict - - def _heaviside_operator_func(self, data_dict: DATA_DICT) -> DATA_DICT: - data_dict[self.key] = self._auxiliary_func(data_dict[self.childs[0]]) - return data_dict - - def _minimum_operator_func(self, data_dict: DATA_DICT) -> DATA_DICT: - data_dict[self.key] = paddle.minimum( - data_dict[self.childs[0]], data_dict[self.childs[1]] - ) - for i in range(2, len(self.childs)): - data_dict[self.key] = paddle.minimum( - data_dict[self.key], - data_dict[self.childs[i]], - ) - return data_dict - - def _maximum_operator_func(self, data_dict: DATA_DICT) -> DATA_DICT: - data_dict[self.key] = paddle.maximum( - data_dict[self.childs[0]], data_dict[self.childs[1]] - ) - for i in range(2, len(self.childs)): - data_dict[self.key] = paddle.maximum( - data_dict[self.key], - data_dict[self.childs[i]], - ) - return data_dict - - def _vanilla_operator_func(self, data_dict: DATA_DICT) -> DATA_DICT: - data_dict[self.key] = self._auxiliary_func( - *tuple(data_dict[child] for child in self.childs) - ) - return data_dict - - -class DerivativeNode(Node): - """Class for operator node in converted expression tree. - - Args: - expr (sp.Derivative): Sympy derivative expression. - create_graph (bool, optional): Whether to create the gradient graphs of - the computing process. When it is True, higher order derivatives are - supported to compute; when it is False, the gradient graphs of the - computing process would be discarded. Defaults to True. - retain_graph (Optional[bool]): Whether to retain the forward graph which - is used to calculate the gradient. When it is True, the graph would - be retained, in which way users can calculate backward twice for the - same graph. When it is False, the graph would be freed. Defaults to None, - which means it is equal to `create_graph`. - """ - - def __init__( - self, - expr: sp.Derivative, - create_graph: bool = True, - retain_graph: Optional[bool] = None, - ): - super().__init__(expr) - # preprocess children's key instead of processing at run-time in forward - # which can reduce considerable overhead of time for calling "_cvt_to_key" - self.childs = [_cvt_to_key(self.expr.args[0])] + [ - (_cvt_to_key(arg), int(order)) for (arg, order) in self.expr.args[1:] - ] - self.create_graph = create_graph - self.retain_graph = retain_graph - self._apply_func = self._derivate_operator_func - self.merged = False - - def forward(self, data_dict: DATA_DICT): - # use cache - if self.key in data_dict: - return data_dict - - return self._apply_func(data_dict) - - def _derivate_operator_func(self, data_dict: DATA_DICT) -> DATA_DICT: - # NOTE: Derivative of 'sdf' function will not be executed here, which is already - # generated in 'data_dict' during points sampling using discrete difference - # method(see also: ppsci/geometry/geometry.py: Geometry.sdf_derivatives), - # such as 'sdf__x', 'sdf__y'. - data_dict[self.key] = data_dict[self.childs[0]] - for child, order in self.childs[1:]: - if order & 1: - data_dict[self.key] = jacobian( - data_dict[self.key], - data_dict[child], - create_graph=self.create_graph, - retain_graph=self.retain_graph, - ) - order -= 1 - for _ in range(0, order, 2): - data_dict[self.key] = hessian( - data_dict[self.key], - data_dict[child], - create_graph=self.create_graph, - retain_graph=self.retain_graph, - ) - order -= 2 - return data_dict - - -class FusedDerivativeNode(nn.Layer): - """Class for fused DerivativeNode. - - Args: - f_x_tuples (List[Tuple[Union[sp.Function, sp.Derivative], sp.Symbol]]): - indicate all derivatives of a function in list of tuples. e.g. - [(func1, var1), (func1, var2), (func1, var3), ...]. - create_graph (bool, optional): Whether to create the gradient graphs of - the computing process. When it is True, higher order derivatives are - supported to compute; when it is False, the gradient graphs of the - computing process would be discarded. Defaults to True. - retain_graph (Optional[bool]): Whether to retain the forward graph which - is used to calculate the gradient. When it is True, the graph would - be retained, in which way users can calculate backward twice for the - same graph. When it is False, the graph would be freed. Defaults to None, - which means it is equal to `create_graph`. - """ - - def __init__( - self, - f_x_tuples: List[Tuple[Union[sp.Function, sp.Derivative], sp.Symbol]], - create_graph: bool = True, - retain_graph: Optional[bool] = None, - ): - super().__init__() - self.expr: List[sp.Derivative] = [f.diff(x) for f, x in f_x_tuples] - self.key: List[str] = [_cvt_to_key(expr) for expr in self.expr] - self.create_graph = create_graph - self.retain_graph = retain_graph - - # preprocess children's key instead of processing at run-time in forward - # which can reduce considerable overhead of time for calling "_cvt_to_key" - self.y_key: str = _cvt_to_key(f_x_tuples[0][0]) - self.childs: List[str] = [_cvt_to_key(x) for _, x in f_x_tuples] - self._apply_func = self._parallel_derivate_operator_func - - def forward(self, data_dict: DATA_DICT): - # use cache - if all([key in data_dict for key in self.key]): - return data_dict - - return self._apply_func(data_dict) - - def _parallel_derivate_operator_func(self, data_dict: DATA_DICT) -> DATA_DICT: - # NOTE: Derivative of 'sdf' function will not be executed here, which is already - # generated in 'data_dict' during points sampling using discrete difference - # method(see also: ppsci/geometry/geometry.py: Geometry.sdf_derivatives), - # such as 'sdf__x', 'sdf__y'. - y_data: paddle.Tensor = data_dict[self.y_key] - xs_data: List[paddle.Tensor] = [data_dict[x_key] for x_key in self.childs] - y_wrt_xs_grad: List[paddle.Tensor] = jacobian( - y_data, - xs_data, - create_graph=self.create_graph, - retain_graph=self.retain_graph, - ) - for i, key in enumerate(self.key): - data_dict[key] = y_wrt_xs_grad[i] - return data_dict - - def __str__(self): - return ( - f"{self.__class__.__name__}(expr: {self.expr}, " - f"expr_type: {type(self.expr)})" - ) - - def __repr__(self): - return f"{self.__class__.__name__}(expr: {self.expr})" - - -class LayerNode(Node): - """Class for layer node in converted expression tree. - - Args: - expr (sp.core.function.UndefinedFunction): Sympy expression. - model (arch.Arch): NN model for computing forward result in this node. - """ - - def __init__( - self, - expr: sp.core.function.UndefinedFunction, - model: arch.Arch, - ): - super().__init__(expr) - self.model = model - - def forward(self, data_dict: DATA_DICT) -> DATA_DICT: - # use cache - if self.key in data_dict: - return data_dict - - output_dict = self.model(data_dict) - data_dict.update(output_dict) - - return data_dict - - -class ConstantNode(Node): - """Class for constant variable node in converted expression tree. - - Args: - expr (Union[sp.Number, sp.NumberSymbol]): Number expression. - """ - - def __init__(self, expr: Union[sp.Number, sp.NumberSymbol]): - super().__init__(expr) - if ( - self.expr.is_Float - or self.expr.is_Integer - or self.expr.is_Boolean - or self.expr.is_Rational - ): - self.expr = float(self.expr) - else: - raise TypeError( - "expr({expr}) should be Float/Integer/Boolean/Rational, " - f"but got {type(self.expr)}" - ) - self.expr = paddle.to_tensor(self.expr) - - def forward(self, data_dict: DATA_DICT) -> DATA_DICT: - # use cache - if self.key in data_dict: - return data_dict - - data_dict[self.key] = self.expr - return data_dict - - def __str__(self): - return ( - f"{self.__class__.__name__}(expr: {float(self.expr)}, " - f"expr_type: {type(self.expr)})" - ) - - -class ParameterNode(Node): - """Class for constant variable node in converted expression tree. - - Args: - expr (sp.Symbol): Parameter expression. - parameter (paddle.framework.io.EagerParamBase): Parameter tensor. - """ - - def __init__(self, expr: sp.Symbol, parameter: paddle.framework.io.EagerParamBase): - super().__init__(expr) - self.parameter = parameter - - def forward(self, data_dict: DATA_DICT) -> DATA_DICT: - data_dict[self.key] = self.parameter - return data_dict - - -class ComposedNode(nn.Layer): - """ - Compose list of several callable objects together. - """ - - def __init__(self, callable_nodes: List[Node]): - super().__init__() - assert len(callable_nodes) - self.callable_nodes = nn.LayerList(callable_nodes) - - def forward(self, data_dict: DATA_DICT) -> paddle.Tensor: - # call all callable_nodes in order - for i, func in enumerate(self.callable_nodes): - data_dict = func(data_dict) - - # return result of last node(root node) for target - return data_dict[self.callable_nodes[-1].key] - - -def _post_traverse(cur_node: sp.Basic, nodes: List[sp.Basic]) -> List[sp.Basic]: - """Traverse sympy expression tree in post-order. - - Args: - cur_node (sp.Basic): Sympy expression of current node. - nodes (List[sp.Basic]): Node list storing all tree nodes in post-order. - - Returns: - List[sp.Basic]: Node list storing all tree nodes in post-order. - """ - # traverse into sub-nodes - if isinstance(cur_node, sp.Function): - for arg in cur_node.args: - nodes = _post_traverse(arg, nodes) - nodes.append(cur_node) - elif isinstance(cur_node, sp.Derivative): - nodes = _post_traverse(cur_node.args[0], nodes) - nodes.append(cur_node) - elif isinstance(cur_node, sp.Symbol): - nodes.append(cur_node) - return nodes - elif isinstance(cur_node, sp.Number): - nodes.append(cur_node) - else: - for arg in cur_node.args: - nodes = _post_traverse(arg, nodes) - nodes.append(cur_node) - return nodes - - -def _visualize_graph(nodes: List[sp.Basic], graph_filename: str): - try: - import pygraphviz - except ModuleNotFoundError: - raise ModuleNotFoundError( - "Please install pygraphviz by steps below:\n" - "1. apt-get install graphviz graphviz-dev\n" - "2. python -m pip install pygraphviz" - ) - - SYMPY_BUILTIN_NAME = { - sp.sin: "sin", - sp.sinh: "sinh", - sp.asin: "asin", - sp.cos: "cos", - sp.acos: "acos", - sp.cosh: "cosh", - sp.tan: "tan", - sp.atan: "atan", - sp.atan2: "atan2", - sp.acosh: "acosh", - sp.asinh: "asinh", - sp.tanh: "tanh", - sp.atanh: "atanh", - sp.erf: "erf", - sp.loggamma: "loggamma", - sp.exp: "exp", - sp.Pow: "Pow", - sp.log: "log", - sp.Max: "Max", - sp.Min: "Min", - sp.Abs: "Abs", - sp.Heaviside: "Heaviside", - sp.sign: "sign", - sp.ceiling: "ceiling", - sp.floor: "floor", - sp.Add: "Add", - sp.Mul: "Mul", - } - naming_counter = {k: 0 for k in SYMPY_BUILTIN_NAME} - - def get_operator_name(node: sp.Function): - ret = f"{SYMPY_BUILTIN_NAME[node.func]}_{naming_counter[node.func]}" - naming_counter[node.func] += 1 - return ret - - graph = pygraphviz.AGraph(directed=True, rankdir="TB") - C_FUNC = "#9196f1" # purple color function node - C_DATA = "#feb64d" # orange color for data node - C_EDGE = "#000000" # black color for edge - - def add_edge(u: str, v: str, u_color: str = C_DATA, v_color: str = C_DATA): - """Add an edge from `u` to `v`. - - Args: - u (str): Name of begin node u. - v (str): Name of end node v. - u_color (str, optional): Color of node u. Defaults to '#feb64d'. - v_color (str, optional): Color of node v. Defaults to '#feb64d'. - """ - graph.add_node(u, style="filled", shape="ellipse", color=u_color) - graph.add_node(v, style="filled", shape="ellipse", color=v_color) - graph.add_edge(u, v, color=C_EDGE, style="solid", penwidth=0.5, arrowsize=0.5) - - for node in nodes: - if isinstance(node, tuple(SYMPY_BUILTIN_NAME.keys())): - operator_str = get_operator_name(node) - for arg in node.args: - add_edge(_cvt_to_key(arg), operator_str, v_color=C_FUNC) - add_edge(operator_str, _cvt_to_key(node), u_color=C_FUNC) - elif isinstance(node, sp.Function): - for arg in node.args: - add_edge(_cvt_to_key(arg), str(node), v_color=C_FUNC) - add_edge(str(node), _cvt_to_key(node), u_color=C_FUNC) - elif isinstance(node, sp.Derivative): - add_edge(str(node), _cvt_to_key(node), u_color=C_FUNC) - add_edge(_cvt_to_key(node.args[0]), str(node), v_color=C_FUNC) - for arg in node.args[1:]: - add_edge(_cvt_to_key(arg[0]), str(node), v_color=C_FUNC) - - # export graph to image - graph.layout() - image_path = f"{graph_filename}.png" - dot_path = f"{graph_filename}.dot" - if len(os.path.dirname(image_path)): - os.makedirs(os.path.dirname(image_path), exist_ok=True) - graph.draw(image_path, prog="dot") - graph.write(dot_path) - logger.message( - f"Computational graph has been written to: {image_path} and {dot_path}, " - "which can be visualized at: https://dreampuf.github.io/GraphvizOnline/" - ) - - -def _fuse_derivative_nodes( - derivative_exprs: List[sp.Derivative], -) -> List[FusedDerivativeNode]: - """Merge derivative nodes and return in list of FusedDerivativeNode after merger. - - Args: - derivative_exprs (List[sp.Derivative]): Derivatives sympy expression of same - function, e.g. [Derivative(u(x,y), x), Derivative(u(x,y), y)] - - Returns: - List[FusedDerivativeNode]: List of FusedDerivativeNode converting from mergeable - derivatives. - """ - - class DerivativeTrie: - """Trie for unrolling derivative.""" - - def __init__(self, expr: sp.Basic): - self.expr: sp.Basic = expr - self.next: Dict["sp.Symbol", "DerivativeTrie"] = {} - - # unroll derivative expressions into a trie structure - trie_root = DerivativeTrie(derivative_exprs[0].args[0]) - for derivative_expr in derivative_exprs: - cur_node = trie_root - for (child, order) in derivative_expr.args[1:]: - for _ in range(order): - if child not in cur_node.next: - cur_node.next[child] = DerivativeTrie(cur_node.expr.diff(child)) - cur_node = cur_node.next[child] - - def dfs_trie( - node: DerivativeTrie, fused_derivative_nodes: List[FusedDerivativeNode] - ) -> None: - if node.next: - fused_derivative_nodes.append( - FusedDerivativeNode( - [(node.expr, name) for name in node.next], - ) - ) - for child in node.next: - dfs_trie(node.next[child], fused_derivative_nodes) - - # walk on derivative trie in pre-order and log fusable nodes - fused_derivative_nodes: List[FusedDerivativeNode] = [] - dfs_trie(trie_root, fused_derivative_nodes) - - return fused_derivative_nodes - - -def lambdify( - expr: Union[sp.Basic, List[sp.Basic]], - models: Optional[Union[arch.Arch, Tuple[arch.Arch, ...]]] = None, - extra_parameters: Optional[Sequence[paddle.Tensor]] = None, - graph_filename: Optional[str] = None, - create_graph: bool = True, - retain_graph: Optional[bool] = None, - fuse_derivative: bool = False, -) -> Union[ComposedNode, List[ComposedNode]]: - """Convert sympy expression to callable function. - - Args: - expr (Union[sp.Basic, List[sp.Basic]]): Sympy expression(s) to be converted. - Will return callable functions in list if multiple expressions are given, - else return one single callable function. - models (Optional[Union[arch.Arch, Tuple[arch.Arch, ...]]]): Model(s) for - computing forward result in `LayerNode`. - extra_parameters (Optional[nn.ParameterList]): Extra learnable parameters. - Defaults to None. - graph_filename (Optional[str]): Save computational graph to `graph_filename.png` - for given `expr`, if `graph_filename` is not None and a valid string, - such as 'momentum_x'. Defaults to None. - create_graph (bool, optional): Whether to create the gradient graphs of - the computing process. When it is True, higher order derivatives are - supported to compute. When it is False, the gradient graphs of the - computing process would be discarded. Defaults to True. - retain_graph (Optional[bool]): Whether to retain the forward graph which - is used to calculate the gradient. When it is True, the graph would - be retained, in which way users can calculate backward twice for the - same graph. When it is False, the graph would be freed. Defaults to None, - which means it is equal to `create_graph`. - fuse_derivative (bool, optional): Whether to fuse the derivative nodes. - For example, if `expr` is 'Derivative(u, x) + Derivative(u, y)' - It will compute grad(u, x) + grad(u, y) if fuse_derivative=False, - else will compute sum(grad(u, [x, y])) if fuse_derivative=True as is more - efficient in backward-graph. Defaults to False, as it is experimental so not - enabled by default if used independently. - - Returns: - Union[ComposedNode, List[ComposedNode]]: Callable object(s) for computing expr - with necessary input(s) data in dict given. - - Examples: - >>> import paddle - >>> import ppsci - >>> import sympy as sp - - >>> a, b, c, x, y = sp.symbols("a b c x y") - >>> u = sp.Function("u")(x, y) - >>> v = sp.Function("v")(x, y) - >>> z = -a + b * (c ** 2) + u * v + 2.3 - - >>> model = ppsci.arch.MLP(("x", "y"), ("u", "v"), 4, 16) - - >>> batch_size = 13 - >>> a_tensor = paddle.randn([batch_size, 1]) - >>> b_tensor = paddle.randn([batch_size, 1]) - >>> c_tensor = paddle.randn([batch_size, 1]) - >>> x_tensor = paddle.randn([batch_size, 1]) - >>> y_tensor = paddle.randn([batch_size, 1]) - - >>> model_output_dict = model({"x": x_tensor, "y": y_tensor}) - >>> u_tensor, v_tensor = model_output_dict["u"], model_output_dict["v"] - - >>> z_tensor_manually = ( - ... -a_tensor + b_tensor * (c_tensor ** 2) - ... + u_tensor * v_tensor + 2.3 - ... ) - >>> z_tensor_sympy = ppsci.lambdify(z, model)( - ... { - ... "a": a_tensor, - ... "b": b_tensor, - ... "c": c_tensor, - ... "x": x_tensor, - ... "y": y_tensor, - ... } - ... ) - - >>> paddle.allclose(z_tensor_manually, z_tensor_sympy).item() - True - """ - if not extra_parameters: - extra_parameters = () - - if isinstance(models, arch.ModelList): - models = tuple(models.model_list[i] for i in range(len(models.model_list))) - if not isinstance(models, (tuple, list)): - models = (models,) - - def _expr_to_callable_nodes( - single_expr: sp.Basic, graph_filename_: Optional[str] = None - ) -> List[Node]: - """Convert sympy expression to a sequence of nodes in topologic order. - - Args: - single_expr (sp.Basic): Single sympy expression, such as "a+b*c". - graph_filename_ (Optional[str]): Save computational graph to - `/path/to/graph_filename.png` for given `expr`, if `graph_filename` is not - None and a valid string, such as 'momentum_x'. Defaults to None. - - Returns: - List[Node]: Sequence of callable nodes. - """ - # NOTE: Those simplify methods may complicate given expr instead, so not use here - # simplify expression to reduce nodes in tree - # expr = sp.nsimplify(expr) - # expr = sp.expand(expr) - # expr = sp.simplify(expr) - - # remove 1.0 from sympy expression tree - single_expr = single_expr.subs(1.0, 1) - - # convert sympy expression tree to list of nodes in post-order - sympy_nodes: List[sp.Basic] = [] - sympy_nodes = _post_traverse(single_expr, sympy_nodes) - - # remove unnecessary symbol nodes already in input dict(except for parameter symbol) - _parameter_names = tuple(param.name for param in extra_parameters) - sympy_nodes = [ - node - for node in sympy_nodes - if (not node.is_Symbol) or (_cvt_to_key(node) in _parameter_names) - ] - - # remove duplicated node(s) with topological order kept - sympy_nodes = list(dict.fromkeys(sympy_nodes)) - - # convert sympy node to callable node - callable_nodes = [] - for i, node in enumerate(sympy_nodes): - if isinstance( - node, tuple(SYMPY_TO_PADDLE.keys()) + (sp.Add, sp.Mul, sp.Derivative) - ): - if isinstance(node, sp.Derivative): - callable_nodes.append( - DerivativeNode(node, create_graph, retain_graph) - ) - else: - callable_nodes.append(OperatorNode(node)) - elif isinstance(node, sp.Function): - if str(node.func) == equation.DETACH_FUNC_NAME: - callable_nodes.append(DetachNode(node)) - logger.debug(f"Detected detach node {node}") - else: - match_index = None - for j, model in enumerate(models): - if str(node.func) in model.output_keys: - callable_nodes.append( - LayerNode( - node, - model, - ) - ) - if match_index is not None: - raise ValueError( - f"Name of function: '{node}' should be unique along given" - f" models, but got same output_key: '{str(node.func)}' " - f"in given models[{match_index}] and models[{j}]." - ) - match_index = j - # NOTE: Skip 'sdf' function, which should be already generated in - # given data_dict - if match_index is None and str(node.func) != "sdf": - raise ValueError( - f"Node {node} can not match any model in given model(s)." - ) - elif node.is_Number or node.is_NumberSymbol: - callable_nodes.append(ConstantNode(node)) - elif isinstance(node, sp.Symbol): - callable_nodes.append( - ParameterNode( - node, - *[ - param - for param in extra_parameters - if param.name == node.name - ], - ) - ) - else: - raise NotImplementedError( - f"The node {node} is not supported in lambdify." - ) - - # NOTE: visualize computational graph using 'pygraphviz' - if isinstance(graph_filename, str): - _visualize_graph(sympy_nodes, os.path.join(graph_filename, graph_filename_)) - - return callable_nodes - - if isinstance(expr, sp.Basic): - callable_nodes_group = [_expr_to_callable_nodes(expr, "expr")] - else: - callable_nodes_group = [ - _expr_to_callable_nodes(expr_i, f"expr_{i}") - for i, expr_i in enumerate(expr) - ] - - # [Optional] Fused derivatives nodes that with same function to be differentiated - while fuse_derivative: - candidate_pos: List[Tuple[int, int]] = [] # [(group_id, node_id), ...] - - # use 4-nested for-loop to find all potential mergeable derivative nodes - for i in range(len(callable_nodes_group)): - for j in range(len(callable_nodes_group[i])): - # skip non-derivative node - if not isinstance(callable_nodes_group[i][j], DerivativeNode): - continue - # skip sdf function since it is always already given in data_dict - if callable_nodes_group[i][j].expr.args[0].name == "sdf": - continue - # skip merged node - if callable_nodes_group[i][j].merged: - continue - - candidate_pos = [[i, j]] - for ii in range(len(callable_nodes_group)): - for jj in range(len(callable_nodes_group[ii])): - # skip non-derivative node - if not isinstance(callable_nodes_group[ii][jj], DerivativeNode): - continue - - # skip same node - if i == ii and j == jj: - continue - # skip merged node - if callable_nodes_group[ii][jj].merged: - continue - - # has same function item - if ( - callable_nodes_group[i][j].expr.args[0] - == callable_nodes_group[ii][jj].expr.args[0] - ): - candidate_pos.append([ii, jj]) - - if len(candidate_pos) > 1: - break - if len(candidate_pos) > 1: - break - - # merge all candidate nodes into one or more FusedDerivativeNode node - if len(candidate_pos) > 1: - fused_node_seq = _fuse_derivative_nodes( - [callable_nodes_group[gid][nid].expr for gid, nid in candidate_pos] - ) - assert isinstance( - fused_node_seq, list - ), "'fused_node_seq' should be list of 'FusedDerivativeNode'" - gid0, nid0 = candidate_pos[0] - logger.debug( - f"Fused {len(candidate_pos)} derivatives nodes: " - f"{[callable_nodes_group[i][j].expr for i, j in candidate_pos]} into" - f" {len(fused_node_seq)} fuse node sequence: {fused_node_seq} at position: ([{gid0}][{nid0}])" - ) - - # mark merged node - for i, (gid, nid) in enumerate(candidate_pos): - assert isinstance(callable_nodes_group[gid][nid], DerivativeNode) - callable_nodes_group[gid][nid].merged = True - - # replace first mergeable node with fused node sequence(packed in list) - # then mask the rest merged node to None(except [gid0, nid0]) - for i, (gid, nid) in enumerate(candidate_pos[1:]): - # keep the end node of each group to avoid generating empty callable - # node sequence, this will not effect performance since cache strategy - # in Node.forward - if nid != len(callable_nodes_group[gid]) - 1: - callable_nodes_group[gid][nid] = None - - if nid0 == len(callable_nodes_group[gid0]) - 1: - callable_nodes_group[gid0].insert(nid0, fused_node_seq) - else: - callable_nodes_group[gid0][nid0] = fused_node_seq - - # re-organize callable_nodes_group, remove None element and unpack list - for i in range(len(callable_nodes_group)): - tmp = [] - for j in range(len(callable_nodes_group[i])): - if isinstance( - callable_nodes_group[i][j], (Node, FusedDerivativeNode) - ): - tmp.append(callable_nodes_group[i][j]) - elif isinstance(callable_nodes_group[i][j], list) and isinstance( - callable_nodes_group[i][j][0], FusedDerivativeNode - ): - tmp.extend(callable_nodes_group[i][j]) - else: - assert ( - callable_nodes_group[i][j] is None - ), f"Unexpected element: {callable_nodes_group[i][j]}" - callable_nodes_group[i] = tmp - else: - # exit while loop if no more fused - break - - # Compose callable nodes into one callable object - if isinstance(expr, sp.Basic): - return ComposedNode(callable_nodes_group[0]) - else: - return [ComposedNode(callable_nodes) for callable_nodes in callable_nodes_group] diff --git a/examples/smc_reac/ppsci/utils/writer.py b/examples/smc_reac/ppsci/utils/writer.py deleted file mode 100644 index d1b4c503f3..0000000000 --- a/examples/smc_reac/ppsci/utils/writer.py +++ /dev/null @@ -1,225 +0,0 @@ -# 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. - -from __future__ import annotations - -import csv -import os -from typing import Dict -from typing import Optional -from typing import Tuple -from typing import Union - -import numpy as np -import paddle - -from ppsci.utils import logger - -__all__ = [ - "save_csv_file", -] - - -def save_csv_file( - filename: str, - data_dict: Dict[str, Union[np.ndarray, "paddle.Tensor"]], - keys: Tuple[str, ...], - alias_dict: Optional[Dict[str, str]] = None, - use_header: bool = True, - delimiter: str = ",", - encoding: str = "utf-8", -): - """Write numpy or tensor data into csv file. - - Args: - filename (str): Dump file path. - data_dict (Dict[str, Union[np.ndarray, paddle.Tensor]]): Numpy or tensor data in dict. - keys (Tuple[str, ...]): Keys for data_dict to be fetched. - alias_dict (Optional[Dict[str, str]], optional): Alias dict for keys, - i.e. {dump_key: dict_key}. Defaults to None. - use_header (bool, optional): Whether save csv with header. Defaults to True. - delimiter (str, optional): Delimiter for splitting different data field. Defaults to ",". - encoding (str, optional): Encoding. Defaults to "utf-8". - - Examples: - >>> import numpy as np - >>> from ppsci.utils import save_csv_file - >>> data_dict = { - ... "a": np.array([[1], [2], [3]]).astype("int64"), # [3, 1] - ... "b": np.array([[4.12], [5.25], [6.3370]]).astype("float32"), # [3, 1] - ... } - >>> save_csv_file( - ... "test.csv", - ... data_dict, - ... ("A", "B"), - ... alias_dict={"A": "a", "B": "b"}, - ... use_header=True, - ... delimiter=",", - ... encoding="utf-8", - ... ) # doctest: +SKIP - - >>> # == test.csv == - >>> # A,B - >>> # 1,4.12 - >>> # 2,5.25 - >>> # 3,6.337 - """ - if alias_dict is None: - alias_dict = {} - - # convert to numpy array - data_fields = [] - header = [] - for key in keys: - fetch_key = alias_dict.get(key, key) - data = data_dict[fetch_key] - if isinstance(data, paddle.Tensor): - data = data.numpy() # [num_of_samples, ] - - if isinstance(data, np.ndarray): - data = data.flatten() - data_fields.append(data) - - header.append(key) - - assert len(header) == len(data_fields) - - if len(os.path.dirname(filename)): - os.makedirs(os.path.dirname(filename), exist_ok=True) - - data_fields = zip(*data_fields) # transpose col data to row data - with open(filename, "w", newline="", encoding=encoding) as file: - writer = csv.writer(file, delimiter=delimiter) - - if use_header: - writer.writerow(header) - - writer.writerows(data_fields) - - logger.message(f"csv file has been dumped to: {filename}") - - -def save_tecplot_file( - filename: str, - data_dict: Dict[str, Union[np.ndarray, "paddle.Tensor"]], - keys: Tuple[str, ...], - num_x: int, - num_y: int, - alias_dict: Optional[Dict[str, str]] = None, - delimiter: str = " ", - encoding: str = "utf-8", - num_timestamps: int = 1, -): - """Write numpy or tensor data into tecplot file(s). - - Args: - filename (str): Tecplot file path. - data_dict (Dict[str, Union[np.ndarray, paddle.Tensor]]): Numpy or Tensor data in dict. - keys (Tuple[str, ...]): Target keys to be dumped. - num_x (int): The number of discrete points of the grid in the X-axis. Assuming - the discrete grid size is 20 x 30, then num_x=20. - num_y (int): The number of discrete points of the grid in the Y-axis. Assuming - the discrete grid size is 20 x 30, then num_y=30. - alias_dict (Optional[Dict[str, str]], optional): Alias dict for keys, - i.e. {dump_key: dict_key}. Defaults to None. - delimiter (str, optional): Delimiter for splitting different data field. Defaults to " ". - encoding (str, optional): Encoding. Defaults to "utf-8". - num_timestamps (int, optional): Number of timestamp over coord and value. Defaults to 1. - - Examples: - >>> import numpy as np - >>> from ppsci.utils import save_tecplot_file - >>> data_dict = { - ... "x": np.array([[-1.0], [-1.0], [-1.0], [-1.0], [-1.0], [-1.0]]), # [6, 1] - ... "y": np.array([[1.0], [2.0], [3.0], [1.0], [2.0], [3.0]]), # [6, 1] - ... "value": np.array([[3], [33], [333], [3333], [33333], [333333]]), # [6, 1] - ... } - >>> save_tecplot_file( - ... "./test.dat", - ... data_dict, - ... ("X", "Y", "value"), - ... num_x=1, - ... num_y=3, - ... alias_dict={"X": "x", "Y": "y"}, - ... num_timestamps=2, - ... ) # doctest: +SKIP - >>> # == test_t-0.dat == - >>> # title = "./test_t-0.dat" - >>> # variables = "X", "Y" - >>> # Zone I = 3, J = 1, F = POINT - >>> # -1.0 1.0 3.0 - >>> # -1.0 2.0 33.0 - >>> # -1.0 3.0 333.0 - - - >>> # == test_t-1.dat == - >>> # title = "./test_t-1.dat" - >>> # variables = "X", "Y" - >>> # Zone I = 3, J = 1, F = POINT - >>> # -1.0 1.0 3333.0 - >>> # -1.0 2.0 33333.0 - >>> # -1.0 3.0 333333.0 - """ - if alias_dict is None: - alias_dict = {} - - ntxy = len(next(iter(data_dict.values()))) - if ntxy % num_timestamps != 0: - raise ValueError( - f"num_points({ntxy}) must be a multiple of " - f"num_timestamps({num_timestamps})." - ) - nxy = ntxy // num_timestamps - - nx, ny = num_x, num_y - assert nx * ny == nxy, f"nx({nx}) * ny({ny}) != nxy({nxy})" - - if len(os.path.dirname(filename)): - os.makedirs(os.path.dirname(filename), exist_ok=True) - - if filename.endswith(".dat"): - filename = filename[:-4] - - for t in range(num_timestamps): - # write 1 tecplot file for each timestep - if num_timestamps > 1: - dump_filename = f"{filename}_t-{t}.dat" - else: - dump_filename = f"{filename}.dat" - - fetch_keys = [alias_dict.get(key, key) for key in keys] - with open(dump_filename, "w", encoding=encoding) as f: - # write meta information of tec - f.write(f'title = "{dump_filename}"\n') - header = ", ".join([f'"{key}"' for key in keys]) - f.write(f"variables = {header}\n") - - # NOTE: Tecplot is column-major, so we need to specify I = ny, J = nx, - # which is in contrast to our habits. - f.write(f"Zone I = {ny}, J = {nx}, F = POINT\n") - - # write points data into file - data_cur_time_step = [ - data_dict[key][t * nxy : (t + 1) * nxy] for key in fetch_keys - ] - - for items in zip(*data_cur_time_step): - f.write(delimiter.join([str(float(x)) for x in items]) + "\n") - - if num_timestamps > 1: - logger.message( - f"tecplot files are saved to: {filename}_t-0.dat ~ {filename}_t-{num_timestamps - 1}.dat" - ) - else: - logger.message(f"tecplot file is saved to: {filename}.dat") diff --git a/examples/smc_reac/ppsci/validate/__init__.py b/examples/smc_reac/ppsci/validate/__init__.py deleted file mode 100644 index 3bc1c9ae4d..0000000000 --- a/examples/smc_reac/ppsci/validate/__init__.py +++ /dev/null @@ -1,81 +0,0 @@ -# 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. - -import copy - -from ppsci.loss import build_loss -from ppsci.metric import build_metric -from ppsci.utils import logger -from ppsci.utils import misc -from ppsci.validate.base import Validator -from ppsci.validate.geo_validator import GeometryValidator -from ppsci.validate.sup_validator import SupervisedValidator - -__all__ = [ - "Validator", - "GeometryValidator", - "SupervisedValidator", -] - - -def build_validator(cfg, equation_dict, geom_dict): - """Build validator(s). - - Args: - cfg (List[DictConfig]): Validator(s) config list. - geom_dict (Dct[str, Geometry]): Geometry(ies) in dict. - equation_dict (Dct[str, Equation]): Equation(s) in dict. - - Returns: - Dict[str, Validator]: Validator(s) in dict. - """ - if cfg is None: - return None - cfg = copy.deepcopy(cfg) - global_dataloader_cfg = cfg["dataloader"] - validator_cfg = cfg["content"] - - validator_dict = misc.PrettyOrderedDict() - for _item in validator_cfg: - validator_cls = next(iter(_item.keys())) - _validator_cfg = _item[validator_cls] - validator_name = _validator_cfg.get("name", validator_cls) - # select geometry - geom_name = _validator_cfg.pop("geom") - _validator_cfg["geom"] = geom_dict[geom_name] - - # update complete dataloader config - local_dataloader_cfg = _validator_cfg["dataloader"] - local_dataloader_cfg.update(global_dataloader_cfg) - - # select equation - for name, expr in _validator_cfg["output_expr"].items(): - if isinstance(expr, str) and expr in equation_dict: - _validator_cfg["output_expr"][name] = equation_dict[expr].equations[ - name - ] - - # build loss - _validator_cfg["loss"] = build_loss(_validator_cfg["loss"]) - - # build metric - _validator_cfg["metric"] = build_metric(_validator_cfg["metric"]) - - # instantiate validator - _validator_cfg["dataloader_cfg"] = _validator_cfg.pop("dataloader") - validator_dict[validator_name] = eval(validator_cls)(**_validator_cfg) - - logger.debug(str(validator_dict[validator_name])) - - return validator_dict diff --git a/examples/smc_reac/ppsci/validate/base.py b/examples/smc_reac/ppsci/validate/base.py deleted file mode 100644 index 84760f25d7..0000000000 --- a/examples/smc_reac/ppsci/validate/base.py +++ /dev/null @@ -1,69 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import TYPE_CHECKING -from typing import Any -from typing import Dict -from typing import Optional - -from paddle import io - -from ppsci import data - -if TYPE_CHECKING: - from ppsci import loss - from ppsci import metric - - -class Validator: - """Base class for validators. - - Args: - dataset (io.Dataset): Dataset for validator. - dataloader_cfg (Dict[str, Any]): Dataloader config. - loss (loss.Loss): Loss functor. - metric (Optional[Dict[str, metric.Metric]]): Named metric functors in dict. - name (str): Name of validator. - """ - - def __init__( - self, - dataset: io.Dataset, - dataloader_cfg: Dict[str, Any], - loss: "loss.Loss", - metric: Optional[Dict[str, "metric.Metric"]], - name: str, - ): - self.data_loader = data.build_dataloader(dataset, dataloader_cfg) - self.data_iter = iter(self.data_loader) - self.loss = loss - self.metric = metric - self.name = name - - def __str__(self): - return ", ".join( - [ - self.__class__.__name__, - f"name = {self.name}", - f"input_keys = {self.input_keys}", - f"output_keys = {self.output_keys}", - f"output_expr = {self.output_expr}", - f"label_dict = {self.label_dict}", - f"len(dataloader) = {len(self.data_loader)}", - f"loss = {self.loss}", - f"metric = {list(self.metric.keys())}", - ] - ) diff --git a/examples/smc_reac/ppsci/validate/geo_validator.py b/examples/smc_reac/ppsci/validate/geo_validator.py deleted file mode 100644 index 08f2e663a4..0000000000 --- a/examples/smc_reac/ppsci/validate/geo_validator.py +++ /dev/null @@ -1,161 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import Any -from typing import Callable -from typing import Dict -from typing import Optional -from typing import Union - -import numpy as np -import paddle -import sympy -from typing_extensions import Literal - -from ppsci import geometry -from ppsci import loss -from ppsci import metric -from ppsci.data import dataset -from ppsci.validate import base - - -class GeometryValidator(base.Validator): - """Validator for geometry. - - Args: - output_expr (Dict[str, Callable]): Function in dict for computing output. - e.g. {"u_mul_v": lambda out: out["u"] * out["v"]} means the model output u - will be multiplied by model output v and the result will be named "u_mul_v". - label_dict (Dict[str, Union[float, Callable]]): Function in dict for computing - label, which will be a reference value to participate in the loss calculation. - geom (geometry.Geometry): Geometry where data sampled from. - dataloader_cfg (Dict[str, Any]): Dataloader config. - loss (loss.Loss): Loss functor. - random (Literal["pseudo", "Halton", "LHS"], optional): Random method for sampling data in - geometry. Defaults to "pseudo". - criteria (Optional[Callable]): Criteria for refining specified domain. Defaults to None. - evenly (bool, optional): Whether to use evenly distribution sampling. Defaults to False. - metric (Optional[Dict[str, metric.Metric]]): Named metric functors in dict. Defaults to None. - with_initial (bool, optional): Whether the data contains time t0. Defaults to False. - name (Optional[str]): Name of validator. Defaults to None. - - Examples: - >>> import ppsci - >>> rect = ppsci.geometry.Rectangle((0, 0), (1, 1)) - >>> geom_validator = ppsci.validate.GeometryValidator( - ... {"u": lambda out: out["u"]}, - ... {"u": 0}, - ... rect, - ... { - ... "dataset": "IterableNamedArrayDataset", - ... "iters_per_epoch": 1, - ... "total_size": 32, - ... "batch_size": 16, - ... }, - ... ppsci.loss.MSELoss("mean"), - ... ) - """ - - def __init__( - self, - output_expr: Dict[str, Callable], - label_dict: Dict[str, Union[float, Callable]], - geom: geometry.Geometry, - dataloader_cfg: Dict[str, Any], - loss: loss.Loss, - random: Literal["pseudo", "Halton", "LHS"] = "pseudo", - criteria: Optional[Callable] = None, - evenly: bool = False, - metric: Optional[Dict[str, metric.Metric]] = None, - with_initial: bool = False, - name: Optional[str] = None, - ): - self.output_expr = output_expr - self.label_dict = label_dict - self.input_keys = geom.dim_keys - self.output_keys = tuple(label_dict.keys()) - - nx = dataloader_cfg["total_size"] - self.num_timestamps = 1 - # TODO(sensen): Simplify code below - if isinstance(geom, geometry.TimeXGeometry): - if geom.timedomain.num_timestamps is not None: - if with_initial: - # include t0 - self.num_timestamps = geom.timedomain.num_timestamps - assert ( - nx % self.num_timestamps == 0 - ), f"{nx} % {self.num_timestamps} != 0" - nx //= self.num_timestamps - input = geom.sample_interior( - nx * (geom.timedomain.num_timestamps - 1), - random, - criteria, - evenly, - ) - initial = geom.sample_initial_interior(nx, random, criteria, evenly) - input = { - key: np.vstack((initial[key], input[key])) for key in input - } - else: - # exclude t0 - self.num_timestamps = geom.timedomain.num_timestamps - 1 - assert ( - nx % self.num_timestamps == 0 - ), f"{nx} % {self.num_timestamps} != 0" - nx //= self.num_timestamps - input = geom.sample_interior( - nx * (geom.timedomain.num_timestamps - 1), - random, - criteria, - evenly, - ) - else: - raise NotImplementedError( - "TimeXGeometry with random timestamp not implemented yet." - ) - else: - input = geom.sample_interior(nx, random, criteria, evenly) - - label = {} - for key, value in label_dict.items(): - if isinstance(value, (int, float)): - label[key] = np.full_like(next(iter(input.values())), value) - elif isinstance(value, sympy.Basic): - func = sympy.lambdify( - sympy.symbols(geom.dim_keys), - value, - [{"amax": lambda xy, _: np.maximum(xy[0], xy[1])}, "numpy"], - ) - label[key] = func( - **{k: v for k, v in input.items() if k in geom.dim_keys} - ) - elif callable(value): - func = value - label[key] = func(input) - if isinstance(label[key], (int, float)): - label[key] = np.full( - (next(iter(input.values())).shape[0], 1), - label[key], - paddle.get_default_dtype(), - ) - else: - raise NotImplementedError(f"type of {type(value)} is invalid yet.") - - weight = {key: np.ones_like(next(iter(label.values()))) for key in label} - - _dataset = getattr(dataset, dataloader_cfg["dataset"])(input, label, weight) - super().__init__(_dataset, dataloader_cfg, loss, metric, name) diff --git a/examples/smc_reac/ppsci/validate/sup_validator.py b/examples/smc_reac/ppsci/validate/sup_validator.py deleted file mode 100644 index a88a02af5e..0000000000 --- a/examples/smc_reac/ppsci/validate/sup_validator.py +++ /dev/null @@ -1,103 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import Any -from typing import Callable -from typing import Dict -from typing import Optional - -from ppsci import loss -from ppsci import metric -from ppsci.data import dataset -from ppsci.validate import base - - -class SupervisedValidator(base.Validator): - """Validator for supervised models. - - Args: - dataloader_cfg (Dict[str, Any]): Config of building a dataloader. - loss (loss.Loss): Loss functor. - output_expr (Optional[Dict[str, Callable]]): List of label expression. - metric (Optional[Dict[str, metric.Metric]]): Named metric functors in dict. Defaults to None. - name (Optional[str]): Name of validator. Defaults to None. - - Examples: - >>> import ppsci - >>> valid_dataloader_cfg = { - ... "dataset": { - ... "name": "MatDataset", - ... "file_path": "/path/to/file.mat", - ... "input_keys": ("t_f",), - ... "label_keys": ("eta", "f"), - ... }, - ... "batch_size": 32, - ... "sampler": { - ... "name": "BatchSampler", - ... "drop_last": False, - ... "shuffle": False, - ... }, - ... } # doctest: +SKIP - >>> eta_mse_validator = ppsci.validate.SupervisedValidator( - ... valid_dataloader_cfg, - ... ppsci.loss.MSELoss("mean"), - ... {"eta": lambda out: out["eta"]}, - ... metric={"MSE": ppsci.metric.MSE()}, - ... name="eta_mse", - ... ) # doctest: +SKIP - """ - - def __init__( - self, - dataloader_cfg: Dict[str, Any], - loss: loss.Loss, - output_expr: Optional[Dict[str, Callable]] = None, - metric: Optional[Dict[str, metric.Metric]] = None, - name: Optional[str] = None, - ): - self.output_expr = output_expr - - # build dataset - _dataset = dataset.build_dataset(dataloader_cfg["dataset"]) - - self.input_keys = _dataset.input_keys - self.output_keys = ( - tuple(output_expr.keys()) - if output_expr is not None - else _dataset.label_keys - ) - - if self.output_expr is None: - self.output_expr = { - key: lambda out, k=key: out[k] for key in self.output_keys - } - - # construct dataloader with dataset and dataloader_cfg - super().__init__(_dataset, dataloader_cfg, loss, metric, name) - - def __str__(self): - return ", ".join( - [ - self.__class__.__name__, - f"name = {self.name}", - f"input_keys = {self.input_keys}", - f"output_keys = {self.output_keys}", - f"output_expr = {self.output_expr}", - f"len(dataloader) = {len(self.data_loader)}", - f"loss = {self.loss}", - f"metric = {list(self.metric.keys())}", - ] - ) diff --git a/examples/smc_reac/ppsci/visualize/__init__.py b/examples/smc_reac/ppsci/visualize/__init__.py deleted file mode 100644 index e6a5d49186..0000000000 --- a/examples/smc_reac/ppsci/visualize/__init__.py +++ /dev/null @@ -1,82 +0,0 @@ -# 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. - -from __future__ import annotations - -import copy - -from ppsci.visualize.vtu import save_vtu_to_mesh - -from ppsci.visualize.base import Visualizer # isort:skip -from ppsci.visualize.visualizer import VisualizerScatter1D # isort:skip -from ppsci.visualize.visualizer import VisualizerScatter3D # isort:skip -from ppsci.visualize.visualizer import VisualizerVtu # isort:skip -from ppsci.visualize.visualizer import Visualizer2D # isort:skip -from ppsci.visualize.visualizer import Visualizer2DPlot # isort:skip -from ppsci.visualize.visualizer import Visualizer3D # isort:skip -from ppsci.visualize.visualizer import VisualizerWeather # isort:skip -from ppsci.visualize.radar import VisualizerRadar # isort:skip -from ppsci.visualize.vtu import save_vtu_from_dict # isort:skip -from ppsci.visualize.vtu import save_vtp_from_dict # isort:skip -from ppsci.visualize.plot import save_plot_from_1d_dict # isort:skip -from ppsci.visualize.plot import save_plot_from_3d_dict # isort:skip -from ppsci.visualize.plot import save_plot_weather_from_dict # isort:skip - - -__all__ = [ - "Visualizer", - "VisualizerScatter1D", - "VisualizerScatter3D", - "VisualizerVtu", - "Visualizer2D", - "Visualizer2DPlot", - "Visualizer3D", - "VisualizerWeather", - "VisualizerRadar", - "save_vtu_from_dict", - "save_vtp_from_dict", - "save_vtu_to_mesh", - "save_plot_from_1d_dict", - "save_plot_from_3d_dict", - "save_plot_weather_from_dict", -] - - -def build_visualizer(cfg): - """Build visualizer(s). - - Args: - cfg (List[DictConfig]): Visualizer(s) config list. - geom_dict (Dct[str, Geometry]): Geometry(ies) in dict. - equation_dict (Dct[str, Equation]): Equation(s) in dict. - - Returns: - Dict[str, Visualizer]: Visualizer(s) in dict. - """ - if cfg is None: - return None - cfg = copy.deepcopy(cfg) - - visualizer_dict = {} - for _item in cfg: - visualizer_cls = next(iter(_item.keys())) - visualizer_cfg = _item[visualizer_cls] - visualizer = eval(visualizer_cls)(**visualizer_cfg) - - visualizer_name = visualizer_cfg.get("name", visualizer_cls) - if visualizer_name in visualizer_dict: - raise ValueError(f"Name of visualizer({visualizer_name}) should be unique") - visualizer_dict[visualizer_name] = visualizer - - return visualizer_dict diff --git a/examples/smc_reac/ppsci/visualize/base.py b/examples/smc_reac/ppsci/visualize/base.py deleted file mode 100644 index b249efcabc..0000000000 --- a/examples/smc_reac/ppsci/visualize/base.py +++ /dev/null @@ -1,65 +0,0 @@ -# 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. - -from __future__ import annotations - -import abc -from typing import Callable -from typing import Dict - -import numpy as np - - -class Visualizer: - """Base class for visualizer. - - Args: - input_dict (Dict[str, np.ndarray]): Input dict. - output_expr (Dict[str, Callable]): Output expression. - batch_size (int): Batch size of data when computing result in visu.py. - num_timestamps (int): Number of timestamps. - prefix (str): Prefix for output file. - """ - - def __init__( - self, - input_dict: Dict[str, np.ndarray], - output_expr: Dict[str, Callable], - batch_size: int, - num_timestamps: int, - prefix: str, - ): - self.input_dict = input_dict - self.input_keys = tuple(input_dict.keys()) - self.output_expr = output_expr - self.output_keys = tuple(output_expr.keys()) - self.batch_size = batch_size - self.num_timestamps = num_timestamps - self.prefix = prefix - - @abc.abstractmethod - def save(self, data_dict): - """Visualize result from data_dict and save as files""" - - def __str__(self): - return ", ".join( - [ - f"input_keys: {self.input_keys}", - f"output_keys: {self.output_keys}", - f"output_expr: {self.output_expr}", - f"batch_size: {self.batch_size}", - f"num_timestamps: {self.num_timestamps}", - f"output file prefix: {self.prefix}", - ] - ) diff --git a/examples/smc_reac/ppsci/visualize/plot.py b/examples/smc_reac/ppsci/visualize/plot.py deleted file mode 100644 index 8d41ac17c6..0000000000 --- a/examples/smc_reac/ppsci/visualize/plot.py +++ /dev/null @@ -1,580 +0,0 @@ -# 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. - -from __future__ import annotations - -import os -from typing import Dict -from typing import Optional -from typing import Tuple -from typing import Union - -import imageio -import matplotlib -import numpy as np -import paddle -from matplotlib import pyplot as plt -from matplotlib.legend_handler import HandlerBase -from matplotlib.patches import Rectangle -from mpl_toolkits.mplot3d.art3d import Line3DCollection - -from ppsci.utils import logger - -cnames = [ - "bisque", - "black", - "blanchedalmond", - "blue", - "blueviolet", - "brown", - "burlywood", - "cadetblue", - "chartreuse", - "orangered", - "orchid", - "palegoldenrod", - "palegreen", -] - -CMAPS = [ - "Reds", - "Blues", - "Greys", - "Purples", - "Greens", - "Oranges", - "YlOrBr", - "YlOrRd", - "OrRd", - "PuRd", - "RdPu", - "BuPu", - "GnBu", - "PuBu", - "YlGnBu", - "PuBuGn", - "BuGn", - "YlGn", -] - - -def _save_plot_from_1d_array(filename, coord, value, value_keys, num_timestamps=1): - """Save plot from given 1D data. - - Args: - filename (str): Filename. - coord (np.ndarray): Coordinate array. - value (Dict[str, np.ndarray]): Dict of value array. - value_keys (Tuple[str, ...]): Value keys. - num_timestamps (int, optional): Number of timestamps coord/value contains. Defaults to 1. - """ - if len(os.path.dirname(filename)): - os.makedirs(os.path.dirname(filename), exist_ok=True) - fig, a = plt.subplots(len(value_keys), num_timestamps, squeeze=False) - fig.subplots_adjust(hspace=0.8) - - len_ts = len(coord) // num_timestamps - for t in range(num_timestamps): - st = t * len_ts - ed = (t + 1) * len_ts - coord_t = coord[st:ed] - - for i, key in enumerate(value_keys): - _value_t: np.ndarray = value[st:ed, i] - a[i][t].scatter( - coord_t, - _value_t, - color=cnames[i], - label=key, - s=2, - ) - if num_timestamps > 1: - a[i][t].set_title(f"{key}(t={t})") - else: - a[i][t].set_title(f"{key}") - a[i][t].grid(color="#c2ccd0", linestyle="--", linewidth=0.5) - a[i][t].legend() - - if num_timestamps == 1: - fig.savefig(filename, dpi=300) - else: - fig.savefig(f"{filename}_{t}", dpi=300) - - if num_timestamps == 1: - logger.message(f"1D result is saved to: {filename}.png") - else: - logger.message( - f"1D result is saved to: {filename}_0.png" - f" ~ {filename}_{num_timestamps - 1}.png" - ) - - -def save_plot_from_1d_dict( - filename, data_dict, coord_keys, value_keys, num_timestamps=1 -): - """Plot dict data as file. - - Args: - filename (str): Output filename. - data_dict (Dict[str, Union[np.ndarray, paddle.Tensor]]): Data in dict. - coord_keys (Tuple[str, ...]): Tuple of coord key. such as ("x", "y"). - value_keys (Tuple[str, ...]): Tuple of value key. such as ("u", "v"). - num_timestamps (int, optional): Number of timestamp in data_dict. Defaults to 1. - - Examples: - >>> import ppsci - >>> import numpy as np - >>> filename = "path/to/file" - >>> data_dict = { - ... "x": np.array([[1], [2], [3],[4]]), - ... "u": np.array([[4], [5], [6],[4]]), - ... } - >>> coord_keys = ("x",) - >>> value_keys = ("u",) - >>> ppsci.visualize.save_plot_from_1d_dict(filename, data_dict, coord_keys, value_keys) # doctest: +SKIP - """ - space_ndim = len(coord_keys) - int("t" in coord_keys) - if space_ndim not in [1, 2, 3]: - raise ValueError(f"ndim of space coord ({space_ndim}) should be 1, 2 or 3") - - coord = [data_dict[k] for k in coord_keys if k != "t"] - value = [data_dict[k] for k in value_keys] if value_keys else None - - if isinstance(coord[0], paddle.Tensor): - coord = [x.numpy() for x in coord] - else: - coord = [x for x in coord] - coord = np.concatenate(coord, axis=1) - - if value is not None: - if isinstance(value[0], paddle.Tensor): - value = [x.numpy() for x in value] - else: - value = [x for x in value] - value = np.concatenate(value, axis=1) - - _save_plot_from_1d_array(filename, coord, value, value_keys, num_timestamps) - - -def _save_plot_from_2d_array( - filename: str, - visu_data: Tuple[np.ndarray, ...], - visu_keys: Tuple[str, ...], - num_timestamps: int = 1, - stride: int = 1, - xticks: Optional[Tuple[float, ...]] = None, - yticks: Optional[Tuple[float, ...]] = None, -): - """Save plot from given 2D data. - - Args: - filename (str): Filename. - visu_data (Tuple[np.ndarray, ...]): Data that requires visualization. - visu_keys (Tuple[str, ...]): Keys for visualizing data. such as ("u", "v"). - num_timestamps (int, optional): Number of timestamps coord/value contains. Defaults to 1. - stride (int, optional): The time stride of visualization. Defaults to 1. - xticks (Optional[Tuple[float, ...]]): Tuple of xtick locations. Defaults to None. - yticks (Optional[Tuple[float, ...]]): Tuple of ytick locations. Defaults to None. - """ - if len(os.path.dirname(filename)): - os.makedirs(os.path.dirname(filename), exist_ok=True) - - plt.close("all") - matplotlib.rcParams["xtick.labelsize"] = 5 - matplotlib.rcParams["ytick.labelsize"] = 5 - - fig, ax = plt.subplots( - len(visu_keys), - num_timestamps, - squeeze=False, - sharey=True, - figsize=(num_timestamps, len(visu_keys)), - ) - fig.subplots_adjust(hspace=0.3) - target_flag = any("target" in key for key in visu_keys) - for i, data in enumerate(visu_data): - if target_flag is False or "target" in visu_keys[i]: - c_max = np.amax(data) - c_min = np.amin(data) - - for t_idx in range(num_timestamps): - t = t_idx * stride - ax[i, t_idx].imshow( - data[t, :, :], - extent=[xticks.min(), xticks.max(), yticks.min(), yticks.max()], - cmap="inferno", - origin="lower", - vmax=c_max, - vmin=c_min, - ) - if xticks is not None: - ax[i, t_idx].set_xticks(xticks) - if yticks is not None: - ax[i, t_idx].set_yticks(yticks) - - ax[i, t_idx].set_title(f"t={t}", fontsize=8) - if t_idx == 0: - ax[i, 0].set_ylabel(visu_keys[i], fontsize=8) - - p0 = ax[i, -1].get_position().get_points().flatten() - ax_cbar = fig.add_axes([p0[2] + 0.005, p0[1], 0.0075, p0[3] - p0[1]]) - ticks = np.linspace(0, 1, 5) - tickLabels = np.linspace(c_min, c_max, 5) - tickLabels = [f"{t0:02.2f}" for t0 in tickLabels] - cbar = matplotlib.colorbar.ColorbarBase( - ax_cbar, cmap=plt.get_cmap("inferno"), orientation="vertical", ticks=ticks - ) - cbar.set_ticklabels(tickLabels, fontsize=5) - plt.savefig(f"{filename}", dpi=300) - - -def save_plot_from_2d_dict( - filename: str, - data_dict: Dict[str, Union[np.ndarray, paddle.Tensor]], - visu_keys: Tuple[str, ...], - num_timestamps: int = 1, - stride: int = 1, - xticks: Optional[Tuple[float, ...]] = None, - yticks: Optional[Tuple[float, ...]] = None, -): - """Plot 2d dict data as file. - - Args: - filename (str): Output filename. - data_dict (Dict[str, Union[np.ndarray, paddle.Tensor]]): Data in dict. - visu_keys (Tuple[str, ...]): Keys for visualizing data. such as ("u", "v"). - num_timestamps (int, optional): Number of timestamp in data_dict. Defaults to 1. - stride (int, optional): The time stride of visualization. Defaults to 1. - xticks (Optional[Tuple[float,...]]): The list of xtick locations. Defaults to None. - yticks (Optional[Tuple[float,...]]): The list of ytick locations. Defaults to None. - """ - visu_data = [data_dict[k] for k in visu_keys] - if isinstance(visu_data[0], paddle.Tensor): - visu_data = [x.numpy() for x in visu_data] - _save_plot_from_2d_array( - filename, visu_data, visu_keys, num_timestamps, stride, xticks, yticks - ) - - -# Interface to LineCollection: -def _colorline3d( - x, y, z, t=None, cmap=plt.get_cmap("viridis"), linewidth=1, alpha=1.0, ax=None -): - """ - Plot a colored line with coordinates x and y - Optionally specify colors in the array z - Optionally specify a colormap, a norm function and a line width - https://stackoverflow.com/questions/52884221/how-to-plot-a-matplotlib-line-plot-using-colormap - """ - # Default colors equally spaced on [0, 1]: - if t is None: - t = np.linspace(0.25, 1.0, len(x)) - if ax is None: - ax = plt.gca() - - points = np.array([x, y, z]).T.reshape(-1, 1, 3) - segments = np.concatenate([points[:-1], points[1:]], axis=1) - - colors = np.array([cmap(i) for i in t]) - lc = Line3DCollection(segments, colors=colors, linewidth=linewidth, alpha=alpha) - ax.add_collection(lc) - ax.scatter(x, y, z, c=colors, marker="*", alpha=alpha) # Adding line markers - - -class HandlerColormap(HandlerBase): - """Class for creating colormap legend rectangles. - - Args: - cmap (matplotlib.cm): Matplotlib colormap. - num_stripes (int, optional): Number of contour levels (strips) in rectangle. Defaults to 8. - """ - - def __init__(self, cmap: matplotlib.cm, num_stripes: int = 8, **kw): - HandlerBase.__init__(self, **kw) - self.cmap = cmap - self.num_stripes = num_stripes - - def create_artists( - self, legend, orig_handle, xdescent, ydescent, width, height, fontsize, trans - ): - stripes = [] - for i in range(self.num_stripes): - s = Rectangle( - [xdescent + i * width / self.num_stripes, ydescent], - width / self.num_stripes, - height, - fc=self.cmap((2 * i + 1) / (2 * self.num_stripes)), - transform=trans, - ) - stripes.append(s) - return stripes - - -def _save_plot_from_3d_array( - filename: str, - visu_data: Tuple[np.ndarray, ...], - visu_keys: Tuple[str, ...], - num_timestamps: int = 1, -): - """Save plot from given 3D data. - - Args: - filename (str): Filename. - visu_data (Tuple[np.ndarray, ...]): Data that requires visualization. - visu_keys (Tuple[str, ...]): Keys for visualizing data. such as ("u", "v"). - num_timestamps (int, optional): Number of timestamps coord/value contains. Defaults to 1. - """ - if len(os.path.dirname(filename)): - os.makedirs(os.path.dirname(filename), exist_ok=True) - - fig = plt.figure(figsize=(10, 10)) - len_ts = len(visu_data[0]) // num_timestamps - for t in range(num_timestamps): - ax = fig.add_subplot(1, num_timestamps, t + 1, projection="3d") - st = t * len_ts - ed = (t + 1) * len_ts - visu_data_t = [data[st:ed] for data in visu_data] - cmaps = [] - for i, data in enumerate(visu_data_t): - cmap = plt.get_cmap(CMAPS[i % len(CMAPS)]) - _colorline3d(data[:, 0], data[:, 1], data[:, 2], cmap=cmap, ax=ax) - cmaps.append(cmap) - cmap_handles = [Rectangle((0, 0), 1, 1) for _ in visu_keys] - handler_map = dict( - zip(cmap_handles, [HandlerColormap(cm, num_stripes=8) for cm in cmaps]) - ) - # Create custom legend with color map rectangles - ax.legend( - handles=cmap_handles, - labels=visu_keys, - handler_map=handler_map, - loc="upper right", - framealpha=0.95, - ) - if num_timestamps == 1: - fig.savefig(filename, dpi=300) - else: - fig.savefig(f"{filename}_{t}", dpi=300) - - if num_timestamps == 1: - logger.message(f"3D result is saved to: {filename}.png") - else: - logger.message( - f"3D result is saved to: {filename}_0.png" - f" ~ {filename}_{num_timestamps - 1}.png" - ) - - -def save_plot_from_3d_dict( - filename: str, - data_dict: Dict[str, Union[np.ndarray, paddle.Tensor]], - visu_keys: Tuple[str, ...], - num_timestamps: int = 1, -): - """Plot dict data as file. - - Args: - filename (str): Output filename. - data_dict (Dict[str, Union[np.ndarray, paddle.Tensor]]): Data in dict. - visu_keys (Tuple[str, ...]): Keys for visualizing data. such as ("u", "v"). - num_timestamps (int, optional): Number of timestamp in data_dict. Defaults to 1. - - Examples: - >>> import numpy as np - >>> import ppsci - - >>> data_dict = { - ... "u": np.array([[[10], [20], [30], [40], [50]]]), - ... "v": np.array([[[5], [15], [25], [35], [45]]]), - ... } - - >>> ppsci.visualize.save_plot_from_3d_dict( - ... "path/to/file", - ... data_dict, - ... ("u", "v"), - ... 1, - ... ) # doctest: +SKIP - """ - visu_data = [data_dict[k] for k in visu_keys] - if isinstance(visu_data[0], paddle.Tensor): - visu_data = [x.numpy() for x in visu_data] - - _save_plot_from_3d_array(filename, visu_data, visu_keys, num_timestamps) - - -def _save_plot_weather_from_array( - filename: str, - pred: np.ndarray, - target: np.ndarray, - pred_key: str, - target_key: str, - xticks: Tuple[float, ...], - xticklabels: Tuple[str, ...], - yticks: Tuple[float, ...], - yticklabels: Tuple[str, ...], - vmin: float, - vmax: float, - colorbar_label: str = "", - log_norm: bool = False, -): - """Plot weather result as file from array data. - - Args: - filename (str): Output file name. - pred (np.ndarray): The predict data. - target (np.ndarray): The target data. - pred_key (str): The key of predict data. - target_key (str): The key of target data. - xticks (Tuple[float, ...]): The list of xtick locations. - xticklabels (Tuple[str, ...]): The x-axis' tick labels. - yticks (Tuple[float, ...]): The list of ytick locations. - yticklabels (Tuple[str, ...]): The y-axis' tick labels. - vmin (float): Minimum value that the colormap covers. - vmax (float): Maximal value that the colormap covers. - colorbar_label (str, optional): The color-bar label. Defaults to "". - log_norm (bool, optional): Whether use log norm. Defaults to False. - """ - - def plot_weather( - ax, - data, - title_text, - xticks, - xticklabels, - yticks, - yticklabels, - vmin, - vmax, - log_norm, - cmap=plt.get_cmap("turbo", 1000), - ): - ax.title.set_text(title_text) - ax.set_yticks(yticks) - ax.set_yticklabels(yticklabels) - ax.set_xticks(xticks) - ax.set_xticklabels(xticklabels) - if not log_norm: - map_ = ax.imshow( - data, - interpolation="nearest", - cmap=cmap, - aspect="auto", - vmin=vmin, - vmax=vmax, - ) - else: - norm = matplotlib.colors.LogNorm(vmin=vmin, vmax=vmax, clip=True) - map_ = ax.imshow( - data, interpolation="nearest", cmap=cmap, aspect="auto", norm=norm - ) - plt.colorbar(mappable=map_, cax=None, ax=None, shrink=0.5, label=colorbar_label) - - if len(os.path.dirname(filename)): - os.makedirs(os.path.dirname(filename), exist_ok=True) - fig = plt.figure(facecolor="w", figsize=(7, 7)) - ax = fig.add_subplot(2, 1, 1) - plot_weather( - ax, - pred, - pred_key, - xticks, - xticklabels, - yticks, - yticklabels, - vmin, - vmax, - log_norm, - ) - bx = fig.add_subplot(2, 1, 2) - plot_weather( - bx, - target, - target_key, - xticks, - xticklabels, - yticks, - yticklabels, - vmin, - vmax, - log_norm, - ) - fig.savefig(filename, dpi=300) - plt.close() - - -def save_plot_weather_from_dict( - foldername: str, - data_dict: Dict[str, Union[np.ndarray, paddle.Tensor]], - visu_keys: Tuple[str, ...], - xticks: Tuple[float, ...], - xticklabels: Tuple[str, ...], - yticks: Tuple[float, ...], - yticklabels: Tuple[str, ...], - vmin: float, - vmax: float, - colorbar_label: str = "", - log_norm: bool = False, - num_timestamps: int = 1, -): - """Plot weather result as file from dict data. - - Args: - foldername (str): Output folder name. - data_dict (Dict[str, Union[np.ndarray, paddle.Tensor]]): Data in dict. - visu_keys (Tuple[str, ...]): Keys for visualizing data. such as ("output_6h", "target_6h"). - xticks (Tuple[float, ...]): The list of xtick locations. - xticklabels (Tuple[str, ...]): The x-axis' tick labels. - yticks (Tuple[float, ...]): The list of ytick locations, - yticklabels (Tuple[str, ...]): The y-axis' tick labels. - vmin (float): Minimum value that the colormap covers. - vmax (float): Maximal value that the colormap covers. - colorbar_label (str, optional): The colorbar label. Defaults to "". - log_norm (bool, optional): Whether use log norm. Defaults to False. - num_timestamps (int): Number of timestamp in data_dict. Defaults to 1. - """ - os.makedirs(foldername, exist_ok=True) - - visu_data = [data_dict[k] for k in visu_keys] - if isinstance(visu_data[0], paddle.Tensor): - visu_data = [x.numpy() for x in visu_data] - - frames = [] - for t in range(num_timestamps): - pred_key, target_key = visu_keys[2 * t], visu_keys[2 * t + 1] - pred_data = visu_data[2 * t] - target_data = visu_data[2 * t + 1] - filename_t = os.path.join(foldername, f"{t}.png") - _save_plot_weather_from_array( - filename_t, - pred_data, - target_data, - pred_key, - target_key, - xticks, - xticklabels, - yticks, - yticklabels, - vmin=vmin, - vmax=vmax, - colorbar_label=colorbar_label, - log_norm=log_norm, - ) - frames.append(imageio.imread(filename_t)) - filename = os.path.join(foldername, "result.gif") - imageio.mimsave( - filename, - frames, - "GIF", - duration=1, - ) diff --git a/examples/smc_reac/ppsci/visualize/radar.py b/examples/smc_reac/ppsci/visualize/radar.py deleted file mode 100644 index abde75b775..0000000000 --- a/examples/smc_reac/ppsci/visualize/radar.py +++ /dev/null @@ -1,124 +0,0 @@ -# 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. - -from __future__ import annotations - -import os -from typing import Callable -from typing import Dict - -import matplotlib.pyplot as plt -import numpy as np - -from ppsci.visualize import base - - -class VisualizerRadar(base.Visualizer): - """Visualizer for NowcastNet Radar Dataset. - - Args: - input_dict (Dict[str, np.ndarray]): Input dict. - output_expr (Dict[str, Callable]): Output expression. - batch_size (int, optional): Batch size of data when computing result in visu.py. Defaults to 64. - num_timestamps (int, optional): Number of timestamps - prefix (str, optional): Prefix for output file. - case_type (str, optional): Case type. - total_length (str, optional): Total length. - - Examples: - >>> import ppsci - >>> import paddle - >>> frames_tensor = paddle.randn([1, 29, 512, 512, 2]) - >>> visualizer = ppsci.visualize.VisualizerRadar( - ... {"input": frames_tensor}, - ... {"output": lambda out: out["output"]}, - ... num_timestamps=1, - ... prefix="v_nowcastnet", - ... ) - """ - - def __init__( - self, - input_dict: Dict[str, np.ndarray], - output_expr: Dict[str, Callable], - batch_size: int = 64, - num_timestamps: int = 1, - prefix: str = "vtu", - case_type: str = "normal", - total_length: int = 29, - ): - super().__init__(input_dict, output_expr, batch_size, num_timestamps, prefix) - self.case_type = case_type - self.total_length = total_length - self.input_dict = input_dict - - def save(self, path, data_dict): - if not os.path.exists(path): - os.makedirs(path) - test_ims = self.input_dict[list(self.input_dict.keys())[0]] - # keys: {"input", "output"} - img_gen = data_dict[list(data_dict.keys())[1]] - vis_info = {"vmin": 1, "vmax": 40} - if self.case_type == "normal": - test_ims_plot = test_ims[0][ - :-2, 256 - 192 : 256 + 192, 256 - 192 : 256 + 192 - ] - img_gen_plot = img_gen[0][:-2, 256 - 192 : 256 + 192, 256 - 192 : 256 + 192] - else: - test_ims_plot = test_ims[0][:-2] - img_gen_plot = img_gen[0][:-2] - save_plots( - test_ims_plot, - labels=[f"gt{i + 1}" for i in range(self.total_length)], - res_path=path, - vmin=vis_info["vmin"], - vmax=vis_info["vmax"], - ) - save_plots( - img_gen_plot, - labels=[f"pd{i + 1}" for i in range(9, self.total_length)], - res_path=path, - vmin=vis_info["vmin"], - vmax=vis_info["vmax"], - ) - - -def save_plots( - field, - labels, - res_path, - figsize=None, - vmin=0, - vmax=10, - cmap="viridis", - npy=False, - **imshow_args, -): - for i, data in enumerate(field): - if i >= len(labels): - break - plt.figure(figsize=figsize) - ax = plt.axes() - ax.set_axis_off() - alpha = data[..., 0] / 1 - alpha[alpha < 1] = 0 - alpha[alpha > 1] = 1 - ax.imshow( - data[..., 0], alpha=alpha, vmin=vmin, vmax=vmax, cmap=cmap, **imshow_args - ) - plt.savefig(os.path.join(res_path, labels[i] + ".png")) - plt.close() - if npy: - with open(os.path.join(res_path, labels[i] + ".npy"), "wb") as f: - np.save(f, data[..., 0]) diff --git a/examples/smc_reac/ppsci/visualize/visualizer.py b/examples/smc_reac/ppsci/visualize/visualizer.py deleted file mode 100644 index e3e602daa5..0000000000 --- a/examples/smc_reac/ppsci/visualize/visualizer.py +++ /dev/null @@ -1,409 +0,0 @@ -# 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. - -from __future__ import annotations - -import os.path as osp -from typing import Callable -from typing import Dict -from typing import Optional -from typing import Tuple - -import numpy as np - -from ppsci.visualize import base -from ppsci.visualize import plot -from ppsci.visualize import vtu - - -class VisualizerScatter1D(base.Visualizer): - """Visualizer for 1d scatter data. - - Args: - input_dict (Dict[str, np.ndarray]): Input dict. - coord_keys (Tuple[str, ...]): Coordinate keys, such as ("x", "y"). - output_expr (Dict[str, Callable]): Output expression. - batch_size (int, optional): Batch size of data when computing result in visu.py. Defaults to 64. - num_timestamps (int, optional): Number of timestamps. Defaults to 1. - prefix (str, optional): Prefix for output file. Defaults to "plot". - - Examples: - >>> import ppsci - >>> visu_mat = {"t_f": np.random.randn(16, 1), "eta": np.random.randn(16, 1)} - >>> visualizer_eta = ppsci.visualize.VisualizerScatter1D( - ... visu_mat, - ... ("t_f",), - ... {"eta": lambda d: d["eta"]}, - ... num_timestamps=1, - ... prefix="viv_pred", - ... ) - """ - - def __init__( - self, - input_dict: Dict[str, np.ndarray], - coord_keys: Tuple[str, ...], - output_expr: Dict[str, Callable], - batch_size: int = 64, - num_timestamps: int = 1, - prefix: str = "plot", - ): - super().__init__(input_dict, output_expr, batch_size, num_timestamps, prefix) - self.coord_keys = coord_keys - - def save(self, filename, data_dict): - plot.save_plot_from_1d_dict( - filename, data_dict, self.coord_keys, self.output_keys, self.num_timestamps - ) - - -class VisualizerScatter3D(base.Visualizer): - """Visualizer for 3d scatter data. - - Args: - input_dict (Dict[str, np.ndarray]): Input dict. - output_expr (Dict[str, Callable]): Output expression. - batch_size (int, optional): Batch size of data when computing result in visu.py. Defaults to 64. - num_timestamps (int, optional): Number of timestamps. Defaults to 1. - prefix (str, optional): Prefix for output file. Defaults to "plot3d_scatter". - - Examples: - >>> import ppsci - >>> vis_data = {"states": np.random.randn(16, 1)} - >>> visualizer = ppsci.visualize.VisualizerScatter3D( - ... vis_data, - ... {"states": lambda d: d["states"]}, - ... num_timestamps=1, - ... prefix="result_states", - ... ) - """ - - def __init__( - self, - input_dict: Dict[str, np.ndarray], - output_expr: Dict[str, Callable], - batch_size: int = 64, - num_timestamps: int = 1, - prefix: str = "plot3d_scatter", - ): - super().__init__(input_dict, output_expr, batch_size, num_timestamps, prefix) - - def save(self, filename, data_dict): - data_dict = { - key: value for key, value in data_dict.items() if key in self.output_keys - } - value = data_dict[self.output_keys[0]] - dim = len(value.shape) - if dim == 3: - # value.shape=(B, T, 3) - for i in range(value.shape[0]): - cur_data_dict = {key: value[i] for key, value in data_dict.items()} - plot.save_plot_from_3d_dict( - filename + str(i), - cur_data_dict, - self.output_keys, - self.num_timestamps, - ) - else: - # value.shape=(T, 3) - plot.save_plot_from_3d_dict( - filename, data_dict, self.output_keys, self.num_timestamps - ) - - -class VisualizerVtu(base.Visualizer): - """Visualizer for 2D points data. - - Args: - input_dict (Dict[str, np.ndarray]): Input dict. - output_expr (Dict[str, Callable]): Output expression. - batch_size (int, optional): Batch size of data when computing result in visu.py. Defaults to 64. - num_timestamps (int, optional): Number of timestamps - prefix (str, optional): Prefix for output file. - - Examples: - >>> import ppsci - >>> vis_points = { - ... "x": np.random.randn(128, 1), - ... "y": np.random.randn(128, 1), - ... "u": np.random.randn(128, 1), - ... "v": np.random.randn(128, 1), - ... } - >>> visualizer_u_v = ppsci.visualize.VisualizerVtu( - ... vis_points, - ... {"u": lambda d: d["u"], "v": lambda d: d["v"]}, - ... num_timestamps=1, - ... prefix="result_u_v", - ... ) - """ - - def __init__( - self, - input_dict: Dict[str, np.ndarray], - output_expr: Dict[str, Callable], - batch_size: int = 64, - num_timestamps: int = 1, - prefix: str = "vtu", - ): - super().__init__(input_dict, output_expr, batch_size, num_timestamps, prefix) - - def save(self, filename, data_dict): - vtu.save_vtu_from_dict( - filename, data_dict, self.input_keys, self.output_keys, self.num_timestamps - ) - - -class Visualizer2D(base.Visualizer): - """Visualizer for 2D data. - - Args: - input_dict (Dict[str, np.ndarray]): Input dict. - output_expr (Dict[str, Callable]): Output expression. - batch_size (int, optional): Batch size of data when computing result in visu.py. Defaults to 64. - num_timestamps (int, optional): Number of timestamps. Defaults to 1. - prefix (str, optional): Prefix for output file. Defaults to "plot2d". - - Examples: - >>> import ppsci - >>> vis_points = { - ... "x": np.random.randn(128, 1), - ... "y": np.random.randn(128, 1), - ... "u": np.random.randn(128, 1), - ... "v": np.random.randn(128, 1), - ... } - >>> visualizer_u_v = ppsci.visualize.Visualizer2D( - ... vis_points, - ... {"u": lambda d: d["u"], "v": lambda d: d["v"]}, - ... num_timestamps=1, - ... prefix="result_u_v", - ... ) - """ - - def __init__( - self, - input_dict: Dict[str, np.ndarray], - output_expr: Dict[str, Callable], - batch_size: int = 64, - num_timestamps: int = 1, - prefix: str = "plot2d", - ): - super().__init__(input_dict, output_expr, batch_size, num_timestamps, prefix) - - -class Visualizer2DPlot(Visualizer2D): - """Visualizer for 2D data use matplotlib. - - Args: - input_dict (Dict[str, np.ndarray]): Input dict. - output_expr (Dict[str, Callable]): Output expression. - batch_size (int, optional): Batch size of data when computing result in visu.py. Defaults to 64. - num_timestamps (int, optional): Number of timestamps. - stride (int, optional): The time stride of visualization. Defaults to 1. - xticks (Optional[Tuple[float,...]]): The list of xtick locations. Defaults to None. - yticks (Optional[Tuple[float,...]]): The list of ytick locations. Defaults to None. - prefix (str, optional): Prefix for output file. Defaults to "plot2d". - - Examples: - >>> import ppsci - >>> vis_data = { - ... "target_ux": np.random.randn(128, 20, 1), - ... "pred_ux": np.random.randn(128, 20, 1), - ... } - >>> visualizer_states = ppsci.visualize.Visualizer2DPlot( - ... vis_data, - ... { - ... "target_ux": lambda d: d["states"][:, :, 0], - ... "pred_ux": lambda d: output_transform(d)[:, :, 0], - ... }, - ... batch_size=1, - ... num_timestamps=10, - ... stride=20, - ... xticks=np.linspace(-2, 14, 9), - ... yticks=np.linspace(-4, 4, 5), - ... prefix="result_states", - ... ) - """ - - def __init__( - self, - input_dict: Dict[str, np.ndarray], - output_expr: Dict[str, Callable], - batch_size: int = 64, - num_timestamps: int = 1, - stride: int = 1, - xticks: Optional[Tuple[float, ...]] = None, - yticks: Optional[Tuple[float, ...]] = None, - prefix: str = "plot2d", - ): - super().__init__(input_dict, output_expr, batch_size, num_timestamps, prefix) - self.stride = stride - self.xticks = xticks - self.yticks = yticks - - def save(self, filename, data_dict): - data_dict = { - key: value for key, value in data_dict.items() if key in self.output_keys - } - value = data_dict[self.output_keys[0]] - dim = len(value.shape) - if dim == 4: - # value.shape=(B, T, H, W) - for i in range(value.shape[0]): - cur_data_dict = {key: value[i] for key, value in data_dict.items()} - plot.save_plot_from_2d_dict( - filename + str(i), - cur_data_dict, - self.output_keys, - self.num_timestamps, - self.stride, - self.xticks, - self.yticks, - ) - else: - # value.shape=(T, H, W) - plot.save_plot_from_2d_dict( - filename, - data_dict, - self.output_keys, - self.num_timestamps, - self.stride, - self.xticks, - self.yticks, - ) - - -class Visualizer3D(base.Visualizer): - """Visualizer for 3D plot data. - - Args: - input_dict (Dict[str, np.ndarray]): Input dict. - output_expr (Dict[str, Callable]): Output expression. - batch_size (int, optional): Batch size of data when computing result in visu.py. Defaults to 64. - label_dict (Dict[str, np.ndarray]): Label dict. - time_list (Optional[Tuple[float, ...]]): Time list. - prefix (str, optional): Prefix for output file. - """ - - def __init__( - self, - input_dict: Dict[str, np.ndarray], - output_expr: Dict[str, Callable], - batch_size: int = 64, - label_dict: Optional[Dict[str, np.ndarray]] = None, - time_list: Optional[Tuple[float, ...]] = None, - prefix: str = "vtu", - ): - self.label = label_dict - self.time_list = time_list - super().__init__(input_dict, output_expr, batch_size, len(time_list), prefix) - - def save(self, filename: str, data_dict: Dict[str, np.ndarray]): - n = int((next(iter(data_dict.values()))).shape[0] / self.num_timestamps) - coord_keys = [x for x in self.input_dict if x != "t"] - for i in range(len(self.time_list)): - vtu.save_vtu_to_mesh( - osp.join(filename, f"predict_{i+1}.vtu"), - {key: (data_dict[key][i * n : (i + 1) * n]) for key in data_dict}, - coord_keys, - self.output_keys, - ) - - -class VisualizerWeather(base.Visualizer): - """Visualizer for weather data use matplotlib. - - Args: - input_dict (Dict[str, np.ndarray]): Input dict. - output_expr (Dict[str, Callable]): Output expression. - xticks (Tuple[float, ...]): The list of xtick locations. - xticklabels (Tuple[str, ...]): The x-axis' tick labels. - yticks (Tuple[float, ...]): The list of ytick locations. - yticklabels (Tuple[str, ...]): The y-axis' tick labels. - vmin (float): Minimum value that the colormap covers. - vmax (float): Maximal value that the colormap covers. - colorbar_label (str, optional): The color-bar label. Defaults to "". - log_norm (bool, optional): Whether use log norm. Defaults to False. - batch_size (int, optional): : Batch size of data when computing result in visu.py. Defaults to 1. - num_timestamps (int, optional): Number of timestamps. Defaults to 1. - prefix (str, optional): Prefix for output file. Defaults to "plot_weather". - - Examples: - >>> import ppsci - >>> import numpy as np - >>> vis_data = { - ... "output_6h": np.random.randn(1, 720, 1440), - ... "target_6h": np.random.randn(1, 720, 1440), - ... } - >>> visualizer_weather = ppsci.visualize.VisualizerWeather( - ... vis_data, - ... { - ... "output_6h": lambda d: d["output_6h"], - ... "target_6h": lambda d: d["target_6h"], - ... }, - ... xticks=np.linspace(0, 1439, 13), - ... xticklabels=[str(i) for i in range(360, -1, -30)], - ... yticks=np.linspace(0, 719, 7), - ... yticklabels=[str(i) for i in range(90, -91, -30)], - ... vmin=0, - ... vmax=25, - ... prefix="result_states", - ... ) - """ - - def __init__( - self, - input_dict: Dict[str, np.ndarray], - output_expr: Dict[str, Callable], - xticks: Tuple[float, ...], - xticklabels: Tuple[str, ...], - yticks: Tuple[float, ...], - yticklabels: Tuple[str, ...], - vmin: float, - vmax: float, - colorbar_label: str = "", - log_norm: bool = False, - batch_size: int = 1, - num_timestamps: int = 1, - prefix: str = "plot_weather", - ): - super().__init__(input_dict, output_expr, batch_size, num_timestamps, prefix) - self.xticks = xticks - self.xticklabels = xticklabels - self.yticks = yticks - self.yticklabels = yticklabels - self.vmin = vmin - self.vmax = vmax - self.colorbar_label = colorbar_label - self.log_norm = log_norm - - def save(self, filename, data_dict): - data_dict = {key: data_dict[key] for key in self.output_keys} - value = data_dict[self.output_keys[0]] - # value.shape=(B, H, W) - for i in range(value.shape[0]): - cur_data_dict = {key: value[i] for key, value in data_dict.items()} - plot.save_plot_weather_from_dict( - filename + str(i), - cur_data_dict, - self.output_keys, - self.xticks, - self.xticklabels, - self.yticks, - self.yticklabels, - self.vmin, - self.vmax, - self.colorbar_label, - self.log_norm, - self.num_timestamps, - ) diff --git a/examples/smc_reac/ppsci/visualize/vtu.py b/examples/smc_reac/ppsci/visualize/vtu.py deleted file mode 100644 index 500c7e2e84..0000000000 --- a/examples/smc_reac/ppsci/visualize/vtu.py +++ /dev/null @@ -1,278 +0,0 @@ -# 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. - -from __future__ import annotations - -import os -from typing import Dict -from typing import Tuple - -import meshio -import numpy as np -import paddle -from pyevtk import hl - -from ppsci.utils import logger - - -def _save_vtu_from_array(filename, coord, value, value_keys, num_timestamps=1): - """Save data to '*.vtu' file(s). - - Args: - filename (str): Output filename. - coord (np.ndarray): Coordinate points with shape of [N, 2] or [N, 3]. - value (np.ndarray): Value of each coord points with shape of [N, M]. - value_keys (Tuple[str, ...]): Names of each dimension of value, such as ("u", "v"). - num_timestamps (int, optional): Number of timestamp over coord and value. - Defaults to 1. - """ - if not isinstance(coord, np.ndarray): - raise ValueError(f"type of coord({type(coord)}) should be ndarray.") - if value is not None and not isinstance(value, np.ndarray): - raise ValueError(f"type of value({type(value)}) should be ndarray.") - if value is not None and len(coord) != len(value): - raise ValueError( - f"coord length({len(coord)}) should be equal to value length({len(value)})" - ) - if len(coord) % num_timestamps != 0: - raise ValueError( - f"coord length({len(coord)}) should be an integer multiple of " - f"num_timestamps({num_timestamps})" - ) - if coord.shape[1] not in [2, 3]: - raise ValueError(f"ndim of coord({coord.shape[1]}) should be 2 or 3.") - - if len(os.path.dirname(filename)): - os.makedirs(os.path.dirname(filename), exist_ok=True) - - # discard extension name - if filename.endswith(".vtu"): - filename = filename[:-4] - npoint = len(coord) - coord_ndim = coord.shape[1] - - if value is None: - value = np.ones([npoint, 1], dtype=coord.dtype) - value_keys = ["dummy_key"] - - data_ndim = value.shape[1] - nx = npoint // num_timestamps - for t in range(num_timestamps): - # NOTE: each array in data_vtu should be 1-dim, i.e. [N, 1] will occur error. - if coord_ndim == 2: - axis_x = np.ascontiguousarray(coord[t * nx : (t + 1) * nx, 0]) - axis_y = np.ascontiguousarray(coord[t * nx : (t + 1) * nx, 1]) - axis_z = np.zeros([nx], dtype=paddle.get_default_dtype()) - elif coord_ndim == 3: - axis_x = np.ascontiguousarray(coord[t * nx : (t + 1) * nx, 0]) - axis_y = np.ascontiguousarray(coord[t * nx : (t + 1) * nx, 1]) - axis_z = np.ascontiguousarray(coord[t * nx : (t + 1) * nx, 2]) - - data_vtu = {} - for j in range(data_ndim): - data_vtu[value_keys[j]] = np.ascontiguousarray( - value[t * nx : (t + 1) * nx, j] - ) - - if num_timestamps > 1: - width = len(str(num_timestamps - 1)) - hl.pointsToVTK( - f"{filename}_t-{t:0{width}}", axis_x, axis_y, axis_z, data=data_vtu - ) - else: - hl.pointsToVTK(filename, axis_x, axis_y, axis_z, data=data_vtu) - - if num_timestamps > 1: - logger.message( - f"Visualization results are saved to: {filename}_t-{0:0{width}}.vtu ~ " - f"{filename}_t-{num_timestamps - 1:0{width}}.vtu" - ) - else: - logger.message(f"Visualization result is saved to: {filename}.vtu") - - -def save_vtu_from_dict( - filename: str, - data_dict: Dict[str, np.ndarray], - coord_keys: Tuple[str, ...], - value_keys: Tuple[str, ...], - num_timestamps: int = 1, -): - """Save dict data to '*.vtu' file. - - Args: - filename (str): Output filename. - data_dict (Dict[str, np.ndarray]): Data in dict. - coord_keys (Tuple[str, ...]): Tuple of coord key. such as ("x", "y"). - value_keys (Tuple[str, ...]): Tuple of value key. such as ("u", "v"). - num_timestamps (int, optional): Number of timestamp in data_dict. Defaults to 1. - - Examples: - >>> import ppsci - >>> import numpy as np - >>> filename = "path/to/file.vtu" - >>> data_dict = { - ... "x": np.array([[1], [2], [3],[4]]), - ... "y": np.array([[2], [3], [4],[4]]), - ... "z": np.array([[3], [4], [5],[4]]), - ... "u": np.array([[4], [5], [6],[4]]), - ... "v": np.array([[5], [6], [7],[4]]), - ... } - >>> coord_keys = ("x","y","z") - >>> value_keys = ("u","v") - >>> ppsci.visualize.save_vtu_from_dict(filename, data_dict, coord_keys, value_keys) # doctest: +SKIP - """ - if len(coord_keys) not in [2, 3, 4]: - raise ValueError(f"ndim of coord ({len(coord_keys)}) should be 2, 3 or 4") - - coord = [data_dict[k] for k in coord_keys if k not in ("t", "sdf")] - value = [data_dict[k] for k in value_keys] if value_keys else None - - coord = np.concatenate(coord, axis=1) - - if value is not None: - value = np.concatenate(value, axis=1) - - _save_vtu_from_array(filename, coord, value, value_keys, num_timestamps) - - -def save_vtp_from_dict( - filename: str, - data_dict: Dict[str, np.ndarray], - coord_keys: Tuple[str, ...], - value_keys: Tuple[str, ...], - num_timestamps: int = 1, -): - """Save dict data to '*.vtp' file. - - Args: - filename (str): Output filename. - data_dict (Dict[str, np.ndarray]): Data in dict. - coord_keys (Tuple[str, ...]): Tuple of coord key. such as ("x", "y"). - value_keys (Tuple[str, ...]): Tuple of value key. such as ("u", "v"). - num_timestamps (int, optional): Number of timestamp in data_dict. Defaults to 1. - - Examples: - >>> import ppsci - >>> import numpy as np - >>> filename = "path/to/file.vtp" - >>> data_dict = { - ... "x": np.array([[1], [2], [3],[4]]), - ... "y": np.array([[2], [3], [4],[4]]), - ... "z": np.array([[3], [4], [5],[4]]), - ... "u": np.array([[4], [5], [6],[4]]), - ... "v": np.array([[5], [6], [7],[4]]), - ... } - >>> coord_keys = ("x","y","z") - >>> value_keys = ("u","v") - >>> ppsci.visualize.save_vtp_from_dict(filename, data_dict, coord_keys, value_keys) # doctest: +SKIP - """ - import pyvista as pv - - if len(coord_keys) not in [3]: - raise ValueError(f"ndim of coord ({len(coord_keys)}) should be 3 in vtp format") - - coord = [data_dict[k] for k in coord_keys if k not in ("t", "sdf")] - assert all([c.ndim == 2 for c in coord]), "array of each axis should be [*, 1]" - coord = np.concatenate(coord, axis=1) - - if not isinstance(coord, np.ndarray): - raise ValueError(f"type of coord({type(coord)}) should be ndarray.") - if len(coord) % num_timestamps != 0: - raise ValueError( - f"coord length({len(coord)}) should be an integer multiple of " - f"num_timestamps({num_timestamps})" - ) - if coord.shape[1] not in [3]: - raise ValueError(f"ndim of coord({coord.shape[1]}) should be 3 in vtp format.") - - if len(os.path.dirname(filename)): - os.makedirs(os.path.dirname(filename), exist_ok=True) - - npoint = len(coord) - nx = npoint // num_timestamps - if filename.endswith(".vtp"): - filename = filename[:-4] - - for t in range(num_timestamps): - coord_ = coord[t * nx : (t + 1) * nx] - point_cloud = pv.PolyData(coord_) - for k in value_keys: - value_ = data_dict[k][t * nx : (t + 1) * nx] - if value_ is not None and not isinstance(value_, np.ndarray): - raise ValueError(f"type of value({type(value_)}) should be ndarray.") - if value_ is not None and len(coord_) != len(value_): - raise ValueError( - f"coord length({len(coord_)}) should be equal to value length({len(value_)})" - ) - point_cloud[k] = value_ - - if num_timestamps > 1: - width = len(str(num_timestamps - 1)) - point_cloud.save(f"{filename}_t-{t:0{width}}.vtp") - else: - point_cloud.save(f"{filename}.vtp") - - if num_timestamps > 1: - logger.message( - f"Visualization results are saved to: {filename}_t-{0:0{width}}.vtp ~ " - f"{filename}_t-{num_timestamps - 1:0{width}}.vtp" - ) - else: - logger.message(f"Visualization result is saved to: {filename}.vtp") - - -def save_vtu_to_mesh( - filename: str, - data_dict: Dict[str, np.ndarray], - coord_keys: Tuple[str, ...], - value_keys: Tuple[str, ...], -): - """Save data into .vtu format by meshio. - - Args: - filename (str): File name. - data_dict (Dict[str, np.ndarray]): Data in dict. - coord_keys (Tuple[str, ...]): Tuple of coord key. such as ("x", "y"). - value_keys (Tuple[str, ...]): Tuple of value key. such as ("u", "v"). - - Examples: - >>> import ppsci - >>> import numpy as np - >>> filename = "path/to/file.vtu" - >>> data_dict = { - ... "x": np.array([[1], [2], [3],[4]]), - ... "y": np.array([[2], [3], [4],[4]]), - ... "z": np.array([[3], [4], [5],[4]]), - ... "u": np.array([[4], [5], [6],[4]]), - ... "v": np.array([[5], [6], [7],[4]]), - ... } - >>> coord_keys = ("x","y","z") - >>> value_keys = ("u","v") - >>> ppsci.visualize.save_vtu_to_mesh(filename, data_dict, coord_keys, value_keys) # doctest: +SKIP - """ - npoint = len(next(iter(data_dict.values()))) - coord_ndim = len(coord_keys) - - # get the list variable transposed - points = np.stack(tuple(data_dict[key] for key in coord_keys)).reshape( - coord_ndim, npoint - ) - mesh = meshio.Mesh( - points=points.T, cells=[("vertex", np.arange(npoint).reshape(npoint, 1))] - ) - mesh.point_data = {key: data_dict[key] for key in value_keys} - if len(os.path.dirname(filename)): - os.makedirs(os.path.dirname(filename), exist_ok=True) - mesh.write(filename)