Coverage for biobb_pytorch / test / unitests / test_mdae / test_mdfeaturizer.py: 100%

1# type: ignore 

2from biobb_common.tools import test_fixtures as fx 

3from biobb_pytorch.mdae.mdfeaturizer import MDFeaturizer 

4import torch 

5 

6 

7class TestMDFeaturizer: 

8 def setup_class(self): 

9 fx.test_setup(self, 'MDFeaturizer') 

10 

11 def teardown_class(self): 

12 fx.test_teardown(self) 

13 

14 def test_MDFeaturizer(self): 

15 MDFeaturizer(properties=self.properties, **self.paths) 

16 assert fx.not_empty(self.paths['output_dataset_pt_path']) 

17 assert fx.not_empty(self.paths['output_stats_pt_path']) 

18 

19 dataset = torch.load(self.paths['output_dataset_pt_path'], weights_only=False) 

20 ref_dataset = torch.load(self.paths['ref_output_dataset_pt_path'], weights_only=False) 

21 assert (dataset['data'] == ref_dataset['data']).all(), "Datasets are not equal!" 

22 

23 stats = torch.load(self.paths['output_stats_pt_path'], weights_only=False) 

24 ref_stats = torch.load(self.paths['ref_output_stats_pt_path'], weights_only=False) 

25 for key in stats.keys(): 

26 if key == 'topology': 

27 continue 

28 if isinstance(stats[key], torch.Tensor): 

29 assert (stats[key] == ref_stats[key]).all(), f"Stats for {key} are not equal!" 

30 else: 

31 assert stats[key] == ref_stats[key], f"Stats for {key} are not equal!"