Untitled

    parser.add_argument('--hidden_dim', type=eval, default=128, help='Hidden dimension(s)')
    parser.add_argument('--basis_dim', type=int, default=256, help='Basis dimension')
    parser.add_argument('--basis_hidden_dim', type=int, default=128, help='Hidden dimension of the basis function MLP')

    parser.add_argument('--layers', type=eval, default=[0,2,2,2,2], help='Layers per scale')
    parser.add_argument('--orientations', type=int, default=8, help='Number of orientations')
    parser.add_argument('--degree', type=int, default=2, help='Polynomial degree')
    parser.add_argument('--edge_types', type=eval, default=["knn-16", "knn-16", "knn-16", "knn-16", "fc"], help='Edge types')

    parser.add_argument('--ratios', type=eval, default=[0.25, 0.25, 0.25, 0.25], help='Pooling ratios')

    parser.add_argument('--widening', type=int, default=4, help='Network widening factor')
    parser.add_argument('--layer_scale', type=eval, default=None, help='Layer scaling factor')
    parser.add_argument('--equivariance', type=str, default="SEn", help='Type of equivariance')
    parser.add_argument('--skip_connections', type=eval, default=False, help='Use U-Net style skip connections')

    # EDM parameters
    parser.add_argument('--S_churn', type=float, default=0., help='Noise level for EDM sampling')
    parser.add_argument('--sigma_max', type=float, default=5.0, help='Maximum sigma for EDM')
    parser.add_argument('--sigma_min', type=float, default=0.002, help='Minimum sigma for EDM')
    parser.add_argument('--num_steps', type=int, default=50, help='Number of diffusion steps')
    parser.add_argument('--sigma_data', type=float, default=1.0, help='Data sigma for EDM')
    parser.add_argument('--rho', type=int, default=7, help='Rho parameter for EDM')

    # Feature configuration
    parser.add_argument('--use_charges', type=eval, default=False, help='Use atomic charges')
    parser.add_argument('--scalar_features', type=eval, default=["coords"], help='Features to use as scalars: ["coords"]')
    
    parser.add_argument('--vector_features', type=eval, default=[], help='Features to use as vectors: ["coords"]')
    
# Training features
    parser.add_argument('--train_augm', type=eval, default=True, help='Use rotation augmentation during training')
    parser.add_argument('--validation_frequency', type=int, default=10, help='Validation frequency in epochs')