generate-test-data.py 1.4 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546
  1. import numpy as np
  2. import json
  3. import os
  4. import pathlib
  5. from pathlib import Path
  6. # Get the script's directory and create data directory relative to it
  7. script_dir = Path(__file__).parent
  8. data_dir = script_dir / "data"
  9. data_dir.mkdir(parents=True, exist_ok=True)
  10. # Load existing records if any
  11. records_file = script_dir / "records.json"
  12. if os.path.exists(records_file):
  13. with open(records_file) as f:
  14. print(f"Loading records from {records_file}")
  15. records = json.load(f)
  16. else:
  17. records = {}
  18. # Generate test data for each combination
  19. for dimensions in [(10,), (65, 65), (100, 100, 100), (4, 4, 4, 4, 4)]:
  20. for dtype in ["int8", "int16", "int64", "float16", "float32", "float64"]:
  21. name = f"./data/{'x'.join(str(i) for i in dimensions)}-{dtype}"
  22. # Skip if file already exists
  23. if name in records:
  24. continue
  25. data = np.random.randint(0, 255, dimensions).astype(dtype)
  26. # Store the last 5 values consistently for all types
  27. records[name] = data.ravel()[-5:].tolist()
  28. # Save file using the correct path
  29. file_path = script_dir / name.lstrip("./")
  30. file_path.parent.mkdir(parents=True, exist_ok=True)
  31. np.save(file_path, data)
  32. # Save records in a pretty, sorted format
  33. with open(records_file, 'w') as f:
  34. json.dump(
  35. records,
  36. f,
  37. indent=4,
  38. sort_keys=True
  39. )