basic.py

import numpy as np
import os
import time
from paper import PaperMatrix, Plan, EagerNode
from paper.config import TILE_SIZE

# --- Configuration ---
DATA_DIR = "data" # Directory to store large matrix files

def create_random_matrix(filepath, shape):
    """Creates and saves a large matrix with random data, tile by tile."""
    print(f"Creating random matrix at '{filepath}' with shape {shape}...")
    
    # Create a new file for writing
    matrix = PaperMatrix(filepath, shape, mode='w+')
    
    # Iterate through the matrix in blocks and fill with random data
    for r_start in range(0, shape[0], TILE_SIZE):
        r_end = min(r_start + TILE_SIZE, shape[0])
        for c_start in range(0, shape[1], TILE_SIZE):
            c_end = min(c_start + TILE_SIZE, shape[1])
            tile_shape = (r_end - r_start, c_end - c_start)
            # Generate a small in-memory tile
            random_tile = np.random.rand(*tile_shape).astype(matrix.dtype)
            # Write the tile to the memory-mapped file
            matrix.data[r_start:r_end, c_start:c_end] = random_tile
            
    matrix.data.flush() # Ensure all changes are written to disk
    print("Creation complete.")
    matrix.close()

def main():
    # Create the data directory if it doesn't exist
    if not os.path.exists(DATA_DIR):
        os.makedirs(DATA_DIR)

    # Define the shape of our first test matrix
    shape_A = (3500, 4200)
    shape_C = (4200, 4000)
    path_A = os.path.join(DATA_DIR, "A.bin")
    path_B = os.path.join(DATA_DIR, "B_add.bin")
    path_C = os.path.join(DATA_DIR, "C_mul.bin")

    # Create the matrix file if it doesn't already exist
    for path, shape in {path_A: shape_A, path_B: shape_A, path_C: shape_C}.items():
        if not os.path.exists(path):
            writer = PaperMatrix(path, shape, mode='w+')
            writer.close()
    
    # lazy execution logic
    print("\n--- Build and execute a lazy addition plan ---")

    A_handle = PaperMatrix(path_A, shape_A, mode='r')
    B_handle = PaperMatrix(path_B, shape_A, mode='r')
    C_handle = PaperMatrix(path_C, shape_C, mode='r')
    A_lazy = Plan(EagerNode(A_handle))
    B_lazy = Plan(EagerNode(B_handle))
    C_lazy = Plan(EagerNode(C_handle))

    # 1. Build the plan using the '+' operator
    #   This calls our __add__ method.
    plan = (A_lazy) * 2
    plan2 = A_lazy @ C_lazy
    print(f"\n plan built: '{plan!r}'")

    # 2. Execute the plan using .compute()
    result_file_path = os.path.join(DATA_DIR, "L_lazy_add.bin")
    start_time = time.time()
    result_matrix = plan.compute(result_file_path)
    end_time = time.time()

    print(f"\nExecution finished in {end_time - start_time:.2f} seconds.")

    result_file_path = os.path.join(DATA_DIR, "M_lazy_add.bin")
    start_time = time.time()
    result_matrix2 = plan.compute(result_file_path)
    end_time = time.time()

    # 3. Clean up the file handle
    A_handle.close()
    B_handle.close()
    C_handle.close()
    result_matrix.close()
    result_matrix2.close()


if __name__ == '__main__':
    main()