Hadley Talk Demo from JSM 2025 using scorcher

hadley-talk.R

######
### Load libraries - you'll need to install these!
#####

library(tidyverse)
library(torch)
library(scorcher)


######
### Visualize Hadley
#####


dat = read_csv("~/Downloads/hadley_widespaced_scatter.csv") |>
  sample_n(500)
dat |>  ggplot() + 
  geom_point(aes(x=x,y=y)) + 
  theme_void()


#####
## Make a dataloader
#####

input = output = torch_tensor(dat |> as.matrix(), dtype = torch_float())

dl = scorch_create_dataloader(input, output, batch_size = 32)



#####
## Set up a simple model
#####

scorch_model <- dl |> 
  initiate_scorch() |> 
  scorch_layer("linear", 128 * 3, 128) |>
  scorch_layer("gelu") |>
  scorch_layer(c("linear", "gelu"), 128, 128, use_residual = TRUE) |>
  scorch_layer(c("linear", "gelu"), 128, 128, use_residual = TRUE) |>
  scorch_layer("linear", 128, 2)


#####
## Compile the model
#####

compiled_scorch_model <- scorch_model |>
  compile_scorch(init_fn = scorch_2d_diffusion_init, 
                 forward_fn = scorch_2d_diffusion_forward, scale = 1000)


#####
## Set up a noise scheduler
#####

noise_scheduler <- NoiseScheduler$new(
  num_timesteps = 50L,
  beta_schedule = "linear"
)


#####
## Train the model (this takes a while, don't run live, Jeff!)
#####


fitted_scorch_model = compiled_scorch_model |> 
  fit_scorch(loss = nn_mse_loss, optim = optim_adamw, num_epochs = 500,
             verbose = F, preprocess_fn = scorch_2d_diffusion_train, 
             clip_grad = "norm", clip_params = list(max_norm = 1),
             noise_scheduler = noise_scheduler)



#####
## Generate a new Hadley
####

fitted_scorch_model$eval()

sample = torch_randn(1000L, 2)

sample |>
  as_array() |> 
  as_data_frame() |>
  ggplot() + 
  geom_point(aes(x=V1,y=V2)) + 
  theme_void()

timesteps = rev(seq_len(length(noise_scheduler)))

for (t in timesteps) {
  t_tensor = torch_tensor(rep(t, 1000L), dtype = torch_long())
  with_no_grad({
    residual = fitted_scorch_model(sample, t_tensor)
  })
  sample = noise_scheduler$step(residual, t, sample)
}

sample |>
  as_array() |> 
  as_data_frame() |>
  ggplot() + 
  geom_point(aes(x=V1,y=V2)) + 
  theme_void()

hadley_widespaced_scatter.csv