{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import pandas as pd\nfrom sklearn.model_selection import train_test_split\n\nimport wandb\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport pytorch_lightning as pl\nfrom pytorch_lightning.loggers import WandbLogger\nfrom torch.utils.data import Dataset, DataLoader","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import numpy as np\n\ndef normalize_targets(y):\n    y_min = np.array([0.0, 0.0], dtype=np.float32)  # azimuth: [0, 2*pi], zenith: [0, pi]\n    y_max = np.array([2 * np.pi, np.pi], dtype=np.float32)\n    return (y - y_min) / (y_max - y_min)\n\ndef denormalize_targets(y_norm):\n    y_min = np.array([0.0, 0.0], dtype=np.float32)\n    y_max = np.array([2 * np.pi, np.pi], dtype=np.float32)\n    return y_norm * (y_max - y_min) + y_min\n","metadata":{"ExecuteTime":{"end_time":"2023-04-17T01:21:55.010725Z","start_time":"2023-04-17T01:21:54.943337Z"},"collapsed":false,"jupyter":{"outputs_hidden":false}},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def angular_dist_score(az_true, zen_true, az_pred, zen_pred):\n    '''\n    calculate the MAE of the angular distance between two directions.\n    The two vectors are first converted to cartesian unit vectors,\n    and then their scalar product is computed, which is equal to\n    the cosine of the angle between the two vectors. The inverse\n    cosine (arccos) thereof is then the angle between the two input vectors\n\n    Parameters:\n    -----------\n\n    az_true : float (or array thereof)\n        true azimuth value(s) in radian\n    zen_true : float (or array thereof)\n        true zenith value(s) in radian\n    az_pred : float (or array thereof)\n        predicted azimuth value(s) in radian\n    zen_pred : float (or array thereof)\n        predicted zenith value(s) in radian\n\n    Returns:\n    --------\n\n    dist : float\n        mean over the angular distance(s) in radian\n    '''\n\n    if not (np.all(np.isfinite(az_true)) and\n            np.all(np.isfinite(zen_true)) and\n            np.all(np.isfinite(az_pred)) and\n            np.all(np.isfinite(zen_pred))):\n        raise ValueError(\"All arguments must be finite\")\n\n    # pre-compute all sine and cosine values\n    sa1 = np.sin(az_true)\n    ca1 = np.cos(az_true)\n    sz1 = np.sin(zen_true)\n    cz1 = np.cos(zen_true)\n\n    sa2 = np.sin(az_pred)\n    ca2 = np.cos(az_pred)\n    sz2 = np.sin(zen_pred)\n    cz2 = np.cos(zen_pred)\n\n    # scalar product of the two cartesian vectors (x = sz*ca, y = sz*sa, z = cz)\n    scalar_prod = sz1*sz2*(ca1*ca2 + sa1*sa2) + (cz1*cz2)\n\n    # scalar product of two unit vectors is always between -1 and 1, this is against nummerical instability\n    # that might otherwise occure from the finite precision of the sine and cosine functions\n    scalar_prod =  np.clip(scalar_prod, -1, 1)\n\n    # convert back to an angle (in radian)\n    return np.average(np.abs(np.arccos(scalar_prod)))","metadata":{"ExecuteTime":{"end_time":"2023-04-17T01:21:55.010926Z","start_time":"2023-04-17T01:21:54.995186Z"},"collapsed":false,"jupyter":{"outputs_hidden":false}},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class NeutrinoModel(pl.LightningModule):\n    def __init__(self, input_dims, num_layers=2, hidden_units=None, learning_rate=0.001):\n        super().__init__()\n        self.save_hyperparameters()\n\n        layers = []\n        layers.append(nn.BatchNorm1d(input_dims))\n\n        if hidden_units is None:\n            hidden_units = [input_dims*2] * (num_layers - 1)\n\n        for i in range(num_layers):\n            in_dim = input_dims if i == 0 else hidden_units\n            out_dim = 2 if i == num_layers - 1 else hidden_units\n            layers.append(nn.Linear(in_dim, out_dim))\n\n            if i < num_layers - 1:\n                layers.append(nn.ReLU())\n            else:\n                layers.append(nn.Sigmoid())\n\n        self.layers = nn.Sequential(*layers)\n\n    def forward(self, x):\n        return self.layers(x)\n\n    def training_step(self, batch, batch_idx):\n        x, y = batch\n        y_hat = self(x)\n        loss = F.l1_loss(y_hat, y, reduction='mean')\n        self.log(\"train_loss\", loss, on_step=True, on_epoch=True, prog_bar=True, logger=True)\n        return loss\n\n    def validation_step(self, batch, batch_idx):\n        x, y_norm = batch\n        y_hat_norm = self(x)\n        loss = F.l1_loss(y_hat_norm, y_norm, reduction='mean')\n\n        # Denormalize target values and predictions\n        y = denormalize_targets(y_norm.cpu().numpy())\n        y_hat = denormalize_targets(y_hat_norm.detach().cpu().numpy())\n\n        # Calculate angular distance score\n        az_true, zen_true = y[:, 0], y[:, 1]\n        az_pred, zen_pred = y_hat[:, 0], y_hat[:, 1]\n        ang_dist = angular_dist_score(az_true, zen_true, az_pred, zen_pred)\n\n        self.log(\"val_loss\", loss, on_step=False, on_epoch=True, prog_bar=True, logger=True)\n        self.log(\"angular_dist_score\", ang_dist, on_step=False, on_epoch=True, prog_bar=True, logger=True)\n\n        return loss\n\n    def configure_optimizers(self):\n        return torch.optim.Adam(self.parameters(), lr=self.hparams.learning_rate)\n","metadata":{"ExecuteTime":{"end_time":"2023-04-17T01:21:56.742130Z","start_time":"2023-04-17T01:21:55.001102Z"},"collapsed":true,"jupyter":{"outputs_hidden":true}},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Dataset class\nclass NeutrinoDataset(Dataset):\n    def __init__(self, data):\n        self.features = data[:, :-2].astype(np.float32)\n        self.targets = data[:, -2:].astype(np.float32)\n\n    def __len__(self):\n        return len(self.features)\n\n    def __getitem__(self, idx):\n        return self.features[idx], self.targets[idx]\n\n# DataLoader function\ndef create_data_loaders(train_data, val_data, batch_size=32):\n    train_dataset = NeutrinoDataset(train_data)\n    val_dataset = NeutrinoDataset(val_data)\n\n    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)\n    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)\n\n    return train_loader, val_loader\n\n# Load your dataset\ndata_path = \"features.csv\"\ndata = pd.read_csv(data_path).values\n\n# Split data into train and validation sets\ntrain_data, val_data = train_test_split(data, test_size=0.2, random_state=42)\n\n# Create data loaders\ntrain_loader, val_loader = create_data_loaders(train_data, val_data)\n","metadata":{"ExecuteTime":{"end_time":"2023-04-17T01:22:00.053132Z","start_time":"2023-04-17T01:21:56.740989Z"},"collapsed":false,"jupyter":{"outputs_hidden":false}},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"\n# Initialize Wandb\n# Define hyperparameter search space\nsweep_config = {\n    \"name\": \"neutrino-hyperparam-sweep\",\n    \"method\": \"random\",\n    \"metric\": {\n        \"name\": \"val_loss\",\n        \"goal\": \"minimize\"\n    },\n    \"parameters\": {\n        \"num_layers\": {\n            \"values\": [2, 3, 4, 5]\n        },\n        \"hidden_units\": {\n            \"values\": [32, 64, 128, 256, 512]\n        },\n        \"learning_rate\": {\n            \"values\": [1e-3, 1e-4, 1e-5, 1e-6]\n        }\n    }\n}\n\n# Create Wandb sweep\nsweep_id = wandb.sweep(sweep=sweep_config, project=\"neutrino\")\n\ndef train():\n    with wandb.init():\n        config = wandb.config\n        model = NeutrinoModel(\n            input_dims=train_data.shape[1] - 2,\n            num_layers=config.num_layers,\n            hidden_units=config.hidden_units,\n            learning_rate=config.learning_rate\n        )\n\n        # Set up Wandb logger for PyTorch Lightning\n        wandb_logger = WandbLogger()\n\n        # Set up PyTorch Lightning Trainer\n        trainer = pl.Trainer(logger=wandb_logger, max_epochs=3, accelerator='mps')\n\n        # Train the model\n        trainer.fit(model, train_loader, val_loader)\n\n# Run the sweep\nwandb.agent(sweep_id, function=train, count=100)\n","metadata":{"ExecuteTime":{"end_time":"2023-04-17T01:22:00.074396Z","start_time":"2023-04-17T01:22:00.069350Z"},"collapsed":false,"jupyter":{"outputs_hidden":false}},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"<img src=\"https://live.staticflickr.com/65535/52837021368_0ca4b8f87b_b.jpg\" width=\"1024\" height=\"435\">","metadata":{}},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}