FilesExpand file tree

main

/

surrogate-cli.test.ts

Copy path

Blame

More file actions

Blame

More file actions

Latest commit

 

History

History

History

173 lines (155 loc) · 5.87 KB

main

/

surrogate-cli.test.ts

Copy path

Top

File metadata and controls

• Code
• Blame

173 lines (155 loc) · 5.87 KB

Raw

Copy raw file

Download raw file

Open symbols panel

Edit and raw actions

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

22

23

24

25

26

27

28

29

30

31

32

33

34

35

36

37

38

39

40

41

42

43

44

45

46

47

48

49

50

51

52

53

54

55

56

57

58

59

60

61

62

63

64

65

66

67

68

69

70

71

72

73

74

75

76

77

78

79

80

81

82

83

84

85

86

87

88

89

90

91

92

93

94

95

96

97

98

99

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

166

167

168

169

170

171

172

173

// SPDX-License-Identifier: AGPL-3.0-or-later

import { exportROMToONNX, trainROM } from "@modelscript/simulate";

import assert from "node:assert/strict";

import fs from "node:fs/promises";

import path from "node:path";

import test, { after, before, describe } from "node:test";

import { Surrogate } from "../src/commands/surrogate.js";

describe("ModelScript CLI Surrogate Command (Export Formats, Guardrails, ONNX Import)", () => {

const tmpDir = path.resolve(process.cwd(), "scratch_test_surrogate_cli");

const onnxPath = path.join(tmpDir, "test_model.onnx.json");

const moPath = path.join(tmpDir, "Cooler_Surrogate.mo");

const cPath = path.join(tmpDir, "Cooler_surrogate.c");

const reportPath = path.join(tmpDir, "surrogate_report.json");

let origLog: typeof console.log;

let origError: typeof console.error;

let origWrite: typeof process.stdout.write;

before(async () => {

origLog = console.log;

origError = console.error;

origWrite = process.stdout.write;

console.log = () => {};

console.error = () => {};

(process.stdout as any).write = () => true;

await fs.mkdir(tmpDir, { recursive: true });

// Create a mock ONNX graph file using exportROMToONNX

const dummyRom = trainROM({

data: {

inputs: [

[290.0, 1.0],

[310.0, 2.0],

[330.0, 3.0],

],

outputs: [[300.0], [315.0], [325.0]],

inputNames: ["inlet_temp", "mass_flow"],

outputNames: ["outlet_temp"],

snapshotTimes: [0.0],

isTransient: false,

wallClockMs: 1,

failedSamples: 0,

totalSamples: 3,

},

architecture: "mlp",

hiddenLayers: [8, 8],

activation: "tanh",

epochs: 50,

learningRate: 0.01,

seed: 42,

});

const onnxData = exportROMToONNX(dummyRom, "Cooler");

await fs.writeFile(onnxPath, JSON.stringify(onnxData, null, 2), "utf-8");

});

after(async () => {

console.log = origLog;

console.error = origError;

process.stdout.write = origWrite;

await fs.rm(tmpDir, { recursive: true, force: true });

});

test("imports ONNX model and exports Modelica block with active extrapolation guardrails", async () => {

await Surrogate.handler({

name: "Cooler",

fromOnnx: onnxPath,

format: "mo",

"input-bounds": "inlet_temp:280.0:350.0,mass_flow:0.5:5.0",

inputBounds: "inlet_temp:280.0:350.0,mass_flow:0.5:5.0",

"extrapolation-guardrails": true,

extrapolationGuardrails: true,

"extrapolation-flag": true,

extrapolationFlag: true,

"extrapolation-level": "warning",

extrapolationLevel: "warning",

"out-file": moPath,

outFile: moPath,

report: reportPath,

strategy: "latin-hypercube",

"num-samples": 10,

numSamples: 10,

architecture: "mlp",

} as any);

const moContent = await fs.readFile(moPath, "utf-8");

assert.ok(moContent.includes("block Cooler_Surrogate"));

assert.ok(moContent.includes("input Real inlet_temp"));

assert.ok(moContent.includes("input Real mass_flow"));

assert.ok(moContent.includes("output Real outlet_temp"));

assert.ok(moContent.includes("output Boolean isExtrapolating"));

assert.ok(moContent.includes("assert(inlet_temp >= 280.0 and inlet_temp <= 350.0"));

assert.ok(moContent.includes("assert(mass_flow >= 0.5 and mass_flow <= 5.0"));

assert.ok(moContent.includes("isExtrapolating ="));

assert.ok(moContent.includes("algorithm"));

assert.ok(moContent.includes("Modelica.Math.tanh"));

const reportContent = JSON.parse(await fs.readFile(reportPath, "utf-8"));

assert.equal(reportContent.model, "Cooler");

assert.deepEqual(reportContent.parameterBounds.inlet_temp, { min: 280, max: 350 });

});

test("imports ONNX model and exports zero-allocation C evaluation kernel", async () => {

await Surrogate.handler({

name: "Cooler",

fromOnnx: onnxPath,

format: "c",

"out-file": cPath,

outFile: cPath,

report: reportPath,

strategy: "latin-hypercube",

"num-samples": 10,

numSamples: 10,

architecture: "mlp",

} as any);

const cContent = await fs.readFile(cPath, "utf-8");

assert.ok(cContent.includes("#include <math.h>"));

assert.ok(cContent.includes("static const double W0") || cContent.includes("static const double W_0"));

assert.ok(cContent.includes("void Cooler_evaluate") || cContent.includes("rom_evaluate"));

});

test("surrogatizes and replaces submodel component in parent Modelica file", async () => {

const systemMoPath = path.join(tmpDir, "System.mo");

const replacedMoPath = path.join(tmpDir, "System_Surrogatized.mo");

const systemMoContent = `model System

Source pump;

HeatExchanger cooler(A = 15.0);

Sink reservoir;

equation

connect(pump.port, cooler.port_a);

connect(cooler.port_b, reservoir.port);

end System;`;

await fs.writeFile(systemMoPath, systemMoContent, "utf-8");

await Surrogate.handler({

name: "Cooler",

paths: [systemMoPath],

fromOnnx: onnxPath,

format: "mo",

"out-file": moPath,

outFile: moPath,

"replace-component": "cooler",

replaceComponent: "cooler",

"replace-out": replacedMoPath,

replaceOut: replacedMoPath,

report: reportPath,

strategy: "latin-hypercube",

"num-samples": 10,

numSamples: 10,

architecture: "mlp",

} as any);

const replacedContent = await fs.readFile(replacedMoPath, "utf-8");

assert.ok(

replacedContent.includes("Cooler_Surrogate cooler(A = 15.0);"),

"Declaration must be replaced with Cooler_Surrogate",

);

assert.ok(

replacedContent.includes("connect(pump.port, cooler.port_a);"),

"Connector wiring port_a must remain valid and intact",

);

assert.ok(

replacedContent.includes("connect(cooler.port_b, reservoir.port);"),

"Connector wiring port_b must remain valid and intact",

);

});

});