[FR] Allow as_draws_dt for data.table

[SteveBronder]

My workflow is pretty data.table based. It would be nice to have a format like draws_df that uses the data.table backend.

I think this would be semi copy-paste heavy wrt draws_df. Though it would let a few function that work over groups be a bit faster by using the data.table backend.

imo I don't care much about performance though, I just like the data.table scheme for parsing data

[mjskay]

Would it make sense to have a backend option for draws_df that applies the draws_df class on top of a data frame? Or would the semantics of data table break something in that case...

[SteveBronder]

Sorry I'm not sure I understand. For reference I'd like to be able to call something like the below which currently errors

# Run pathfinder
fit <- mod$pathfinder(data=data_list, seed=1234, refresh = 0)
# create a set of draws that is a `draws_dt` object.
path_draws_dt = init$draws(format = "dt")
path_draws_dt[, lw := lp__ - lp_approx__]
path_draws_dt[, pareto_weight := pareto_smooth(exp(lw - max(lw)), tail = "right")[["x"]]]
# This line not possible
path_draws_dt[, posterior_draws := 
  weight_draws(.SD, weights = pareto_weight, log = FALSE)]

I think to cast the data.table I'm using to a draws_df like in the below. But it would be nice if there were a draws_dt type that would just operate like a data.table. I think most of the draws_df functions could be reused. And we could write specializations for the functions that are computationally expensive

weight_init_draws_df = path_draws_dt[, posterior::weight_draws(posterior::as_draws_df(.SD), weights = pareto_weight, log = FALSE)]

[mjskay]

Ah, sorry I wasn't clear. The current draws_df class is:

class(as_draws_df(example_draws()))
## [1] "draws_df"   "draws"      "tbl_df"     "tbl"        "data.frame"

i.e., it uses tibble as its base class. However, S3 is not like classical OOP in that there's no reason we cannot also create objects with a class of c("draws_df", "draws", "data.table", "data.frame"). We would just have to make sure existing draws_df code is compatible with data table's slightly different semantics. If that works then I'd suggest rather than a whole new format, something like as_draws_df(..., backend = "data.table") would make sense.

However, if that leads to a bunch of corner cases, then probably a separate draws_dt() format would be better. A class like c("draws_dt", "draws_df", "draws", "data.table", "data.frame") would allow us to take as much advantage of existing code for draws_df but also make it easy to handle corner cases in draws_dt. A reasonable question would also be if functions like mutate_variables() would have reference semantics for the draws_dt format, or create copies like every other format does.

If we went ahead with this I suspect we would only want data.table as a Suggests level dependency.

Curious what @paul-buerkner thinks.

[paul-buerkner]

I would not be opposed to it. After all, one core principle of posterior is that we allow and support different formats because that represents the reality of its users wanting to interact with their draws differently. Why not have a data.table kind format then. And if we can reuse code from draws_df that this is of course great too.

[mjskay]

Cool. I made a quick attempt at implementing this and quickly realized that if we want the format to allow data.table's specialized selection syntax with [, it probably won't be able to share a lot of code with draws_df unless we factor out a bunch of common stuff into helper functions. The problem is that if we make {posterior} "data.table-aware" (which involves putting .datatable.aware = TRUE somewhere in our namespace and which allows us to take advantage of the specialized [ data.table syntax; e.g. x[, mu] instead of x[, "mu"] to select column mu), normal data.frame selections on the data.table break (because, e.g., if you have a variable called columns containing column names you have to do x[, ..columns] not x[, columns]).

The upshot is there may not be much code-sharing with draws_df unless we put work into some refactoring, so it may take a bit of work to implement.

[paul-buerkner]

In other words, we would just need another draws format to make this work? That would be okay for me :-)

[mjskay]

Yeah, exactly :). That comment was just a lament that the implementation will take more than a half hour of work, heh.

[mjskay]

I put the "quick" implementation on the draws_dt branch here: https://github.com/stan-dev/posterior/tree/draws_dt

This version doesn't allow data.table selection semantics, which we presumably want (else why the format). A starting point for such an implementation could be to change this value to TRUE:

posterior/R/as_draws_dt.R

Lines 1 to 4 in 80daa13

	# set this to TRUE so that data.table:::cedta() will treat {posterior} as a
	# data.table-aware package, then re-implement methods as necessary to use
	# data.table-style indexing with `[`
	.datatable.aware = FALSE

And then begin re-implementing stuff ;). I probably don't have the time to do that at the moment since I'd have to go and learn data.table semantics first. Maybe someone already familiar with data.table would be the right person? (@SteveBronder? ;) )

[SteveBronder]

I'm working on a few things in cmdstanr right now but I think eventually I can find time for this