mo-seph
mo-seph

Reputation: 6223

Extract values from a transition matrix using columns of a data.frame in R

I have a transition matrix, with the cost of going from one state to another, e.g.

cost <- data.frame( a=c("aa","ab"),b=c("ba","bb"))

(pretending that the string "aa" is the cost of moving from a to a)

I've got a data.frame, with states in:

transitions <- data.frame( from=c("a","a","b"), to=c("a","b","b") )

I'd like to be able to add a column to transitions, with the cost of each transition in, so it ends up being:

  from to cost
1    a  a   aa
2    a  b   ab
3    b  b   bb

I'm sure there is an R-ish way to do this. I've ended up using a for loop:

n <- dim(data)[1]
v <- vector("numeric",n)
for( i in 1:n ) 
{ 
    z<-data[i,c(col1,col2),with=FALSE]
    za <- z[[col1]]
    zb <- z[[col2]]
    v[i] <- dist[za,zb]
}
data <- cbind(data,d=v)
names(data)[dim(data)[2]] <- colName
data

But this feels pretty ugly, and it's incredibly slow - it takes about 20 minutes on a 2M row data.frame (and an operation to compute distances between elements of the same table takes less than a second).

Is there a simple, fast, one or two line command that would get me the cost column above?

Upvotes: 2

Views: 1137

Answers (3)

redmode
redmode

Reputation: 4941

UPDATE: Consider known states

data.table solution:

require(utils)
require(data.table)

## Data generation
N <- 2e6
set.seed(1)
states <- c("a","b")
cost <- data.frame(a=c("aa","ab"),b=c("ba","bb"))
transitions <- data.frame(from=sample(states, N, replace=T), 
                            to=sample(states, N, replace=T))

## Expanded cost matrix construction
f <- expand.grid(states, states)
f <- f[order(f$Var1, f$Var2),]
f$cost <- unlist(cost)

## Prepare data.table
dt <- data.table(transitions)
setkey(dt, from, to)

## Routine itself  
dt[,cost:=as.character("")] # You don't need this line if cost is numeric
apply(f, 1, function(x) dt[J(x[1],x[2]),cost:=x[3]])

With 2M rows in transitions it takes about 0.3sec to proceed.

Upvotes: 3

mo-seph
mo-seph

Reputation: 6223

Starting from Arun's answer, I went with:

library(reshape)
cost <- data.frame( a = c("aa","ab"), b = c("ba","bb") )
transitions <- data.frame(from = c("a","a","b"), to = c("a","b","b") )
row.names(cost) <- c("a","b") #Normally get this from the csv file
cost$from <- row.names(cost)
m <- melt(cost, id.vars=c("from"))
m$transition = paste(m$from,m$variable)
transitions$transition=paste(transitions$from,transitions$to)
merge(m, transitions, by.x="transition",by.y="transition")

It's a few more lines, but I'm a bit untrusting of factor orderings as indexes. It also means that when they are data.tables, I can do:

setkey(m,transition)
setkey(transitions,transition)
m[transitions]

I haven't benchmarked, but on large datasets, I'm pretty confident the data.table merge will be faster than the merge or vector scan approaches.

Upvotes: 2

Arun
Arun

Reputation: 118869

Here's one way: (At least this works on this example and I believe it'll work on larger data as well. Please write back with an example if it doesn't)

# load both cost and transition with stringsAsFactors = FALSE
# so that strings are NOT by default loaded as factors
cost <- data.frame( a = c("aa","ab"), b = c("ba","bb"), stringsAsFactors=F)
transitions <- data.frame(from = c("a","a","b"), to = c("a","b","b"), 
                                      stringsAsFactors = FALSE)

# convert cost to vector: it'll have names a1, a2, b1, b2. we'll exploit that.
cost.vec <- unlist(cost)
# convert "to" to factor and create id column with "from" and as.integer(to)
# the as.integer(to) will convert it into its levels
transitions$to <- as.factor(transitions$to)
transitions$id <- paste0(transitions$from, as.integer(transitions$to))

# now, you'll have a1, a2 etc.. here as well, just match it with the vector
transitions$val <- cost.vec[!is.na(match(names(cost.vec), transitions$id))]

#   from to id val
# 1    a  a a1  aa
# 2    a  b a2  ab
# 3    b  b b2  bb

You can of course remove the id. If this wouldn't work in any case, let me know. I'll try to fix it.

Upvotes: 2

Related Questions