Skip to content

Commit

Permalink
Update test_jit.py
Browse files Browse the repository at this point in the history
  • Loading branch information
@iProzd
iProzd committed Oct 23, 2024
1 parent 332d7d9 commit 35a3acc
Showing 1 changed file with 95 additions and 97 deletions.
192 changes: 95 additions & 97 deletions source/tests/pt/model/test_jit.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -47,101 +47,101 @@ def tearDown(self):
os.remove(f)


# class TestEnergyModelSeA(unittest.TestCase, JITTest):
# def setUp(self):
# input_json = str(Path(__file__).parent / "water/se_atten.json")
# with open(input_json) as f:
# self.config = json.load(f)
# data_file = [str(Path(__file__).parent / "water/data/data_0")]
# self.config["training"]["training_data"]["systems"] = data_file
# self.config["training"]["validation_data"]["systems"] = data_file
# self.config["model"] = deepcopy(model_se_e2_a)
# self.config["training"]["numb_steps"] = 10
# self.config["training"]["save_freq"] = 10
#
# def tearDown(self):
# JITTest.tearDown(self)
#
#
# class TestDOSModelSeA(unittest.TestCase, JITTest):
# def setUp(self):
# input_json = str(Path(__file__).parent.parent / "dos/input.json")
# with open(input_json) as f:
# self.config = json.load(f)
# data_file = [str(Path(__file__).parent.parent / "dos/data/global_system")]
# self.config["training"]["training_data"]["systems"] = data_file
# self.config["training"]["validation_data"]["systems"] = data_file
# self.config["model"] = deepcopy(model_dos)
# self.config["training"]["numb_steps"] = 10
# self.config["training"]["save_freq"] = 10
#
# def tearDown(self):
# JITTest.tearDown(self)
#
#
# class TestEnergyModelDPA1(unittest.TestCase, JITTest):
# def setUp(self):
# input_json = str(Path(__file__).parent / "water/se_atten.json")
# with open(input_json) as f:
# self.config = json.load(f)
# data_file = [str(Path(__file__).parent / "water/data/data_0")]
# self.config["training"]["training_data"]["systems"] = data_file
# self.config["training"]["validation_data"]["systems"] = data_file
# self.config["model"] = deepcopy(model_dpa1)
# self.config["training"]["numb_steps"] = 10
# self.config["training"]["save_freq"] = 10
#
# def tearDown(self):
# JITTest.tearDown(self)
#
#
# class TestEnergyModelDPA2(unittest.TestCase, JITTest):
# def setUp(self):
# input_json = str(Path(__file__).parent / "water/se_atten.json")
# with open(input_json) as f:
# self.config = json.load(f)
# data_file = [str(Path(__file__).parent / "water/data/data_0")]
# self.config["training"]["training_data"]["systems"] = data_file
# self.config["training"]["validation_data"]["systems"] = data_file
# self.config["model"] = deepcopy(model_dpa2)
# self.config["training"]["numb_steps"] = 10
# self.config["training"]["save_freq"] = 10
#
# def tearDown(self):
# JITTest.tearDown(self)
#
#
# class TestEnergyModelHybrid(unittest.TestCase, JITTest):
# def setUp(self):
# input_json = str(Path(__file__).parent / "water/se_atten.json")
# with open(input_json) as f:
# self.config = json.load(f)
# data_file = [str(Path(__file__).parent / "water/data/data_0")]
# self.config["training"]["training_data"]["systems"] = data_file
# self.config["training"]["validation_data"]["systems"] = data_file
# self.config["model"] = deepcopy(model_hybrid)
# self.config["training"]["numb_steps"] = 10
# self.config["training"]["save_freq"] = 10
#
# def tearDown(self):
# JITTest.tearDown(self)
#
#
# class TestEnergyModelHybrid2(unittest.TestCase, JITTest):
# def setUp(self):
# input_json = str(Path(__file__).parent / "water/se_atten.json")
# with open(input_json) as f:
# self.config = json.load(f)
# data_file = [str(Path(__file__).parent / "water/data/data_0")]
# self.config["training"]["training_data"]["systems"] = data_file
# self.config["training"]["validation_data"]["systems"] = data_file
# self.config["model"] = deepcopy(model_hybrid)
# # self.config["model"]["descriptor"]["hybrid_mode"] = "sequential"
# self.config["training"]["numb_steps"] = 10
# self.config["training"]["save_freq"] = 10
#
# def tearDown(self):
# JITTest.tearDown(self)
class TestEnergyModelSeA(unittest.TestCase, JITTest):
def setUp(self):
input_json = str(Path(__file__).parent / "water/se_atten.json")
with open(input_json) as f:
self.config = json.load(f)
data_file = [str(Path(__file__).parent / "water/data/data_0")]
self.config["training"]["training_data"]["systems"] = data_file
self.config["training"]["validation_data"]["systems"] = data_file
self.config["model"] = deepcopy(model_se_e2_a)
self.config["training"]["numb_steps"] = 10
self.config["training"]["save_freq"] = 10

