fix cmloss num param issue (fc726995) · Commits · HyGNN / capabilities / reflectometry

scripts/train_ntrace.py

+1 −1

Original line number	Diff line number	Diff line
		@@ -23,7 +23,7 @@ if __name__ == "__main__":
		"optimizer": "SGD",
		"loss": "composite",
		"cache_dir": "data",
		"experiment_name": "nTrace_mo_film_5layer",
		"experiment_name": "nTrace_cu_film_6layer",
		"run_name": "d2048_h32_l6",
		"datadir": "expdata",
		}

src/tgreft/train/generic.py

+2 −2

Original line number	Diff line number	Diff line
		@@ -215,7 +215,7 @@ def get_optimizer(
		raise ValueError(f"Unknown optimizer: {name}")


		def get_loss(name: str) -> torch.nn.Module:
		def get_loss(name: str, num_param: int = 17) -> torch.nn.Module:
		"""Get the loss function.

		Parameters
		@@ -235,6 +235,6 @@ def get_loss(name: str) -> torch.nn.Module:
		elif name.lower() == "huber":
		return torch.nn.SmoothL1Loss()
		elif name.lower() == "composite":
		return CompositeLoss()
		return CompositeLoss(num_param=num_param)
		else:
		raise ValueError(f"Unknown loss function: {name}")

src/tgreft/train/train_ntrace.py

+2 −3

Original line number	Diff line number	Diff line
		#!/usr/bin/env python3
		"""Train function for the transformer model."""
		import os
		import shutil
		import torch
		import mlflow
		import logging
		@@ -127,7 +126,7 @@ def train(

		# prepare loss function
		logger.info("Preparing loss function...")
		loss_fn = get_loss(name=loss)
		loss_fn = get_loss(name=loss, num_param=output_dim)

		# prepare optimizer
		logger.info("Preparing optimizer...")

src/tgreft/utils/data/data_loader.py

+38 −43

Original line number	Diff line number	Diff line
		@@ -85,26 +85,45 @@ def generate_data(
		# NOTE: the n-Trace model is more expressive with relaxed bounds

		# cu_film, 6 layer
		parameters_ref = np.column_stack(
		[
		np.random.uniform(5.0, 7.0, n_dataset), # electolyte_sld,
		np.random.uniform(5, 120, n_dataset), # electolyte_roughness,
		np.random.uniform(-5.0, 6.5, n_dataset), # sei_sld,
		np.random.uniform(10, 500, n_dataset), # sei_thickness,
		np.random.uniform(1, 80, n_dataset), # sei_roughness,
		np.random.uniform(-2, 6, n_dataset), # material_sld,
		np.random.uniform(10, 200, n_dataset), # material_thickness,
		np.random.uniform(1, 35, n_dataset), # material_roughness,
		np.random.uniform(6, 7, n_dataset), # cu_sld,
		np.random.uniform(20, 700, n_dataset), # cu_thickness,
		np.random.uniform(1, 35, n_dataset), # cu_roughness,
		np.random.uniform(-3.5, 0, n_dataset), # ti_sld,
		np.random.uniform(10, 100, n_dataset), # ti_thickness,
		np.random.uniform(1, 35, n_dataset), # ti_roughness,
		np.random.uniform(1, 4.2, n_dataset), # oxide_sld,
		np.random.uniform(5, 50, n_dataset), # oxide_thickness,
		np.random.uniform(1, 10, n_dataset), # oxide_roughness,
		]
		)

		# mo_film, 5 layer
		# parameters_ref = np.column_stack(
		# [
		# np.random.uniform(5.0, 7.0, n_dataset), # electolyte_sld,
		# np.random.uniform(-1.0, 7.0, n_dataset), # electolyte_sld, I changed the lower bound to -1 to include air and H2O
		# np.random.uniform(5, 120, n_dataset), # electolyte_roughness,

		# np.random.uniform(-5.0, 6.5, n_dataset), # sei_sld,
		# np.random.uniform(10, 500, n_dataset), # sei_thickness,
		# np.random.uniform(1, 80, n_dataset), # sei_roughness,

		# np.random.uniform(-2, 6, n_dataset), # material_sld,
		# np.random.uniform(-2, 7, n_dataset), # bulk_3_sld,
		# np.random.uniform(10, 200, n_dataset), # material_thickness,
		# np.random.uniform(1, 35, n_dataset), # material_roughness,
		# np.random.uniform(1, 55, n_dataset), # material_roughness,

		# np.random.uniform(6, 7, n_dataset), # cu_sld,
		# np.random.uniform(2, 6, n_dataset), # bulk_2_sld,
		# np.random.uniform(20, 700, n_dataset), # cu_thickness,
		# np.random.uniform(1, 35, n_dataset), # cu_roughness,

		# np.random.uniform(-3.5, 0, n_dataset), # ti_sld,
		# np.random.uniform(10, 100, n_dataset), # ti_thickness,
		# np.random.uniform(1, 35, n_dataset), # ti_roughness,
		# np.random.uniform(1, 55, n_dataset), # cu_roughness,

		# np.random.uniform(1, 4.2, n_dataset), # oxide_sld,
		# np.random.uniform(5, 50, n_dataset), # oxide_thickness,
		@@ -112,53 +131,29 @@ def generate_data(
		# ]
		# )

		# mo_film, 5 layer
		# xx_film, 6 layer model
		# parameters_ref = np.column_stack(
		# [
		# np.random.uniform(-1.0, 7.0, n_dataset), # electolyte_sld, I changed the lower bound to -1 to include air and H2O
		# np.random.uniform(-1.0, 7.0, n_dataset), # electolyte_sld, including air and H2O
		# np.random.uniform(5, 120, n_dataset), # electolyte_roughness,

		# np.random.uniform(-5.0, 6.5, n_dataset), # sei_sld,
		# np.random.uniform(10, 500, n_dataset), # sei_thickness,
		# np.random.uniform(1, 80, n_dataset), # sei_roughness,

		# np.random.uniform(-2, 7, n_dataset), # bulk_3_sld,
		# np.random.uniform(10, 200, n_dataset), # material_thickness,
		# np.random.uniform(1, 55, n_dataset), # material_roughness,

		# np.random.uniform(2, 6, n_dataset), # bulk_2_sld,
		# np.random.uniform(20, 700, n_dataset), # cu_thickness,
		# np.random.uniform(1, 55, n_dataset), # cu_roughness,

		# np.random.uniform(10, 200, n_dataset), # bulk_3_thickness,
		# np.random.uniform(1, 55, n_dataset), # bulk_3_roughness,
		# np.random.uniform(2, 7, n_dataset), # bulk_2_sld,
		# np.random.uniform(20, 700, n_dataset), # bulk_2_thickness (cu_thickness),
		# np.random.uniform(1, 55, n_dataset), # bulk_2_roughness (cu_roughness),
		# np.random.uniform(-3.5, 7, n_dataset), # bulk_1_sld,
		# np.random.uniform(10, 200, n_dataset), # bulk_1_thickness,
		# np.random.uniform(1, 55, n_dataset), # bulk_1_roughness,
		# np.random.uniform(1, 4.2, n_dataset), # oxide_sld,
		# np.random.uniform(5, 50, n_dataset), # oxide_thickness,
		# np.random.uniform(1, 10, n_dataset), # oxide_roughness,
		# ]
		# )

		# xx_film, 6 layer model
		parameters_ref = np.column_stack(
		[
		np.random.uniform(-1.0, 7.0, n_dataset), # electolyte_sld, including air and H2O
		np.random.uniform(5, 120, n_dataset), # electolyte_roughness,
		np.random.uniform(-5.0, 6.5, n_dataset), # sei_sld,
		np.random.uniform(10, 500, n_dataset), # sei_thickness,
		np.random.uniform(1, 80, n_dataset), # sei_roughness,
		np.random.uniform(-2, 7, n_dataset), # bulk_3_sld,
		np.random.uniform(10, 200, n_dataset), # bulk_3_thickness,
		np.random.uniform(1, 55, n_dataset), # bulk_3_roughness,
		np.random.uniform(2, 7, n_dataset), # bulk_2_sld,
		np.random.uniform(20, 700, n_dataset), # bulk_2_thickness (cu_thickness),
		np.random.uniform(1, 55, n_dataset), # bulk_2_roughness (cu_roughness),
		np.random.uniform(-3.5, 7, n_dataset), # bulk_1_sld,
		np.random.uniform(10, 200, n_dataset), # bulk_1_thickness,
		np.random.uniform(1, 55, n_dataset), # bulk_1_roughness,
		np.random.uniform(1, 4.2, n_dataset), # oxide_sld,
		np.random.uniform(5, 50, n_dataset), # oxide_thickness,
		np.random.uniform(1, 10, n_dataset), # oxide_roughness,
		]
		)

		# generate the reference rcurves
		r_curves = np.apply_along_axis(param_to_rcurve, 1, parameters_ref)
		logger.debug(f"r_curves.shape: {r_curves.shape}")