Tuneando funciones en R

Tuneando funciones en R Hace un rato necesitaba poder remuestrear filas de un data frame para testear un modelo (vendrá un próximo post al respecto, esperemos...). El tema es que necesitaba que el muestreo cumpliera tres condiciones:
  1. respetara los tamaños de los estratos (el muestreo original había sido estratificado)
  2. que al interior de cada estrato se pudiera tener la opción de samplear con resposición
  3. que se pudieran elegir la proporción del tamaño de remuestra en cada uno de los estratos
Estuve un rato pensando cómo construir una función que lo hiciera... y en ese interín se me ocurrió que quizás sería mejor usar una función que ya hiciera algo parecido y evaluar si era posible modificarla un poco para que se adaptara a esos tres requisitos.
En un búsqueda de Google me acordé del paquete "survey" para R (gran paquete para trabajar con muestras complejas, tiene algoritmos de raking, post-estratificación, calibración, etc.; acá, hay un tutorial rápido). El paquete tiene una función -stratsample- que hace muestreo estratificado y que es relativamente simple. Acá está el script de la función original:
stratsample<-function (strata, counts){
        strata<-as.character(strata)
        n<-length(strata)
        rval<-integer(sum(counts))
        allrows<-1:n
        j<-0
        for (i in 1:length(counts)){
                thisstrat<-names(counts)[i]
                rval[j+(1:counts[i])]<-sample(allrows[strata == thisstrat],
                                                  counts[i])
                j<-j+counts[i]
        }
        rval
}
 
El argumento "strata" es un vector que identfica a las unidades y sus respectivos estratos y el argumento "counts" es un vector con nombres ("named num") con los tamaños deseados de la muestra en cada estrato. La función devuelve un vector de índices dentro del vector "strata" que sirve para identificar los casos que entran en la muestra.

Cambiando dos líneas y agregando dos argumentos en cinco minutos tuve la función que necesitaba:
#### FUNCION PARA BOOSTRAPEO POR ESTRATOS ####
# strata = Vector que identifica los estratos
# counts = Cantidad de casos en cada uno de los estratos
# rep = Argumento lógico: TRUE si el muestreo es con reposición; FALSE si no
# prop = Argumento numérico: indica la proporción de casos que se muestrean al interior del estrato

stratify<-function(strata, counts, rep=FALSE, prop=1){
        strata<-as.character(strata)
        n<-length(strata)
        rval<-integer(sum(counts))
        allrows<-1:n
        j<-0
        for (i in 1:length(counts)){
                c<-counts[i]
                thisstrat<-names(counts)[i]
                rows<-allrows[strata == thisstrat]
                rval[j+(1:round(c*prop,0))]<-sample(rows,round(c*prop,0),replace=rep)
                j<-j+c
        }
        rval
}
"rep" es un argumento lógico: rep=TRUE implica que el muestreo es con reposición: rep=FALSE, que es sin reposición.
"prop" es un argumento númerico que afecta el tamaño (vía la proporción de casos) en cada estrato que se quiere muestrear. Dado que quería generar un test set del 20%, el training set debía ser de 0.8. Cuando prop=1 se muestrean los tamaños definidos en "counts".

Quizás a alguno le sirva.

Reflexión final: abandonémonos al principio de Zipf ("least effort principle") y no tratemos de crear de cero lo que ya está creado. En este caso, ya existía una función programada (muy bien, además) que hacía algo muuuy similar a lo que necesitaba y solamente requería 15 minutos de cabeza. En fin, una de las ventajas de la comunidad abierta.

PS: En algún momento me voy a acodar cómo hacer para que el código quede lindo...
PS 20/05/2016: actualicé la función con las sugerencias de @javierburroni para que quede más bonita.

Comentarios

Entradas populares de este blog

Sobre la multicausalidad

Número efectivo de partidos en elecciones presidenciales 1983-2011

El ruido de las capitales (vol. 1)