def tearDown(self):
JITTest.tearDown(self)


class TestDOSModelSeA(unittest.TestCase, JITTest):
def setUp(self):
input_json = str(Path(__file__).parent.parent / "dos/input.json")
with open(input_json) as f:
self.config = json.load(f)
data_file = [str(Path(__file__).parent.parent / "dos/data/global_system")]
self.config["training"]["training_data"]["systems"] = data_file
self.config["training"]["validation_data"]["systems"] = data_file
self.config["model"] = deepcopy(model_dos)
self.config["training"]["numb_steps"] = 10
self.config["training"]["save_freq"] = 10

def tearDown(self):
JITTest.tearDown(self)


class TestEnergyModelDPA1(unittest.TestCase, JITTest):
def setUp(self):
input_json = str(Path(__file__).parent / "water/se_atten.json")
with open(input_json) as f:
self.config = json.load(f)
data_file = [str(Path(__file__).parent / "water/data/data_0")]
self.config["training"]["training_data"]["systems"] = data_file
self.config["training"]["validation_data"]["systems"] = data_file
self.config["model"] = deepcopy(model_dpa1)
self.config["training"]["numb_steps"] = 10
self.config["training"]["save_freq"] = 10

def tearDown(self):
JITTest.tearDown(self)


class TestEnergyModelDPA2(unittest.TestCase, JITTest):
def setUp(self):
input_json = str(Path(__file__).parent / "water/se_atten.json")
with open(input_json) as f:
self.config = json.load(f)
data_file = [str(Path(__file__).parent / "water/data/data_0")]
self.config["training"]["training_data"]["systems"] = data_file
self.config["training"]["validation_data"]["systems"] = data_file
self.config["model"] = deepcopy(model_dpa2)
self.config["training"]["numb_steps"] = 10
self.config["training"]["save_freq"] = 10

def tearDown(self):
JITTest.tearDown(self)


class TestEnergyModelHybrid(unittest.TestCase, JITTest):
def setUp(self):
input_json = str(Path(__file__).parent / "water/se_atten.json")
with open(input_json) as f:
self.config = json.load(f)
data_file = [str(Path(__file__).parent / "water/data/data_0")]
self.config["training"]["training_data"]["systems"] = data_file
self.config["training"]["validation_data"]["systems"] = data_file
self.config["model"] = deepcopy(model_hybrid)
self.config["training"]["numb_steps"] = 10
self.config["training"]["save_freq"] = 10

def tearDown(self):
JITTest.tearDown(self)


class TestEnergyModelHybrid2(unittest.TestCase, JITTest):
def setUp(self):
input_json = str(Path(__file__).parent / "water/se_atten.json")
with open(input_json) as f:
self.config = json.load(f)
data_file = [str(Path(__file__).parent / "water/data/data_0")]
self.config["training"]["training_data"]["systems"] = data_file
self.config["training"]["validation_data"]["systems"] = data_file
self.config["model"] = deepcopy(model_hybrid)
# self.config["model"]["descriptor"]["hybrid_mode"] = "sequential"
self.config["training"]["numb_steps"] = 10
self.config["training"]["save_freq"] = 10

def tearDown(self):
JITTest.tearDown(self)

class TestEnergyModelDPA2IntRcut(unittest.TestCase, JITTest):
def setUp(self):
Expand All @@ -154,8 +154,6 @@ def setUp(self):
self.config["model"] = deepcopy(model_dpa2)
self.config["model"]["descriptor"]["repinit"]["rcut"] = int(self.config["model"]["descriptor"]["repinit"]["rcut"])
self.config["model"]["descriptor"]["repinit"]["rcut_smth"] = int(self.config["model"]["descriptor"]["repinit"]["rcut_smth"])
# from IPython import embed
# embed()
self.config["training"]["numb_steps"] = 10
self.config["training"]["save_freq"] = 10

Expand Down

0 comments on commit 35a3acc

Please sign in to comment.